feat: adapt Parquet storage for Variant projection

native/core/Cargo.toml

@@ -38,7 +38,7 @@ publish = false
 arrow = { workspace = true }
 base64 = "0.23.0"
 bytes = { workspace = true }
-parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
+parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "arrow_canonical_extension_types", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
 futures = { workspace = true }
 mimalloc = { version = "*", default-features = false, optional = true }
 tikv-jemallocator = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls"] }

native/core/src/parquet/cast_column/variant.rs

@@ -21,8 +21,8 @@ use arrow::{
         StructArray,
     },
     buffer::NullBuffer,
-    compute::cast,
-    datatypes::{DataType, FieldRef},
+    compute::{cast, cast_with_options},
+    datatypes::{DataType, FieldRef, TimeUnit, DECIMAL128_MAX_PRECISION},
     error::ArrowError,
 };
 use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -59,6 +59,7 @@ pub(super) fn normalize_variant_array(
     // VariantArray resolves metadata/value/typed_value by name, so the reader's child order is
     // irrelevant. Legacy Spark residuals must be put in Arrow order before the single upstream
     // unshred call; the whole output is then put back in the order expected by released Spark 4.
+    let array = normalize_variant_storage(array)?;
     let variant = VariantArray::try_new(array.as_ref())?;
     let prepared = prepare_variant_for_unshredding(&variant)?;
     let unshredded = unshred_variant(&prepared)?;
@@ -74,6 +75,102 @@ pub(super) fn normalize_variant_array(
     )?))
 }
 
+/// Arrow Variant compute rejects some storage types that Spark's Parquet reader accepts.
+/// Choose supported types recursively for encoded, unsigned, decimal, timestamp, and fixed
+/// binary/list children before reconstructing the whole value.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn normalize_variant_type(data_type: &DataType) -> Option<DataType> {
+    fn normalize_field(field: &FieldRef) -> Option<FieldRef> {
+        normalize_variant_type(field.data_type())
+            .map(|data_type| Arc::new(field.as_ref().clone().with_data_type(data_type)))
+    }
+
+    match data_type {
+        DataType::Dictionary(_, value_type) => {
+            Some(normalize_variant_type(value_type).unwrap_or_else(|| value_type.as_ref().clone()))
+        }
+        DataType::UInt8 => Some(DataType::Int16),
+        DataType::UInt16 => Some(DataType::Int32),
+        DataType::UInt32 => Some(DataType::Int64),
+        // Spark reads Parquet UINT_64 as Decimal(20, 0). This is lossless for the full range and
+        // preserves values larger than i64::MAX for Variant decimal encoding.
+        DataType::UInt64 => Some(DataType::Decimal128(20, 0)),
+        // Arrow chooses Decimal256 from the physical byte width, but Spark's DecimalType is
+        // precision-based and stores every supported precision (<= 38) in 128 bits.
+        DataType::Decimal256(precision, scale) if *precision <= DECIMAL128_MAX_PRECISION => {
+            Some(DataType::Decimal128(*precision, *scale))
+        }
+        DataType::Timestamp(TimeUnit::Millisecond, timezone) => {
+            Some(DataType::Timestamp(TimeUnit::Microsecond, timezone.clone()))
+        }
+        DataType::FixedSizeBinary(_) => Some(DataType::Binary),
+        DataType::FixedSizeList(field, _) => Some(DataType::List(
+            normalize_field(field).unwrap_or_else(|| Arc::clone(field)),
+        )),
+        DataType::List(field) => normalize_field(field).map(DataType::List),
+        DataType::LargeList(field) => normalize_field(field).map(DataType::LargeList),
+        DataType::ListView(field) => normalize_field(field).map(DataType::ListView),
+        DataType::LargeListView(field) => normalize_field(field).map(DataType::LargeListView),
+        DataType::Struct(fields) => {
+            let mut changed = false;
+            let fields = fields
+                .iter()
+                .map(|field| match normalize_field(field) {
+                    Some(field) => {
+                        changed = true;
+                        field
+                    }
+                    None => Arc::clone(field),
+                })
+                .collect::<Vec<_>>();
+            changed.then(|| DataType::Struct(fields.into()))
+        }
+        _ => None,
+    }
+}
+
+fn contains_uuid_extension(data_type: &DataType) -> bool {
+    fn field_contains_uuid(field: &FieldRef) -> bool {
+        (field.data_type() == &DataType::FixedSizeBinary(16)
+            && field.extension_type_name() == Some("arrow.uuid"))
+            || contains_uuid_extension(field.data_type())
+    }
+
+    match data_type {
+        DataType::Struct(fields) => fields.iter().any(field_contains_uuid),
+        DataType::List(field)
+        | DataType::LargeList(field)
+        | DataType::ListView(field)
+        | DataType::LargeListView(field)
+        | DataType::FixedSizeList(field, _)
+        | DataType::Map(field, _) => field_contains_uuid(field),
+        DataType::Dictionary(_, value_type) => contains_uuid_extension(value_type),
+        _ => false,
+    }
+}
+
+/// Arrow Variant compute cannot consume every storage type Spark reads. Decode and cast those
+/// children before validation, preserving UUID rejection and reporting conversion overflow.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn normalize_variant_storage(array: &ArrayRef) -> DataFusionResult<ArrayRef> {
+    if contains_uuid_extension(array.data_type()) {
+        return Err(DataFusionError::Execution(
+            "Parquet UUID is not supported as a shredded Variant child".to_string(),
+        ));
+    }
+    let Some(data_type) = normalize_variant_type(array.data_type()) else {
+        return Ok(Arc::clone(array));
+    };
+    Ok(cast_with_options(
+        array.as_ref(),
+        &data_type,
+        &arrow::compute::CastOptions {
+            safe: false,
+            ..Default::default()
+        },
+    )?)
+}
+
 /// Arrow validates every residual `value` while unshredding. Spark versions before SPARK-58949
 /// wrote object keys in Java UTF-16 order, so rewrite every reachable legacy residual to Arrow's
 /// UTF-8 order before unshredding. `metadata_rows` carries each root metadata row through nested

