chore: upgrade to DataFusion main + arrow/parquet 59.1.0 [WIP]

native/Cargo.toml

@@ -34,14 +34,14 @@ edition = "2021"
 rust-version = "1.88"
 
 [workspace.dependencies]
-arrow = { version = "58.3.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow = { version = "59.1.0", features = ["prettyprint", "ffi", "chrono-tz"] }
 async-trait = { version = "0.1" }
 bytes = { version = "1.11.1" }
-parquet = { version = "58.3.0", default-features = false, features = ["experimental"] }
-datafusion = { version = "54.0.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
-datafusion-datasource = { version = "54.0.0" }
-datafusion-physical-expr-adapter = { version = "54.0.0" }
-datafusion-spark = { version = "54.0.0", features = ["core"] }
+parquet = { version = "59.1.0", default-features = false, features = ["experimental"] }
+datafusion = { git = "https://github.com/apache/datafusion", rev = "63e25c15bc72e731c174ba03d1143b243cc1aaae", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
+datafusion-datasource = { git = "https://github.com/apache/datafusion", rev = "63e25c15bc72e731c174ba03d1143b243cc1aaae" }
+datafusion-physical-expr-adapter = { git = "https://github.com/apache/datafusion", rev = "63e25c15bc72e731c174ba03d1143b243cc1aaae" }
+datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "63e25c15bc72e731c174ba03d1143b243cc1aaae", features = ["core"] }
 datafusion-comet-spark-expr = { path = "spark-expr" }
 datafusion-comet-common = { path = "common" }
 datafusion-comet-jni-bridge = { path = "jni-bridge" }

native/core/Cargo.toml

@@ -91,7 +91,7 @@ jni = { version = "0.22.4", features = ["invocation"] }
 lazy_static = "1.4"
 assertables = "10"
 hex = "0.4.3"
-datafusion-functions-nested = { version = "54.0.0" }
+datafusion-functions-nested = { git = "https://github.com/apache/datafusion", rev = "63e25c15bc72e731c174ba03d1143b243cc1aaae" }
 
 [features]
 backtrace = ["datafusion/backtrace"]

native/core/src/execution/columnar_to_row.rs

@@ -2499,11 +2499,14 @@ mod tests {
         let schema = vec![DataType::FixedSizeBinary(3)];
         let mut ctx = ColumnarToRowContext::new(schema, 100);
 
-        let array: ArrayRef = Arc::new(FixedSizeBinaryArray::from(vec![
-            Some(&[1u8, 2, 3][..]),
-            Some(&[4u8, 5, 6][..]),
-            None, // Test null handling
-        ]));
+        let array: ArrayRef = Arc::new(
+            FixedSizeBinaryArray::try_from(vec![
+                Some(&[1u8, 2, 3][..]),
+                Some(&[4u8, 5, 6][..]),
+                None, // Test null handling
+            ])
+            .unwrap(),
+        );
         let arrays = vec![array];
 
         let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap();

native/core/src/execution/merge_as_partial.rs

@@ -207,23 +207,22 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
+        _opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         // Redirect update to merge — this is the key trick.
         self.inner
-            .merge_batch(values, group_indices, opt_filter, total_num_groups)
+            .merge_batch(values, group_indices, total_num_groups)
     }
 
     fn merge_batch(
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         self.inner
-            .merge_batch(values, group_indices, opt_filter, total_num_groups)
+            .merge_batch(values, group_indices, total_num_groups)
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {

native/core/src/execution/operators/copy.rs

@@ -37,7 +37,9 @@ pub(crate) fn copy_array(array: &dyn Array) -> ArrayRef {
 
     let mut mutable = MutableArrayData::new(vec![&data], false, capacity);
 
-    mutable.extend(0, 0, capacity);
+    mutable
+        .try_extend(0, 0, capacity)
+        .expect("copy_array: extend within existing array cannot overflow");
 
     if matches!(array.data_type(), DataType::Dictionary(_, _)) {
         let copied_dict = make_array(mutable.freeze());
@@ -50,7 +52,9 @@ pub(crate) fn copy_array(array: &dyn Array) -> ArrayRef {
                 let data = values.to_data();
 
                 let mut mutable = MutableArrayData::new(vec![&data], false, values.len());
-                mutable.extend(0, 0, values.len());
+                mutable
+                    .try_extend(0, 0, values.len())
+                    .expect("copy_array: extend within existing array cannot overflow");
 
                 let copied_dict = ref_copied_dict.with_values(make_array(mutable.freeze()));
                 Arc::new(copied_dict)

native/core/src/parquet/parquet_exec.rs

@@ -19,7 +19,7 @@ use crate::execution::operators::ExecutionError;
 use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID};
 use crate::parquet::parquet_support::SparkParquetOptions;
 use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory;
-use arrow::datatypes::{Field, SchemaRef};
+use arrow::datatypes::{Field, FieldRef, SchemaRef};
 use datafusion::config::TableParquetOptions;
 use datafusion::datasource::listing::PartitionedFile;
 use datafusion::datasource::physical_plan::parquet::CachedParquetFileReaderFactory;
@@ -122,13 +122,14 @@ pub(crate) fn init_datasource_exec(
         }
         _ => (Arc::clone(&required_schema), None),
     };
-    let partition_fields: Vec<_> = partition_schema
+    let partition_fields: Vec<FieldRef> = partition_schema
         .iter()
         .flat_map(|s| s.fields().iter())
-        .map(|f| Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable())) as _)
+        .map(|f| Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable())))
         .collect();
-    let table_schema =
-        TableSchema::from_file_schema(base_schema).with_table_partition_cols(partition_fields);
+    let table_schema = TableSchema::builder(base_schema)
+        .with_table_partition_cols(partition_fields)
+        .build();
 
     let mut parquet_source = ParquetSource::new(table_schema)
         .with_table_parquet_options(table_parquet_options)

native/shuffle/src/writers/local/spill.rs

@@ -20,12 +20,13 @@ use crate::writers::BufBatchWriter;
 use crate::ShuffleBlockWriter;
 use arrow::record_batch::RecordBatch;
 use datafusion::common::DataFusionError;
-use datafusion::execution::disk_manager::RefCountedTempFile;
 use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::execution::SpillFile as DfSpillFile;
 use std::fs::{File, OpenOptions};
+use std::sync::Arc;
 
 struct SpillFile {
-    temp_file: RefCountedTempFile,
+    temp_file: Arc<dyn DfSpillFile>,
     file: File,
 }
 
@@ -97,7 +98,11 @@ impl SpillWriter {
                 .write(true)
                 .create(true)
                 .truncate(true)
-                .open(spill_file.path())
+                .open(spill_file.path().ok_or_else(|| {
+                    DataFusionError::Execution(
+                        "Spill file backend does not expose a local path".to_string(),
+                    )
+                })?)
                 .map_err(|e| {
                     DataFusionError::Execution(format!("Error occurred while spilling {e}"))
                 })?;
@@ -112,7 +117,7 @@ impl SpillWriter {
     pub(crate) fn path(&self) -> Option<&std::path::Path> {
         self.spill_file
             .as_ref()
-            .map(|spill_file| spill_file.temp_file.path())
+            .and_then(|spill_file| spill_file.temp_file.path())
     }
 
     #[cfg(test)]

native/spark-expr/src/agg_funcs/avg.rs

@@ -283,7 +283,6 @@ where
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        _opt_filter: Option<&arrow::array::BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 2, "two arguments to merge_batch");

native/spark-expr/src/agg_funcs/avg_decimal.rs

@@ -536,7 +536,6 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        _opt_filter: Option<&arrow::array::BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 2, "two arguments to merge_batch");

native/spark-expr/src/agg_funcs/correlation.rs

@@ -335,7 +335,6 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 6, "six state columns to merge_batch");
@@ -358,11 +357,11 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         ];
 
         self.covar
