hilbert recluster support block stream write

Author: zhyass

src/query/expression/src/utils/block_thresholds.rs

@@ -152,8 +152,8 @@ impl BlockThresholds {
 
         let bytes_per_block = total_bytes.div_ceil(block_num_by_compressed);
         // Adjust the number of blocks based on block size thresholds.
-        let max_bytes_per_block = (4 * self.min_bytes_per_block).min(400 * 1024 * 1024);
-        let min_bytes_per_block = (self.min_bytes_per_block / 2).min(50 * 1024 * 1024);
+        let max_bytes_per_block = self.max_bytes_per_block.min(400 * 1024 * 1024);
+        let min_bytes_per_block = self.min_bytes_per_block.min(100 * 1024 * 1024);
         let block_nums = if bytes_per_block > max_bytes_per_block {
             // Case 1: If the block size is too bigger.
             total_bytes.div_ceil(max_bytes_per_block)

src/query/service/src/interpreters/interpreter_table_recluster.rs

@@ -28,6 +28,7 @@ use databend_common_catalog::table::TableExt;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::type_check::check_function;
+use databend_common_expression::types::NumberScalar;
 use databend_common_expression::DataBlock;
 use databend_common_expression::Scalar;
 use databend_common_functions::BUILTIN_FUNCTIONS;
@@ -52,6 +53,8 @@ use databend_common_sql::plans::plan_hilbert_sql;
 use databend_common_sql::plans::replace_with_constant;
 use databend_common_sql::plans::set_update_stream_columns;
 use databend_common_sql::plans::BoundColumnRef;
+use databend_common_sql::plans::ConstantExpr;
+use databend_common_sql::plans::FunctionCall;
 use databend_common_sql::plans::Plan;
 use databend_common_sql::plans::ReclusterPlan;
 use databend_common_sql::IdentifierNormalizer;
@@ -325,19 +328,7 @@ impl ReclusterTableInterpreter {
             block_thresholds.calc_rows_for_recluster(total_rows, total_bytes, total_compressed);
 
         // Calculate initial partition count based on data volume and block size
-        let mut total_partitions = std::cmp::max(total_rows / rows_per_block, 1);
-
-        // Adjust number of partitions according to the block size thresholds
-        if total_partitions < block_thresholds.block_per_segment
-            && block_thresholds.check_perfect_segment(
-                block_thresholds.block_per_segment, // this effectively by-pass the total_blocks criteria
-                total_rows,
-                total_bytes,
-                total_compressed,
-            )
-        {
-            total_partitions = block_thresholds.block_per_segment;
-        }
+        let total_partitions = std::cmp::max(total_rows / rows_per_block, 1);
 
         warn!(
             "Do hilbert recluster, total_bytes: {}, total_rows: {}, total_partitions: {}",
@@ -439,15 +430,37 @@ impl ReclusterTableInterpreter {
 
         // For distributed execution, add an exchange operator to distribute work
         if is_distributed {
+            let nodes_num = cluster.nodes.len() as u64;
+            let scalar_expr = ScalarExpr::FunctionCall(FunctionCall {
+                span: None,
+                func_name: "div".to_string(),
+                params: vec![],
+                arguments: vec![
+                    ScalarExpr::FunctionCall(FunctionCall {
+                        span: None,
+                        func_name: "multiply".to_string(),
+                        params: vec![],
+                        arguments: vec![
+                            ScalarExpr::BoundColumnRef(BoundColumnRef {
+                                span: None,
+                                column: bind_context.columns.last().unwrap().clone(),
+                            }),
+                            ScalarExpr::ConstantExpr(ConstantExpr {
+                                span: None,
+                                value: Scalar::Number(NumberScalar::UInt64(nodes_num)),
+                            }),
+                        ],
+                    }),
+                    ScalarExpr::ConstantExpr(ConstantExpr {
+                        span: None,
+                        value: Scalar::Number(NumberScalar::UInt64(total_partitions as u64)),
+                    }),
+                ],
+            });
+
             // Create an expression for the partition column,
             // i.e.`range_partition_id(hilbert_range_index({hilbert_keys_str}), [...]) AS _predicate`
-            let expr = scalar_expr_to_remote_expr(
-                &ScalarExpr::BoundColumnRef(BoundColumnRef {
-                    span: None,
-                    column: bind_context.columns.last().unwrap().clone(),
-                }),
-                plan.output_schema()?.as_ref(),
-            )?;
+            let expr = scalar_expr_to_remote_expr(&scalar_expr, plan.output_schema()?.as_ref())?;
 
             // Add exchange operator for data distribution,
             // shuffling data based on the hash of range partition IDs derived from the Hilbert index.
@@ -471,7 +484,8 @@ impl ReclusterTableInterpreter {
             plan_id: 0,
             input: plan,
             table_info: table_info.clone(),
-            num_partitions: total_partitions,
+            range_start: 0,
+            range_width: total_partitions,
             table_meta_timestamps,
             rows_per_block,
         }));

src/query/service/src/pipelines/builders/builder_hilbert_partition.rs

@@ -15,13 +15,15 @@
 use std::sync::atomic;
 use std::sync::atomic::AtomicUsize;
 
+use databend_common_catalog::table::Table;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
 use databend_common_pipeline_core::processors::ProcessorPtr;
 use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_sql::executor::physical_plans::HilbertPartition;
 use databend_common_sql::executor::physical_plans::MutationKind;
+use databend_common_storages_fuse::operations::TransformBlockWriter;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
 use databend_common_storages_fuse::statistics::ClusterStatsGenerator;
 use databend_common_storages_fuse::FuseTable;
@@ -43,10 +45,12 @@ impl PipelineBuilder {
             .ctx
             .build_table_by_table_info(&partition.table_info, None)?;
         let table = FuseTable::try_from_table(table.as_ref())?;
+        let enable_stream_writer = self.ctx.get_settings().get_enable_block_stream_write()?
+            && table.storage_format_as_parquet();
 
         self.main_pipeline.exchange(
             num_processors,
-            HilbertPartitionExchange::create(partition.num_partitions),
+            HilbertPartitionExchange::create(partition.range_start, partition.range_width),
         );
 
         let settings = self.ctx.get_settings();
@@ -77,26 +81,43 @@ impl PipelineBuilder {
                     &settings,
                     processor_id.fetch_add(1, atomic::Ordering::AcqRel),
                     num_processors,
-                    partition.num_partitions,
+                    partition.range_width,
                     window_spill_settings.clone(),
                     disk_spill.clone(),
-                    CompactStrategy::new(partition.rows_per_block, max_bytes_per_block),
+                    CompactStrategy::new(
+                        partition.rows_per_block,
+                        max_bytes_per_block,
+                        enable_stream_writer,
+                    ),
                 )?,
             )))
         })?;
 
