DataFusion 52 migration

Author: comphead

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

@@ -42,9 +42,9 @@ use iceberg::io::FileIO;
 
 use crate::execution::operators::ExecutionError;
 use crate::parquet::parquet_support::SparkParquetOptions;
-use crate::parquet::schema_adapter::SparkSchemaMapperFactory;
 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.
 ///
@@ -298,19 +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 mapper_factory = SparkSchemaMapperFactory::new(spark_options, None);
-                    let file_schema = batch.schema();
-
-                    let result = mapper_factory
-                        .create_mapper(Arc::clone(&target_schema), file_schema.as_ref())
-                        .and_then(|mapper| {
-                            mapper.map_batch(batch).map_err(|e| {
-                                DataFusionError::Execution(format!("Batch mapping failed: {}", e))
-                            })
-                        })
-                        .map_err(|e| {
-                            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

@@ -1034,6 +1034,42 @@ impl PhysicalPlanner {
                     .map(|expr| self.create_expr(expr, Arc::clone(&required_schema)))
                     .collect();
 
+                let default_values: Option<HashMap<usize, ScalarValue>> = if !scan
+                    .default_values
+                    .is_empty()
+                {
+                    // We have default values. Extract the two lists (same length) of values and
+                    // indexes in the schema, and then create a HashMap to use in the SchemaMapper.
+                    let default_values: Result<Vec<ScalarValue>, DataFusionError> = scan
+                        .default_values
+                        .iter()
+                        .map(|expr| {
+                            let literal = self.create_expr(expr, Arc::clone(&required_schema))?;
+                            let df_literal = literal
+                                .as_any()
+                                .downcast_ref::<DataFusionLiteral>()
+                                .ok_or_else(|| {
+                                    GeneralError("Expected literal of default value.".to_string())
+                                })?;
+                            Ok(df_literal.value().clone())
+                        })
+                        .collect();
+                    let default_values = default_values?;
+                    let default_values_indexes: Vec<usize> = scan
+                        .default_values_indexes
+                        .iter()
+                        .map(|offset| *offset as usize)
+                        .collect();
+                    Some(
+                        default_values_indexes
+                            .into_iter()
+                            .zip(default_values)
+                            .collect(),
+                    )
+                } else {
+                    None
+                };
+
                 // Get one file from this partition (we know it's not empty due to early return above)
                 let one_file = partition_files
                     .partitioned_file
@@ -1066,6 +1102,7 @@ impl PhysicalPlanner {
                     file_groups,
                     Some(projection_vector),
                     Some(data_filters?),
+                    default_values,
                     scan.session_timezone.as_str(),
                     scan.case_sensitive,
                     self.session_ctx(),

native/core/src/parquet/mod.rs

@@ -769,6 +769,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             file_groups,
             None,
             data_filters,
+            None,
             session_timezone.as_str(),
             case_sensitive != JNI_FALSE,
             session_ctx,

native/core/src/parquet/parquet_exec.rs

@@ -15,6 +15,7 @@
 // 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;
@@ -32,6 +33,9 @@ use datafusion::prelude::SessionContext;
 use datafusion_comet_spark_expr::EvalMode;
 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.
@@ -61,12 +65,13 @@ pub(crate) fn init_datasource_exec(
     file_groups: Vec<Vec<PartitionedFile>>,
     projection_vector: Option<Vec<usize>>,
     data_filters: Option<Vec<Arc<dyn PhysicalExpr>>>,
+    default_values: Option<HashMap<usize, ScalarValue>>,
     session_timezone: &str,
     case_sensitive: bool,
     session_ctx: &Arc<SessionContext>,
     encryption_enabled: bool,
 ) -> Result<Arc<DataSourceExec>, ExecutionError> {
-    let (table_parquet_options, _) = get_options(
+    let (table_parquet_options, spark_parquet_options) = get_options(
         session_timezone,
         case_sensitive,
         &object_store_url,
@@ -118,7 +123,11 @@ pub(crate) fn init_datasource_exec(
         );
     }
 
-    let file_source = Arc::new(parquet_source) as Arc<dyn FileSource>;
+    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()
@@ -133,7 +142,7 @@ pub(crate) fn init_datasource_exec(
             file_scan_config_builder.with_projection_indices(Some(projection_vector))?;
     }
 
-    let file_scan_config = file_scan_config_builder.build();
+    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))))
 }