diff --git a/src/common/query/src/stream.rs b/src/common/query/src/stream.rs index 99e02dbc41..1af777ab15 100644 --- a/src/common/query/src/stream.rs +++ b/src/common/query/src/stream.rs @@ -27,9 +27,17 @@ use datafusion_common::DataFusionError; use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datatypes::schema::SchemaRef; +/// Opens a new scan stream equivalent to the one this plan was built from. +pub type StreamFactoryRef = + Arc datafusion_common::Result + Send + Sync>; + /// Adapts greptime's [SendableRecordBatchStream] to DataFusion's [ExecutionPlan]. pub struct StreamScanAdapter { stream: Mutex>, + /// Opens a replacement stream once `stream` has been handed out. DataFusion + /// executes a plan more than once in a recursive CTE, where every iteration + /// re-executes the recursive term. + stream_factory: Option, schema: SchemaRef, arrow_schema: ArrowSchemaRef, properties: Arc, @@ -58,6 +66,7 @@ impl StreamScanAdapter { Self { stream: Mutex::new(Some(stream)), + stream_factory: None, schema, arrow_schema, properties, @@ -69,6 +78,13 @@ impl StreamScanAdapter { self.output_ordering = output_ordering; self } + + /// Makes this plan re-executable. The factory must open a stream over the + /// same scan request. + pub fn with_stream_factory(mut self, stream_factory: StreamFactoryRef) -> Self { + self.stream_factory = Some(stream_factory); + self + } } impl DisplayAs for StreamScanAdapter { @@ -113,10 +129,16 @@ impl ExecutionPlan for StreamScanAdapter { _partition: usize, _context: Arc, ) -> datafusion_common::Result { - let mut stream = self.stream.lock().unwrap(); - let stream = stream - .take() - .ok_or_else(|| DataFusionError::Execution("Stream already exhausted".to_string()))?; + let stream = self.stream.lock().unwrap().take(); + let stream = match stream { + Some(stream) => stream, + None => { + let factory = self.stream_factory.as_ref().ok_or_else(|| { + DataFusionError::Execution("Stream already exhausted".to_string()) + })?; + factory()? + } + }; Ok(Box::pin(DfRecordBatchStreamAdapter::new(stream))) } @@ -176,4 +198,41 @@ mod test { _ => unreachable!(), } } + + #[tokio::test] + async fn test_re_execute_with_stream_factory() { + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1, 2])) as _], + ) + .unwrap(); + + let factory_schema = schema.clone(); + let factory_batch = batch.clone(); + let scan = StreamScanAdapter::new( + RecordBatches::try_new(schema.clone(), vec![batch.clone()]) + .unwrap() + .as_stream(), + ) + .with_stream_factory(Arc::new(move || { + Ok( + RecordBatches::try_new(factory_schema.clone(), vec![factory_batch.clone()]) + .unwrap() + .as_stream(), + ) + })); + + for _ in 0..3 { + let stream = scan.execute(0, ctx.task_ctx()).unwrap(); + let recordbatches = stream.try_collect::>().await.unwrap(); + assert_eq!(recordbatches, vec![batch.clone().into_df_record_batch()]); + } + } } diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index ac71b448e9..a8ec6ee93e 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -16,7 +16,7 @@ use std::any::Any; use std::sync::{Arc, Mutex}; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; -use common_query::stream::StreamScanAdapter; +use common_query::stream::{StreamFactoryRef, StreamScanAdapter}; use common_recordbatch::OrderOption; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef as DfSchemaRef}; use datafusion::catalog::Session; @@ -29,8 +29,10 @@ use datafusion_expr::expr::Expr; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datatypes::types::json_type::JsonNativeType; +use snafu::ResultExt; use store_api::storage::{ScanRequest, VectorSearchRequest}; +use crate::error::TablesRecordBatchSnafu; use crate::table::{TableRef, TableType}; /// Adapt greptime's [TableRef] to DataFusion's [TableProvider]. @@ -169,7 +171,7 @@ impl TableProvider for DfTableProviderAdapter { return Ok(plan); } - let stream = self.table.scan_to_stream(request).await?; + let stream = self.table.scan_to_stream(request.clone()).await?; // build sort physical expr let schema = stream.schema(); @@ -187,8 +189,20 @@ impl TableProvider for DfTableProviderAdapter { .collect::>() }); + // The stream above is single-use, and a recursive CTE re-executes its + // recursive term on every iteration. + let data_source = self.table.data_source(); + let stream_factory: StreamFactoryRef = Arc::new(move || { + data_source + .get_stream(request.clone()) + .context(TablesRecordBatchSnafu) + .map_err(Into::into) + }); + Ok(Arc::new( - StreamScanAdapter::new(stream).with_output_ordering(sort_expr), + StreamScanAdapter::new(stream) + .with_output_ordering(sort_expr) + .with_stream_factory(stream_factory), )) } diff --git a/tests/cases/standalone/common/cte/cte.result b/tests/cases/standalone/common/cte/cte.result index 0a796b68b7..1ea77cae68 100644 --- a/tests/cases/standalone/common/cte/cte.result +++ b/tests/cases/standalone/common/cte/cte.result @@ -115,6 +115,27 @@ SELECT max(d) FROM cte; | 1 | +------------+ +-- the recursive term is re-executed once per iteration, so its scans must be re-runnable +WITH RECURSIVE t AS ( + SELECT 0 AS depth + UNION ALL + SELECT t.depth + 1 AS depth + FROM t, information_schema.tables + WHERE table_catalog = 'greptime' + AND table_schema = 'information_schema' + AND table_name = 'tables' + AND t.depth < 2 +) +SELECT depth FROM t ORDER BY depth; + ++-------+ +| depth | ++-------+ +| 0 | +| 1 | +| 2 | ++-------+ + -- Nested aliases is not supported in datafusion with cte (a) as ( select 1 diff --git a/tests/cases/standalone/common/cte/cte.sql b/tests/cases/standalone/common/cte/cte.sql index ebd517e31b..a9e63ae845 100644 --- a/tests/cases/standalone/common/cte/cte.sql +++ b/tests/cases/standalone/common/cte/cte.sql @@ -39,6 +39,19 @@ WITH RECURSIVE cte(d) AS ( ) SELECT max(d) FROM cte; +-- the recursive term is re-executed once per iteration, so its scans must be re-runnable +WITH RECURSIVE t AS ( + SELECT 0 AS depth + UNION ALL + SELECT t.depth + 1 AS depth + FROM t, information_schema.tables + WHERE table_catalog = 'greptime' + AND table_schema = 'information_schema' + AND table_name = 'tables' + AND t.depth < 2 +) +SELECT depth FROM t ORDER BY depth; + -- Nested aliases is not supported in datafusion with cte (a) as ( select 1