-        self.main_pipeline
-            .add_transform(|transform_input_port, transform_output_port| {
-                let proc = TransformSerializeBlock::try_create(
+        if enable_stream_writer {
+            self.main_pipeline.add_transform(|input, output| {
+                TransformBlockWriter::try_create(
                     self.ctx.clone(),
-                    transform_input_port,
-                    transform_output_port,
+                    input,
+                    output,
                     table,
-                    ClusterStatsGenerator::default(),
-                    MutationKind::Recluster,
                     partition.table_meta_timestamps,
-                )?;
-                proc.into_processor()
+                    false,
+                )
             })
+        } else {
+            self.main_pipeline
+                .add_transform(|transform_input_port, transform_output_port| {
+                    let proc = TransformSerializeBlock::try_create(
+                        self.ctx.clone(),
+                        transform_input_port,
+                        transform_output_port,
+                        table,
+                        ClusterStatsGenerator::default(),
+                        MutationKind::Recluster,
+                        partition.table_meta_timestamps,
+                    )?;
+                    proc.into_processor()
+                })
+        }
     }
 }

src/query/service/src/pipelines/processors/transforms/window/partition/data_processor_strategy.rs

@@ -27,13 +27,19 @@ pub trait DataProcessorStrategy: Send + Sync + 'static {
 pub struct CompactStrategy {
     max_bytes_per_block: usize,
     max_rows_per_block: usize,
+    enable_stream_writer: bool,
 }
 
 impl CompactStrategy {
-    pub fn new(max_rows_per_block: usize, max_bytes_per_block: usize) -> Self {
+    pub fn new(
+        max_rows_per_block: usize,
+        max_bytes_per_block: usize,
+        enable_stream_writer: bool,
+    ) -> Self {
         Self {
             max_bytes_per_block,
             max_rows_per_block,
+            enable_stream_writer,
         }
     }
 
@@ -50,6 +56,10 @@ impl DataProcessorStrategy for CompactStrategy {
     const NAME: &'static str = "Compact";
 
     fn process_data_blocks(&self, data_blocks: Vec<DataBlock>) -> Result<Vec<DataBlock>> {
+        if self.enable_stream_writer {
+            return Ok(data_blocks);
+        }
+
         let blocks_num = data_blocks.len();
         if blocks_num < 2 {
             return Ok(data_blocks);

src/query/service/src/pipelines/processors/transforms/window/partition/hilbert_partition_exchange.rs

@@ -25,12 +25,13 @@ use databend_common_pipeline_core::processors::Exchange;
 use crate::pipelines::processors::transforms::WindowPartitionMeta;
 
 pub struct HilbertPartitionExchange {
-    num_partitions: usize,
+    start: u64,
+    width: usize,
 }
 
 impl HilbertPartitionExchange {
-    pub fn create(num_partitions: usize) -> Arc<HilbertPartitionExchange> {
-        Arc::new(HilbertPartitionExchange { num_partitions })
+    pub fn create(start: u64, width: usize) -> Arc<HilbertPartitionExchange> {
+        Arc::new(HilbertPartitionExchange { start, width })
     }
 }
 
@@ -48,20 +49,25 @@ impl Exchange for HilbertPartitionExchange {
         // Scatter the data block to different partitions.
         let indices = range_ids
             .iter()
-            .map(|&id| (id % self.num_partitions as u64) as u16)
+            .map(|&id| (id - self.start) as u16)
             .collect::<Vec<_>>();
         data_block.pop_columns(1);
-        let scatter_indices =
-            DataBlock::divide_indices_by_scatter_size(&indices, self.num_partitions);
+
+        let scatter_indices = DataBlock::divide_indices_by_scatter_size(&indices, self.width);
         // Partition the data blocks to different processors.
+        let base = self.width / n;
+        let remainder = self.width % n;
         let mut output_data_blocks = vec![vec![]; n];
-        for (partition_id, indices) in scatter_indices.iter().take(self.num_partitions).enumerate()
-        {
-            if indices.is_empty() {
-                continue;
+        for (partition_id, indices) in scatter_indices.into_iter().take(self.width).enumerate() {
+            if !indices.is_empty() {
+                let target = if partition_id < remainder * (base + 1) {
+                    partition_id / (base + 1)
+                } else {
+                    (partition_id - remainder) / base
+                };
+                let block = data_block.take_with_optimize_size(&indices)?;
+                output_data_blocks[target].push((partition_id, block));
             }
-            let block = data_block.take_with_optimize_size(indices)?;
-            output_data_blocks[partition_id % n].push((partition_id, block));
         }
 
         // Union data blocks for each processor.

src/query/service/src/schedulers/fragments/fragmenter.rs

@@ -26,6 +26,7 @@ use databend_common_sql::executor::physical_plans::ExchangeSink;
 use databend_common_sql::executor::physical_plans::ExchangeSource;
 use databend_common_sql::executor::physical_plans::FragmentKind;
 use databend_common_sql::executor::physical_plans::HashJoin;
+use databend_common_sql::executor::physical_plans::HilbertPartition;
 use databend_common_sql::executor::physical_plans::MutationSource;
 use databend_common_sql::executor::physical_plans::Recluster;
 use databend_common_sql::executor::physical_plans::ReplaceInto;
@@ -41,7 +42,6 @@ use crate::servers::flight::v1::exchange::DataExchange;
 use crate::servers::flight::v1::exchange::MergeExchange;
 use crate::servers::flight::v1::exchange::ShuffleDataExchange;
 use crate::sessions::QueryContext;
-use crate::sql::executor::physical_plans::Mutation;
 use crate::sql::executor::PhysicalPlan;
 
 /// Visitor to split a `PhysicalPlan` into fragments.
@@ -67,6 +67,7 @@ enum State {
     Compact,
     Recluster,
     Other,
+    HilbertRecluster,
 }
 
 impl Fragmenter {
@@ -170,14 +171,6 @@ impl PhysicalPlanReplacer for Fragmenter {
         Ok(PhysicalPlan::MutationSource(plan.clone()))
     }
 
-    fn replace_mutation(&mut self, plan: &Mutation) -> Result<PhysicalPlan> {
-        let input = self.replace(&plan.input)?;
-        Ok(PhysicalPlan::Mutation(Box::new(Mutation {
-            input: Box::new(input),
-            ..plan.clone()
-        })))
-    }
-
     fn replace_replace_into(&mut self, plan: &ReplaceInto) -> Result<PhysicalPlan> {
         let input = self.replace(&plan.input)?;
         self.state = State::ReplaceInto;
@@ -209,6 +202,11 @@ impl PhysicalPlanReplacer for Fragmenter {
         Ok(PhysicalPlan::Recluster(Box::new(plan.clone())))
     }
 
+    fn replace_hilbert_serialize(&mut self, plan: &HilbertPartition) -> Result<PhysicalPlan> {
+        self.state = State::HilbertRecluster;
+        Ok(PhysicalPlan::HilbertPartition(Box::new(plan.clone())))
+    }
+
     fn replace_compact_source(&mut self, plan: &CompactSource) -> Result<PhysicalPlan> {
         self.state = State::Compact;
         Ok(PhysicalPlan::CompactSource(Box::new(plan.clone())))
@@ -310,6 +308,7 @@ impl PhysicalPlanReplacer for Fragmenter {
             State::ReplaceInto => FragmentType::ReplaceInto,
             State::Compact => FragmentType::Compact,
             State::Recluster => FragmentType::Recluster,
+            State::HilbertRecluster => FragmentType::HilbertRecluster,
         };
         self.state = State::Other;
         let exchange = Self::get_exchange(self.ctx.clone(), &plan)?;