feat: project Spark 4 VARIANT columns in native Parquet scans

native/core/src/execution/jni_api.rs

@@ -686,6 +686,7 @@ fn prepare_output(
     let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? };
     let schema_addrs = &*schema_addrs;
 
+    let output_schema = output_batch.schema();
     let results = output_batch.columns();
     let num_rows = output_batch.num_rows();
 
@@ -712,6 +713,7 @@ fn prepare_output(
         let mut i = 0;
         while i < results.len() {
             let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?;
+            let field = output_schema.field(i);
 
             if array_ref.offset() != 0 {
                 // https://github.com/apache/datafusion-comet/issues/2051
@@ -728,11 +730,11 @@ fn prepare_output(
 
                 new_array
                     .to_data()
-                    .move_to_spark(array_addrs[i], schema_addrs[i])?;
+                    .move_to_spark(field, array_addrs[i], schema_addrs[i])?;
             } else {
                 array_ref
                     .to_data()
-                    .move_to_spark(array_addrs[i], schema_addrs[i])?;
+                    .move_to_spark(field, array_addrs[i], schema_addrs[i])?;
             }
             i += 1;
         }

native/core/src/execution/planner.rs

@@ -43,7 +43,7 @@ use crate::execution::{
     },
     planner::expression_registry::ExpressionRegistry,
     planner::operator_registry::OperatorRegistry,
-    serde::to_arrow_datatype,
+    serde::{to_arrow_datatype, to_arrow_field},
     shuffle::{SchemaAlignExec, ShuffleWriterExec},
 };
 use crate::jvm_bridge::{jni_call, JVMClasses};
@@ -3772,15 +3772,17 @@ pub(crate) fn convert_spark_types_to_arrow_schema(
     let arrow_fields = spark_types
         .iter()
         .map(|spark_type| {
-            let field = Field::new(
+            let field = to_arrow_field(
                 String::clone(&spark_type.name),
-                to_arrow_datatype(spark_type.data_type.as_ref().unwrap()),
+                spark_type.data_type.as_ref().unwrap(),
                 spark_type.nullable,
             );
             if spark_type.metadata.is_empty() {
                 field
             } else {
-                field.with_metadata(spark_type.metadata.clone())
+                let mut metadata = spark_type.metadata.clone();
+                metadata.extend(field.metadata().clone());
+                field.with_metadata(metadata)
             }
         })
         .collect_vec();

native/core/src/execution/serde.rs

@@ -31,9 +31,9 @@ use datafusion_comet_proto::{
     spark_expression::DataType,
     spark_operator,
 };
-use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+use parquet::{arrow::PARQUET_FIELD_ID_META_KEY, variant::VariantType};
 use prost::Message;
-use std::{collections::HashMap, io::Cursor, sync::Arc};
+use std::{io::Cursor, sync::Arc};
 
 /// Deserialize bytes to protobuf type of expression
 pub fn deserialize_expr(buf: &[u8]) -> Result<spark_expression::Expr, ExpressionError> {
@@ -106,6 +106,10 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
         // Spark's CalendarIntervalType stores months, days, and microseconds. Arrow stores the
         // same components with nanosecond precision.
         DataTypeId::CalendarInterval => ArrowDataType::Interval(IntervalUnit::MonthDayNano),
+        DataTypeId::Variant => ArrowDataType::Struct(Fields::from(vec![
+            Field::new("value", ArrowDataType::Binary, false),
+            Field::new("metadata", ArrowDataType::Binary, false),
+        ])),
         DataTypeId::Null => ArrowDataType::Null,
         DataTypeId::List => match dt_value
             .type_info
@@ -117,9 +121,9 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
         {
             DatatypeStruct::List(info) => {
                 let field = with_parquet_field_id(
-                    Field::new(
+                    to_arrow_field(
                         "item",
-                        to_arrow_datatype(info.element_type.as_ref().unwrap()),
+                        info.element_type.as_ref().unwrap(),
                         info.contains_null,
                     ),
                     info.element_field_id,
@@ -138,17 +142,13 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
         {
             DatatypeStruct::Map(info) => {
                 let key_field = with_parquet_field_id(
-                    Field::new(
-                        "key",
-                        to_arrow_datatype(info.key_type.as_ref().unwrap()),
-                        false,
-                    ),
+                    to_arrow_field("key", info.key_type.as_ref().unwrap(), false),
                     info.key_field_id,
                 );
                 let value_field = with_parquet_field_id(
-                    Field::new(
+                    to_arrow_field(
                         "value",
-                        to_arrow_datatype(info.value_type.as_ref().unwrap()),
+                        info.value_type.as_ref().unwrap(),
                         info.value_contains_null,
                     ),
                     info.value_field_id,
@@ -176,16 +176,18 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
                     .iter()
                     .enumerate()
                     .map(|(idx, name)| {
-                        let field = Field::new(
+                        let field = to_arrow_field(
                             name,
-                            to_arrow_datatype(&info.field_datatypes[idx]),
+                            &info.field_datatypes[idx],
                             info.field_nullable[idx],
                         );
                         // Attach Spark field metadata (currently parquet.field.id) when present.
                         // field_metadata is parallel to field_names; either empty or full length.
                         if let Some(meta) = info.field_metadata.get(idx) {
                             if !meta.metadata.is_empty() {
-                                return field.with_metadata(meta.metadata.clone());
+                                let mut metadata = meta.metadata.clone();
+                                metadata.extend(field.metadata().clone());
+                                return field.with_metadata(metadata);
                             }
                         }
                         field
@@ -198,13 +200,32 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
     }
 }
 
+/// Converts a protobuf type to an Arrow field, preserving logical extension identity.
+pub fn to_arrow_field(
+    name: impl Into<std::string::String>,
+    data_type: &DataType,
+    nullable: bool,
+) -> Field {
+    let field = Field::new(name, to_arrow_datatype(data_type), nullable);
+    if DataTypeId::try_from(data_type.type_id).unwrap() == DataTypeId::Variant {
+        field.with_extension_type(VariantType)
+    } else {
+        field
+    }
+}
+
+pub fn is_variant_field(field: &Field) -> bool {
+    field.has_valid_extension_type::<VariantType>()
+}
+
 /// Attach a Parquet field ID without changing synthetic fields when Catalyst did not supply one.
 fn with_parquet_field_id(field: Field, field_id: Option<i32>) -> Field {
     match field_id {
-        Some(id) => field.with_metadata(HashMap::from([(
-            PARQUET_FIELD_ID_META_KEY.to_string(),
-            id.to_string(),
-        )])),
+        Some(id) => {
+            let mut metadata = field.metadata().clone();
+            metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string());
+            field.with_metadata(metadata)
+        }
         None => field,
     }
 }

native/core/src/execution/utils.rs

@@ -19,17 +19,18 @@
 use crate::execution::operators::ExecutionError;
 use arrow::{
     array::ArrayData,
+    datatypes::Field,
     ffi::{FFI_ArrowArray, FFI_ArrowSchema},
 };
 
 pub trait SparkArrowConvert {
     /// Move Arrow Arrays to C data interface.
-    fn move_to_spark(&self, array: i64, schema: i64) -> Result<(), ExecutionError>;
+    fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError>;
 }
 
 impl SparkArrowConvert for ArrayData {
     /// Move this ArrowData to pointers of Arrow C data interface.
-    fn move_to_spark(&self, array: i64, schema: i64) -> Result<(), ExecutionError> {
+    fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError> {
         let array_ptr = array as *mut FFI_ArrowArray;
         let schema_ptr = schema as *mut FFI_ArrowSchema;
 
@@ -40,7 +41,7 @@ impl SparkArrowConvert for ArrayData {
         if array_ptr.align_offset(array_align) != 0 || schema_ptr.align_offset(schema_align) != 0 {
             unsafe {
                 std::ptr::write_unaligned(array_ptr, FFI_ArrowArray::new(self));
-                std::ptr::write_unaligned(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);
+                std::ptr::write_unaligned(schema_ptr, FFI_ArrowSchema::try_from(field)?);
             }
         } else {
             // SAFETY: `array_ptr` and `schema_ptr` are aligned correctly.
@@ -56,7 +57,7 @@ impl SparkArrowConvert for ArrayData {
             );
             unsafe {
                 std::ptr::write(array_ptr, FFI_ArrowArray::new(self));
-                std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);
+                std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(field)?);
             }
         }
 

native/core/src/parquet/cast_column.rs

@@ -19,17 +19,21 @@ use arrow::{
         make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray,
         TimestampMicrosecondArray, TimestampMillisecondArray,
     },
-    compute::CastOptions,
+    compute::{cast, CastOptions},
     datatypes::{DataType, FieldRef, Schema, TimeUnit},
     record_batch::RecordBatch,
 };
 
-use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions};
+use crate::{
+    execution::serde::is_variant_field,
+    parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions},
+};
 use datafusion::common::format::DEFAULT_CAST_OPTIONS;
-use datafusion::common::Result as DataFusionResult;
 use datafusion::common::ScalarValue;
+use datafusion::common::{DataFusionError, Result as DataFusionResult};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
+use parquet::variant::{unshred_variant, VariantArray};
 use std::{
     fmt::{self, Display},
     hash::Hash,
@@ -176,6 +180,42 @@ fn cast_timestamp_micros_to_millis_scalar(
     ScalarValue::TimestampMillisecond(new_val, target_tz)
 }
 
+fn normalize_variant_array(
+    array: &ArrayRef,
+    target_field: &FieldRef,
+) -> DataFusionResult<ArrayRef> {
+    let DataType::Struct(fields) = target_field.data_type() else {
+        return Err(DataFusionError::Execution(
+            "Variant extension field must use Struct storage".to_string(),
+        ));
+    };
+    if fields.len() != 2
+        || fields[0].name() != "value"
+        || fields[1].name() != "metadata"
+        || fields
+            .iter()
+            .any(|field| field.data_type() != &DataType::Binary)
+    {
+        return Err(DataFusionError::Execution(
+            "Variant output must contain Binary children [value, metadata]".to_string(),
+        ));
+    }
+
+    let variant = VariantArray::try_new(array.as_ref())?;
+    let unshredded = unshred_variant(&variant)?;
+    let value = unshredded.value_field().ok_or_else(|| {
+        DataFusionError::Execution("Unshredded Variant is missing its value field".to_string())
+    })?;
+    let value = cast(value.as_ref(), &DataType::Binary)?;
+    let metadata = cast(unshredded.metadata_field().as_ref(), &DataType::Binary)?;
+    let output = StructArray::try_new(
+        fields.clone(),
+        vec![value, metadata],
+        unshredded.inner().nulls().cloned(),
+    )?;
+    Ok(Arc::new(output))
+}
+
 #[derive(Debug, Clone, Eq)]
 pub struct CometCastColumnExpr {
     /// The physical expression producing the value to cast.
@@ -260,6 +300,18 @@ impl PhysicalExpr for CometCastColumnExpr {
     fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
         let value = self.expr.evaluate(batch)?;
 
+        if is_variant_field(&self.target_field) {
+            return match value {
+                ColumnarValue::Array(array) => Ok(ColumnarValue::Array(normalize_variant_array(
+                    &array,
+                    &self.target_field,
+                )?)),
+                ColumnarValue::Scalar(_) => Err(DataFusionError::Execution(
+                    "Variant Parquet projection requires an array".to_string(),
+                )),
+            };
+        }
+
         // Use == (PartialEq) instead of equals_datatype because equals_datatype
         // ignores field names in nested types (Struct, List, Map). We need to detect
         // when field names differ (e.g., Struct("a","b") vs Struct("c","d")) so that
@@ -349,9 +401,70 @@ impl PhysicalExpr for CometCastColumnExpr {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{Array, Int32Array, StringArray};
+    use arrow::array::{Array, AsArray, Int32Array, Int64Array, StringArray};
     use arrow::datatypes::{Field, Fields};
     use datafusion::physical_expr::expressions::Column;
+    use parquet::variant::{Variant, VariantArrayBuilder, VariantType};
+
+    #[test]
+    fn test_normalize_shredded_variant_for_spark() {
+        let mut builder = VariantArrayBuilder::new(3);
+        builder.append_variant(Variant::from(1_i64));
+        builder.append_null();
+        builder.append_variant(Variant::from(3_i64));
+        let base = builder.build();
+        let metadata = Arc::clone(base.metadata_field());
+        let typed_value: ArrayRef = Arc::new(Int64Array::from(vec![Some(10), None, Some(30)]));
+        let physical_fields = Fields::from(vec![
+            Field::new("typed_value", DataType::Int64, true),
+            Field::new("metadata", metadata.data_type().clone(), false),
+        ]);
+        let physical = StructArray::try_new(
+            physical_fields,
+            vec![typed_value, metadata],
+            base.inner().nulls().cloned(),
+        )
+        .unwrap();
+
+        let input_field = Arc::new(Field::new("v", physical.data_type().clone(), true));
+        let target_fields = Fields::from(vec![
+            Field::new("value", DataType::Binary, false),
+            Field::new("metadata", DataType::Binary, false),
+        ]);
+        let target_field = Arc::new(
+            Field::new("v", DataType::Struct(target_fields), true).with_extension_type(VariantType),
+        );
+        let schema = Schema::new(vec![Arc::clone(&input_field)]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(physical)]).unwrap();
+        let expr = CometCastColumnExpr::new(
+            Arc::new(Column::new("v", 0)),
+            input_field,
+            target_field,
+            None,
+        );
+
+        let ColumnarValue::Array(output) = expr.evaluate(&batch).unwrap() else {
+            panic!("expected array")
+        };
+        let output = output.as_struct();
+        assert_eq!(
+            output
+                .fields()
+                .iter()
+                .map(|field| field.name().as_str())
+                .collect::<Vec<_>>(),
+            vec!["value", "metadata"]
+        );
+        assert!(output
+            .columns()
+            .iter()
+            .all(|column| column.data_type() == &DataType::Binary));
+        assert!(output.is_null(1));
+
+        let variant = VariantArray::try_new(output).unwrap();
+        assert_eq!(variant.value(0), Variant::from(10_i64));
+        assert_eq!(variant.value(2), Variant::from(30_i64));
+    }
 
     #[test]
     fn test_cast_timestamp_micros_to_millis_array() {

native/core/src/parquet/mod.rs

@@ -308,8 +308,10 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_currentColumnBatch(
             .ok_or_else(|| CometError::Execution {
                 source: ExecutionError::GeneralError("There is no more data to read".to_string()),
             });
-        let data = batch_reader?.column(column_idx as usize).into_data();
-        data.move_to_spark(array_addr, schema_addr)
+        let batch = batch_reader?;
+        let field = batch.schema().field(column_idx as usize).clone();
+        let data = batch.column(column_idx as usize).into_data();
+        data.move_to_spark(&field, array_addr, schema_addr)
             .map_err(|e| e.into())
     })
 }