-            .merge_batch(&covar_state, group_indices, opt_filter, total_num_groups)?;
+            .merge_batch(&covar_state, group_indices, total_num_groups)?;
         self.var1
-            .merge_batch(&var1_state, group_indices, opt_filter, total_num_groups)?;
+            .merge_batch(&var1_state, group_indices, total_num_groups)?;
         self.var2
-            .merge_batch(&var2_state, group_indices, opt_filter, total_num_groups)?;
+            .merge_batch(&var2_state, group_indices, total_num_groups)?;
         Ok(())
     }
 

native/spark-expr/src/agg_funcs/covariance.rs

@@ -414,7 +414,6 @@ impl GroupsAccumulator for CovarianceGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        _opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 4, "four arguments to merge_batch");
@@ -563,8 +562,8 @@ mod groups_tests {
         let rstate = right.state(EmitTo::All).unwrap();
 
         let mut merged = pop();
-        merged.merge_batch(&lstate, &[0], None, 1).unwrap();
-        merged.merge_batch(&rstate, &[0], None, 1).unwrap();
+        merged.merge_batch(&lstate, &[0], 1).unwrap();
+        merged.merge_batch(&rstate, &[0], 1).unwrap();
         let merged_result = evaluate(&mut merged)[0].unwrap();
 
         assert!((single - merged_result).abs() < 1e-12);

native/spark-expr/src/agg_funcs/percentile.rs

@@ -221,7 +221,6 @@ impl GroupsAccumulator for SparkPercentileGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        _opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         let input_group_values = values[0].as_list::<i32>();

native/spark-expr/src/agg_funcs/stddev.rs

@@ -231,11 +231,10 @@ impl GroupsAccumulator for StddevGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         self.inner
-            .merge_batch(values, group_indices, opt_filter, total_num_groups)
+            .merge_batch(values, group_indices, total_num_groups)
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {

native/spark-expr/src/agg_funcs/sum_decimal.rs

@@ -541,14 +541,8 @@ impl GroupsAccumulator for SumDecimalGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> DFResult<()> {
-        debug_assert!(
-            opt_filter.is_none(),
-            "opt_filter is not supported in merge_batch"
-        );
-
         self.resize_helper(total_num_groups);
 
         // For decimal sum, always expect 2 arrays regardless of eval_mode

native/spark-expr/src/agg_funcs/sum_int.rs

@@ -495,14 +495,8 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorLegacy {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> DFResult<()> {
-        debug_assert!(
-            opt_filter.is_none(),
-            "opt_filter is not supported in merge_batch"
-        );
-
         if values.len() != 1 {
             return Err(DataFusionError::Internal(format!(
                 "Invalid state while merging batch. Expected 1 element but found {}",
@@ -642,14 +636,8 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorAnsi {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> DFResult<()> {
-        debug_assert!(
-            opt_filter.is_none(),
-            "opt_filter is not supported in merge_batch"
-        );
-
         if values.len() != 1 {
             return Err(DataFusionError::Internal(format!(
                 "Invalid state while merging batch. Expected 1 element but found {}",
@@ -826,14 +814,8 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> DFResult<()> {
-        debug_assert!(
-            opt_filter.is_none(),
-            "opt_filter is not supported in merge_batch"
-        );
-
         if values.len() != 2 {
             return Err(DataFusionError::Internal(format!(
                 "Invalid state while merging batch. Expected 2 elements but found {}",

native/spark-expr/src/agg_funcs/variance.rs

@@ -340,7 +340,6 @@ impl GroupsAccumulator for VarianceGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        _opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 3, "three arguments to merge_batch");
@@ -512,8 +511,8 @@ mod groups_tests {
         let right_state = right.state(EmitTo::All).unwrap();
 
         let mut merged = pop_acc();
-        merged.merge_batch(&left_state, &[0], None, 1).unwrap();
-        merged.merge_batch(&right_state, &[0], None, 1).unwrap();
+        merged.merge_batch(&left_state, &[0], 1).unwrap();
+        merged.merge_batch(&right_state, &[0], 1).unwrap();
         let merged_result = evaluate(&mut merged)[0].unwrap();
 
         assert!((single_result - merged_result).abs() < 1e-12);

native/spark-expr/src/array_funcs/array_compact.rs

@@ -146,7 +146,7 @@ fn compact_list<OffsetSize: OffsetSizeTrait>(
         for i in start..end {
             let is_null = value_nulls.as_ref().map(|n| n.is_null(i)).unwrap_or(false);
             if !is_null {
-                mutable.extend(0, i, i + 1);
+                mutable.try_extend(0, i, i + 1)?;
                 copied += 1;
             }
         }

native/spark-expr/src/array_funcs/array_insert.rs

@@ -266,17 +266,17 @@ fn array_insert<O: OffsetSizeTrait>(
             if pos1 <= len + 1 {
                 // In-range insertion (including appending to end)
                 let corrected = pos1 - 1; // 0-based insertion point
-                mutable_values.extend(0, start, start + corrected);
-                mutable_values.extend(1, row_index, row_index + 1);
-                mutable_values.extend(0, start + corrected, end);
+                mutable_values.try_extend(0, start, start + corrected)?;
+                mutable_values.try_extend(1, row_index, row_index + 1)?;
+                mutable_values.try_extend(0, start + corrected, end)?;
                 final_len = len + 1;
             } else {
                 // Beyond end: pad with nulls then insert
                 let corrected = pos1 - 1;
                 let padding = corrected - len;
-                mutable_values.extend(0, start, end);
-                mutable_values.extend_nulls(padding);
-                mutable_values.extend(1, row_index, row_index + 1);
+                mutable_values.try_extend(0, start, end)?;
+                mutable_values.try_extend_nulls(padding)?;
+                mutable_values.try_extend(1, row_index, row_index + 1)?;
                 final_len = corrected + 1; // equals pos1
             }
         } else {
@@ -289,9 +289,9 @@ fn array_insert<O: OffsetSizeTrait>(
                 // Legacy:     -1 behaves like insert before the last element (corrected = len - k)
                 let base_offset = if legacy_mode { 0 } else { 1 };
                 let corrected = len - k + base_offset;
-                mutable_values.extend(0, start, start + corrected);
-                mutable_values.extend(1, row_index, row_index + 1);
-                mutable_values.extend(0, start + corrected, end);
+                mutable_values.try_extend(0, start, start + corrected)?;
+                mutable_values.try_extend(1, row_index, row_index + 1)?;
+                mutable_values.try_extend(0, start + corrected, end)?;
                 final_len = len + 1;
             } else {
                 // Negative index beyond the start (Spark-specific behavior):
@@ -300,9 +300,9 @@ fn array_insert<O: OffsetSizeTrait>(
                 let base_offset = if legacy_mode { 1 } else { 0 };
                 let target_len = k + base_offset;
                 let padding = target_len.saturating_sub(len + 1);
-                mutable_values.extend(1, row_index, row_index + 1); // insert item first
-                mutable_values.extend_nulls(padding); // pad nulls
-                mutable_values.extend(0, start, end); // append original values
+                mutable_values.try_extend(1, row_index, row_index + 1)?; // insert item first
+                mutable_values.try_extend_nulls(padding)?; // pad nulls
+                mutable_values.try_extend(0, start, end)?; // append original values
                 final_len = target_len;
             }
         }

native/spark-expr/src/array_funcs/array_slice.rs

@@ -163,7 +163,7 @@ fn slice_list<O: OffsetSizeTrait>(
         } else {
             let take = std::cmp::min(length_value, arr_len - zero_based_start) as usize;
             let begin = row_start + zero_based_start as usize;
-            mutable.extend(0, begin, begin + take);
+            mutable.try_extend(0, begin, begin + take)?;
             take
         };