feat: Implement an unordered scanner for append mode

src/mito2/Cargo.toml

@@ -70,6 +70,7 @@ common-procedure-test.workspace = true
 common-test-util.workspace = true
 criterion = "0.4"
 log-store.workspace = true
+object-store = { workspace = true, features = ["services-memory"] }
 toml.workspace = true
 
 [[bench]]

src/mito2/src/compaction/twcs.rs

@@ -35,6 +35,7 @@ use crate::config::MitoConfig;
 use crate::error::{self, CompactRegionSnafu};
 use crate::metrics::{COMPACTION_FAILURE_COUNT, COMPACTION_STAGE_ELAPSED};
 use crate::read::projection::ProjectionMapper;
+use crate::read::scan_region::ScanInput;
 use crate::read::seq_scan::SeqScan;
 use crate::read::{BoxedBatchReader, Source};
 use crate::region::options::IndexOptions;
@@ -577,13 +578,12 @@ async fn build_sst_reader(
     inputs: &[FileHandle],
     append_mode: bool,
 ) -> error::Result<BoxedBatchReader> {
-    SeqScan::new(sst_layer, ProjectionMapper::all(&metadata)?)
+    let scan_input = ScanInput::new(sst_layer, ProjectionMapper::all(&metadata)?)
         .with_files(inputs.to_vec())
         .with_append_mode(append_mode)
         // We ignore file not found error during compaction.
-        .with_ignore_file_not_found(true)
-        .build_reader()
-        .await
+        .with_ignore_file_not_found(true);
+    SeqScan::new(scan_input).build_reader().await
 }
 
 #[cfg(test)]

src/mito2/src/engine.rs

