feat: Adds RegionScanner trait

Cargo.toml

@@ -109,6 +109,7 @@ datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev
 datafusion-functions = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
 datafusion-optimizer = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
 datafusion-physical-expr = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
+datafusion-physical-plan = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
 datafusion-sql = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
 datafusion-substrait = { git = "https://github.com/apache/arrow-datafusion.git", rev = "34eda15b73a9e278af8844b30ed2f1c21c10359c" }
 derive_builder = "0.12"

src/datanode/src/tests.rs

@@ -23,7 +23,6 @@ use common_function::function::FunctionRef;
 use common_function::scalars::aggregate::AggregateFunctionMetaRef;
 use common_query::prelude::ScalarUdf;
 use common_query::Output;
-use common_recordbatch::SendableRecordBatchStream;
 use common_runtime::Runtime;
 use query::dataframe::DataFrame;
 use query::plan::LogicalPlan;
@@ -32,7 +31,7 @@ use query::query_engine::DescribeResult;
 use query::{QueryEngine, QueryEngineContext};
 use session::context::QueryContextRef;
 use store_api::metadata::RegionMetadataRef;
-use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
+use store_api::region_engine::{RegionEngine, RegionRole, RegionScannerRef, SetReadonlyResponse};
 use store_api::region_request::{AffectedRows, RegionRequest};
 use store_api::storage::{RegionId, ScanRequest};
 use table::TableRef;
@@ -189,11 +188,11 @@ impl RegionEngine for MockRegionEngine {
         Ok(RegionResponse::new(0))
     }
 
-    async fn handle_query(
+    async fn handle_partitioned_query(
         &self,
         _region_id: RegionId,
         _request: ScanRequest,
-    ) -> Result<SendableRecordBatchStream, BoxedError> {
+    ) -> Result<RegionScannerRef, BoxedError> {
         unimplemented!()
     }
 

src/file-engine/src/engine.rs

@@ -25,7 +25,9 @@ use common_telemetry::{error, info};
 use object_store::ObjectStore;
 use snafu::{ensure, OptionExt};
 use store_api::metadata::RegionMetadataRef;
-use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
+use store_api::region_engine::{
+    RegionEngine, RegionRole, RegionScannerRef, SetReadonlyResponse, SinglePartitionScanner,
+};
 use store_api::region_request::{
     AffectedRows, RegionCloseRequest, RegionCreateRequest, RegionDropRequest, RegionOpenRequest,
     RegionRequest,
@@ -49,6 +51,20 @@ impl FileRegionEngine {
             inner: Arc::new(EngineInner::new(object_store)),
         }
     }
+
+    async fn handle_query(
+        &self,
+        region_id: RegionId,
+        request: ScanRequest,
+    ) -> Result<SendableRecordBatchStream, BoxedError> {
+        self.inner
+            .get_region(region_id)
+            .await
+            .context(RegionNotFoundSnafu { region_id })
+            .map_err(BoxedError::new)?
+            .query(request)
+            .map_err(BoxedError::new)
+    }
 }
 
 #[async_trait]
@@ -68,18 +84,14 @@ impl RegionEngine for FileRegionEngine {
             .map_err(BoxedError::new)
     }
 
-    async fn handle_query(
+    async fn handle_partitioned_query(
         &self,
         region_id: RegionId,
         request: ScanRequest,
-    ) -> Result<SendableRecordBatchStream, BoxedError> {
-        self.inner
-            .get_region(region_id)
-            .await
-            .context(RegionNotFoundSnafu { region_id })
-            .map_err(BoxedError::new)?
-            .query(request)
-            .map_err(BoxedError::new)
+    ) -> Result<RegionScannerRef, BoxedError> {
+        let stream = self.handle_query(region_id, request).await?;
+        let scanner = Arc::new(SinglePartitionScanner::new(stream));
+        Ok(scanner)
     }
 
     async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError> {

src/metric-engine/src/engine.rs

@@ -35,7 +35,9 @@ use common_recordbatch::SendableRecordBatchStream;
 use mito2::engine::MitoEngine;
 use store_api::metadata::RegionMetadataRef;
 use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
-use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
+use store_api::region_engine::{
+    RegionEngine, RegionRole, RegionScannerRef, SetReadonlyResponse, SinglePartitionScanner,
+};
 use store_api::region_request::RegionRequest;
 use store_api::storage::{RegionId, ScanRequest};
 
@@ -155,16 +157,14 @@ impl RegionEngine for MetricEngine {
         })
     }
 
