chore: Bump DataFusion to git rev 2c2f225

native/Cargo.toml

@@ -38,8 +38,8 @@ arrow = { version = "55.1.0", features = ["prettyprint", "ffi", "chrono-tz"] }
 async-trait = { version = "0.1" }
 bytes = { version = "1.10.0" }
 parquet = { version = "55.1.0", default-features = false, features = ["experimental"] }
-datafusion = { git = "https://github.com/apache/datafusion", rev = "efd9587", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
-datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "efd9587" }
+datafusion = { git = "https://github.com/apache/datafusion", rev = "2c2f225", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
+datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "2c2f225" }
 datafusion-comet-spark-expr = { path = "spark-expr" }
 datafusion-comet-proto = { path = "proto" }
 chrono = { version = "0.4", default-features = false, features = ["clock"] }

native/core/Cargo.toml

@@ -82,7 +82,7 @@ jni = { version = "0.21", features = ["invocation"] }
 lazy_static = "1.4"
 assertables = "9"
 hex = "0.4.3"
-datafusion-functions-nested = { git = "https://github.com/apache/datafusion", rev = "efd9587" }
+datafusion-functions-nested = { git = "https://github.com/apache/datafusion", rev = "2c2f225" }
 
 [features]
 default = []

native/core/src/execution/expressions/bloom_filter_agg.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::datatypes::Field;
+use arrow::datatypes::{Field, FieldRef};
 use datafusion::{arrow::datatypes::DataType, logical_expr::Volatility};
 use std::{any::Any, sync::Arc};
 
@@ -97,8 +97,8 @@ impl AggregateUDFImpl for BloomFilterAgg {
         ))))
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
-        Ok(vec![Field::new("bits", DataType::Binary, false)])
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+        Ok(vec![Arc::new(Field::new("bits", DataType::Binary, false))])
     }
 
     fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {

native/core/src/execution/jni_api.rs

@@ -22,7 +22,7 @@ use arrow::array::RecordBatch;
 use arrow::datatypes::DataType as ArrowDataType;
 use datafusion::execution::memory_pool::MemoryPool;
 use datafusion::{
-    execution::{disk_manager::DiskManagerConfig, runtime_env::RuntimeEnv},
+    execution::{disk_manager::DiskManagerBuilder, runtime_env::RuntimeEnv},
     physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream},
     prelude::{SessionConfig, SessionContext},
 };
@@ -49,6 +49,7 @@ use crate::{
     jvm_bridge::{jni_new_global_ref, JVMClasses},
 };
 use datafusion::common::ScalarValue;
+use datafusion::execution::disk_manager::DiskManagerMode;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_comet_proto::spark_operator::Operator;
 use futures::stream::StreamExt;
@@ -255,9 +256,9 @@ fn prepare_datafusion_session_context(
     memory_pool: Arc<dyn MemoryPool>,
     local_dirs: Vec<String>,
 ) -> CometResult<SessionContext> {
-    let disk_manager_config =
-        DiskManagerConfig::NewSpecified(local_dirs.into_iter().map(PathBuf::from).collect());
-    let mut rt_config = RuntimeEnvBuilder::new().with_disk_manager(disk_manager_config);
+    let paths = local_dirs.into_iter().map(PathBuf::from).collect();
+    let disk_manager = DiskManagerBuilder::default().with_mode(DiskManagerMode::Directories(paths));
+    let mut rt_config = RuntimeEnvBuilder::new().with_disk_manager_builder(disk_manager);
     rt_config = rt_config.with_memory_pool(memory_pool);
 
     // Get Datafusion configuration from Spark Execution context

native/core/src/execution/planner.rs

@@ -893,7 +893,7 @@ impl PhysicalPlanner {
                     func_name,
                     fun_expr,
                     vec![left, right],
-                    Field::new(func_name, data_type, true),
+                    Arc::new(Field::new(func_name, data_type, true)),
                 )))
             }
             _ => Ok(Arc::new(BinaryExpr::new(left, op, right))),
@@ -2239,7 +2239,7 @@ impl PhysicalPlanner {
                     let arg_fields = coerced_types
                         .iter()
                         .enumerate()
-                        .map(|(i, dt)| Field::new(format!("arg{i}"), dt.clone(), true))
+                        .map(|(i, dt)| Arc::new(Field::new(format!("arg{i}"), dt.clone(), true)))
                         .collect::<Vec<_>>();
 
                     // TODO this should try and find scalar
@@ -2258,10 +2258,7 @@ impl PhysicalPlanner {
                         scalar_arguments: &arguments,
                     };
 
-                    let data_type = func
-                        .inner()
-                        .return_field_from_args(args)?
-                        .clone()
+                    let data_type = Arc::clone(&func.inner().return_field_from_args(args)?)
                         .data_type()
                         .clone();
 
@@ -2295,7 +2292,7 @@ impl PhysicalPlanner {
             fun_name,
             fun_expr,
             args.to_vec(),
-            Field::new(fun_name, data_type, true),
+            Arc::new(Field::new(fun_name, data_type, true)),
         ));
 
         Ok(scalar_expr)

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

