chore(cubestore): Upgrade DF: Regroup batches in compaction to avoid i32 offset overflow

Author: srh

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -2073,24 +2073,33 @@ fn combine_filters(filters: &[Expr]) -> Option<Expr> {
     Some(combined_filter)
 }
 
+pub fn regroup_batch_onto(
+    b: RecordBatch,
+    max_rows: usize,
+    onto: &mut Vec<RecordBatch>,
+) -> Result<(), CubeError> {
+    let mut row = 0;
+    while row != b.num_rows() {
+        let slice_len = min(b.num_rows() - row, max_rows);
+        onto.push(RecordBatch::try_new(
+            b.schema(),
+            b.columns()
+                .iter()
+                .map(|c| slice_copy(c.as_ref(), row, slice_len))
+                .collect(),
+        )?);
+        row += slice_len;
+    }
+    Ok(())
+}
+
 fn regroup_batches(
     batches: Vec<RecordBatch>,
     max_rows: usize,
 ) -> Result<Vec<RecordBatch>, CubeError> {
     let mut r = Vec::with_capacity(batches.len());
     for b in batches {
-        let mut row = 0;
-        while row != b.num_rows() {
-            let slice_len = min(b.num_rows() - row, max_rows);
-            r.push(RecordBatch::try_new(
-                b.schema(),
-                b.columns()
-                    .iter()
-                    .map(|c| slice_copy(c.as_ref(), row, slice_len))
-                    .collect(),
-            )?);
-            row += slice_len
-        }
+        regroup_batch_onto(b, max_rows, &mut r)?;
     }
     Ok(r)
 }

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

@@ -11,6 +11,7 @@ use crate::metastore::{
 };
 use crate::queryplanner::merge_sort::LastRowByUniqueKeyExec;
 use crate::queryplanner::metadata_cache::MetadataCacheFactory;
+use crate::queryplanner::query_executor::regroup_batch_onto;
 use crate::queryplanner::trace_data_loaded::{DataLoadedSize, TraceDataLoadedExec};
 use crate::queryplanner::{try_make_memory_data_source, QueryPlannerImpl};
 use crate::remotefs::{ensure_temp_file_is_dropped, RemoteFs};
@@ -668,6 +669,7 @@ impl CompactionService for CompactionServiceImpl {
                     None,
                 )?)
             }
+
             Ok((store, new))
         })
         .await??;
@@ -1425,7 +1427,11 @@ pub async fn merge_chunks(
     task_context: Arc<TaskContext>,
 ) -> Result<SendableRecordBatchStream, CubeError> {
     let schema = l.schema();
-    let r = RecordBatch::try_new(schema.clone(), r)?;
+    let r_batch = RecordBatch::try_new(schema.clone(), r)?;
+    let mut r = Vec::<RecordBatch>::new();
+    // Regroup batches -- which had been concatenated and sorted -- so that SortPreservingMergeExec
+    // doesn't overflow i32 in interleaving or building a Utf8Array.
+    regroup_batch_onto(r_batch, 8192, &mut r)?;
 
     let mut key = Vec::with_capacity(key_size);
     for i in 0..key_size {
@@ -1436,10 +1442,7 @@ pub async fn merge_chunks(
         ));
     }
 
-    let inputs = UnionExec::new(vec![
-        l,
-        try_make_memory_data_source(&[vec![r]], schema, None)?,
-    ]);
+    let inputs = UnionExec::new(vec![l, try_make_memory_data_source(&[r], schema, None)?]);
     let mut res: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(
         LexOrdering::new(key),
         Arc::new(inputs),