Preserve VariantType for missing variant fields

Import VariantArrayBuilder and change shredded path handling to gracefully handle non-Struct fields by falling back to the value column instead of erroring.

Author: codephage2020

parquet-variant-compute/src/variant_get.rs

@@ -24,6 +24,7 @@ use arrow_schema::{ArrowError, DataType, FieldRef};
 use parquet_variant::{VariantPath, VariantPathElement};
 
 use crate::VariantArray;
+use crate::VariantArrayBuilder;
 use crate::VariantType;
 use crate::variant_array::BorrowedShreddingState;
 use crate::variant_to_arrow::make_variant_to_arrow_row_builder;
@@ -90,17 +91,10 @@ pub(crate) fn follow_shredded_path_element<'a>(
             // The field might be a VariantArray (StructArray) if shredded,
             // or it might be a primitive array. Only proceed if it's a StructArray.
             let Some(struct_array) = field.as_struct_opt() else {
-                // Field exists but is not a StructArray (VariantArray),
-                // which means it's not shredded further.
-                if !cast_options.safe {
-                    return Err(ArrowError::CastError(format!(
-                        "Expected Struct array while following path, got {}",
-                        field.data_type(),
-                    )));
-                }
-                // In safe mode, return NotShredded to let the caller
-                // handle it via value column.
-                return Ok(ShreddedPathStep::NotShredded);
+                // Field exists but is not a StructArray, so it cannot be
+                // followed further. Fall back to the value column if present,
+                // otherwise the path is missing.
+                return Ok(missing_path_step());
             };
 
             let state = BorrowedShreddingState::try_from(struct_array)?;
@@ -224,47 +218,48 @@ fn shredded_get_path(
     // further, and build up the final struct from those individually shredded results.
     if let DataType::Struct(fields) = as_field.data_type() {
         let mut updated_fields = Vec::with_capacity(fields.len());
-        let children: Result<Vec<_>> = fields
-            .iter()
-            .map(|field| {
-                // If the field has VariantType extension metadata, extract it as a
-                // VariantArray instead of casting to the declared data type. This allows
-                // callers to request structs where some fields remain as variants.
-                // See test_struct_extraction_with_variant_fields for usage example.
-                let is_strongly_typed = field.try_extension_type::<VariantType>().is_err();
-                let field_as_type = is_strongly_typed.then(|| field.as_ref());
-                let child = shredded_get_path(
-                    &target,
-                    &[VariantPathElement::from(field.name().as_str())],
-                    field_as_type,
-                    cast_options,
-                )?;
-
-                // Update field type if it was a Variant marker (extracted as VariantArray).
-                // The actual data type will be the internal structure of VariantArray.
-                // Preserve VariantType extension metadata so downstream consumers
-                // can recognize this field as a Variant column.
-                //
-                // When the field is entirely absent in the data, shredded_get_path
-                // returns a NullArray (DataType::Null). VariantType only supports
-                // Struct storage, so we must skip the extension in that case.
-                let is_variant_field = !is_strongly_typed;
-                let new_field = field
-                    .as_ref()
-                    .clone()
-                    .with_data_type(child.data_type().clone());
-                let updated_field =
-                    if is_variant_field && matches!(child.data_type(), DataType::Struct(_)) {
-                        new_field.with_extension_type(VariantType)
-                    } else {
-                        new_field
-                    };
-                updated_fields.push(updated_field);
+        let mut children = Vec::with_capacity(fields.len());
+        for field in fields.iter() {
+            // If the field has VariantType extension metadata, extract it as a
+            // VariantArray instead of casting to the declared data type. This allows
+            // callers to request structs where some fields remain as variants.
+            // See test_struct_extraction_with_variant_fields for usage example.
+            let is_variant_field = field.try_extension_type::<VariantType>().is_ok();
+            let field_as_type = (!is_variant_field).then(|| field.as_ref());
+            let child = shredded_get_path(
+                &target,
+                &[VariantPathElement::from(field.name().as_str())],
+                field_as_type,
+                cast_options,
+            )?;
+
+            // When the field is entirely absent in the data, shredded_get_path
+            // returns a NullArray. For variant fields, construct an all-null
+            // VariantArray so the extension metadata is preserved.
+            let child = if is_variant_field && child.data_type() == &DataType::Null {
+                let mut builder = VariantArrayBuilder::new(child.len());
+                for _ in 0..child.len() {
+                    builder.append_null();
+                }
+                let null_variant = builder.build();
+                Arc::new(null_variant.into_inner()) as ArrayRef
+            } else {
+                child
+            };
 
-                Ok(child)
-            })
-            .collect();
-        let children = children?;
+            // Update field data type to match the actual child array.
+            // Preserve VariantType extension metadata for variant fields so
+            // downstream consumers can recognize them as Variant columns.
+            let mut new_field = field
+                .as_ref()
+                .clone()
+                .with_data_type(child.data_type().clone());
+            if is_variant_field {
+                new_field = new_field.with_extension_type(VariantType);
+            }
+            updated_fields.push(new_field);
+            children.push(child);
+        }
 
         let struct_nulls = target.nulls().cloned();
 
@@ -4383,7 +4378,9 @@ mod test {
 
         // Row 0: data = {"score": 95}
         let data0 = data_variant_array.value(0);
-        assert!(matches!(data0, Variant::Object(_)));
+        let obj0 = data0.as_object().expect("row 0 data should be an object");
+        let score = obj0.get("score").expect("row 0 data should have 'score'");
+        assert_eq!(score.as_int16(), Some(95));
 
         // Row 1: data = null
         assert!(
@@ -4392,7 +4389,9 @@ mod test {
 
         // Row 2: data = {"level": 5}
         let data2 = data_variant_array.value(2);
-        assert!(matches!(data2, Variant::Object(_)));
+        let obj2 = data2.as_object().expect("row 2 data should be an object");
+        let level = obj2.get("level").expect("row 2 data should have 'level'");
+        assert_eq!(level.as_int8(), Some(5));
     }
 
     /// Test that requesting a variant field absent in all rows does not panic.
@@ -4429,5 +4428,25 @@ mod test {
         // The missing variant field should be all nulls
         let missing_col = struct_result.column(1);
         assert_eq!(missing_col.null_count(), missing_col.len());
+
+        // The missing variant field should preserve VariantType extension metadata
+        let missing_schema_field = struct_result
+            .fields()
+            .iter()
+            .find(|f| f.name() == "missing_field")
+            .unwrap();
+        assert!(
+            missing_schema_field
+                .try_extension_type::<VariantType>()
+                .is_ok(),
+            "missing variant field should preserve VariantType extension metadata"
+        );
+
+        // The missing variant field should be a valid VariantArray
+        let missing_variant = VariantArray::try_new(missing_col);
+        assert!(
+            missing_variant.is_ok(),
+            "missing variant field should be a valid VariantArray"
+        );
     }
 }