DataFusion 52 migration

native/Cargo.toml

@@ -38,9 +38,9 @@ arrow = { version = "57.0.0", features = ["prettyprint", "ffi", "chrono-tz"] }
 async-trait = { version = "0.1" }
 bytes = { version = "1.10.0" }
 parquet = { version = "57.0.0", default-features = false, features = ["experimental"] }
-datafusion = { version = "51.0.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
-datafusion-datasource = { version = "51.0.0" }
-datafusion-spark = { version = "51.0.0" }
+datafusion = { git = "https://github.com/apache/datafusion", branch = "branch-52", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
+datafusion-datasource = { git = "https://github.com/apache/datafusion", branch = "branch-52" }
+datafusion-spark = { git = "https://github.com/apache/datafusion", branch = "branch-52" }
 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

@@ -95,7 +95,7 @@ jni = { version = "0.21", features = ["invocation"] }
 lazy_static = "1.4"
 assertables = "9"
 hex = "0.4.3"
-datafusion-functions-nested = { version = "51.0.0" }
+datafusion-functions-nested = { git = "https://github.com/apache/datafusion", branch = "branch-52" }
 
 [features]
 backtrace = ["datafusion/backtrace"]

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

@@ -42,10 +42,9 @@ use iceberg::io::FileIO;
 
 use crate::execution::operators::ExecutionError;
 use crate::parquet::parquet_support::SparkParquetOptions;
-use crate::parquet::schema_adapter::SparkSchemaAdapterFactory;
-use datafusion::datasource::schema_adapter::SchemaAdapterFactory;
 use datafusion_comet_spark_expr::EvalMode;
 use datafusion_datasource::file_stream::FileStreamMetrics;
+use crate::parquet::schema_adapter::adapt_batch_with_expressions;
 
 /// Iceberg table scan operator that uses iceberg-rust to read Iceberg tables.
 ///
@@ -299,22 +298,14 @@ impl IcebergFileStream {
                 .map_err(|e| DataFusionError::Execution(format!("Iceberg scan error: {}", e)))
                 .and_then(move |batch| {
                     let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);
-                    let adapter_factory = SparkSchemaAdapterFactory::new(spark_options, None);
-                    let file_schema = batch.schema();
-                    let adapter = adapter_factory
-                        .create(Arc::clone(&target_schema), Arc::clone(&file_schema));
-
-                    let result = match adapter.map_schema(file_schema.as_ref()) {
-                        Ok((schema_mapper, _projection)) => {
-                            schema_mapper.map_batch(batch).map_err(|e| {
-                                DataFusionError::Execution(format!("Batch mapping failed: {}", e))
-                            })
-                        }
-                        Err(e) => Err(DataFusionError::Execution(format!(
-                            "Schema mapping failed: {}",
-                            e
-                        ))),
-                    };
+                    let result =
+                        adapt_batch_with_expressions(batch, &target_schema, &spark_options)
+                            .map_err(|e| {
+                                DataFusionError::Execution(format!(
+                                    "Batch adaptation failed: {}",
+                                    e
+                                ))
+                            });
                     futures::future::ready(result)
                 });
 

native/core/src/execution/planner.rs

@@ -1049,8 +1049,8 @@ impl PhysicalPlanner {
                                 .as_any()
                                 .downcast_ref::<DataFusionLiteral>()
                                 .ok_or_else(|| {
-                                GeneralError("Expected literal of default value.".to_string())
-                            })?;
+                                    GeneralError("Expected literal of default value.".to_string())
+                                })?;
                             Ok(df_literal.value().clone())
                         })
                         .collect();
