perf(cubestore): early compaction split (#11079)

Author: waralexrom

rust/cubestore/cubestore/src/config/mod.rs

@@ -393,6 +393,8 @@ pub trait ConfigObj: DIService {
 
     fn wal_split_threshold(&self) -> u64;
 
+    fn wal_split_size_threshold_bytes(&self) -> u64;
+
     fn select_worker_pool_size(&self) -> usize;
 
     fn select_worker_idle_timeout(&self) -> u64;
@@ -587,6 +589,7 @@ pub struct ConfigObjImpl {
     pub max_partition_split_threshold: u64,
     pub compaction_chunks_total_size_threshold: u64,
     pub compaction_chunks_count_threshold: u64,
+    pub compaction_chunks_threshold_multiplier: f64,
     pub compaction_chunks_in_memory_size_threshold: u64,
     pub compaction_chunks_max_lifetime_threshold: u64,
     pub compaction_in_memory_chunks_max_lifetime_threshold: u64,
@@ -597,6 +600,7 @@ pub struct ConfigObjImpl {
     pub compaction_in_memory_chunks_ratio_check_threshold: u64,
     pub compaction_in_memory_chunks_schedule_period_secs: u64,
     pub wal_split_threshold: u64,
+    pub wal_split_size_threshold_bytes: u64,
     pub data_dir: PathBuf,
     pub dump_dir: Option<PathBuf>,
     pub store_provider: FileStoreProvider,
@@ -709,11 +713,15 @@ impl ConfigObj for ConfigObjImpl {
     }
 
     fn compaction_chunks_total_size_threshold(&self) -> u64 {
-        self.compaction_chunks_total_size_threshold
+        (self.compaction_chunks_total_size_threshold as f64
+            * self.compaction_chunks_threshold_multiplier)
+            .floor() as u64
     }
 
     fn compaction_chunks_count_threshold(&self) -> u64 {
-        self.compaction_chunks_count_threshold
+        (self.compaction_chunks_count_threshold as f64
+            * self.compaction_chunks_threshold_multiplier)
+            .floor() as u64
     }
 
     fn compaction_chunks_in_memory_size_threshold(&self) -> u64 {
@@ -756,6 +764,10 @@ impl ConfigObj for ConfigObjImpl {
         self.wal_split_threshold
     }
 
+    fn wal_split_size_threshold_bytes(&self) -> u64 {
+        self.wal_split_size_threshold_bytes
+    }
+
     fn select_worker_pool_size(&self) -> usize {
         self.select_worker_pool_size
     }
@@ -1330,6 +1342,10 @@ impl Config {
                     1048576 * 8,
                 ),
                 compaction_chunks_count_threshold: env_parse("CUBESTORE_CHUNKS_COUNT_THRESHOLD", 4),
+                compaction_chunks_threshold_multiplier: env_parse(
+                    "CUBESTORE_COMPACTION_CHUNKS_THRESHOLD_MULTIPLIER",
+                    1.0_f64,
+                ),
                 compaction_chunks_in_memory_size_threshold: env_parse_size(
                     "CUBESTORE_COMPACTION_CHUNKS_IN_MEMORY_SIZE_THRESHOLD",
                     1 * 1024 * 1024 * 1024,
@@ -1533,6 +1549,12 @@ impl Config {
                 download_concurrency: env_parse("CUBESTORE_MAX_ACTIVE_DOWNLOADS", 8),
                 max_ingestion_data_frames: env_parse("CUBESTORE_MAX_DATA_FRAMES", 4),
                 wal_split_threshold: env_parse("CUBESTORE_WAL_SPLIT_THRESHOLD", 1048576 / 2),
+                wal_split_size_threshold_bytes: env_parse_size(
+                    "CUBESTORE_WAL_SPLIT_SIZE_THRESHOLD",
+                    100 * 1024 * 1024,
+                    None,
+                    None,
+                ) as u64,
                 job_runners_count: env_parse("CUBESTORE_JOB_RUNNERS", 4),
                 long_term_job_runners_count: env_parse("CUBESTORE_LONG_TERM_JOB_RUNNERS", 32),
                 csv_import_job_runners_count: env_parse("CUBESTORE_CSV_IMPORT_JOB_RUNNERS", 0),
@@ -1720,6 +1742,7 @@ impl Config {
                 partition_size_split_threshold_bytes: 2 * 1024,
                 max_partition_split_threshold: 20,
                 compaction_chunks_count_threshold: 1,
+                compaction_chunks_threshold_multiplier: 1.0,
                 compaction_chunks_in_memory_size_threshold: 3 * 1024 * 1024 * 1024,
                 compaction_chunks_total_size_threshold: 10,
                 compaction_chunks_max_lifetime_threshold: 600,
@@ -1778,6 +1801,7 @@ impl Config {
                 download_concurrency: 8,
                 max_ingestion_data_frames: 4,
                 wal_split_threshold: 262144,
+                wal_split_size_threshold_bytes: 100 * 1024 * 1024,
                 connection_timeout: 60,
                 server_name: "localhost".to_string(),
                 upload_to_remote: true,

rust/cubestore/cubestore/src/import/mod.rs

@@ -87,6 +87,20 @@ impl ImportFormat {
         Ok(rows.boxed())
     }
 
+    /// Estimated in-memory Arrow buffer footprint of a single appended value, used to
+    /// cap a data frame by size during import. Mirrors `columns_vec_buffer_size` accounting
+    /// closely enough to drive the split; exactness is unnecessary for a threshold.
+    fn estimate_arrow_value_size(column_type: &ColumnType, value: Option<&str>) -> usize {
+        match column_type {
+            ColumnType::String | ColumnType::Bytes | ColumnType::HyperLogLog(_) => {
+                value.map_or(0, |v| v.len()) + 4
+            }
+            ColumnType::Int | ColumnType::Timestamp | ColumnType::Float => 8,
+            ColumnType::Int96 | ColumnType::Decimal { .. } | ColumnType::Decimal96 { .. } => 16,
+            ColumnType::Boolean => 1,
+        }
+    }
+
     pub fn parse_column_value_str(column: &Column, value: &str) -> Result<TableValue, CubeError> {
         Ok(match column.get_column_type() {
             ColumnType::String => TableValue::String(value.to_string()),
@@ -777,12 +791,17 @@ impl ImportServiceImpl {
         };
 
         let table_cols = table.get_row().get_columns().as_slice();
+        let row_threshold = self.config_obj.wal_split_threshold() as usize;
+        let size_threshold = self.config_obj.wal_split_size_threshold_bytes() as usize;
         let mut builders = create_array_builders(table_cols);
         let mut num_rows = 0;
+        let mut estimated_bytes = 0;
         while let Some(line) = lines.next().await {
             let line = line?;
             let is_data_row = parser.visit_line(line.as_str(), |insert_pos, column, value| {
                 let builder = builders[insert_pos].as_mut();
+                estimated_bytes +=
+                    ImportFormat::estimate_arrow_value_size(column.get_column_type(), value);
                 match value {
                     None => {
                         append_value(builder, column.get_column_type(), &TableValue::Null);
@@ -799,10 +818,11 @@ impl ImportServiceImpl {
             }
             num_rows += 1;
 
-            if num_rows >= self.config_obj.wal_split_threshold() as usize {
+            if num_rows >= row_threshold || estimated_bytes >= size_threshold {
                 let mut to_add = create_array_builders(table_cols);
                 mem::swap(&mut builders, &mut to_add);
                 num_rows = 0;
+                estimated_bytes = 0;
 
                 let builded_rows = finish(to_add);
 

rust/cubestore/cubestore/src/sql/mod.rs

@@ -3315,7 +3315,9 @@ mod tests {
         Config::test("over_10k_join").update_config(|mut c| {
             c.partition_split_threshold = 1000000;
             c.compaction_chunks_count_threshold = 50;
-            c.max_joined_partitions = 10;
+            // Eager split-by-file-size fragments the right table further; lift the join cap
+            // above the resulting partition count (this test asserts join correctness, not the cap).
+            c.max_joined_partitions = 50;
             c
         }).start_test(async move |services| {
             let service = services.sql_service;
@@ -4048,6 +4050,8 @@ mod tests {
         Config::test("inactive_partitions_cleanup")
             .update_config(|mut c| {
                 c.partition_split_threshold = 1000000;
+                // Keep a single active partition: this test covers inactive-file GC, not size split.
+                c.partition_size_split_threshold_bytes = 1024 * 1024;
                 c.compaction_chunks_count_threshold = 0;
                 c.not_used_timeout = 0;
                 c.meta_store_log_upload_interval = 1;

rust/cubestore/cubestore/src/store/compaction.rs

@@ -583,13 +583,20 @@ impl CompactionService for CompactionServiceImpl {
             ) as usize)
                 // Do not allow to much of new partitions to limit partition accuracy trade off
                 .min(16);
-            let new_partitions_count_by_file_size =
-                if let Some(partition_file_size) = partition.get_row().file_size() {
+            // Size the split by the bytes actually being written: the existing main table plus
+            // the pending chunks merged in this pass. A partition with a small (or empty) main
+            // table but large accumulated chunks must still split by size in a single pass,
+            // otherwise it under-splits and re-splits on the next round.
+            let new_partitions_count_by_file_size = {
+                let total_file_size =
+                    partition.get_row().file_size().unwrap_or(0) + chunks_total_file_size;
+                if total_file_size > 0 {
                     let threshold = self.config.partition_size_split_threshold_bytes();
-                    (div_ceil(partition_file_size, threshold) as usize).min(16)
+                    (div_ceil(total_file_size, threshold) as usize).min(16)
                 } else {
                     1
-                };
+                }
+            };
 
             let new_partitions_count =
                 new_partitions_count_by_rows.max(new_partitions_count_by_file_size);
@@ -2397,6 +2404,8 @@ mod tests {
         Config::test("partition_compaction_decimal96")
             .update_config(|mut c| {
                 c.partition_split_threshold = 20;
+                // Keep the split row-based: this test covers decimal handling, not size split.
+                c.partition_size_split_threshold_bytes = 1024 * 1024;
                 c
             })
             .start_test(async move |services| {
@@ -2497,7 +2506,9 @@ mod tests {
                     .get_active_partitions_by_index_id(1)
                     .await
                     .unwrap();
-                assert_eq!(partitions.len(), 1);
+                // Eager split-by-file-size: the first compaction already splits because the
+                // pending chunks exceed partition_size_split_threshold_bytes.
+                assert!(partitions.len() > 1);
                 let values = (0..10)
                     .map(|_| format!("('{}', '{}')", "a".repeat(10), "b".repeat(10)))
                     .collect::<Vec<_>>()