chore: calc rows per block for recluster (#17639)

* fix calc rows per block

* fix

* enable compact when block_size_threshold changed

* fix

* fix hilbert recluster get small blocks

* add unit test

* fix test

Author: zhyass

src/common/io/src/constants.rs

@@ -25,15 +25,13 @@ pub const INF_BYTES_LOWER: &str = "inf";
 pub const INF_BYTES_LONG: &str = "Infinity";
 
 // The size of the I/O read/write block buffer by default.
-pub const DEFAULT_BLOCK_BUFFER_SIZE: usize = 100 * 1024 * 1024;
+pub const DEFAULT_BLOCK_BUFFER_SIZE: usize = 125 * 1024 * 1024;
 // The size of the block compressed by default.
-pub const DEFAULT_BLOCK_COMPRESSED_SIZE: usize = 10 * 1024 * 1024;
+pub const DEFAULT_BLOCK_COMPRESSED_SIZE: usize = 16 * 1024 * 1024;
 // The size of the I/O read/write block index buffer by default.
 pub const DEFAULT_BLOCK_INDEX_BUFFER_SIZE: usize = 300 * 1024;
 // The max number of a block by default.
-pub const DEFAULT_BLOCK_MAX_ROWS: usize = 1000 * 1000;
-// The min number of a block by default.
-pub const DEFAULT_BLOCK_MIN_ROWS: usize = 800 * 1000;
+pub const DEFAULT_BLOCK_ROW_COUNT: usize = 1000 * 1000;
 // The number of blocks in a segment by default.
 pub const DEFAULT_BLOCK_PER_SEGMENT: usize = 1000;
 /// The number of bytes read at the end of the file on first read

src/query/catalog/src/table.rs

@@ -28,11 +28,6 @@ use databend_common_expression::BlockThresholds;
 use databend_common_expression::ColumnId;
 use databend_common_expression::Scalar;
 use databend_common_expression::TableSchema;
-use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
-use databend_common_io::constants::DEFAULT_BLOCK_COMPRESSED_SIZE;
-use databend_common_io::constants::DEFAULT_BLOCK_MAX_ROWS;
-use databend_common_io::constants::DEFAULT_BLOCK_MIN_ROWS;
-use databend_common_io::constants::DEFAULT_BLOCK_PER_SEGMENT;
 use databend_common_meta_app::app_error::AppError;
 use databend_common_meta_app::app_error::UnknownTableId;
 use databend_common_meta_app::schema::TableIdent;
@@ -382,13 +377,7 @@ pub trait Table: Sync + Send {
     }
 
     fn get_block_thresholds(&self) -> BlockThresholds {
-        BlockThresholds {
-            max_rows_per_block: DEFAULT_BLOCK_MAX_ROWS,
-            min_rows_per_block: DEFAULT_BLOCK_MIN_ROWS,
-            max_bytes_per_block: DEFAULT_BLOCK_BUFFER_SIZE,
-            max_bytes_per_file: DEFAULT_BLOCK_COMPRESSED_SIZE,
-            block_per_segment: DEFAULT_BLOCK_PER_SEGMENT,
-        }
+        BlockThresholds::default()
     }
 
     #[async_backtrace::framed]

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

@@ -14,26 +14,32 @@
 
 use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
 use databend_common_io::constants::DEFAULT_BLOCK_COMPRESSED_SIZE;
-use databend_common_io::constants::DEFAULT_BLOCK_MAX_ROWS;
-use databend_common_io::constants::DEFAULT_BLOCK_MIN_ROWS;
 use databend_common_io::constants::DEFAULT_BLOCK_PER_SEGMENT;
+use databend_common_io::constants::DEFAULT_BLOCK_ROW_COUNT;
 
 #[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
 pub struct BlockThresholds {
     pub max_rows_per_block: usize,
     pub min_rows_per_block: usize,
+
     pub max_bytes_per_block: usize,
-    pub max_bytes_per_file: usize,
+    pub min_bytes_per_block: usize,
+
+    pub max_compressed_per_block: usize,
+    pub min_compressed_per_block: usize,
+
     pub block_per_segment: usize,
 }
 
 impl Default for BlockThresholds {
     fn default() -> BlockThresholds {
         BlockThresholds {
-            max_rows_per_block: DEFAULT_BLOCK_MAX_ROWS,
-            min_rows_per_block: DEFAULT_BLOCK_MIN_ROWS,
+            max_rows_per_block: DEFAULT_BLOCK_ROW_COUNT,
+            min_rows_per_block: (DEFAULT_BLOCK_ROW_COUNT * 4).div_ceil(5),
             max_bytes_per_block: DEFAULT_BLOCK_BUFFER_SIZE,
-            max_bytes_per_file: DEFAULT_BLOCK_COMPRESSED_SIZE,
+            min_bytes_per_block: (DEFAULT_BLOCK_BUFFER_SIZE * 4).div_ceil(5),
+            max_compressed_per_block: DEFAULT_BLOCK_COMPRESSED_SIZE,
+            min_compressed_per_block: (DEFAULT_BLOCK_COMPRESSED_SIZE * 4).div_ceil(5),
             block_per_segment: DEFAULT_BLOCK_PER_SEGMENT,
         }
     }
@@ -42,16 +48,17 @@ impl Default for BlockThresholds {
 impl BlockThresholds {
     pub fn new(
         max_rows_per_block: usize,
-        min_rows_per_block: usize,
         max_bytes_per_block: usize,
-        max_bytes_per_file: usize,
+        max_compressed_per_block: usize,
         block_per_segment: usize,
     ) -> Self {
         BlockThresholds {
             max_rows_per_block,
-            min_rows_per_block,
+            min_rows_per_block: (max_rows_per_block * 4).div_ceil(5),
             max_bytes_per_block,
-            max_bytes_per_file,
+            min_bytes_per_block: (max_bytes_per_block * 4).div_ceil(5),
+            max_compressed_per_block,
+            min_compressed_per_block: (max_compressed_per_block * 4).div_ceil(5),
             block_per_segment,
         }
     }
@@ -64,8 +71,8 @@ impl BlockThresholds {
         file_size: usize,
     ) -> bool {
         row_count >= self.min_rows_per_block
-            || block_size >= self.max_bytes_per_block
-            || file_size >= self.max_bytes_per_file
+            || block_size >= self.min_bytes_per_block
+            || file_size >= self.min_compressed_per_block
     }
 
     #[inline]
@@ -78,53 +85,85 @@ impl BlockThresholds {
     ) -> bool {
         total_blocks >= self.block_per_segment
             && (total_rows >= self.min_rows_per_block * self.block_per_segment
-                || total_bytes >= self.max_bytes_per_block * self.block_per_segment
-                || total_compressed >= self.max_bytes_per_file * self.block_per_segment)
+                || total_bytes >= self.min_bytes_per_block * self.block_per_segment
+                || total_compressed >= self.min_compressed_per_block * self.block_per_segment)
     }
 
     #[inline]
     pub fn check_large_enough(&self, row_count: usize, block_size: usize) -> bool {
-        row_count >= self.min_rows_per_block || block_size >= self.max_bytes_per_block
+        row_count >= self.min_rows_per_block || block_size >= self.min_bytes_per_block
     }
 
     #[inline]
     pub fn check_for_compact(&self, row_count: usize, block_size: usize) -> bool {
-        row_count < 2 * self.min_rows_per_block && block_size < 2 * self.max_bytes_per_block
+        row_count < 2 * self.min_rows_per_block && block_size < 2 * self.min_bytes_per_block
     }
 
     #[inline]
     pub fn check_too_small(&self, row_count: usize, block_size: usize, file_size: usize) -> bool {
         row_count < self.min_rows_per_block / 2
-            && block_size < self.max_bytes_per_block / 2
-            && file_size <= self.max_bytes_per_file / 2
+            && block_size < self.min_bytes_per_block / 2
+            && file_size < self.min_compressed_per_block / 2
     }
 
     #[inline]
-    pub fn calc_rows_per_block(
+    pub fn calc_rows_for_compact(&self, total_bytes: usize, total_rows: usize) -> usize {
+        if self.check_for_compact(total_rows, total_bytes) {
+            return total_rows;
+        }
+
+        let block_num_by_rows = std::cmp::max(total_rows / self.min_rows_per_block, 1);
+        let block_num_by_size = total_bytes / self.min_bytes_per_block;
+        if block_num_by_rows >= block_num_by_size {
+            return self.max_rows_per_block;
+        }
+        total_rows.div_ceil(block_num_by_size)
+    }
+
+    /// Calculates the optimal number of rows per block based on total data size and row count.
+    ///
+    /// # Parameters
+    /// - `total_bytes`: The total size of the data in bytes.
+    /// - `total_rows`: The total number of rows in the data.
+    /// - `total_compressed`: The total compressed size of the data in bytes.
+    ///
+    /// # Returns
+    /// - The calculated number of rows per block that satisfies the thresholds.
+    #[inline]
+    pub fn calc_rows_for_recluster(
         &self,
-        total_bytes: usize,
         total_rows: usize,
+        total_bytes: usize,
         total_compressed: usize,
     ) -> usize {
-        if self.check_for_compact(total_rows, total_bytes) {
+        // Check if the data is compact enough to skip further calculations.
+        if self.check_for_compact(total_rows, total_bytes)
+            && total_compressed < 2 * self.min_compressed_per_block
+        {
             return total_rows;
         }
 
         let block_num_by_rows = std::cmp::max(total_rows / self.min_rows_per_block, 1);
-        let block_num_by_size = std::cmp::max(
-            total_bytes / self.max_bytes_per_block,
-            total_compressed / self.max_bytes_per_file,
-        );
-        if block_num_by_rows >= block_num_by_size {
+        let block_num_by_compressed = total_compressed.div_ceil(self.max_compressed_per_block);
+        // If row-based block count exceeds compressed-based block count, use max rows per block.
+        if block_num_by_rows >= block_num_by_compressed {
             return self.max_rows_per_block;
         }
 
-        let mut rows_per_block = total_rows.div_ceil(block_num_by_size);
-        if rows_per_block < self.max_rows_per_block / 2 {
-            // If block rows < 500_000, max_bytes_per_block set to 125M
-            let block_num_by_size = (4 * block_num_by_size / 5).max(1);
-            rows_per_block = total_rows.div_ceil(block_num_by_size);
-        }
-        rows_per_block
+        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 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)
+        } else if bytes_per_block < min_bytes_per_block {
+            // Case 2: If the block size is too smaller.
+            total_bytes / min_bytes_per_block
+        } else {
+            // Case 3: Otherwise, use the compressed-based block count.
+            block_num_by_compressed
+        };
+        total_rows.div_ceil(block_nums.max(1)).max(1)
     }
 }

src/query/expression/tests/it/block.rs

@@ -1,3 +1,17 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 use databend_common_column::buffer::Buffer;
 use databend_common_expression::block_debug::box_render;
 use databend_common_expression::types::number::NumberScalar;

src/query/expression/tests/it/block_thresholds.rs

@@ -0,0 +1,114 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_expression::BlockThresholds;
+
+fn default_thresholds() -> BlockThresholds {
+    BlockThresholds::new(1000, 1_000_000, 100_000, 4)
+}
+
+#[test]
+fn test_check_perfect_block() {
+    let t = default_thresholds();
+
+    // All below threshold
+    assert!(!t.check_perfect_block(100, 1000, 1000));
+
+    // Any one above threshold
+    assert!(t.check_perfect_block(800, 1000, 1000));
+    assert!(t.check_perfect_block(100, 800_000, 1000));
+    assert!(t.check_perfect_block(100, 1000, 80_000));
+}
+
+#[test]
+fn test_check_perfect_segment() {
+    let t = default_thresholds();
+    let total_blocks = 4;
+
+    // Below threshold
+    assert!(!t.check_perfect_segment(total_blocks, 100, 1000, 1000));
+
+    // One condition meets threshold
+    assert!(t.check_perfect_segment(total_blocks, 4000, 1000, 1000));
+    assert!(t.check_perfect_segment(total_blocks, 100, 4_000_000, 1000));
+    assert!(t.check_perfect_segment(total_blocks, 100, 1000, 320_000));
+}
+
+#[test]
+fn test_check_large_enough() {
+    let t = default_thresholds();
+
+    assert!(!t.check_large_enough(100, 1000));
+    assert!(t.check_large_enough(800, 1000));
+    assert!(t.check_large_enough(100, 800_000));
+}
+
+#[test]
+fn test_check_for_compact() {
+    let t = default_thresholds();
+
+    assert!(t.check_for_compact(1500, 1_500_000)); // just under 2x min
+    assert!(!t.check_for_compact(3000, 3_000_000)); // too large
+}
+
+#[test]
+fn test_check_too_small() {
+    let t = default_thresholds();
+
+    assert!(t.check_too_small(50, 10_000, 10_000)); // All very small
+    assert!(!t.check_too_small(800, 10_000, 10_000)); // Row count not too small
+    assert!(!t.check_too_small(50, 800_000, 10_000)); // Block size not too small
+}
+
+#[test]
+fn test_calc_rows_for_compact() {
+    let t = default_thresholds();
+
+    assert_eq!(t.calc_rows_for_compact(500_000, 1000), 1000);
+
+    // Block number by rows wins → max_rows_per_block
+    assert_eq!(
+        t.calc_rows_for_compact(2_000_000, 10_000),
+        t.max_rows_per_block
+    );
+
+    // Size-based block number wins
+    assert_eq!(t.calc_rows_for_compact(4_000_000, 2000), 400);
+}
+
+#[test]
+fn test_calc_rows_for_recluster() {
+    let t = default_thresholds();
+
+    // compact enough to skip further calculations
+    assert_eq!(t.calc_rows_for_recluster(1000, 500_000, 100_000), 1000);
+
+    // row-based block count exceeds compressed-based block count, use max rows per block.
+    assert_eq!(
+        t.calc_rows_for_recluster(10_000, 2_000_000, 100_000),
+        t.max_rows_per_block
+    );
+
+    // 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);
+
+    // 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);
+
+    // 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);
+}

src/query/expression/tests/it/main.rs

@@ -26,6 +26,7 @@ extern crate core;
 
 mod arrow;
 mod block;
+mod block_thresholds;
 mod common;
 mod decimal;
 mod fill_field_default_value;

src/query/pipeline/transforms/src/processors/transforms/transform_compact_builder.rs

@@ -69,7 +69,7 @@ impl AccumulatingTransform for BlockCompactBuilder {
             // holding slices of blocks to merge later may lead to oom, so
             // 1. we expect blocks from file formats are not slice.
             // 2. if block is split here, cut evenly and emit them at once.
-            let rows_per_block = self.thresholds.calc_rows_per_block(num_bytes, num_rows, 0);
+            let rows_per_block = self.thresholds.calc_rows_for_compact(num_bytes, num_rows);
             Ok(vec![DataBlock::empty_with_meta(Box::new(
                 BlockCompactMeta::Split {
                     block: data,

src/query/service/src/interpreters/common/table_option_validation.rs

@@ -20,7 +20,7 @@ use chrono::Duration;
 use databend_common_ast::ast::Engine;
 use databend_common_exception::ErrorCode;
 use databend_common_expression::TableSchemaRef;
-use databend_common_io::constants::DEFAULT_BLOCK_MAX_ROWS;
+use databend_common_io::constants::DEFAULT_BLOCK_ROW_COUNT;
 use databend_common_settings::Settings;
 use databend_common_sql::BloomIndexColumns;
 use databend_common_storages_fuse::FUSE_OPT_KEY_BLOCK_IN_MEM_SIZE_THRESHOLD;
@@ -154,7 +154,7 @@ pub fn is_valid_row_per_block(
         let row_per_block = value.parse::<u64>()?;
         let error_str = "invalid row_per_block option, can't be over 1000000";
 
-        if row_per_block > DEFAULT_BLOCK_MAX_ROWS as u64 {
+        if row_per_block > DEFAULT_BLOCK_ROW_COUNT as u64 {
             error!("{}", error_str);
             return Err(ErrorCode::TableOptionInvalid(error_str));
         }