@@ -1093,18 +1093,11 @@ impl PhysicalPlanner {
                 let files =
                     self.get_partitioned_files(&scan.file_partitions[self.partition as usize])?;
                 let file_groups: Vec<Vec<PartitionedFile>> = vec![files];
-                let partition_fields: Vec<Field> = partition_schema
-                    .fields()
-                    .iter()
-                    .map(|field| {
-                        Field::new(field.name(), field.data_type().clone(), field.is_nullable())
-                    })
-                    .collect_vec();
+
                 let scan = init_datasource_exec(
                     required_schema,
                     Some(data_schema),
                     Some(partition_schema),
-                    Some(partition_fields),
                     object_store_url,
                     file_groups,
                     Some(projection_vector),
@@ -3385,6 +3378,7 @@ mod tests {
     use arrow::array::{Array, DictionaryArray, Int32Array, ListArray, RecordBatch, StringArray};
     use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema};
     use datafusion::catalog::memory::DataSourceExec;
+    use datafusion::config::TableParquetOptions;
     use datafusion::datasource::listing::PartitionedFile;
     use datafusion::datasource::object_store::ObjectStoreUrl;
     use datafusion::datasource::physical_plan::{
@@ -3401,8 +3395,6 @@ mod tests {
 
     use crate::execution::operators::ExecutionError;
     use crate::execution::planner::literal_to_array_ref;
-    use crate::parquet::parquet_support::SparkParquetOptions;
-    use crate::parquet::schema_adapter::SparkSchemaAdapterFactory;
     use datafusion_comet_proto::spark_expression::expr::ExprStruct;
     use datafusion_comet_proto::spark_expression::ListLiteral;
     use datafusion_comet_proto::{
@@ -3412,7 +3404,6 @@ mod tests {
         spark_operator,
         spark_operator::{operator::OpStruct, Operator},
     };
-    use datafusion_comet_spark_expr::EvalMode;
 
     #[test]
     fn test_unpack_dictionary_primitive() {
@@ -4004,18 +3995,15 @@ mod tests {
             }
         }
 
-        let source = ParquetSource::default().with_schema_adapter_factory(Arc::new(
-            SparkSchemaAdapterFactory::new(
-                SparkParquetOptions::new(EvalMode::Ansi, "", false),
-                None,
-            ),
-        ))?;
+        let source = Arc::new(
+            ParquetSource::new(Arc::new(read_schema.clone()))
+                .with_table_parquet_options(TableParquetOptions::new()),
+        ) as Arc<dyn FileSource>;
 
         let object_store_url = ObjectStoreUrl::local_filesystem();
-        let file_scan_config =
-            FileScanConfigBuilder::new(object_store_url, read_schema.into(), source)
-                .with_file_groups(file_groups)
-                .build();
+        let file_scan_config = FileScanConfigBuilder::new(object_store_url, source)
+            .with_file_groups(file_groups)
+            .build();
 
         // Run native read
         let scan = Arc::new(DataSourceExec::new(Arc::new(file_scan_config.clone())));

native/core/src/parquet/mod.rs

@@ -765,7 +765,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             required_schema,
             Some(data_schema),
             None,
-            None,
             object_store_url,
             file_groups,
             None,

native/core/src/parquet/parquet_exec.rs

@@ -15,10 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::collections::HashMap;
 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::SparkSchemaAdapterFactory;
 use arrow::datatypes::{Field, SchemaRef};
 use datafusion::config::TableParquetOptions;
 use datafusion::datasource::listing::PartitionedFile;
@@ -30,11 +30,12 @@ use datafusion::execution::object_store::ObjectStoreUrl;
 use datafusion::physical_expr::expressions::BinaryExpr;
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::prelude::SessionContext;
-use datafusion::scalar::ScalarValue;
 use datafusion_comet_spark_expr::EvalMode;
-use itertools::Itertools;
-use std::collections::HashMap;
+use datafusion_datasource::TableSchema;
 use std::sync::Arc;
+use datafusion::physical_expr_adapter::PhysicalExprAdapterFactory;
+use datafusion::scalar::ScalarValue;
+use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory;
 
 /// Initializes a DataSourceExec plan with a ParquetSource. This may be used by either the
 /// `native_datafusion` scan or the `native_iceberg_compat` scan.
@@ -60,7 +61,6 @@ pub(crate) fn init_datasource_exec(
     required_schema: SchemaRef,
     data_schema: Option<SchemaRef>,
     partition_schema: Option<SchemaRef>,
-    partition_fields: Option<Vec<Field>>,
     object_store_url: ObjectStoreUrl,
     file_groups: Vec<Vec<PartitionedFile>>,
     projection_vector: Option<Vec<usize>>,
@@ -78,7 +78,26 @@ pub(crate) fn init_datasource_exec(
         encryption_enabled,
     );
 
-    let mut parquet_source = ParquetSource::new(table_parquet_options);
+    // Determine the schema to use for ParquetSource
+    let table_schema = if let Some(ref data_schema) = data_schema {
+        if let Some(ref partition_schema) = partition_schema {
+            let partition_fields: Vec<_> = partition_schema
+                .fields()
+                .iter()
+                .map(|f| {
+                    Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable())) as _
+                })
+                .collect();
+            TableSchema::new(Arc::clone(data_schema), partition_fields)
+        } else {
+            TableSchema::from_file_schema(Arc::clone(data_schema))
+        }
+    } else {
+        TableSchema::from_file_schema(Arc::clone(&required_schema))
+    };
+
+    let mut parquet_source =
+        ParquetSource::new(table_schema).with_table_parquet_options(table_parquet_options);
 
     // Create a conjunctive form of the vector because ParquetExecBuilder takes
     // a single expression
@@ -104,37 +123,26 @@ pub(crate) fn init_datasource_exec(
         );
     }
 
-    let file_source = parquet_source.with_schema_adapter_factory(Arc::new(
-        SparkSchemaAdapterFactory::new(spark_parquet_options, default_values),
-    ))?;
+    let expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory> = Arc::new(
+        SparkPhysicalExprAdapterFactory::new(spark_parquet_options, default_values),
+    );
+
+    let file_source: Arc<dyn FileSource> = Arc::new(parquet_source);
 
     let file_groups = file_groups
         .iter()
         .map(|files| FileGroup::new(files.clone()))
         .collect();
 
-    let file_scan_config = match (data_schema, projection_vector, partition_fields) {
-        (Some(data_schema), Some(projection_vector), Some(partition_fields)) => {
-            get_file_config_builder(
-                data_schema,
-                partition_schema,
-                file_groups,
-                object_store_url,
-                file_source,
-            )
-            .with_projection_indices(Some(projection_vector))
-            .with_table_partition_cols(partition_fields)
-            .build()
-        }
-        _ => get_file_config_builder(
-            required_schema,
-            partition_schema,
-            file_groups,
-            object_store_url,
-            file_source,
-        )
-        .build(),
-    };
+    let mut file_scan_config_builder =
+        FileScanConfigBuilder::new(object_store_url, file_source).with_file_groups(file_groups);
+
+    if let Some(projection_vector) = projection_vector {
+        file_scan_config_builder =
+            file_scan_config_builder.with_projection_indices(Some(projection_vector))?;
+    }
+
+    let file_scan_config = file_scan_config_builder.with_expr_adapter(Some(expr_adapter_factory)).build();
 
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
@@ -165,28 +173,3 @@ fn get_options(
 
     (table_parquet_options, spark_parquet_options)
 }
-
-fn get_file_config_builder(
-    schema: SchemaRef,
-    partition_schema: Option<SchemaRef>,
-    file_groups: Vec<FileGroup>,
-    object_store_url: ObjectStoreUrl,
-    file_source: Arc<dyn FileSource>,
-) -> FileScanConfigBuilder {
-    match partition_schema {
-        Some(partition_schema) => {
-            let partition_fields: Vec<Field> = partition_schema
-                .fields()
-                .iter()
-                .map(|field| {
-                    Field::new(field.name(), field.data_type().clone(), field.is_nullable())
-                })
-                .collect_vec();
-            FileScanConfigBuilder::new(object_store_url, Arc::clone(&schema), file_source)
-                .with_file_groups(file_groups)
-                .with_table_partition_cols(partition_fields)
-        }
-        _ => FileScanConfigBuilder::new(object_store_url, Arc::clone(&schema), file_source)
-            .with_file_groups(file_groups),
-    }
-}