DataFusion 52 migration

Author: comphead

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

@@ -94,7 +94,7 @@ impl ScanExec {
 
         // Build schema directly from data types since get_next now always unpacks dictionaries
         let schema = schema_from_data_types(&data_types);
-        dbg!(&schema);
+        // dbg!(&schema);
 
         let cache = PlanProperties::new(
             EquivalenceProperties::new(Arc::clone(&schema)),
@@ -210,7 +210,7 @@ impl ScanExec {
 
             let array = make_array(array_data);
 
-            dbg!(&array, &selection_indices_arrays);
+            // dbg!(&array, &selection_indices_arrays);
 
             // Apply selection if selection vectors exist (applies to all columns)
             let array = if let Some(ref selection_arrays) = selection_indices_arrays {
@@ -490,7 +490,7 @@ impl ScanStream<'_> {
     ) -> DataFusionResult<RecordBatch, DataFusionError> {
         let schema_fields = self.schema.fields();
         assert_eq!(columns.len(), schema_fields.len());
-        dbg!(&columns, &self.schema);
+        // dbg!(&columns, &self.schema);
         // Cast dictionary-encoded primitive arrays to regular arrays and cast
         // Utf8/LargeUtf8/Binary arrays to dictionary-encoded if the schema is
         // defined as dictionary-encoded and the data in this batch is not
@@ -510,7 +510,7 @@ impl ScanStream<'_> {
             })
             .collect::<Result<Vec<_>, _>>()?;
         let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
-        dbg!(&new_columns, &self.schema);
+        // dbg!(&new_columns, &self.schema);
         RecordBatch::try_new_with_options(Arc::clone(&self.schema), new_columns, &options)
             .map_err(|e| arrow_datafusion_err!(e))
     }
@@ -521,7 +521,7 @@ impl Stream for ScanStream<'_> {
 
     fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let mut timer = self.baseline_metrics.elapsed_compute().timer();
-        dbg!(&self.scan);
+        // dbg!(&self.scan);
         let mut scan_batch = self.scan.batch.try_lock().unwrap();
 
         let input_batch = &*scan_batch;

native/core/src/execution/planner.rs

@@ -946,7 +946,7 @@ impl PhysicalPlanner {
                 ))
             }
             OpStruct::NativeScan(scan) => {
-                dbg!(&scan);
+                // dbg!(&scan);
                 let data_schema = convert_spark_types_to_arrow_schema(scan.data_schema.as_slice());
                 let required_schema: SchemaRef =
                     convert_spark_types_to_arrow_schema(scan.required_schema.as_slice());
@@ -1094,7 +1094,7 @@ impl PhysicalPlanner {
                 ))
             }
             OpStruct::Scan(scan) => {
-                dbg!(&scan);
+                // dbg!(&scan);
                 let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec();
 
                 // If it is not test execution context for unit test, we should have at least one
@@ -1121,7 +1121,7 @@ impl PhysicalPlanner {
                     scan.arrow_ffi_safe,
                 )?;
 
-                dbg!(&scan);
+                // dbg!(&scan);
 
                 Ok((
                     vec![scan.clone()],
@@ -4449,7 +4449,8 @@ mod tests {
         let int8_array = Int8Array::from(vec![Some(1i8), Some(2i8), Some(3i8)]);
 
         // Set input batch for the scan
-        let input_batch = InputBatch::Batch(vec![Arc::new(date_array), Arc::new(int8_array)], row_count);
+        let input_batch =
+            InputBatch::Batch(vec![Arc::new(date_array), Arc::new(int8_array)], row_count);
         scans[0].set_input_batch(input_batch);
 
         let session_ctx = SessionContext::new();
@@ -4464,7 +4465,8 @@ mod tests {
             // Create test data again for the second batch
             let date_array = Date32Array::from(vec![Some(19000), Some(19001), Some(19002)]);
             let int8_array = Int8Array::from(vec![Some(1i8), Some(2i8), Some(3i8)]);
-            let input_batch1 = InputBatch::Batch(vec![Arc::new(date_array), Arc::new(int8_array)], row_count);
+            let input_batch1 =
+                InputBatch::Batch(vec![Arc::new(date_array), Arc::new(int8_array)], row_count);
             let input_batch2 = InputBatch::EOF;
 
             let batches = vec![input_batch1, input_batch2];

native/core/src/parquet/mod.rs

@@ -703,7 +703,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
     key_unwrapper_obj: JObject,
     metrics_node: JObject,
 ) -> jlong {
-    dbg!("Java_org_apache_comet_parquet_Native_initRecordBatchReader");
+    // dbg!("Java_org_apache_comet_parquet_Native_initRecordBatchReader");
     try_unwrap_or_throw(&e, |mut env| unsafe {
         JVMClasses::init(&mut env);
         let session_config = SessionConfig::new().with_batch_size(batch_size as usize);
@@ -777,7 +777,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             encryption_enabled,
         )?;
 
-        dbg!(&scan);
+        // dbg!(&scan);
 
         let partition_index: usize = 0;
         let batch_stream = Some(scan.execute(partition_index, session_ctx.task_ctx())?);
@@ -791,7 +791,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
         };
         let res = Box::new(ctx);
 
-        dbg!("end Java_org_apache_comet_parquet_Native_initRecordBatchReader");
+        // dbg!("end Java_org_apache_comet_parquet_Native_initRecordBatchReader");
 
         Ok(Box::into_raw(res) as i64)
     })

native/core/src/parquet/parquet_exec.rs

@@ -36,9 +36,6 @@ use datafusion_comet_spark_expr::EvalMode;
 use datafusion_datasource::TableSchema;
 use std::collections::HashMap;
 use std::sync::Arc;
-use arrow::util::pretty::print_batches;
-use datafusion::physical_plan::ExecutionPlan;
-use futures::StreamExt;
 
 /// Initializes a DataSourceExec plan with a ParquetSource. This may be used by either the
 /// `native_datafusion` scan or the `native_iceberg_compat` scan.
@@ -81,10 +78,14 @@ pub(crate) fn init_datasource_exec(
         encryption_enabled,
     );
 
-    dbg!(&required_schema, &data_schema);
+    // dbg!(&required_schema, &data_schema);
 
     // Determine the schema to use for ParquetSource
-    let base_schema = required_schema.clone();
+    // Use data_schema only if both data_schema and data_filters are set
+    let base_schema = match (&data_schema, &data_filters) {
+        (Some(schema), Some(_)) => Arc::clone(schema),
+        _ => Arc::clone(&required_schema),
+    };
     let partition_fields: Vec<_> = partition_schema
         .iter()
         .flat_map(|s| s.fields().iter())
@@ -98,7 +99,7 @@ pub(crate) fn init_datasource_exec(
     let mut parquet_source =
         ParquetSource::new(table_schema).with_table_parquet_options(table_parquet_options);
 
-    dbg!(&parquet_source);
+    // dbg!(&parquet_source);
 
     // Create a conjunctive form of the vector because ParquetExecBuilder takes
     // a single expression
@@ -135,12 +136,13 @@ pub(crate) fn init_datasource_exec(
         .map(|files| FileGroup::new(files.clone()))
         .collect();
 
-    let mut file_scan_config_builder =
-        FileScanConfigBuilder::new(object_store_url, file_source).with_file_groups(file_groups);//.with_expr_adapter(Some(expr_adapter_factory));
+    let mut file_scan_config_builder = FileScanConfigBuilder::new(object_store_url, file_source)
+        .with_file_groups(file_groups)
+        .with_expr_adapter(Some(expr_adapter_factory));
 
     if let Some(projection_vector) = projection_vector {
-        file_scan_config_builder = file_scan_config_builder
-            .with_projection_indices(Some(projection_vector))?;
+        file_scan_config_builder =
+            file_scan_config_builder.with_projection_indices(Some(projection_vector))?;
     }
 
     let file_scan_config = file_scan_config_builder.build();

native/core/src/parquet/schema_adapter.rs

@@ -182,15 +182,15 @@ impl SparkPhysicalExprAdapter {
                 // Get the physical datatype (actual file schema)
                 let physical_field = self.physical_file_schema.fields().get(col_idx);
 
-                dbg!(&logical_field, &physical_field);
+                // dbg!(&logical_field, &physical_field);
 
                 if let (Some(logical_field), Some(physical_field)) = (logical_field, physical_field)
                 {
                     let logical_type = logical_field.data_type();
                     let physical_type = physical_field.data_type();
 
                     // If datatypes differ, insert a CastColumnExpr
-                    if logical_type != physical_type || 1==1 {
+                    if logical_type != physical_type {
                         let input_field = Arc::new(Field::new(
                             physical_field.name(),
                             physical_type.clone(),
@@ -203,12 +203,12 @@ impl SparkPhysicalExprAdapter {
                         ));
 
                         let cast_expr: Arc<dyn PhysicalExpr> = Arc::new(CastColumnExpr::new(
-                            e.clone(),
+                            Arc::clone(&e),
                             input_field,
                             target_field,
                             None,
                         ));
-                        dbg!(&cast_expr);
+                        // dbg!(&cast_expr);
                         return Ok(Transformed::yes(cast_expr));
                     }
                 }

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -347,8 +347,8 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   }
 
   test("date_sub with int scalars") {
-    Seq(false).foreach { dictionaryEnabled =>
-      Seq("INT").foreach { intType =>
+    Seq(true, false).foreach { dictionaryEnabled =>
+      Seq("TINYINT", "SHORT", "INT").foreach { intType =>
         withTempDir { dir =>
           val path = new Path(dir.toURI.toString, "test.parquet")
           makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = dictionaryEnabled, 10000)