-    /// Handles substrait query and return a stream of record batches
-    async fn handle_query(
+    async fn handle_partitioned_query(
         &self,
         region_id: RegionId,
         request: ScanRequest,
-    ) -> Result<SendableRecordBatchStream, BoxedError> {
-        self.inner
-            .read_region(region_id, request)
-            .await
-            .map_err(BoxedError::new)
+    ) -> Result<RegionScannerRef, BoxedError> {
+        let stream = self.handle_query(region_id, request).await?;
+        let scanner = Arc::new(SinglePartitionScanner::new(stream));
+        Ok(scanner)
     }
 
     /// Retrieves region's metadata.
@@ -251,6 +251,18 @@ impl MetricEngine {
             .logical_regions(physical_region_id)
             .await
     }
+
+    /// Handles substrait query and return a stream of record batches
+    async fn handle_query(
+        &self,
+        region_id: RegionId,
+        request: ScanRequest,
+    ) -> Result<SendableRecordBatchStream, BoxedError> {
+        self.inner
+            .read_region(region_id, request)
+            .await
+            .map_err(BoxedError::new)
+    }
 }
 
 struct MetricEngineInner {

src/mito2/src/engine.rs

@@ -62,7 +62,7 @@ use object_store::manager::ObjectStoreManagerRef;
 use snafu::{ensure, OptionExt, ResultExt};
 use store_api::logstore::LogStore;
 use store_api::metadata::RegionMetadataRef;
-use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
+use store_api::region_engine::{RegionEngine, RegionRole, RegionScannerRef, SetReadonlyResponse};
 use store_api::region_request::{AffectedRows, RegionRequest};
 use store_api::storage::{RegionId, ScanRequest};
 use tokio::sync::oneshot;
@@ -115,11 +115,35 @@ impl MitoEngine {
         Ok(region.region_usage().await)
     }
 
+    /// Handle substrait query and return a stream of record batches
+    #[tracing::instrument(skip_all)]
+    pub async fn handle_query(
+        &self,
+        region_id: RegionId,
+        request: ScanRequest,
+    ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
+        self.scanner(region_id, request)
+            .map_err(BoxedError::new)?
+            .scan()
+            .await
+            .map_err(BoxedError::new)
+    }
+
     /// Returns a scanner to scan for `request`.
     fn scanner(&self, region_id: RegionId, request: ScanRequest) -> Result<Scanner> {
         self.scan_region(region_id, request)?.scanner()
     }
 
+    /// Returns a region scanner to scan the region for `request`.
+    async fn region_scanner(
+        &self,
+        region_id: RegionId,
+        request: ScanRequest,
+    ) -> Result<RegionScannerRef> {
+        let scanner = self.scanner(region_id, request)?;
+        scanner.region_scanner().await
+    }
+
     /// Scans a region.
     fn scan_region(&self, region_id: RegionId, request: ScanRequest) -> Result<ScanRegion> {
         self.inner.handle_query(region_id, request)
@@ -312,16 +336,13 @@ impl RegionEngine for MitoEngine {
             .map_err(BoxedError::new)
     }
 
-    /// Handle substrait query and return a stream of record batches
     #[tracing::instrument(skip_all)]
-    async fn handle_query(
+    async fn handle_partitioned_query(
         &self,
         region_id: RegionId,
         request: ScanRequest,
-    ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
-        self.scanner(region_id, request)
-            .map_err(BoxedError::new)?
-            .scan()
+    ) -> Result<RegionScannerRef, BoxedError> {
+        self.region_scanner(region_id, request)
             .await
             .map_err(BoxedError::new)
     }

src/mito2/src/read/scan_region.rs

@@ -20,6 +20,7 @@ use std::time::Instant;
 use common_recordbatch::SendableRecordBatchStream;
 use common_telemetry::{debug, error, warn};
 use common_time::range::TimestampRange;
+use store_api::region_engine::{RegionScannerRef, SinglePartitionScanner};
 use store_api::storage::ScanRequest;
 use table::predicate::{Predicate, TimeRangePredicateBuilder};
 use tokio::sync::{mpsc, Semaphore};
@@ -57,6 +58,14 @@ impl Scanner {
             Scanner::Unordered(unordered_scan) => unordered_scan.build_stream().await,
         }
     }
+
+    /// Returns a [RegionScanner] to scan the region.
+    pub(crate) async fn region_scanner(&self) -> Result<RegionScannerRef> {
+        let stream = self.scan().await?;
+        let scanner = SinglePartitionScanner::new(stream);
+
+        Ok(Arc::new(scanner))
+    }
 }
 
 #[cfg(test)]

src/mito2/src/read/unordered_scan.rs

@@ -198,6 +198,14 @@ impl UnorderedScan {
     }
 }
 
