chore: upgrade to DataFusion 50.0.0, Arrow 56.1.0, Parquet 56.0.0 among others (#2286)

mbutrovich · web-flow · commit 24f520975bd5 · 2025-09-17T07:30:11.000-04:00
diff --git a/native/Cargo.lock b/native/Cargo.lock
diff --git a/native/Cargo.toml b/native/Cargo.toml
@@ -31,15 +31,15 @@ license = "Apache-2.0"
 edition = "2021"
 
 # Comet uses the same minimum Rust version as DataFusion
-rust-version = "1.85"
+rust-version = "1.86"
 
 [workspace.dependencies]
-arrow = { version = "55.2.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow = { version = "56.0.0", features = ["prettyprint", "ffi", "chrono-tz"] }
 async-trait = { version = "0.1" }
 bytes = { version = "1.10.0" }
-parquet = { version = "55.2.0", default-features = false, features = ["experimental"] }
-datafusion = { version = "49.0.2", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
-datafusion-spark = { version = "49.0.2" }
+parquet = { version = "=56.0.0", default-features = false, features = ["experimental"] }
+datafusion = { version = "50.0.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
+datafusion-spark = { version = "50.0.0" }
 datafusion-comet-spark-expr = { path = "spark-expr" }
 datafusion-comet-proto = { path = "proto" }
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml
@@ -36,7 +36,7 @@ publish = false
 
 [dependencies]
 arrow = { workspace = true }
-parquet = { workspace = true, default-features = false, features = ["experimental"] }
+parquet = { workspace = true, default-features = false, features = ["experimental", "arrow"] }
 futures = { workspace = true }
 mimalloc = { version = "*", default-features = false, optional = true }
 tikv-jemallocator = { version = "0.6.0", optional = true, features = ["disable_initial_exec_tls"] }
@@ -91,7 +91,7 @@ jni = { version = "0.21", features = ["invocation"] }
 lazy_static = "1.4"
 assertables = "9"
 hex = "0.4.3"
-datafusion-functions-nested = { version = "49.0.2" }
+datafusion-functions-nested = { version = "50.0.0" }
 
 [features]
 default = []
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
@@ -35,14 +35,14 @@ use datafusion::execution::memory_pool::MemoryPool;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::logical_expr::ScalarUDF;
 use datafusion::{
-    execution::{disk_manager::DiskManagerBuilder, runtime_env::RuntimeEnv},
+    execution::disk_manager::DiskManagerBuilder,
     physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream},
     prelude::{SessionConfig, SessionContext},
 };
 use datafusion_comet_proto::spark_operator::Operator;
 use datafusion_spark::function::hash::sha2::SparkSha2;
 use datafusion_spark::function::math::expm1::SparkExpm1;
-use datafusion_spark::function::string::char::SparkChar;
+use datafusion_spark::function::string::char::CharFunc;
 use futures::poll;
 use futures::stream::StreamExt;
 use jni::objects::JByteBuffer;
@@ -291,8 +291,7 @@ fn prepare_datafusion_session_context(
             &ScalarValue::Float64(Some(1.1)),
         );
 
-    #[allow(deprecated)]
-    let runtime = RuntimeEnv::try_new(rt_config)?;
+    let runtime = rt_config.build()?;
 
     let mut session_ctx = SessionContext::new_with_config_rt(session_config, Arc::new(runtime));
 
@@ -301,7 +300,7 @@ fn prepare_datafusion_session_context(
     // register UDFs from datafusion-spark crate
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default()));
-    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkChar::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(CharFunc::default()));
 
     // Must be the last one to override existing functions with the same name
     datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
@@ -40,6 +40,7 @@ use datafusion::physical_plan::InputOrderMode;
 use datafusion::{
     arrow::{compute::SortOptions, datatypes::SchemaRef},
     common::DataFusionError,
+    config::ConfigOptions,
     execution::FunctionRegistry,
     functions_aggregate::first_last::{FirstValue, LastValue},
     logical_expr::Operator as DataFusionOperator,
@@ -623,8 +624,13 @@ impl PhysicalPlanner {
                 let args = vec![child];
                 let comet_hour = Arc::new(ScalarUDF::new_from_impl(SparkHour::new(timezone)));
                 let field_ref = Arc::new(Field::new("hour", DataType::Int32, true));
-                let expr: ScalarFunctionExpr =
-                    ScalarFunctionExpr::new("hour", comet_hour, args, field_ref);
+                let expr: ScalarFunctionExpr = ScalarFunctionExpr::new(
+                    "hour",
+                    comet_hour,
+                    args,
+                    field_ref,
+                    Arc::new(ConfigOptions::default()),
+                );
 
                 Ok(Arc::new(expr))
             }
@@ -635,8 +641,13 @@ impl PhysicalPlanner {
                 let args = vec![child];
                 let comet_minute = Arc::new(ScalarUDF::new_from_impl(SparkMinute::new(timezone)));
                 let field_ref = Arc::new(Field::new("minute", DataType::Int32, true));
-                let expr: ScalarFunctionExpr =
-                    ScalarFunctionExpr::new("minute", comet_minute, args, field_ref);
+                let expr: ScalarFunctionExpr = ScalarFunctionExpr::new(
+                    "minute",
+                    comet_minute,
+                    args,
+                    field_ref,
+                    Arc::new(ConfigOptions::default()),
+                );
 
                 Ok(Arc::new(expr))
             }
@@ -647,8 +658,13 @@ impl PhysicalPlanner {
                 let args = vec![child];
                 let comet_second = Arc::new(ScalarUDF::new_from_impl(SparkSecond::new(timezone)));
                 let field_ref = Arc::new(Field::new("second", DataType::Int32, true));
-                let expr: ScalarFunctionExpr =
-                    ScalarFunctionExpr::new("second", comet_second, args, field_ref);
+                let expr: ScalarFunctionExpr = ScalarFunctionExpr::new(
+                    "second",
+                    comet_second,
+                    args,
+                    field_ref,
+                    Arc::new(ConfigOptions::default()),
+                );
 
                 Ok(Arc::new(expr))
             }
@@ -870,8 +886,13 @@ impl PhysicalPlanner {
                     ScalarUDF::new_from_impl(BloomFilterMightContain::try_new(bloom_filter_expr)?);
 
                 let field_ref = Arc::new(Field::new("might_contain", DataType::Boolean, true));
-                let expr: ScalarFunctionExpr =
-                    ScalarFunctionExpr::new("might_contain", Arc::new(udf), args, field_ref);
+                let expr: ScalarFunctionExpr = ScalarFunctionExpr::new(
+                    "might_contain",
+                    Arc::new(udf),
+                    args,
+                    field_ref,
+                    Arc::new(ConfigOptions::default()),
+                );
                 Ok(Arc::new(expr))
             }
             ExprStruct::CreateNamedStruct(expr) => {
@@ -1090,6 +1111,7 @@ impl PhysicalPlanner {
                     fun_expr,
                     vec![left, right],
                     Arc::new(Field::new(func_name, data_type, true)),
+                    Arc::new(ConfigOptions::default()),
                 )))
             }
             _ => {
@@ -1115,6 +1137,7 @@ impl PhysicalPlanner {
                         fun_expr,
                         vec![left, right],
                         Arc::new(Field::new(op_str, data_type, true)),
+                        Arc::new(ConfigOptions::default()),
                     )))
                 } else {
                     Ok(Arc::new(BinaryExpr::new(left, op, right)))
@@ -2354,6 +2377,8 @@ impl PhysicalPlanner {
             window_frame.into(),
             input_schema.as_ref(),
             false, // TODO: Ignore nulls
+            false, // TODO: Spark does not support DISTINCT ... OVER
+            None,
         )
         .map_err(|e| ExecutionError::DataFusionError(e.to_string()))
     }
