From f465040acc44c48267e83e137a3e52d3447e0c74 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Fri, 18 Nov 2022 21:42:10 +0800 Subject: [PATCH] feat: lazy evaluated record batch stream (#573) Signed-off-by: Ruihang Xia --- src/catalog/src/system.rs | 1 - src/catalog/src/tables.rs | 1 - src/common/query/src/physical_plan.rs | 24 +++----- src/common/recordbatch/src/adapter.rs | 79 +++++++++++++++++++++++++++ src/common/recordbatch/src/lib.rs | 14 +++++ src/frontend/src/table.rs | 40 ++++++++------ src/mito/src/engine.rs | 30 ++-------- src/query/src/datafusion.rs | 1 - src/table/src/table/scan.rs | 8 +-- src/table/src/test_util/memtable.rs | 2 - 10 files changed, 135 insertions(+), 65 deletions(-) diff --git a/src/catalog/src/system.rs b/src/catalog/src/system.rs index 09980b7482c..a55f0d06f79 100644 --- a/src/catalog/src/system.rs +++ b/src/catalog/src/system.rs @@ -134,7 +134,6 @@ impl SystemCatalogTable { .context(error::SystemCatalogTableScanSnafu)?; let stream = scan .execute(0, Arc::new(RuntimeEnv::default())) - .await .context(error::SystemCatalogTableScanExecSnafu)?; Ok(stream) } diff --git a/src/catalog/src/tables.rs b/src/catalog/src/tables.rs index 3379943af65..b11fc870dea 100644 --- a/src/catalog/src/tables.rs +++ b/src/catalog/src/tables.rs @@ -368,7 +368,6 @@ mod tests { let tables_stream = tables.scan(&None, &[], None).await.unwrap(); let mut tables_stream = tables_stream .execute(0, Arc::new(RuntimeEnv::default())) - .await .unwrap(); if let Some(t) = tables_stream.next().await { diff --git a/src/common/query/src/physical_plan.rs b/src/common/query/src/physical_plan.rs index fa566b634f6..fae04438972 100644 --- a/src/common/query/src/physical_plan.rs +++ b/src/common/query/src/physical_plan.rs @@ -17,7 +17,7 @@ use std::fmt::Debug; use std::sync::Arc; use async_trait::async_trait; -use common_recordbatch::adapter::{DfRecordBatchStreamAdapter, RecordBatchStreamAdapter}; +use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter}; use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream}; use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef; use datafusion::error::Result as DfResult; @@ -39,7 +39,6 @@ pub type PhysicalPlanRef = Arc; /// creating the actual `async` [`SendableRecordBatchStream`]s /// of [`RecordBatch`] that incrementally compute the operator's /// output from its input partition. -#[async_trait] pub trait PhysicalPlan: Debug + Send + Sync { /// Returns the physical plan as [`Any`](std::any::Any) so that it can be /// downcast to a specific implementation. @@ -61,7 +60,7 @@ pub trait PhysicalPlan: Debug + Send + Sync { fn with_new_children(&self, children: Vec) -> Result; /// Creates an RecordBatch stream. - async fn execute( + fn execute( &self, partition: usize, runtime: Arc, @@ -84,7 +83,6 @@ impl PhysicalPlanAdapter { } } -#[async_trait] impl PhysicalPlan for PhysicalPlanAdapter { fn as_any(&self) -> &dyn Any { self @@ -118,18 +116,15 @@ impl PhysicalPlan for PhysicalPlanAdapter { Ok(Arc::new(PhysicalPlanAdapter::new(self.schema(), plan))) } - async fn execute( + fn execute( &self, partition: usize, runtime: Arc, ) -> Result { - let stream = self - .df_plan - .execute(partition, runtime) - .await - .context(error::DataFusionExecutionPlanSnafu)?; - let stream = RecordBatchStreamAdapter::try_new(stream) - .context(error::ConvertDfRecordBatchStreamSnafu)?; + let df_plan = self.df_plan.clone(); + let stream = Box::pin(async move { df_plan.execute(partition, runtime).await }); + let stream = AsyncRecordBatchStreamAdapter::new(self.schema(), stream); + Ok(Box::pin(stream)) } } @@ -187,7 +182,7 @@ impl DfPhysicalPlan for DfPhysicalPlanAdapter { partition: usize, runtime: Arc, ) -> DfResult { - let stream = self.0.execute(partition, runtime).await?; + let stream = self.0.execute(partition, runtime)?; Ok(Box::pin(DfRecordBatchStreamAdapter::new(stream))) } @@ -250,7 +245,6 @@ mod test { schema: SchemaRef, } - #[async_trait] impl PhysicalPlan for MyExecutionPlan { fn as_any(&self) -> &dyn Any { self @@ -272,7 +266,7 @@ mod test { unimplemented!() } - async fn execute( + fn execute( &self, _partition: usize, _runtime: Arc, diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index 700ba329612..2994d8f0788 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -19,8 +20,10 @@ use std::task::{Context, Poll}; use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef; use datafusion::physical_plan::RecordBatchStream as DfRecordBatchStream; use datafusion_common::record_batch::RecordBatch as DfRecordBatch; +use datafusion_common::DataFusionError; use datatypes::arrow::error::{ArrowError, Result as ArrowResult}; use datatypes::schema::{Schema, SchemaRef}; +use futures::ready; use snafu::ResultExt; use crate::error::{self, Result}; @@ -28,6 +31,14 @@ use crate::{ DfSendableRecordBatchStream, RecordBatch, RecordBatchStream, SendableRecordBatchStream, Stream, }; +type FutureStream = Pin< + Box< + dyn std::future::Future< + Output = std::result::Result, + > + Send, + >, +>; + /// Greptime SendableRecordBatchStream -> DataFusion RecordBatchStream pub struct DfRecordBatchStreamAdapter { stream: SendableRecordBatchStream, @@ -104,3 +115,71 @@ impl Stream for RecordBatchStreamAdapter { self.stream.size_hint() } } + +enum AsyncRecordBatchStreamAdapterState { + Uninit(FutureStream), + Inited(std::result::Result), +} + +pub struct AsyncRecordBatchStreamAdapter { + schema: SchemaRef, + state: AsyncRecordBatchStreamAdapterState, +} + +impl AsyncRecordBatchStreamAdapter { + pub fn new(schema: SchemaRef, stream: FutureStream) -> Self { + Self { + schema, + state: AsyncRecordBatchStreamAdapterState::Uninit(stream), + } + } +} + +impl RecordBatchStream for AsyncRecordBatchStreamAdapter { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Stream for AsyncRecordBatchStreamAdapter { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + AsyncRecordBatchStreamAdapterState::Uninit(stream_future) => { + self.state = AsyncRecordBatchStreamAdapterState::Inited(ready!(Pin::new( + stream_future + ) + .poll(cx))); + continue; + } + AsyncRecordBatchStreamAdapterState::Inited(stream) => match stream { + Ok(stream) => { + return Poll::Ready(ready!(Pin::new(stream).poll_next(cx)).map(|df| { + Ok(RecordBatch { + schema: self.schema(), + df_recordbatch: df.context(error::PollStreamSnafu)?, + }) + })); + } + Err(e) => { + return Poll::Ready(Some( + error::CreateRecordBatchesSnafu { + reason: format!("Read error {:?} from stream", e), + } + .fail() + .map_err(|e| e.into()), + )) + } + }, + } + } + } + + // This is not supported for lazy stream. + #[inline] + fn size_hint(&self) -> (usize, Option) { + (0, None) + } +} diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 71bf32bb5d4..ce2c2f1e5a9 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -21,6 +21,7 @@ use std::pin::Pin; use std::sync::Arc; use datafusion::arrow_print; +use datafusion::physical_plan::memory::MemoryStream; pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream; use datatypes::prelude::VectorRef; use datatypes::schema::{Schema, SchemaRef}; @@ -132,6 +133,19 @@ impl RecordBatches { index: 0, }) } + + pub fn into_df_stream(self) -> DfSendableRecordBatchStream { + let df_record_batches = self + .batches + .into_iter() + .map(|batch| batch.df_recordbatch) + .collect(); + // unwrap safety: `MemoryStream::try_new` won't fail + Box::pin( + MemoryStream::try_new(df_record_batches, self.schema.arrow_schema().clone(), None) + .unwrap(), + ) + } } pub struct SimpleRecordBatchStream { diff --git a/src/frontend/src/table.rs b/src/frontend/src/table.rs index 8499972cb44..ea4d70e3874 100644 --- a/src/frontend/src/table.rs +++ b/src/frontend/src/table.rs @@ -20,14 +20,17 @@ use std::sync::Arc; use async_trait::async_trait; use client::Database; -use common_error::prelude::BoxedError; use common_query::error::Result as QueryResult; use common_query::logical_plan::Expr; use common_query::physical_plan::{PhysicalPlan, PhysicalPlanRef}; +use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::logical_plan::Expr as DfExpr; -use datafusion::physical_plan::Partitioning; +use datafusion::physical_plan::{ + Partitioning, SendableRecordBatchStream as DfSendableRecordBatchStream, +}; +use datafusion_common::DataFusionError; use datatypes::prelude::Value; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use meta_client::rpc::{Peer, TableName}; @@ -110,14 +113,14 @@ impl Table for DistTable { let datanode_instance = DatanodeInstance::new(Arc::new(self.clone()) as _, db); // TODO(LFC): Pass in "regions" when Datanode supports multi regions for a table. - partition_execs.push(PartitionExec { + partition_execs.push(Arc::new(PartitionExec { table_name: self.table_name.clone(), datanode_instance, projection: projection.clone(), filters: filters.to_vec(), limit, batches: Arc::new(RwLock::new(None)), - }) + })); } let dist_scan = DistTableScan { @@ -385,10 +388,9 @@ fn reverse_operator(op: &Operator) -> Operator { #[derive(Debug)] struct DistTableScan { schema: SchemaRef, - partition_execs: Vec, + partition_execs: Vec>, } -#[async_trait] impl PhysicalPlan for DistTableScan { fn as_any(&self) -> &dyn Any { self @@ -410,14 +412,20 @@ impl PhysicalPlan for DistTableScan { unimplemented!() } - async fn execute( + fn execute( &self, partition: usize, _runtime: Arc, ) -> QueryResult { - let exec = &self.partition_execs[partition]; - exec.maybe_init().await.map_err(BoxedError::new)?; - Ok(exec.as_stream().await) + let exec = self.partition_execs[partition].clone(); + let stream = Box::pin(async move { + exec.maybe_init() + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + exec.as_stream().await + }); + let stream = AsyncRecordBatchStreamAdapter::new(self.schema(), stream); + Ok(Box::pin(stream)) } } @@ -453,12 +461,13 @@ impl PartitionExec { Ok(()) } - async fn as_stream(&self) -> SendableRecordBatchStream { - let batches = self.batches.read().await; - batches - .as_ref() + /// Notice: the record batch will be consumed. + async fn as_stream(&self) -> std::result::Result { + let mut batches = self.batches.write().await; + Ok(batches + .take() .expect("should have been initialized in \"maybe_init\"") - .as_stream() + .into_df_stream()) } } @@ -759,7 +768,6 @@ mod test { for partition in 0..table_scan.output_partitioning().partition_count() { let result = table_scan .execute(partition, Arc::new(RuntimeEnv::default())) - .await .unwrap(); let recordbatches = util::collect(result).await.unwrap(); diff --git a/src/mito/src/engine.rs b/src/mito/src/engine.rs index 468b0013e52..7fbc6479fa6 100644 --- a/src/mito/src/engine.rs +++ b/src/mito/src/engine.rs @@ -569,10 +569,7 @@ mod tests { assert_eq!(2, table.insert(insert_req).await.unwrap()); let stream = table.scan(&None, &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); assert_eq!(1, batches.len()); @@ -605,10 +602,7 @@ mod tests { assert_eq!(2, table.insert(insert_req).await.unwrap()); let stream = table.scan(&None, &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); assert_eq!(1, batches.len()); @@ -660,10 +654,7 @@ mod tests { assert_eq!(2, table.insert(insert_req).await.unwrap()); let stream = table.scan(&None, &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); assert_eq!(1, batches.len()); assert_eq!(batches[0].df_recordbatch.num_columns(), 4); @@ -685,10 +676,7 @@ mod tests { // Scan with projections: cpu and memory let stream = table.scan(&Some(vec![1, 2]), &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); assert_eq!(1, batches.len()); assert_eq!(batches[0].df_recordbatch.num_columns(), 2); @@ -706,10 +694,7 @@ mod tests { // Scan with projections: only ts let stream = table.scan(&Some(vec![3]), &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); assert_eq!(1, batches.len()); assert_eq!(batches[0].df_recordbatch.num_columns(), 1); @@ -749,10 +734,7 @@ mod tests { assert_eq!(test_batch_size, table.insert(insert_req).await.unwrap()); let stream = table.scan(&None, &[], None).await.unwrap(); - let stream = stream - .execute(0, Arc::new(RuntimeEnv::default())) - .await - .unwrap(); + let stream = stream.execute(0, Arc::new(RuntimeEnv::default())).unwrap(); let batches = util::collect(stream).await.unwrap(); let mut total = 0; for batch in batches { diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 1139416803b..7c65c758b3b 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -219,7 +219,6 @@ impl QueryExecutor for DatafusionQueryEngine { 0 => Ok(Box::pin(EmptyRecordBatchStream::new(plan.schema()))), 1 => Ok(plan .execute(0, ctx.state().runtime()) - .await .context(error::ExecutePhysicalPlanSnafu)?), _ => { // merge into a single partition diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 73d5dd81812..4e1ef884e74 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -16,7 +16,6 @@ use std::any::Any; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, Mutex}; -use async_trait::async_trait; use common_query::error as query_error; use common_query::error::Result as QueryResult; use common_query::physical_plan::{Partitioning, PhysicalPlan, PhysicalPlanRef, RuntimeEnv}; @@ -48,7 +47,6 @@ impl SimpleTableScan { } } -#[async_trait] impl PhysicalPlan for SimpleTableScan { fn as_any(&self) -> &dyn Any { self @@ -70,7 +68,7 @@ impl PhysicalPlan for SimpleTableScan { unimplemented!() } - async fn execute( + fn execute( &self, _partition: usize, _runtime: Arc, @@ -117,12 +115,12 @@ mod test { assert_eq!(scan.schema(), schema); let runtime = Arc::new(RuntimeEnv::default()); - let stream = scan.execute(0, runtime.clone()).await.unwrap(); + let stream = scan.execute(0, runtime.clone()).unwrap(); let recordbatches = util::collect(stream).await.unwrap(); assert_eq!(recordbatches[0], batch1); assert_eq!(recordbatches[1], batch2); - let result = scan.execute(0, runtime).await; + let result = scan.execute(0, runtime); assert!(result.is_err()); match result { Err(e) => assert!(e diff --git a/src/table/src/test_util/memtable.rs b/src/table/src/test_util/memtable.rs index aa16d15b794..5f35e73c82b 100644 --- a/src/table/src/test_util/memtable.rs +++ b/src/table/src/test_util/memtable.rs @@ -212,7 +212,6 @@ mod test { let scan_stream = table.scan(&Some(vec![1]), &[], None).await.unwrap(); let scan_stream = scan_stream .execute(0, Arc::new(RuntimeEnv::default())) - .await .unwrap(); let recordbatch = util::collect(scan_stream).await.unwrap(); assert_eq!(1, recordbatch.len()); @@ -235,7 +234,6 @@ mod test { let scan_stream = table.scan(&None, &[], Some(2)).await.unwrap(); let scan_stream = scan_stream .execute(0, Arc::new(RuntimeEnv::default())) - .await .unwrap(); let recordbatch = util::collect(scan_stream).await.unwrap(); assert_eq!(1, recordbatch.len());