+#[cfg(test)]
+impl UnorderedScan {
+    /// Returns the input.
+    pub(crate) fn input(&self) -> &ScanInput {
+        &self.input
+    }
+}
+
 /// Metrics for [UnorderedScan].
 #[derive(Debug, Default)]
 struct Metrics {
@@ -216,11 +224,3 @@ struct Metrics {
     /// Number of rows returned.
     num_rows: usize,
 }
-
-#[cfg(test)]
-impl UnorderedScan {
-    /// Returns the input.
-    pub(crate) fn input(&self) -> &ScanInput {
-        &self.input
-    }
-}

src/query/src/dummy_catalog.rs

@@ -31,7 +31,7 @@ use snafu::ResultExt;
 use store_api::metadata::RegionMetadataRef;
 use store_api::region_engine::RegionEngineRef;
 use store_api::storage::{RegionId, ScanRequest};
-use table::table::scan::StreamScanAdapter;
+use table::table::scan::ReadFromRegion;
 
 use crate::error::{GetRegionMetadataSnafu, Result};
 
@@ -168,12 +168,12 @@ impl TableProvider for DummyTableProvider {
             .collect();
         request.limit = limit;
 
-        let stream = self
+        let scanner = self
             .engine
-            .handle_query(self.region_id, request)
+            .handle_partitioned_query(self.region_id, request)
             .await
             .map_err(|e| DataFusionError::External(Box::new(e)))?;
-        Ok(Arc::new(StreamScanAdapter::new(stream)))
+        Ok(Arc::new(ReadFromRegion::new(scanner)))
     }
 
     fn supports_filters_pushdown(

src/query/src/optimizer/test_util.rs

@@ -23,12 +23,11 @@ use api::v1::SemanticType;
 use async_trait::async_trait;
 use common_error::ext::{BoxedError, PlainError};
 use common_error::status_code::StatusCode;
-use common_recordbatch::SendableRecordBatchStream;
 use datatypes::schema::ColumnSchema;
 use store_api::metadata::{
     ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef,
 };
-use store_api::region_engine::{RegionEngine, RegionRole, SetReadonlyResponse};
+use store_api::region_engine::{RegionEngine, RegionRole, RegionScannerRef, SetReadonlyResponse};
 use store_api::region_request::RegionRequest;
 use store_api::storage::{ConcreteDataType, RegionId, ScanRequest};
 
@@ -63,11 +62,11 @@ impl RegionEngine for MetaRegionEngine {
         unimplemented!()
     }
 
-    async fn handle_query(
+    async fn handle_partitioned_query(
         &self,
         _region_id: RegionId,
         _request: ScanRequest,
-    ) -> Result<SendableRecordBatchStream, BoxedError> {
+    ) -> Result<RegionScannerRef, BoxedError> {
         unimplemented!()
     }
 

src/store-api/Cargo.toml

@@ -17,6 +17,7 @@ common-macro.workspace = true
 common-query.workspace = true
 common-recordbatch.workspace = true
 common-wal.workspace = true
+datafusion-physical-plan.workspace = true
 datatypes.workspace = true
 derive_builder.workspace = true
 futures.workspace = true