@@ -2533,6 +2558,7 @@ impl PhysicalPlanner {
             fun_expr,
             args.to_vec(),
             Arc::new(Field::new(fun_name, data_type, true)),
+            Arc::new(ConfigOptions::default()),
         ));
 
         Ok(scalar_expr)
diff --git a/native/core/src/execution/shuffle/shuffle_writer.rs b/native/core/src/execution/shuffle/shuffle_writer.rs
@@ -1350,7 +1350,7 @@ mod test {
     #[tokio::test]
     async fn shuffle_repartitioner_memory() {
         let batch = create_batch(900);
-        assert_eq!(8376, batch.get_array_memory_size());
+        assert_eq!(8316, batch.get_array_memory_size()); // Not stable across Arrow versions
 
         let memory_limit = 512 * 1024;
         let num_partitions = 2;
diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs
@@ -37,7 +37,7 @@ use datafusion::logical_expr::Volatility::Immutable;
 use DataType::*;
 
 /// AVG aggregate expression
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Avg {
     name: String,
     signature: Signature,
diff --git a/native/spark-expr/src/agg_funcs/avg_decimal.rs b/native/spark-expr/src/agg_funcs/avg_decimal.rs
@@ -40,7 +40,7 @@ use num::{integer::div_ceil, Integer};
 use DataType::*;
 
 /// AVG aggregate expression
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct AvgDecimal {
     signature: Signature,
     sum_data_type: DataType,
diff --git a/native/spark-expr/src/agg_funcs/correlation.rs b/native/spark-expr/src/agg_funcs/correlation.rs
@@ -38,7 +38,7 @@ use datafusion::physical_expr::expressions::StatsType;
 /// we have our own implementation is that DataFusion has UInt64 for state_field `count`,
 /// while Spark has Double for count. Also we have added `null_on_divide_by_zero`
 /// to be consistent with Spark's implementation.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Correlation {
     name: String,
     signature: Signature,
diff --git a/native/spark-expr/src/agg_funcs/covariance.rs b/native/spark-expr/src/agg_funcs/covariance.rs
@@ -38,14 +38,23 @@ use std::sync::Arc;
 /// The implementation mostly is the same as the DataFusion's implementation. The reason
 /// we have our own implementation is that DataFusion has UInt64 for state_field count,
 /// while Spark has Double for count.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Covariance {
     name: String,
     signature: Signature,
     stats_type: StatsType,
     null_on_divide_by_zero: bool,
 }
 
+impl std::hash::Hash for Covariance {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        (self.stats_type as u8).hash(state);
+        self.null_on_divide_by_zero.hash(state);
+    }
+}
+
 impl Covariance {
     /// Create a new COVAR aggregate function
     pub fn new(
diff --git a/native/spark-expr/src/agg_funcs/stddev.rs b/native/spark-expr/src/agg_funcs/stddev.rs
@@ -36,14 +36,23 @@ use datafusion::physical_expr::expressions::StatsType;
 /// we have our own implementation is that DataFusion has UInt64 for state_field `count`,
 /// while Spark has Double for count. Also we have added `null_on_divide_by_zero`
 /// to be consistent with Spark's implementation.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub struct Stddev {
     name: String,
     signature: Signature,
     stats_type: StatsType,
     null_on_divide_by_zero: bool,
 }
 
+impl std::hash::Hash for Stddev {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        (self.stats_type as u8).hash(state);
+        self.null_on_divide_by_zero.hash(state);
+    }
+}
+
 impl Stddev {
     /// Create a new STDDEV aggregate function
     pub fn new(
diff --git a/native/spark-expr/src/agg_funcs/sum_decimal.rs b/native/spark-expr/src/agg_funcs/sum_decimal.rs
@@ -29,7 +29,7 @@ use datafusion::logical_expr::{
 };
 use std::{any::Any, ops::BitAnd, sync::Arc};
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct SumDecimal {
     /// Aggregate function signature
     signature: Signature,
diff --git a/native/spark-expr/src/agg_funcs/variance.rs b/native/spark-expr/src/agg_funcs/variance.rs
@@ -34,14 +34,23 @@ use std::sync::Arc;
 /// we have our own implementation is that DataFusion has UInt64 for state_field `count`,
 /// while Spark has Double for count. Also we have added `null_on_divide_by_zero`
 /// to be consistent with Spark's implementation.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub struct Variance {
     name: String,
     signature: Signature,
     stats_type: StatsType,
     null_on_divide_by_zero: bool,
 }
 
+impl std::hash::Hash for Variance {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        (self.stats_type as u8).hash(state);
+        self.null_on_divide_by_zero.hash(state);
+    }
+}
+
 impl Variance {
     /// Create a new VARIANCE aggregate function
     pub fn new(
diff --git a/native/spark-expr/src/bitwise_funcs/bitwise_count.rs b/native/spark-expr/src/bitwise_funcs/bitwise_count.rs
@@ -22,7 +22,7 @@ use datafusion::{error::DataFusionError, logical_expr::ColumnarValue};
 use std::any::Any;
 use std::sync::Arc;
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct SparkBitwiseCount {
     signature: Signature,
     aliases: Vec<String>,
diff --git a/native/spark-expr/src/bitwise_funcs/bitwise_get.rs b/native/spark-expr/src/bitwise_funcs/bitwise_get.rs
@@ -22,7 +22,7 @@ use datafusion::logical_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Vol
 use std::any::Any;
 use std::sync::Arc;
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct SparkBitwiseGet {
     signature: Signature,
     aliases: Vec<String>,
diff --git a/native/spark-expr/src/bitwise_funcs/bitwise_not.rs b/native/spark-expr/src/bitwise_funcs/bitwise_not.rs
@@ -23,7 +23,7 @@ use datafusion::logical_expr::{ColumnarValue, Volatility};
 use datafusion::logical_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
 use std::{any::Any, sync::Arc};
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct SparkBitwiseNot {
     signature: Signature,
     aliases: Vec<String>,
diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs b/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs
@@ -32,7 +32,7 @@ use datafusion::physical_expr::expressions::Literal;
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::Accumulator;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct BloomFilterAgg {
     signature: Signature,
     num_items: i32,
diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs b/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs
@@ -27,7 +27,7 @@ use std::sync::Arc;
 
 use crate::bloom_filter::spark_bloom_filter::SparkBloomFilter;
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct BloomFilterMightContain {
     signature: Signature,
     bloom_filter: Option<SparkBloomFilter>,
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -192,6 +192,26 @@ struct CometScalarFunction {
     func: ScalarFunctionImplementation,
 }
 
+impl PartialEq for CometScalarFunction {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+            && self.signature == other.signature
+            && self.data_type == other.data_type
+        // Note: we do not test ScalarFunctionImplementation equality, relying on function metadata.
+    }
+}
+
+impl Eq for CometScalarFunction {}
+
+impl std::hash::Hash for CometScalarFunction {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+        self.signature.hash(state);
+        self.data_type.hash(state);
+        // Note: we do not hash ScalarFunctionImplementation, relying on function metadata.
+    }
+}
+
 impl Debug for CometScalarFunction {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CometScalarFunction")
diff --git a/native/spark-expr/src/datetime_funcs/date_trunc.rs b/native/spark-expr/src/datetime_funcs/date_trunc.rs
@@ -24,7 +24,7 @@ use std::any::Any;
 
 use crate::kernels::temporal::{date_trunc_array_fmt_dyn, date_trunc_dyn};
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct SparkDateTrunc {
     signature: Signature,
     aliases: Vec<String>,
diff --git a/native/spark-expr/src/datetime_funcs/extract_date_part.rs b/native/spark-expr/src/datetime_funcs/extract_date_part.rs
@@ -26,7 +26,7 @@ use std::{any::Any, fmt::Debug};
 
 macro_rules! extract_date_part {
     ($struct_name:ident, $fn_name:expr, $date_part_variant:ident) => {
-        #[derive(Debug)]
+        #[derive(Debug, PartialEq, Eq, Hash)]
         pub struct $struct_name {
             signature: Signature,
             aliases: Vec<String>,
diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs
diff --git a/native/spark-expr/src/predicate_funcs/rlike.rs b/native/spark-expr/src/predicate_funcs/rlike.rs
diff --git a/native/spark-expr/src/string_funcs/string_space.rs b/native/spark-expr/src/string_funcs/string_space.rs
