fix

Author: zhyass

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

@@ -39,22 +39,20 @@ use databend_common_sql::executor::physical_plans::Recluster;
 use databend_common_sql::StreamContext;
 use databend_common_storages_factory::Table;
 use databend_common_storages_fuse::io::StreamBlockProperties;
-use databend_common_storages_fuse::operations::TransformBlockBuilder;
 use databend_common_storages_fuse::operations::TransformBlockWriter;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
 use databend_common_storages_fuse::FuseTable;
 use databend_common_storages_fuse::TableContext;
 
 use crate::pipelines::builders::SortPipelineBuilder;
 use crate::pipelines::processors::transforms::ReclusterPartitionExchange;
-use crate::pipelines::processors::transforms::ReclusterPartitionStrategys;
+use crate::pipelines::processors::transforms::ReclusterPartitionStrategy;
 use crate::pipelines::processors::transforms::SampleState;
 use crate::pipelines::processors::transforms::TransformAddOrderColumn;
 use crate::pipelines::processors::transforms::TransformAddStreamColumns;
 use crate::pipelines::processors::transforms::TransformPartitionCollect;
 use crate::pipelines::processors::transforms::TransformRangePartitionIndexer;
 use crate::pipelines::processors::transforms::TransformReclusterCollect;
-use crate::pipelines::processors::transforms::TransformReclusterPartition;
 use crate::pipelines::PipelineBuilder;
 
 impl PipelineBuilder {
@@ -173,7 +171,6 @@ impl PipelineBuilder {
                     let fields_with_cluster_key = properties.fields_with_cluster_key();
                     let schema = DataSchemaRefExt::create(fields_with_cluster_key);
                     let schema = add_order_field(schema, &sort_desc);
-                    let order_offset = schema.fields.len() - 1;
 
                     let num_processors = self.main_pipeline.output_len();
                     let sample_size = self
@@ -203,44 +200,24 @@ impl PipelineBuilder {
                     let processor_id = AtomicUsize::new(0);
 
                     let settings = self.ctx.get_settings();
-                    let enable_writings = settings.get_enable_block_stream_writes()?;
-                    if enable_writings {
-                        let memory_settings = MemorySettings::disable_spill();
-                        self.main_pipeline.add_transform(|input, output| {
-                            let strategy =
-                                ReclusterPartitionStrategys::new(properties.clone(), order_offset);
-
-                            Ok(ProcessorPtr::create(Box::new(
-                                TransformPartitionCollect::new(
-                                    self.ctx.clone(),
-                                    input,
-                                    output,
-                                    &settings,
-                                    processor_id.fetch_add(1, atomic::Ordering::AcqRel),
-                                    num_processors,
-                                    partitions,
-                                    memory_settings.clone(),
-                                    None,
-                                    strategy,
-                                )?,
-                            )))
-                        })?;
-
-                        self.main_pipeline.add_transform(|input, output| {
-                            TransformBlockBuilder::try_create(input, output, properties.clone())
-                        })?;
-                    } else {
-                        self.main_pipeline.add_transform(|input, output| {
-                            TransformReclusterPartition::try_create(
+                    let memory_settings = MemorySettings::disable_spill();
+                    self.main_pipeline.add_transform(|input, output| {
+                        let strategy = ReclusterPartitionStrategy::new(properties.clone());
+                        Ok(ProcessorPtr::create(Box::new(
+                            TransformPartitionCollect::new(
+                                self.ctx.clone(),
                                 input,
                                 output,
-                                properties.clone(),
+                                &settings,
                                 processor_id.fetch_add(1, atomic::Ordering::AcqRel),
                                 num_processors,
                                 partitions,
-                            )
-                        })?;
-                    }
+                                memory_settings.clone(),
+                                None,
+                                strategy,
+                            )?,
+                        )))
+                    })?;
 
                     self.main_pipeline.add_async_accumulating_transformer(|| {
                         TransformBlockWriter::create(

src/query/service/src/pipelines/processors/transforms/recluster/mod.rs

@@ -25,7 +25,6 @@ pub use range_bound_sampler::RangeBoundSampler;
 pub use recluster_partition_exchange::ReclusterPartitionExchange;
 pub use recluster_partition_strategy::CompactPartitionStrategy;
 pub use recluster_partition_strategy::ReclusterPartitionStrategy;
-pub use recluster_partition_strategy::ReclusterPartitionStrategys;
 pub use recluster_sample_state::SampleState;
 pub use transform_add_order_column::TransformAddOrderColumn;
 pub use transform_range_partition_indexer::TransformRangePartitionIndexer;

src/query/service/src/pipelines/processors/transforms/recluster/recluster_partition_strategy.rs

@@ -16,8 +16,6 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
-use databend_common_expression::LimitType;
-use databend_common_expression::SortColumnDescription;
 use databend_common_storages_fuse::io::StreamBlockBuilder;
 use databend_common_storages_fuse::io::StreamBlockProperties;
 
@@ -34,10 +32,6 @@ impl ReclusterPartitionStrategy {
     pub fn new(properties: Arc<StreamBlockProperties>) -> Self {
         Self { properties }
     }
-
-    fn concat_blocks(blocks: Vec<DataBlock>) -> Result<DataBlock> {
-        DataBlock::concat(&blocks)
-    }
 }
 
 impl PartitionProcessStrategy for ReclusterPartitionStrategy {
@@ -74,21 +68,23 @@ impl PartitionProcessStrategy for ReclusterPartitionStrategy {
                 continue;
             }
             if !staged_blocks.is_empty() {
-                compacted.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
+                compacted.push(std::mem::take(&mut staged_blocks));
             }
             std::mem::swap(&mut staged_blocks, &mut pending_blocks);
             accumulated_rows = 0;
             accumulated_bytes = 0;
         }
         staged_blocks.append(&mut pending_blocks);
         if !staged_blocks.is_empty() {
-            compacted.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
+            compacted.push(std::mem::take(&mut staged_blocks));
         }
 
         let mut result = Vec::new();
         let mut builder = StreamBlockBuilder::try_new_with_config(self.properties.clone())?;
-        for block in compacted {
-            builder.write(block)?;
+        for blocks in compacted {
+            for block in blocks {
+                builder.write(block)?;
+            }
             if builder.need_flush() {
                 let serialized = builder.finish()?;
                 result.push(DataBlock::empty_with_meta(Box::new(serialized)));
@@ -177,89 +173,3 @@ impl PartitionProcessStrategy for CompactPartitionStrategy {
         Ok(result)
     }
 }
-
-pub struct ReclusterPartitionStrategys {
-    properties: Arc<StreamBlockProperties>,
-    sort_desc: Vec<SortColumnDescription>,
-}
-
-impl ReclusterPartitionStrategys {
-    pub fn new(properties: Arc<StreamBlockProperties>, offset: usize) -> Self {
-        Self {
-            properties,
-            sort_desc: vec![SortColumnDescription {
-                offset,
-                asc: true,
-                nulls_first: false,
-            }],
-        }
-    }
-
-    fn concat_blocks(blocks: Vec<DataBlock>) -> Result<DataBlock> {
-        DataBlock::concat(&blocks)
-    }
-}
-
-impl PartitionProcessStrategy for ReclusterPartitionStrategys {
-    const NAME: &'static str = "Recluster";
-
-    fn calc_partitions(
-        &self,
-        processor_id: usize,
-        num_processors: usize,
-        num_partitions: usize,
-    ) -> Vec<usize> {
-        (0..num_partitions)
-            .filter(|&partition| (partition * num_processors) / num_partitions == processor_id)
-            .collect()
-    }
-
-    /// Stream write each block, and flush it conditionally based on builder status
-    /// and input size estimation.
-    fn process_data_blocks(&self, data_blocks: Vec<DataBlock>) -> Result<Vec<DataBlock>> {
-        let blocks_num = data_blocks.len();
-        let mut accumulated_rows = 0;
-        let mut accumulated_bytes = 0;
-        let mut pending_blocks = Vec::with_capacity(blocks_num);
-        let mut staged_blocks = Vec::with_capacity(blocks_num);
-        let mut compacted = Vec::with_capacity(blocks_num);
-        for block in data_blocks {
-            accumulated_rows += block.num_rows();
-            accumulated_bytes += block.estimate_block_size();
-            pending_blocks.push(block);
-            if !self
-                .properties
-                .check_large_enough(accumulated_rows, accumulated_bytes)
-            {
-                continue;
-            }
-            if !staged_blocks.is_empty() {
-                compacted.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
-            }
-            std::mem::swap(&mut staged_blocks, &mut pending_blocks);
-            accumulated_rows = 0;
-            accumulated_bytes = 0;
-        }
-        staged_blocks.append(&mut pending_blocks);
-        if !staged_blocks.is_empty() {
-            compacted.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
-        }
-
-        let mut result = Vec::new();
-        let mut builder = StreamBlockBuilder::try_new_with_config(self.properties.clone())?;
-        for block in compacted {
-            let block = DataBlock::sort_with_type(&block, &self.sort_desc, LimitType::None)?;
-            builder.write(block)?;
-            if builder.need_flush() {
-                let serialized = builder.finish()?;
-                result.push(DataBlock::empty_with_meta(Box::new(serialized)));
-                builder = StreamBlockBuilder::try_new_with_config(self.properties.clone())?;
-            }
-        }
-        if !builder.is_empty() {
-            let serialized = builder.finish()?;
-            result.push(DataBlock::empty_with_meta(Box::new(serialized)));
-        }
-        Ok(result)
-    }
-}

src/query/service/src/pipelines/processors/transforms/recluster/recluster_sample_state.rs

@@ -114,7 +114,7 @@ impl SampleStateInner {
             let weight = weights[idx];
             cum_weight += weight;
 
-            if cum_weight >= target && previous_bound.map_or(true, |prev| value > prev) {
+            if cum_weight >= target && previous_bound.is_none_or(|prev| value > prev) {
                 if unlikely(value == max_val) {
                     self.max_value = Some(max_val.clone());
                     break;

src/query/service/src/pipelines/processors/transforms/recluster/transform_range_partition_indexer.rs

@@ -118,28 +118,40 @@ impl Processor for TransformRangePartitionIndexer {
     fn process(&mut self) -> Result<()> {
         let start = Instant::now();
         if let Some(mut block) = self.input_data.pop() {
-            let bound_len = self.bounds.len();
             let num_rows = block.num_rows();
             let mut builder = Vec::with_capacity(num_rows);
             let last_col = block.get_last_column().as_binary().unwrap();
-            for index in 0..num_rows {
-                let val = unsafe { last_col.index_unchecked(index) };
-                if unlikely(self.max_value.as_ref().is_some_and(|v| val >= v.as_slice())) {
-                    let range_id = bound_len + 1;
-                    builder.push(range_id as u64);
-                    continue;
+            if let Some(max_value) = self.max_value.as_ref() {
+                let bound_len = self.bounds.len();
+                for index in 0..num_rows {
+                    let val = unsafe { last_col.index_unchecked(index) };
+                    if unlikely(val >= max_value.as_slice()) {
+                        let range_id = bound_len + 1;
+                        builder.push(range_id as u64);
+                        continue;
+                    }
+
+                    let idx = self
+                        .bounds
+                        .binary_search_by(|b| b.as_slice().cmp(val))
+                        .unwrap_or_else(|i| i);
+                    builder.push(idx as u64);
+                }
+            } else {
+                for index in 0..num_rows {
+                    let val = unsafe { last_col.index_unchecked(index) };
+                    let idx = self
+                        .bounds
+                        .binary_search_by(|b| b.as_slice().cmp(val))
+                        .unwrap_or_else(|i| i);
+                    builder.push(idx as u64);
                 }
-
-                let idx = self
-                    .bounds
-                    .binary_search_by(|b| b.as_slice().cmp(val))
-                    .unwrap_or_else(|i| i);
-                builder.push(idx as u64);
             }
-
+            block.pop_columns(1);
             block.add_column(UInt64Type::from_data(builder));
             self.output_data.push_back(block);
         }
+
         log::info!("Recluster range output: {:?}", start.elapsed());
         Ok(())
     }

src/query/settings/src/settings_default.rs

@@ -1312,13 +1312,6 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
-                ("enable_block_stream_writes", DefaultSettingValue {
-                    value: UserSettingValue::UInt64(0),
-                    desc: "Enables block stream write",
-                    mode: SettingMode::Both,
-                    scope: SettingScope::Both,
-                    range: Some(SettingRange::Numeric(0..=1)),
-                }),
                 ("trace_sample_rate", DefaultSettingValue {
                     value: UserSettingValue::UInt64(1),
                     desc: "Setting the trace sample rate. The value should be between '0' and '100'",

src/query/settings/src/settings_getter_setter.rs

@@ -964,10 +964,6 @@ impl Settings {
         Ok(self.try_get_u64("enable_block_stream_write")? == 1)
     }
 
-    pub fn get_enable_block_stream_writes(&self) -> Result<bool> {
-        Ok(self.try_get_u64("enable_block_stream_writes")? == 1)
-    }
-
     pub fn get_statement_queue_ttl_in_seconds(&self) -> Result<u64> {
         self.try_get_u64("statement_queue_ttl_in_seconds")
     }