add test

Author: zhyass

src/query/ee/src/hilbert_clustering/handler.rs

@@ -63,7 +63,7 @@ impl HilbertClusteringHandler for RealHilbertClusteringHandler {
         let max_bytes_per_block = fuse_table.get_option(
             FUSE_OPT_KEY_BLOCK_IN_MEM_SIZE_THRESHOLD,
             DEFAULT_BLOCK_BUFFER_SIZE,
-        );
+        ) * 2;
         let hilbert_min_bytes = std::cmp::max(
             hilbert_clustering_min_bytes,
             max_bytes_per_block * block_per_seg,
@@ -76,6 +76,7 @@ impl HilbertClusteringHandler for RealHilbertClusteringHandler {
         let mut checker = ReclusterChecker::new(
             cluster_key_id,
             hilbert_min_bytes,
+            block_per_seg,
             push_downs.as_ref().is_none_or(|v| v.filters.is_none()),
         );
         'FOR: for chunk in segment_locations.chunks(chunk_size) {
@@ -139,19 +140,29 @@ struct ReclusterChecker {
     hilbert_min_bytes: usize,
     total_bytes: usize,
 
+    hilbert_min_blocks: usize,
+    total_blocks: usize,
+
     finished: bool,
     // Whether the target segments is at the head of snapshot.
     head_of_snapshot: bool,
 }
 
 impl ReclusterChecker {
-    fn new(default_cluster_id: u32, hilbert_min_bytes: usize, head_of_snapshot: bool) -> Self {
+    fn new(
+        default_cluster_id: u32,
+        hilbert_min_bytes: usize,
+        hilbert_min_blocks: usize,
+        head_of_snapshot: bool,
+    ) -> Self {
         Self {
             segments: vec![],
             last_segment: None,
             default_cluster_id,
+            hilbert_min_blocks,
             hilbert_min_bytes,
             total_bytes: 0,
+            total_blocks: 0,
             finished: false,
             head_of_snapshot,
         }
@@ -164,10 +175,14 @@ impl ReclusterChecker {
 
         if segment_should_recluster || !self.head_of_snapshot {
             self.total_bytes += segment.summary.uncompressed_byte_size as usize;
+            self.total_blocks += segment.summary.block_count as usize;
             self.segments.push((location.clone(), segment.clone()));
         }
 
-        if !segment_should_recluster || self.total_bytes >= self.hilbert_min_bytes {
+        if !segment_should_recluster
+            || (self.total_bytes >= self.hilbert_min_bytes
+                && self.total_blocks >= self.hilbert_min_blocks)
+        {
             if self.check_for_recluster() {
                 self.finished = true;
                 return true;
@@ -208,6 +223,7 @@ impl ReclusterChecker {
 
     fn reset(&mut self) {
         self.total_bytes = 0;
+        self.total_blocks = 0;
         self.head_of_snapshot = false;
         self.segments.clear();
     }

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

@@ -153,7 +153,7 @@ 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 = self.max_bytes_per_block.min(400 * 1024 * 1024);
-        let min_bytes_per_block = self.min_bytes_per_block.min(100 * 1024 * 1024);
+        let min_bytes_per_block = max_bytes_per_block / 2;
         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/expression/tests/it/block_thresholds.rs

@@ -15,7 +15,7 @@
 use databend_common_expression::BlockThresholds;
 
 fn default_thresholds() -> BlockThresholds {
-    BlockThresholds::new(1000, 1_000_000, 100_000, 4)
+    BlockThresholds::new(1_000, 1_000_000, 100_000, 4)
 }
 
 #[test]
@@ -101,14 +101,41 @@ fn test_calc_rows_for_recluster() {
     );
 
     // Case 1: If the block size is too bigger.
-    let result = t.calc_rows_for_recluster(4_000, 30_000_000, 600_000);
-    assert_eq!(result, 400);
+    let result = t.calc_rows_for_recluster(4_500, 30_000_000, 600_000);
+    assert_eq!(result, 300);
 
     // Case 2: If the block size is too smaller.
-    let result = t.calc_rows_for_recluster(4_000, 2_000_000, 600_000);
-    assert_eq!(result, 800);
+    let result = t.calc_rows_for_recluster(4_000, 4_000_000, 600_000);
+    assert_eq!(result, 1000);
 
     // Case 3: use the compressed-based block count.
     let result = t.calc_rows_for_recluster(4_000, 10_000_000, 600_000);
     assert_eq!(result, 667);
 }
+
+#[test]
+fn test_calc_partitions_for_recluster() {
+    let t = default_thresholds();
+
+    // compact enough to skip further calculations
+    assert_eq!(t.calc_partitions_for_recluster(1000, 500_000, 100_000), 1);
+
+    // row-based block count exceeds compressed-based block count, use max rows per block.
+    assert_eq!(
+        t.calc_partitions_for_recluster(10_000, 2_000_000, 100_000),
+        10
+    );
+
+    // Case 1: If the block size is too bigger.
+    let result = t.calc_partitions_for_recluster(4_500, 30_000_000, 600_000);
+    assert_eq!(result, 15);
+
+    // Case 2: If the block size is too smaller.
+    let result = t.calc_partitions_for_recluster(4_000, 4_000_000, 600_000);
+    assert_eq!(result, 4);
+
+    // Case 3: use the compressed-based block count.
+    let result = t.calc_partitions_for_recluster(4_000, 10_000_000, 600_000);
+    assert_eq!(result, 6);
+}
+

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

@@ -78,6 +78,7 @@ impl PipelineBuilder {
                         partition.range_width,
                         window_spill_settings.clone(),
                         disk_spill.clone(),
+                        partition.rows_per_block,
                         partition.bytes_per_block,
                     )?,
                 )))

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

@@ -73,6 +73,7 @@ impl TransformHilbertCollect {
         num_partitions: usize,
         memory_settings: MemorySettings,
         disk_spill: Option<SpillerDiskConfig>,
+        max_block_rows: usize,
         max_block_size: usize,
     ) -> Result<Self> {
         // Calculate the partition ids collected by the processor.
@@ -99,13 +100,10 @@ impl TransformHilbertCollect {
         let spiller = Spiller::create(ctx, operator, spill_config)?;
 
         // Create the window partition buffer.
-        let sort_block_size = settings.get_window_partition_sort_block_size()? as usize;
-        let buffer = WindowPartitionBuffer::new(
-            spiller,
-            partitions.len(),
-            sort_block_size,
-            memory_settings,
-        )?;
+        let max_block_rows =
+            max_block_rows.min(settings.get_window_partition_sort_block_size()? as usize);
+        let buffer =
+            WindowPartitionBuffer::new(spiller, partitions.len(), max_block_rows, memory_settings)?;
 
         Ok(Self {
             input,