fix: re-scan stream-backed tables in recursive CTEs (#9039)

* fix: re-scan stream-backed tables in recursive CTEs

A recursive CTE re-executes its recursive term on every iteration, but
DfTableProviderAdapter hands StreamScanAdapter a single-use stream built at
planning time. The second iteration failed with "Stream already exhausted"
for every table served through DataSource::get_stream — information_schema,
pg_catalog, the computed entity-graph tables and numbers.

Keep that stream for the first execution and open a new one over the same
scan request for later executions.

Closes #9037

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* refactor: drop redundant binding in stream factory

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

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<dyn Fn() -> datafusion_common::Result<SendableRecordBatchStream> + Send + Sync>;
+
 /// Adapts greptime's [SendableRecordBatchStream] to DataFusion's [ExecutionPlan].
 pub struct StreamScanAdapter {
     stream: Mutex<Option<SendableRecordBatchStream>>,
+    /// 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<StreamFactoryRef>,
     schema: SchemaRef,
     arrow_schema: ArrowSchemaRef,
     properties: Arc<PlanProperties>,
@@ -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<TaskContext>,
     ) -> datafusion_common::Result<DfSendableRecordBatchStream> {
-        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::<Vec<_>>().await.unwrap();
+            assert_eq!(recordbatches, vec![batch.clone().into_df_record_batch()]);
+        }
+    }
 }

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::<Vec<_>>()
         });
 
+        // 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),
         ))
     }
 

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

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