feat: [comet-parquet-exec] Schema adapter fixes

native/core/src/execution/datafusion/planner.rs

@@ -19,7 +19,7 @@
 
 use super::expressions::EvalMode;
 use crate::execution::datafusion::expressions::comet_scalar_funcs::create_comet_physical_fun;
-use crate::execution::operators::{CopyMode, FilterExec};
+use crate::execution::operators::{CopyMode, FilterExec as CometFilterExec};
 use crate::{
     errors::ExpressionError,
     execution::{
@@ -55,6 +55,7 @@ use datafusion::functions_aggregate::bit_and_or_xor::{bit_and_udaf, bit_or_udaf,
 use datafusion::functions_aggregate::min_max::max_udaf;
 use datafusion::functions_aggregate::min_max::min_udaf;
 use datafusion::functions_aggregate::sum::sum_udaf;
+use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::windows::BoundedWindowAggExec;
 use datafusion::physical_plan::InputOrderMode;
 use datafusion::{
@@ -851,7 +852,11 @@ impl PhysicalPlanner {
                 let predicate =
                     self.create_expr(filter.predicate.as_ref().unwrap(), child.schema())?;
 
-                Ok((scans, Arc::new(FilterExec::try_new(predicate, child)?)))
+                if can_reuse_input_batch(&child) {
+                    Ok((scans, Arc::new(CometFilterExec::try_new(predicate, child)?)))
+                } else {
+                    Ok((scans, Arc::new(FilterExec::try_new(predicate, child)?)))
+                }
             }
             OpStruct::HashAgg(agg) => {
                 assert!(children.len() == 1);

native/core/src/execution/datafusion/schema_adapter.rs

@@ -19,10 +19,10 @@
 
 use arrow::compute::can_cast_types;
 use arrow_array::{new_null_array, Array, RecordBatch, RecordBatchOptions};
-use arrow_schema::{DataType, Schema, SchemaRef};
+use arrow_schema::{DataType, Schema, SchemaRef, TimeUnit};
 use datafusion::datasource::schema_adapter::{SchemaAdapter, SchemaAdapterFactory, SchemaMapper};
 use datafusion_comet_spark_expr::{spark_cast, EvalMode};
-use datafusion_common::plan_err;
+use datafusion_common::{plan_err, DataFusionError};
 use datafusion_expr::ColumnarValue;
 use std::sync::Arc;
 
@@ -38,11 +38,11 @@ impl SchemaAdapterFactory for CometSchemaAdapterFactory {
     /// schema.
     fn create(
         &self,
-        projected_table_schema: SchemaRef,
+        required_schema: SchemaRef,
         table_schema: SchemaRef,
     ) -> Box<dyn SchemaAdapter> {
         Box::new(CometSchemaAdapter {
-            projected_table_schema,
+            required_schema,
             table_schema,
         })
     }
@@ -54,7 +54,7 @@ impl SchemaAdapterFactory for CometSchemaAdapterFactory {
 pub struct CometSchemaAdapter {
     /// The schema for the table, projected to include only the fields being output (projected) by the
     /// associated ParquetExec
-    projected_table_schema: SchemaRef,
+    required_schema: SchemaRef,
     /// The entire table schema for the table we're using this to adapt.
     ///
     /// This is used to evaluate any filters pushed down into the scan
@@ -69,7 +69,7 @@ impl SchemaAdapter for CometSchemaAdapter {
     ///
     /// Panics if index is not in range for the table schema
     fn map_column_index(&self, index: usize, file_schema: &Schema) -> Option<usize> {
-        let field = self.projected_table_schema.field(index);
+        let field = self.required_schema.field(index);
         Some(file_schema.fields.find(field.name())?.0)
     }
 
@@ -87,40 +87,29 @@ impl SchemaAdapter for CometSchemaAdapter {
         file_schema: &Schema,
     ) -> datafusion_common::Result<(Arc<dyn SchemaMapper>, Vec<usize>)> {
         let mut projection = Vec::with_capacity(file_schema.fields().len());
-        let mut field_mappings = vec![None; self.projected_table_schema.fields().len()];
+        let mut field_mappings = vec![None; self.required_schema.fields().len()];
 
         for (file_idx, file_field) in file_schema.fields.iter().enumerate() {
             if let Some((table_idx, table_field)) =
-                self.projected_table_schema.fields().find(file_field.name())
+                self.required_schema.fields().find(file_field.name())
             {
-                // workaround for struct casting
-                match (file_field.data_type(), table_field.data_type()) {
-                    // TODO need to use Comet cast logic to determine which casts are supported,
-                    // but for now just add a hack to support casting between struct types
-                    (DataType::Struct(_), DataType::Struct(_)) => {
-                        field_mappings[table_idx] = Some(projection.len());
-                        projection.push(file_idx);
-                    }
-                    _ => {
-                        if can_cast_types(file_field.data_type(), table_field.data_type()) {
-                            field_mappings[table_idx] = Some(projection.len());
-                            projection.push(file_idx);
-                        } else {
-                            return plan_err!(
-                                "Cannot cast file schema field {} of type {:?} to table schema field of type {:?}",
-                                file_field.name(),
-                                file_field.data_type(),
-                                table_field.data_type()
-                            );
-                        }
-                    }
+                if comet_can_cast_types(file_field.data_type(), table_field.data_type()) {
+                    field_mappings[table_idx] = Some(projection.len());
+                    projection.push(file_idx);
+                } else {
+                    return plan_err!(
+                        "Cannot cast file schema field {} of type {:?} to required schema field of type {:?}",
+                        file_field.name(),
+                        file_field.data_type(),
+                        table_field.data_type()
+                    );
                 }
             }
         }
 
         Ok((
             Arc::new(SchemaMapping {
-                projected_table_schema: self.projected_table_schema.clone(),
+                required_schema: self.required_schema.clone(),
                 field_mappings,
                 table_schema: self.table_schema.clone(),
             }),
@@ -161,7 +150,7 @@ impl SchemaAdapter for CometSchemaAdapter {
 pub struct SchemaMapping {
     /// The schema of the table. This is the expected schema after conversion
     /// and it should match the schema of the query result.
-    projected_table_schema: SchemaRef,
+    required_schema: SchemaRef,
     /// Mapping from field index in `projected_table_schema` to index in
     /// projected file_schema.
     ///
@@ -185,7 +174,7 @@ impl SchemaMapper for SchemaMapping {
         let batch_cols = batch.columns().to_vec();
 
         let cols = self
-            .projected_table_schema
+            .required_schema
             // go through each field in the projected schema
             .fields()
             .iter()
@@ -208,6 +197,7 @@ impl SchemaMapper for SchemaMapping {
                             EvalMode::Legacy,
                             "UTC",
                             false,
+                            true,
                         )?
                         .into_array(batch_rows)
                     },
@@ -218,7 +208,7 @@ impl SchemaMapper for SchemaMapping {
         // Necessary to handle empty batches
         let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
 
-        let schema = self.projected_table_schema.clone();
+        let schema = self.required_schema.clone();
         let record_batch = RecordBatch::try_new_with_options(schema, cols, &options)?;
         Ok(record_batch)
     }
@@ -259,6 +249,7 @@ impl SchemaMapper for SchemaMapping {
                             EvalMode::Legacy,
                             "UTC",
                             false,
+                            true,
                         )?
                         .into_array(batch_col.len())
                         // and if that works, return the field and column.
@@ -277,3 +268,31 @@ impl SchemaMapper for SchemaMapping {
         Ok(record_batch)
     }
 }
+
+
+
+fn comet_can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
+    // TODO this is just a quick hack to get tests passing
+    match (from_type, to_type) {
+        (DataType::Struct(_), DataType::Struct(_)) => {
+            // workaround for struct casting
+            true
+        }
+        // TODO this is maybe no longer needed
+        (_, DataType::Timestamp(TimeUnit::Nanosecond, _)) => false,
+        _ => can_cast_types(from_type, to_type),
+    }
+}
+
+pub fn comet_cast(
+    arg: ColumnarValue,
+    data_type: &DataType,
+    eval_mode: EvalMode,
+    timezone: &str,
+    allow_incompat: bool,
+) -> Result<ColumnarValue, DataFusionError> {
+    // TODO for now we are re-using the spark cast rules, with a hack to override
+    // unsupported cases and let those fall through to arrow. This is just a short term
+    // hack and we need to implement specific Parquet to Spark conversions here instead
+    spark_cast(arg, data_type, eval_mode, timezone, allow_incompat, true)
+}

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

@@ -227,7 +227,7 @@ impl DisplayAs for FilterExec {
 
 impl ExecutionPlan for FilterExec {
     fn name(&self) -> &'static str {
-        "FilterExec"
+        "CometFilterExec"
     }
 
     /// Return a reference to Any that can be used for downcasting

native/spark-expr/src/cast.rs

@@ -585,6 +585,7 @@ pub fn spark_cast(
     eval_mode: EvalMode,
     timezone: &str,
     allow_incompat: bool,
+    ugly_hack_for_poc: bool, // TODO we definitely don't want to do this
 ) -> DataFusionResult<ColumnarValue> {
     match arg {
         ColumnarValue::Array(array) => Ok(ColumnarValue::Array(cast_array(
@@ -593,6 +594,7 @@ pub fn spark_cast(
             eval_mode,
             timezone.to_owned(),
             allow_incompat,
+            ugly_hack_for_poc,
         )?)),
         ColumnarValue::Scalar(scalar) => {
             // Note that normally CAST(scalar) should be fold in Spark JVM side. However, for
@@ -606,6 +608,7 @@ pub fn spark_cast(
                     eval_mode,
                     timezone.to_owned(),
                     allow_incompat,
+                    ugly_hack_for_poc,
                 )?,
                 0,
             )?;
@@ -620,6 +623,7 @@ fn cast_array(
     eval_mode: EvalMode,
     timezone: String,
     allow_incompat: bool,
+    ugly_hack_for_poc: bool,
 ) -> DataFusionResult<ArrayRef> {
     let array = array_with_timezone(array, timezone.clone(), Some(to_type))?;
     let from_type = array.data_type().clone();
@@ -643,6 +647,7 @@ fn cast_array(
                     eval_mode,
                     timezone,
                     allow_incompat,
+                    ugly_hack_for_poc,
                 )?,
             );
 
@@ -723,7 +728,9 @@ fn cast_array(
             timezone,
             allow_incompat,
         )?),
-        _ if is_datafusion_spark_compatible(from_type, to_type, allow_incompat) => {
+        _ if ugly_hack_for_poc
+            || is_datafusion_spark_compatible(from_type, to_type, allow_incompat) =>
+        {
             // use DataFusion cast only when we know that it is compatible with Spark
             Ok(cast_with_options(&array, to_type, &CAST_OPTIONS)?)
         }
@@ -840,6 +847,7 @@ fn cast_struct_to_struct(
                     eval_mode,
                     timezone.clone(),
                     allow_incompat,
+                    false,
                 )?;
                 cast_fields.push((Arc::clone(&to_fields[i]), cast_field));
             }
@@ -861,6 +869,7 @@ fn casts_struct_to_string(array: &StructArray, timezone: &str) -> DataFusionResu
                 EvalMode::Legacy,
                 timezone,
                 true,
+                false,
             )
             .and_then(|cv| cv.into_array(arr.len()))
         })
@@ -1505,6 +1514,7 @@ impl PhysicalExpr for Cast {
             self.eval_mode,
             &self.timezone,
             self.allow_incompat,
+            false,
         )
     }
 
@@ -2117,6 +2127,7 @@ mod tests {
             EvalMode::Legacy,
             timezone.clone(),
             false,
+            false,
         )?;
         assert_eq!(
             *result.data_type(),
@@ -2327,6 +2338,7 @@ mod tests {
             EvalMode::Legacy,
             "UTC".to_owned(),
             false,
+            false,
         );
         assert!(result.is_err())
     }
@@ -2340,6 +2352,7 @@ mod tests {
             EvalMode::Legacy,
             "Not a valid timezone".to_owned(),
             false,
+            false,
         );
         assert!(result.is_err())
     }
@@ -2364,6 +2377,7 @@ mod tests {
             EvalMode::Legacy,
             "UTC".to_owned(),
             false,
+            false,
         )
         .unwrap();
         let string_array = string_array.as_string::<i32>();
@@ -2400,6 +2414,7 @@ mod tests {
             EvalMode::Legacy,
             "UTC",
             false,
+            false,
         )
         .unwrap();
         if let ColumnarValue::Array(cast_array) = cast_array {
@@ -2433,6 +2448,7 @@ mod tests {
             EvalMode::Legacy,
             "UTC",
             false,
+            false,
         )
         .unwrap();
         if let ColumnarValue::Array(cast_array) = cast_array {

native/spark-expr/src/to_json.rs

@@ -120,6 +120,7 @@ fn array_to_json_string(arr: &Arc<dyn Array>, timezone: &str) -> Result<ArrayRef
             EvalMode::Legacy,
             timezone,
             false,
+            false,
         )?
         .into_array(arr.len())
     }