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;
@@ -193,7 +192,7 @@ impl RegionEngine for MockRegionEngine {
         &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]
@@ -72,14 +88,10 @@ impl RegionEngine for FileRegionEngine {
         &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(
         &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/metric-engine/src/engine/read.rs

@@ -62,7 +62,7 @@ impl MetricEngineInner {
             .start_timer();
 
         self.mito
-            .handle_query(region_id, request)
+            .scan_to_stream(region_id, request)
             .await
             .context(MitoReadOperationSnafu)
     }
@@ -82,7 +82,7 @@ impl MetricEngineInner {
             .transform_request(physical_region_id, logical_region_id, request)
             .await?;
         self.mito
-            .handle_query(data_region_id, request)
+            .scan_to_stream(data_region_id, request)
             .await
             .context(MitoReadOperationSnafu)
     }

src/metric-engine/src/metadata_region.rs

@@ -300,7 +300,7 @@ impl MetadataRegion {
         let scan_req = Self::build_read_request(key);
         let record_batch_stream = self
             .mito
-            .handle_query(region_id, scan_req)
+            .scan_to_stream(region_id, scan_req)
             .await
             .context(MitoReadOperationSnafu)?;
         let scan_result = collect(record_batch_stream)
@@ -317,7 +317,7 @@ impl MetadataRegion {
         let scan_req = Self::build_read_request(key);
         let record_batch_stream = self
             .mito
-            .handle_query(region_id, scan_req)
+            .scan_to_stream(region_id, scan_req)
             .await
             .context(MitoReadOperationSnafu)?;
         let scan_result = collect(record_batch_stream)
@@ -351,7 +351,7 @@ impl MetadataRegion {
         };
         let record_batch_stream = self
             .mito
-            .handle_query(region_id, scan_req)
+            .scan_to_stream(region_id, scan_req)
             .await
             .context(MitoReadOperationSnafu)?;
         let scan_result = collect(record_batch_stream)
@@ -590,7 +590,7 @@ mod test {
         let scan_req = MetadataRegion::build_read_request("test_key");
         let record_batch_stream = metadata_region
             .mito
-            .handle_query(region_id, scan_req)
+            .scan_to_stream(region_id, scan_req)
             .await
             .unwrap();
         let scan_result = collect(record_batch_stream).await.unwrap();

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 scan_to_stream(
+        &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(
         &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/engine/alter_test.rs

@@ -245,7 +245,7 @@ async fn test_put_after_alter() {
 |       | b     | 2.0     | 1970-01-01T00:00:02 |
 +-------+-------+---------+---------------------+";
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     assert_eq!(expected, batches.pretty_print().unwrap());
 }

src/mito2/src/engine/append_mode_test.rs

@@ -63,7 +63,7 @@ async fn test_append_mode_write_query() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
@@ -183,7 +183,7 @@ async fn test_append_mode_compaction() {
     // Reopens the region.
     reopen_region(&engine, region_id, region_dir, false, region_opts).await;
     let stream = engine
-        .handle_query(region_id, ScanRequest::default())
+        .scan_to_stream(region_id, ScanRequest::default())
         .await
         .unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();

src/mito2/src/engine/basic_test.rs

@@ -128,7 +128,7 @@ async fn test_region_replay() {
     assert_eq!(0, result.affected_rows);
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     assert_eq!(42, batches.iter().map(|b| b.num_rows()).sum::<usize>());
 
@@ -166,7 +166,7 @@ async fn test_write_query_region() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
@@ -227,7 +227,7 @@ async fn test_different_order() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+-------+---------+---------+---------------------+
@@ -289,7 +289,7 @@ async fn test_different_order_and_type() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+-------+---------+---------+---------------------+
@@ -341,7 +341,7 @@ async fn test_put_delete() {
     delete_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
@@ -383,7 +383,7 @@ async fn test_delete_not_null_fields() {
     delete_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
@@ -398,7 +398,7 @@ async fn test_delete_not_null_fields() {
     // Reopen and scan again.
     reopen_region(&engine, region_id, region_dir, false, HashMap::new()).await;
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     assert_eq!(expected, batches.pretty_print().unwrap());
 }
@@ -447,7 +447,7 @@ async fn test_put_overwrite() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
@@ -688,7 +688,7 @@ async fn test_cache_null_primary_key() {
     put_rows(&engine, region_id, rows).await;
 
     let request = ScanRequest::default();
-    let stream = engine.handle_query(region_id, request).await.unwrap();
+    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+-------+---------+---------------------+

src/mito2/src/engine/catchup_test.rs

@@ -104,7 +104,7 @@ async fn test_catchup_with_last_entry_id() {
     // Scans
     let request = ScanRequest::default();
     let stream = follower_engine
-        .handle_query(region_id, request)
+        .scan_to_stream(region_id, request)
         .await
         .unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
@@ -264,7 +264,7 @@ async fn test_catchup_without_last_entry_id() {
 
     let request = ScanRequest::default();
     let stream = follower_engine
-        .handle_query(region_id, request)
+        .scan_to_stream(region_id, request)
         .await
         .unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
@@ -367,7 +367,7 @@ async fn test_catchup_with_manifest_update() {
 
     let request = ScanRequest::default();
     let stream = follower_engine
-        .handle_query(region_id, request)
+        .scan_to_stream(region_id, request)
         .await
         .unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();