native/core/src/parquet/cast_column/variant/tests.rs

@@ -39,6 +39,174 @@ fn target_field(nullable: bool) -> FieldRef {
     )
 }
 
+#[test]
+fn normalize_encoded_storage_and_unsigned_extremes() {
+    use arrow::array::{
+        DictionaryArray, Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+    };
+    use arrow::datatypes::Int8Type;
+    let mut builder = VariantBuilder::new();
+    builder.append_value(Variant::Null);
+    let (metadata, _) = builder.finish();
+    let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![metadata.as_slice()]));
+    let encoded_metadata: ArrayRef =
+        Arc::new(DictionaryArray::<Int8Type>::try_new(Int8Array::from(vec![0]), metadata).unwrap());
+    let cases: Vec<(ArrayRef, i128)> = vec![
+        (
+            Arc::new(UInt8Array::from(vec![u8::MAX])),
+            i128::from(u8::MAX),
+        ),
+        (
+            Arc::new(UInt16Array::from(vec![u16::MAX])),
+            i128::from(u16::MAX),
+        ),
+        (
+            Arc::new(UInt32Array::from(vec![u32::MAX])),
+            i128::from(u32::MAX),
+        ),
+        (
+            Arc::new(UInt64Array::from(vec![u64::MAX])),
+            i128::from(u64::MAX),
+        ),
+    ];
+    for (typed, expected) in cases {
+        let typed: ArrayRef = Arc::new(
+            DictionaryArray::<Int8Type>::try_new(Int8Array::from(vec![0]), typed).unwrap(),
+        );
+        let input: ArrayRef = Arc::new(StructArray::new(
+            vec![
+                Field::new("typed_value", typed.data_type().clone(), false),
+                Field::new("metadata", encoded_metadata.data_type().clone(), false),
+            ]
+            .into(),
+            vec![typed, Arc::clone(&encoded_metadata)],
+            None,
+        ));
+        let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+        let output = VariantArray::try_new(output.as_ref()).unwrap();
+        let value = output.value(0);
+        if let Variant::Decimal16(decimal) = value {
+            assert_eq!(decimal.integer(), expected);
+            assert_eq!(decimal.scale(), 0);
+        } else {
+            assert_eq!(i128::from(value.as_int64().unwrap()), expected);
+        }
+    }
+
+    let mut builder = VariantArrayBuilder::new(1);
+    builder.append_variant(Variant::from(42_i64));
+    let base = builder.build();
+    let encoded_value: ArrayRef = Arc::new(
+        DictionaryArray::<Int8Type>::try_new(
+            Int8Array::from(vec![0]),
+            Arc::clone(base.value_column()),
+        )
+        .unwrap(),
+    );
+    let input: ArrayRef = Arc::new(StructArray::new(
+        vec![
+            Field::new("value", encoded_value.data_type().clone(), false),
+            Field::new("metadata", encoded_metadata.data_type().clone(), false),
+        ]
+        .into(),
+        vec![encoded_value, encoded_metadata],
+        None,
+    ));
+    let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+    assert_eq!(
+        VariantArray::try_new(output.as_ref())
+            .unwrap()
+            .value(0)
+            .as_int64(),
+        Some(42)
+    );
+}
+
+#[test]
+fn normalize_fixed_storage_and_checked_timestamps() {
+    use arrow::array::{FixedSizeBinaryArray, TimestampMillisecondArray};
+    let mut builder = VariantBuilder::new();
+    builder.append_value(Variant::Null);
+    let (metadata, _) = builder.finish();
+    let normalize = |typed: ArrayRef| {
+        let input: ArrayRef = Arc::new(StructArray::new(
+            vec![
+                Field::new("metadata", DataType::Binary, false),
+                Field::new("typed_value", typed.data_type().clone(), false),
+            ]
+            .into(),
+            vec![
+                Arc::new(BinaryArray::from(vec![metadata.as_slice()])),
+                typed,
+            ],
+            None,
+        ));
+        normalize_variant_array(&input, &target_field(false))
+    };
+    for width in [3, 16] {
+        let bytes = vec![5; width];
+        let typed = FixedSizeBinaryArray::try_from_iter([bytes.as_slice()].into_iter()).unwrap();
+        let output = normalize(Arc::new(typed)).unwrap();
+        assert_eq!(
+            VariantArray::try_new(output.as_ref()).unwrap().value(0),
+            Variant::Binary(&bytes)
+        );
+    }
+    assert!(normalize(Arc::new(TimestampMillisecondArray::from(vec![i64::MAX]))).is_err());
+    assert!(normalize(Arc::new(TimestampMillisecondArray::from(vec![123]))).is_ok());
+}
+
+#[test]
+fn normalize_fixed_size_list_and_reject_uuid() {
+    use arrow::array::{FixedSizeBinaryArray, FixedSizeListArray, UInt16Array};
+    let mut builder = VariantBuilder::new();
+    builder.append_value(Variant::Null);
+    let (metadata, _) = builder.finish();
+    let wrap = |field: Field, typed: ArrayRef| -> ArrayRef {
+        Arc::new(StructArray::new(
+            vec![Field::new("metadata", DataType::Binary, false), field].into(),
+            vec![
+                Arc::new(BinaryArray::from(vec![metadata.as_slice()])),
+                typed,
+            ],
+            None,
+        ))
+    };
+    let elements: ArrayRef = Arc::new(StructArray::new(
+        vec![Field::new("typed_value", DataType::UInt16, false)].into(),
+        vec![Arc::new(UInt16Array::from(vec![1, u16::MAX]))],
+        None,
+    ));
+    let list: ArrayRef = Arc::new(FixedSizeListArray::new(
+        Arc::new(Field::new("item", elements.data_type().clone(), false)),
+        2,
+        elements,
+        None,
+    ));
+    let input = wrap(
+        Field::new("typed_value", list.data_type().clone(), false),
+        list,
+    );
+    let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+    let output = VariantArray::try_new(output.as_ref()).unwrap();
+    let Variant::List(list) = output.value(0) else {
+        panic!("expected list")
+    };
+    assert_eq!(
+        list.iter()
+            .map(|value| value.as_int64().unwrap())
+            .collect::<Vec<_>>(),
+        vec![1, 65535]
+    );
+
+    let uuid: ArrayRef =
+        Arc::new(FixedSizeBinaryArray::try_from_iter([[0_u8; 16]].into_iter()).unwrap());
+    let field = Field::new("typed_value", uuid.data_type().clone(), false)
+        .with_metadata([("ARROW:extension:name".to_string(), "arrow.uuid".to_string())].into());
+    let error = normalize_variant_array(&wrap(field, uuid), &target_field(false)).unwrap_err();
+    assert!(error.to_string().contains("Parquet UUID"));
+}
+
 fn unicode_object_keys() -> Vec<String> {
     let mut keys = (0..30).map(|i| format!("k{i:02}")).collect::<Vec<_>>();
     keys.push("\u{e000}".to_string());