refactor: optimize data block processing in MergeIntoOperationAggregator (#11571)

* refactor: use dedicate thread while decoding block data

* Revert experimental changes

No significant performance improvement observed, back to try_join_all.

* cleanup

* adjust deserialization path so that metrics could be recorded

* shffule segments while build replace into pipeline

* fix ut

* refactor: (MergeIntoOperationAggregator) load datablocks in parallel

Author: dantengsky

Cargo.lock

@@ -14,6 +14,7 @@
 
 use common_exception::Result;
 use common_storages_fuse::FuseTable;
+use itertools::Itertools;
 
 #[test]
 fn test_partition() -> Result<()> {
@@ -36,10 +37,22 @@ fn test_partition() -> Result<()> {
                 rng.gen_range(1..number_segment)
             };
 
-            let chunks = FuseTable::partition_segments(&segments, num_partition);
-            assert_eq!(chunks.len(), num_partition);
-            for (idx, (segment_idx, _)) in chunks.clone().into_iter().flatten().enumerate() {
-                assert_eq!(idx, segment_idx)
+            let partitions = FuseTable::partition_segments(&segments, num_partition);
+            // check number of partitions are as expected
+            assert_eq!(partitions.len(), num_partition);
+
+            // check segments
+            let origin = segments.iter().enumerate();
+            let segment_of_chunks = partitions
+                .iter()
+                .flatten()
+                .sorted_by(|a, b| a.0.cmp(&b.0))
+                .collect::<Vec<_>>();
+
+            for (origin_idx, origin_location) in origin {
+                let (seg_idx, seg_location) = segment_of_chunks[origin_idx];
+                assert_eq!(origin_idx, *seg_idx);
+                assert_eq!(origin_location, seg_location);
             }
         }
     }

src/query/service/tests/it/storages/fuse/operations/replace_into.rs

@@ -50,6 +50,7 @@ futures-util = "0.3.24"
 metrics = "0.20.1"
 opendal = { workspace = true }
 parquet-format-safe = "0.2"
+rand = "0.8.5"
 serde = { workspace = true }
 serde_json = { workspace = true }
 sha2 = "0.10.6"

src/query/storages/fuse/Cargo.toml

@@ -29,6 +29,7 @@ use super::BlockReader;
 use crate::io::read::block::block_reader_merge_io::DataItem;
 use crate::io::ReadSettings;
 use crate::io::UncompressedBuffer;
+use crate::FusePartInfo;
 use crate::FuseStorageFormat;
 
 pub enum DeserializedArray<'a> {
@@ -47,15 +48,47 @@ pub struct FieldDeserializationContext<'a> {
 
 impl BlockReader {
     /// Deserialize column chunks data from parquet format to DataBlock.
-    pub fn deserialize_chunks(
+    pub fn deserialize_chunks_with_part_info(
         &self,
         part: PartInfoPtr,
         chunks: HashMap<ColumnId, DataItem>,
         storage_format: &FuseStorageFormat,
+    ) -> Result<DataBlock> {
+        let part = FusePartInfo::from_part(&part)?;
+        self.deserialize_chunks(
+            &part.location,
+            part.nums_rows,
+            &part.compression,
+            &part.columns_meta,
+            chunks,
+            storage_format,
+        )
+    }
+
+    pub fn deserialize_chunks(
+        &self,
+        block_path: &str,
+        num_rows: usize,
+        compression: &Compression,
+        column_metas: &HashMap<ColumnId, ColumnMeta>,
+        column_chunks: HashMap<ColumnId, DataItem>,
+        storage_format: &FuseStorageFormat,
     ) -> Result<DataBlock> {
         match storage_format {
-            FuseStorageFormat::Parquet => self.deserialize_parquet_chunks(part, chunks),
-            FuseStorageFormat::Native => self.deserialize_native_chunks(part, chunks),
+            FuseStorageFormat::Parquet => self.deserialize_parquet_chunks(
+                block_path,
+                num_rows,
+                compression,
+                column_metas,
+                column_chunks,
+            ),
+            FuseStorageFormat::Native => self.deserialize_native_chunks(
+                block_path,
+                num_rows,
+                compression,
+                column_metas,
+                column_chunks,
+            ),
         }
     }
 

src/query/storages/fuse/src/io/read/block/block_reader_deserialize.rs

@@ -174,7 +174,7 @@ impl BlockReader {
 
         let mut merge_io_read_res =
             Self::merge_io_read(settings, self.operator.clone(), location, ranges).await?;
-        // TODO set
+
         merge_io_read_res.cached_column_data = cached_column_data;
         merge_io_read_res.cached_column_array = cached_column_array;
         Ok(merge_io_read_res)

src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs

@@ -21,7 +21,6 @@ use common_arrow::arrow::chunk::Chunk;
 use common_arrow::arrow::datatypes::Field;
 use common_arrow::native::read::batch_read::batch_read_array;
 use common_arrow::parquet::metadata::ColumnDescriptor;
-use common_catalog::plan::PartInfoPtr;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::ColumnId;
@@ -35,7 +34,6 @@ use storages_common_table_meta::meta::Compression;
 
 use super::block_reader_deserialize::DeserializedArray;
 use super::block_reader_deserialize::FieldDeserializationContext;
-use crate::fuse_part::FusePartInfo;
 use crate::io::read::block::block_reader_merge_io::DataItem;
 use crate::io::BlockReader;
 use crate::io::UncompressedBuffer;
@@ -45,22 +43,24 @@ impl BlockReader {
     /// Deserialize column chunks data from native format to DataBlock.
     pub(super) fn deserialize_native_chunks(
         &self,
-        part: PartInfoPtr,
-        chunks: HashMap<ColumnId, DataItem>,
+        block_path: &str,
+        num_rows: usize,
+        compression: &Compression,
+        column_metas: &HashMap<ColumnId, ColumnMeta>,
+        column_chunks: HashMap<ColumnId, DataItem>,
     ) -> Result<DataBlock> {
-        let part = FusePartInfo::from_part(&part)?;
         let start = Instant::now();
 
-        if chunks.is_empty() {
-            return self.build_default_values_block(part.nums_rows);
+        if column_chunks.is_empty() {
+            return self.build_default_values_block(num_rows);
         }
 
         let deserialized_res = self.deserialize_native_chunks_with_buffer(
-            &part.location,
-            part.nums_rows,
-            &part.compression,
-            &part.columns_meta,
-            chunks,
+            block_path,
+            num_rows,
+            compression,
+            column_metas,
+            column_chunks,
             None,
         );
 

src/query/storages/fuse/src/io/read/block/block_reader_native_deserialize.rs

@@ -24,7 +24,6 @@ use common_arrow::parquet::compression::Compression as ParquetCompression;
 use common_arrow::parquet::metadata::ColumnDescriptor;
 use common_arrow::parquet::read::PageMetaData;
 use common_arrow::parquet::read::PageReader;
-use common_catalog::plan::PartInfoPtr;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::ColumnId;
@@ -38,7 +37,6 @@ use storages_common_table_meta::meta::Compression;
 
 use super::block_reader_deserialize::DeserializedArray;
 use super::block_reader_deserialize::FieldDeserializationContext;
-use crate::fuse_part::FusePartInfo;
 use crate::io::read::block::block_reader_merge_io::DataItem;
 use crate::io::read::block::decompressor::BuffedBasicDecompressor;
 use crate::io::BlockReader;
@@ -49,22 +47,24 @@ impl BlockReader {
     /// Deserialize column chunks data from parquet format to DataBlock.
     pub(super) fn deserialize_parquet_chunks(
         &self,
-        part: PartInfoPtr,
-        chunks: HashMap<ColumnId, DataItem>,
+        block_path: &str,
+        num_rows: usize,
+        compression: &Compression,
+        column_metas: &HashMap<ColumnId, ColumnMeta>,
+        column_chunks: HashMap<ColumnId, DataItem>,
     ) -> Result<DataBlock> {
-        let part = FusePartInfo::from_part(&part)?;
-        let start = Instant::now();
-
-        if chunks.is_empty() {
-            return self.build_default_values_block(part.nums_rows);
+        if column_chunks.is_empty() {
+            return self.build_default_values_block(num_rows);
         }
 
+        let start = Instant::now();
+
         let deserialized_res = self.deserialize_parquet_chunks_with_buffer(
-            &part.location,
-            part.nums_rows,
-            &part.compression,
-            &part.columns_meta,
-            chunks,
+            block_path,
+            num_rows,
+            compression,
+            column_metas,
+            column_chunks,
             None,
         );
 

src/query/storages/fuse/src/io/read/block/block_reader_parquet_deserialize.rs

@@ -35,12 +35,3 @@ impl ReadSettings {
         })
     }
 }
-
-impl Default for ReadSettings {
-    fn default() -> Self {
-        ReadSettings {
-            storage_io_min_bytes_for_seek: 1024,
-            storage_io_max_page_bytes_for_read: 1024 * 1024,
-        }
-    }
-}