@@ -114,6 +114,11 @@ impl MitoEngine {
 
     /// Returns a scanner to scan for `request`.
     fn scanner(&self, region_id: RegionId, request: ScanRequest) -> Result<Scanner> {
+        self.scan_region(region_id, request)?.scanner()
+    }
+
+    /// Scans a region.
+    fn scan_region(&self, region_id: RegionId, request: ScanRequest) -> Result<ScanRegion> {
         self.inner.handle_query(region_id, request)
     }
 
@@ -220,8 +225,8 @@ impl EngineInner {
         receiver.await.context(RecvSnafu)?
     }
 
-    /// Handles the scan `request` and returns a [Scanner] for the `request`.
-    fn handle_query(&self, region_id: RegionId, request: ScanRequest) -> Result<Scanner> {
+    /// Handles the scan `request` and returns a [ScanRegion].
+    fn handle_query(&self, region_id: RegionId, request: ScanRequest) -> Result<ScanRegion> {
         let query_start = Instant::now();
         // Reading a region doesn't need to go through the region worker thread.
         let region = self
@@ -246,7 +251,7 @@ impl EngineInner {
         .with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled())
         .with_start_time(query_start);
 
-        scan_region.scanner()
+        Ok(scan_region)
     }
 
     /// Set writable mode for a region.

src/mito2/src/engine/append_mode_test.rs

@@ -16,14 +16,17 @@
 
 use api::v1::Rows;
 use common_recordbatch::RecordBatches;
+use datatypes::arrow::compute::{self, SortColumn};
+use datatypes::arrow::record_batch::RecordBatch;
+use datatypes::arrow::util::pretty;
 use store_api::region_engine::RegionEngine;
 use store_api::region_request::{RegionCompactRequest, RegionRequest};
 use store_api::storage::{RegionId, ScanRequest};
 
 use crate::config::MitoConfig;
 use crate::test_util::{
-    build_rows, build_rows_for_key, flush_region, put_rows, rows_schema, CreateRequestBuilder,
-    TestEnv,
+    build_rows, build_rows_for_key, flush_region, put_rows, reopen_region, rows_schema,
+    CreateRequestBuilder, TestEnv,
 };
 
 #[tokio::test]
@@ -74,21 +77,37 @@ async fn test_append_mode_write_query() {
 | 1     | 1.0     | 1970-01-01T00:00:01 |
 | 2     | 2.0     | 1970-01-01T00:00:02 |
 +-------+---------+---------------------+";
+    assert_eq!(expected, sort_batches_and_print(&batches, &["tag_0", "ts"]));
+
+    // Tries to use seq scan to test it under append mode.
+    let scan = engine
+        .scan_region(region_id, ScanRequest::default())
+        .unwrap();
+    let seq_scan = scan.seq_scan().unwrap();
+    let stream = seq_scan.build_stream().await.unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
     assert_eq!(expected, batches.pretty_print().unwrap());
 }
 
 #[tokio::test]
 async fn test_append_mode_compaction() {
     let mut env = TestEnv::new();
-    let engine = env.create_engine(MitoConfig::default()).await;
-
+    let engine = env
+        .create_engine(MitoConfig {
+            scan_parallelism: 2,
+            ..Default::default()
+        })
+        .await;
     let region_id = RegionId::new(1, 1);
+
     let request = CreateRequestBuilder::new()
         .insert_option("compaction.type", "twcs")
         .insert_option("compaction.twcs.max_active_window_files", "2")
         .insert_option("compaction.twcs.max_inactive_window_files", "2")
         .insert_option("append_mode", "true")
         .build();
+    let region_dir = request.region_dir.clone();
+    let region_opts = request.options.clone();
 
     let column_schemas = rows_schema(&request);
     engine
@@ -132,10 +151,6 @@ async fn test_append_mode_compaction() {
     };
     put_rows(&engine, region_id, rows).await;
 
-    let scanner = engine.scanner(region_id, ScanRequest::default()).unwrap();
-    assert_eq!(1, scanner.num_files());
-    let stream = scanner.scan().await.unwrap();
-    let batches = RecordBatches::try_collect(stream).await.unwrap();
     let expected = "\
 +-------+---------+---------------------+
 | tag_0 | field_0 | ts                  |
@@ -149,5 +164,58 @@ async fn test_append_mode_compaction() {
 | b     | 0.0     | 1970-01-01T00:00:00 |
 | b     | 1.0     | 1970-01-01T00:00:01 |
 +-------+---------+---------------------+";
-    assert_eq!(expected, batches.pretty_print().unwrap());
+    // Scans in parallel.
+    let scanner = engine.scanner(region_id, ScanRequest::default()).unwrap();
+    assert_eq!(1, scanner.num_files());
+    assert_eq!(1, scanner.num_memtables());
+    let stream = scanner.scan().await.unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    assert_eq!(expected, sort_batches_and_print(&batches, &["tag_0", "ts"]));
+
+    // Reopens engine with parallelism 1.
+    let engine = env
+        .reopen_engine(
+            engine,
+            MitoConfig {
+                scan_parallelism: 1,
+                ..Default::default()
+            },
+        )
+        .await;
+    // Reopens the region.
+    reopen_region(&engine, region_id, region_dir, false, region_opts).await;
+    let stream = engine
+        .handle_query(region_id, ScanRequest::default())
+        .await
+        .unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    assert_eq!(expected, sort_batches_and_print(&batches, &["tag_0", "ts"]));
+}
+
+/// Sorts `batches` by column `names`.
+fn sort_batches_and_print(batches: &RecordBatches, names: &[&str]) -> String {
+    let schema = batches.schema();
+    let record_batches = batches.iter().map(|batch| batch.df_record_batch());
+    let record_batch = compute::concat_batches(schema.arrow_schema(), record_batches).unwrap();
+    let columns: Vec<_> = names
+        .iter()
+        .map(|name| {
+            let array = record_batch.column_by_name(name).unwrap();
+            SortColumn {
+                values: array.clone(),
+                options: None,
+            }
+        })
+        .collect();
+    let indices = compute::lexsort_to_indices(&columns, None).unwrap();
+    let columns = record_batch
+        .columns()
+        .iter()
+        .map(|array| compute::take(&array, &indices, None).unwrap())
+        .collect();
+    let record_batch = RecordBatch::try_new(record_batch.schema(), columns).unwrap();
+
+    pretty::pretty_format_batches(&[record_batch])
+        .unwrap()
+        .to_string()
 }

src/mito2/src/engine/basic_test.rs

@@ -394,7 +394,7 @@ async fn test_delete_not_null_fields() {
     assert_eq!(expected, batches.pretty_print().unwrap());
 
     // Reopen and scan again.
-    reopen_region(&engine, region_id, region_dir, false).await;
+    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 batches = RecordBatches::try_collect(stream).await.unwrap();

src/mito2/src/engine/flush_test.rs

@@ -263,7 +263,7 @@ async fn test_flush_reopen_region() {
     };
     check_region();
 
-    reopen_region(&engine, region_id, region_dir, true).await;
+    reopen_region(&engine, region_id, region_dir, true, Default::default()).await;
     check_region();
 
     // Puts again.

src/mito2/src/engine/open_test.rs

@@ -95,7 +95,7 @@ async fn test_engine_reopen_region() {
         .await
         .unwrap();
 
-    reopen_region(&engine, region_id, region_dir, false).await;
+    reopen_region(&engine, region_id, region_dir, false, Default::default()).await;
     assert!(engine.is_region_exists(region_id));
 }
 
@@ -113,7 +113,7 @@ async fn test_engine_open_readonly() {
         .await
         .unwrap();
 
-    reopen_region(&engine, region_id, region_dir, false).await;
+    reopen_region(&engine, region_id, region_dir, false, Default::default()).await;
 
     // Region is readonly.
     let rows = Rows {

src/mito2/src/read.rs

@@ -19,6 +19,7 @@ pub mod merge;
 pub mod projection;
 pub(crate) mod scan_region;
 pub(crate) mod seq_scan;
+pub(crate) mod unordered_scan;
 
 use std::collections::HashSet;
 use std::sync::Arc;