@@ -22,7 +22,7 @@ use arrow::array::{
     Array, ArrayRef, ArrowNumericType, Int64Array, PrimitiveArray,
 };
 use arrow::compute::sum;
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion::common::{not_impl_err, Result, ScalarValue};
 use datafusion::logical_expr::{
     type_coercion::aggregates::avg_return_type, Accumulator, AggregateUDFImpl, EmitTo,
@@ -78,18 +78,18 @@ impl AggregateUDFImpl for Avg {
         }
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new(
+            Arc::new(Field::new(
                 format_state_name(&self.name, "sum"),
                 self.input_data_type.clone(),
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "count"),
                 DataType::Int64,
                 true,
-            ),
+            )),
         ])
     }
 

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

@@ -21,7 +21,7 @@ use arrow::array::{
     types::{Decimal128Type, Int64Type},
     Array, ArrayRef, Decimal128Array, Int64Array, PrimitiveArray,
 };
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use arrow::{array::BooleanBufferBuilder, buffer::NullBuffer, compute::sum};
 use datafusion::common::{not_impl_err, Result, ScalarValue};
 use datafusion::logical_expr::{
@@ -82,18 +82,18 @@ impl AggregateUDFImpl for AvgDecimal {
         }
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new(
+            Arc::new(Field::new(
                 format_state_name(self.name(), "sum"),
                 self.sum_data_type.clone(),
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(self.name(), "count"),
                 DataType::Int64,
                 true,
-            ),
+            )),
         ])
     }
 

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

@@ -21,6 +21,7 @@ use std::{any::Any, sync::Arc};
 
 use crate::agg_funcs::covariance::CovarianceAccumulator;
 use crate::agg_funcs::stddev::StddevAccumulator;
+use arrow::datatypes::FieldRef;
 use arrow::{
     array::ArrayRef,
     datatypes::{DataType, Field},
@@ -83,38 +84,38 @@ impl AggregateUDFImpl for Correlation {
         )?))
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new(
+            Arc::new(Field::new(
                 format_state_name(&self.name, "count"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "mean1"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "mean2"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "algo_const"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "m2_1"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "m2_2"),
                 DataType::Float64,
                 true,
-            ),
+            )),
         ])
     }
 }

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

@@ -17,8 +17,7 @@
  * under the License.
  */
 
-use std::any::Any;
-
+use arrow::datatypes::FieldRef;
 use arrow::{
     array::{ArrayRef, Float64Array},
     compute::cast,
@@ -32,6 +31,8 @@ use datafusion::logical_expr::type_coercion::aggregates::NUMERICS;
 use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
 use datafusion::physical_expr::expressions::format_state_name;
 use datafusion::physical_expr::expressions::StatsType;
+use std::any::Any;
+use std::sync::Arc;
 
 /// COVAR_SAMP and COVAR_POP aggregate expression
 /// The implementation mostly is the same as the DataFusion's implementation. The reason
@@ -92,28 +93,28 @@ impl AggregateUDFImpl for Covariance {
         )?))
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new(
+            Arc::new(Field::new(
                 format_state_name(&self.name, "count"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "mean1"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "mean2"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "algo_const"),
                 DataType::Float64,
                 true,
-            ),
+            )),
         ])
     }
 }

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

@@ -18,6 +18,7 @@
 use std::{any::Any, sync::Arc};
 
 use crate::agg_funcs::variance::VarianceAccumulator;
+use arrow::datatypes::FieldRef;
 use arrow::{
     array::ArrayRef,
     datatypes::{DataType, Field},
@@ -102,19 +103,23 @@ impl AggregateUDFImpl for Stddev {
         )?))
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
         Ok(vec![
-            Field::new(
+            Arc::new(Field::new(
                 format_state_name(&self.name, "count"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(
+            )),
+            Arc::new(Field::new(
                 format_state_name(&self.name, "mean"),
                 DataType::Float64,
                 true,
-            ),
-            Field::new(format_state_name(&self.name, "m2"), DataType::Float64, true),
+            )),
+            Arc::new(Field::new(
+                format_state_name(&self.name, "m2"),
+                DataType::Float64,
+                true,
+            )),
         ])
     }
 

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

@@ -19,7 +19,7 @@ use crate::utils::{is_valid_decimal_precision, unlikely};
 use arrow::array::{
     cast::AsArray, types::Decimal128Type, Array, ArrayRef, BooleanArray, Decimal128Array,
 };
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use arrow::{
     array::BooleanBufferBuilder,
     buffer::{BooleanBuffer, NullBuffer},
@@ -77,10 +77,14 @@ impl AggregateUDFImpl for SumDecimal {
         )))
     }
 
-    fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<Field>> {
+    fn state_fields(&self, _args: StateFieldsArgs) -> DFResult<Vec<FieldRef>> {
         let fields = vec![
-            Field::new(self.name(), self.result_type.clone(), self.is_nullable()),
-            Field::new("is_empty", DataType::Boolean, false),
+            Arc::new(Field::new(
+                self.name(),
+                self.result_type.clone(),
+                self.is_nullable(),
+            )),
+            Arc::new(Field::new("is_empty", DataType::Boolean, false)),
         ];
         Ok(fields)
     }