Add support for FixedSizeList to variant_to_arrow

rishvin · rishvin · commit 808b1b543afa · 2026-04-05T19:16:19.000-07:00
diff --git a/parquet-variant-compute/src/shred_variant.rs b/parquet-variant-compute/src/shred_variant.rs
@@ -689,9 +689,9 @@ mod tests {
     use super::*;
     use crate::VariantArrayBuilder;
     use arrow::array::{
-        Array, BinaryViewArray, FixedSizeBinaryArray, Float64Array, GenericListArray,
-        GenericListViewArray, Int64Array, LargeBinaryArray, LargeStringArray, ListArray,
-        ListLikeArray, OffsetSizeTrait, PrimitiveArray, StringArray,
+        Array, BinaryViewArray, FixedSizeBinaryArray, FixedSizeListArray, Float64Array,
+        GenericListArray, GenericListViewArray, Int64Array, LargeBinaryArray, LargeStringArray,
+        ListArray, ListLikeArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray,
     };
     use arrow::datatypes::{
         ArrowPrimitiveType, DataType, Field, Fields, Int64Type, TimeUnit, UnionFields, UnionMode,
@@ -1607,6 +1607,67 @@ mod tests {
 
     #[test]
     fn test_array_shredding_as_fixed_size_list() {
+        let input = build_variant_array(vec![
+            VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
+            VariantRow::Value(VariantValue::from("This should not be shredded")),
+            VariantRow::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]),
+        ]);
+
+        let list_schema =
+            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
+        let result = shred_variant(&input, &list_schema).unwrap();
+        assert_eq!(result.len(), 3);
+
+        // The first row should be shredded, so the `value` field should be null and the
+        // `typed_value` field should contain the list
+        assert!(result.is_valid(0));
+        assert!(result.value_field().unwrap().is_null(0));
+        assert!(result.typed_value_field().unwrap().is_valid(0));
+
+        // The second row should not be shredded because the provided schema for shredding did not
+        // match. Hence, the `value` field should contain the raw value and the `typed_value` field
+        // should be null.
+        assert!(result.is_valid(1));
+        assert!(result.value_field().unwrap().is_valid(1));
+        assert!(result.typed_value_field().unwrap().is_null(1));
+
+        // The third row should be shredded, so the `value` field should be null and the
+        // `typed_value` field should contain the list
+        assert!(result.is_valid(2));
+        assert!(result.value_field().unwrap().is_null(2));
+        assert!(result.typed_value_field().unwrap().is_valid(2));
+
+        let typed_value = result.typed_value_field().unwrap();
+        let fixed_size_list = typed_value
+            .as_any()
+            .downcast_ref::<FixedSizeListArray>()
+            .expect("Expected FixedSizeListArray");
+
+        // Verify that typed value is `FixedSizeList`.
+        assert_eq!(fixed_size_list.len(), 3);
+        assert_eq!(fixed_size_list.value_length(), 2);
+
+        // Verify that the first entry in the `FixedSizeList` contains the expected value.
+        let val0 = fixed_size_list.value(0);
+        let val0_struct = val0.as_any().downcast_ref::<StructArray>().unwrap();
+        let val0_typed = val0_struct.column_by_name("typed_value").unwrap();
+        let val0_ints = val0_typed.as_any().downcast_ref::<Int64Array>().unwrap();
+        assert_eq!(val0_ints.values(), &[1i64, 2i64]);
+
+        // Verify that second entry in the `FixedSizeList` cannot be shredded hence the value is
+        // invalid.
+        assert!(fixed_size_list.is_null(1));
+
+        // Verify that the third entry in the `FixedSizeList` contains the expected value.
+        let val2 = fixed_size_list.value(2);
+        let val2_struct = val2.as_any().downcast_ref::<StructArray>().unwrap();
+        let val2_typed = val2_struct.column_by_name("typed_value").unwrap();
+        let val2_ints = val2_typed.as_any().downcast_ref::<Int64Array>().unwrap();
+        assert_eq!(val2_ints.values(), &[3i64, 4i64]);
+    }
+
+    #[test]
+    fn test_array_shredding_as_fixed_size_list_wrong_size() {
         let input = build_variant_array(vec![VariantRow::List(vec![
             VariantValue::from(1i64),
             VariantValue::from(2i64),
@@ -1615,9 +1676,12 @@ mod tests {
         let list_schema =
             DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
         let err = shred_variant(&input, &list_schema).unwrap_err();
-        assert_eq!(
-            err.to_string(),
-            "Not yet implemented: Converting unshredded variant arrays to arrow fixed-size lists"
+        println!("{}", err);
+        assert!(
+            err.to_string()
+                .contains("Expected fixed size list of size 2, got size 3"),
+            "got: {}",
+            err
         );
     }
 
diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs
@@ -351,9 +351,9 @@ mod test {
     use arrow::array::{
         Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Date32Array,
         Date64Array, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array,
-        Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
-        LargeBinaryArray, LargeListArray, LargeListViewArray, LargeStringArray, ListArray,
-        ListViewArray, NullBuilder, StringArray, StringViewArray, StructArray,
+        FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
+        Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray, LargeStringArray,
+        ListArray, ListViewArray, NullBuilder, StringArray, StringViewArray, StructArray,
         Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
     };
     use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
@@ -4336,7 +4336,8 @@ mod test {
             DataType::List(item_field.clone()),
             DataType::LargeList(item_field.clone()),
             DataType::ListView(item_field.clone()),
-            DataType::LargeListView(item_field),
+            DataType::LargeListView(item_field.clone()),
+            DataType::FixedSizeList(item_field, 2),
         ];
 
         for data_type in data_types {
@@ -4353,10 +4354,65 @@ mod test {
     }
 
     #[test]
-    fn test_variant_get_fixed_size_list_not_implemented() {
-        let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2]", "\"not a list\""]));
+    fn test_variant_get_fixed_size_list_with_safe_option() {
+        let string_array: ArrayRef = Arc::new(StringArray::from(vec![
+            "[1, 2]",
+            "[3, 4]",
+            "\"not a list\"",
+        ]));
         let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
         let item_field = Arc::new(Field::new("item", Int64, true));
+
+        // Request shredding on `FixedSizeList` with `safe` set to true, such that `variant_get`
+        // does not raise error on type mismatch.
+        let options = GetOptions::new()
+            .with_as_type(Some(FieldRef::from(Field::new(
+                "result",
+                DataType::FixedSizeList(item_field.clone(), 2),
+                true,
+            ))))
+            .with_cast_options(CastOptions {
+                safe: true,
+                ..Default::default()
+            });
+
+        // Verify that the shredded value is a `FixedSizeList`.
+        let result = variant_get(&variant_array, options).unwrap();
+        let fixed_size_list = result
+            .as_any()
+            .downcast_ref::<FixedSizeListArray>()
+            .expect("Expected FixedSizeListArray");
+        assert_eq!(fixed_size_list.len(), 3);
+        assert_eq!(fixed_size_list.value_length(), 2);
+
+        // Verify that the first entry in the `FixedSizeList` contains the expected value.
+        assert!(fixed_size_list.is_valid(0));
+        let val0 = fixed_size_list.value(0);
+        let val0_struct = val0.as_any().downcast_ref::<StructArray>().unwrap();
+        let val0_typed = val0_struct.column_by_name("typed_value").unwrap();
+        let val0_ints = val0_typed.as_any().downcast_ref::<Int64Array>().unwrap();
+        assert_eq!(val0_ints.values(), &[1i64, 2i64]);
+
+        // Verify that the second entry in the `FixedSizeList` contains the expected value.
+        assert!(fixed_size_list.is_valid(1));
+        let val1 = fixed_size_list.value(1);
+        let val1_struct = val1.as_any().downcast_ref::<StructArray>().unwrap();
+        let val1_typed = val1_struct.column_by_name("typed_value").unwrap();
+        let val1_ints = val1_typed.as_any().downcast_ref::<Int64Array>().unwrap();
+        assert_eq!(val1_ints.values(), &[3i64, 4i64]);
+
+        // Verify that the third entry is null due to type mismatch.
+        assert!(fixed_size_list.is_null(2));
+    }
+
+    #[test]
+    fn test_variant_get_fixed_size_list_wrong_size() {
+        let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2, 3]"]));
+        let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
+        let item_field = Arc::new(Field::new("item", Int64, true));
+
+        // Set the safe flag to both true and false and verify that size mismatch raises an error
+        // for `FixedSizeList`, regardless.
         for safe in [true, false] {
             let options = GetOptions::new()
                 .with_as_type(Some(FieldRef::from(Field::new(
@@ -4368,11 +4424,11 @@ mod test {
                     safe,
                     ..Default::default()
                 });
-
             let err = variant_get(&variant_array, options).unwrap_err();
             assert!(
                 err.to_string()
-                    .contains("Converting unshredded variant arrays to arrow fixed-size lists")
+                    .contains("Expected fixed size list of size 2, got size 3"),
+                "safe={safe}, got: {err}",
             );
         }
     }
diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs
@@ -27,10 +27,10 @@ use crate::variant_array::ShreddedVariantFieldArray;
 use crate::{VariantArray, VariantValueArrayBuilder};
 use arrow::array::{
     ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder, BinaryViewArray,
-    BinaryViewBuilder, BooleanBuilder, FixedSizeBinaryBuilder, GenericListArray,
-    GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, NullArray, NullBufferBuilder,
-    OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder, StringViewBuilder,
-    StructArray,
+    BinaryViewBuilder, BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray,
+    GenericListArray, GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, NullArray,
+    NullBufferBuilder, OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder,
+    StringViewBuilder, StructArray,
 };
 use arrow::buffer::{OffsetBuffer, ScalarBuffer};
 use arrow::compute::{CastOptions, DecimalCast};
@@ -507,6 +507,7 @@ pub(crate) enum ArrayVariantToArrowRowBuilder<'a> {
     LargeList(VariantToListArrowRowBuilder<'a, i64, false>),
     ListView(VariantToListArrowRowBuilder<'a, i32, true>),
     LargeListView(VariantToListArrowRowBuilder<'a, i64, true>),
+    FixedSizeList(VariantToFixedSizeListArrowRowBuilder<'a>),
 }
 
 pub(crate) struct StructVariantToArrowRowBuilder<'a> {
@@ -611,10 +612,14 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> {
             DataType::LargeList(field) => make_list_builder!(LargeList, i64, false, field),
             DataType::ListView(field) => make_list_builder!(ListView, i32, true, field),
             DataType::LargeListView(field) => make_list_builder!(LargeListView, i64, true, field),
-            DataType::FixedSizeList(..) => {
-                return Err(ArrowError::NotYetImplemented(
-                    "Converting unshredded variant arrays to arrow fixed-size lists".to_string(),
-                ));
+            DataType::FixedSizeList(field, size) => {
+                FixedSizeList(VariantToFixedSizeListArrowRowBuilder::try_new(
+                    field.clone(),
+                    field.data_type(),
+                    *size,
+                    cast_options,
+                    capacity,
+                )?)
             }
             other => {
                 return Err(ArrowError::InvalidArgumentError(format!(
@@ -631,6 +636,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> {
             Self::LargeList(builder) => builder.append_null(),
             Self::ListView(builder) => builder.append_null(),
             Self::LargeListView(builder) => builder.append_null(),
+            Self::FixedSizeList(builder) => builder.append_null(),
         }
     }
 
@@ -640,6 +646,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> {
             Self::LargeList(builder) => builder.append_value(value),
             Self::ListView(builder) => builder.append_value(value),
             Self::LargeListView(builder) => builder.append_value(value),
+            Self::FixedSizeList(builder) => builder.append_value(value),
         }
     }
 
@@ -649,6 +656,7 @@ impl<'a> ArrayVariantToArrowRowBuilder<'a> {
             Self::LargeList(builder) => builder.finish(),
             Self::ListView(builder) => builder.finish(),
             Self::LargeListView(builder) => builder.finish(),
+            Self::FixedSizeList(builder) => builder.finish(),
         }
     }
 }
@@ -1003,6 +1011,92 @@ where
     }
 }
 
+pub(crate) struct VariantToFixedSizeListArrowRowBuilder<'a> {
+    field: FieldRef,
+    list_size: i32,
+    element_builder: Box<VariantToShreddedVariantRowBuilder<'a>>,
+    nulls: NullBufferBuilder,
+    cast_options: &'a CastOptions<'a>,
+}
+
+impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> {
+    fn try_new(
+        field: FieldRef,
+        element_data_type: &'a DataType,
+        list_size: i32,
+        cast_options: &'a CastOptions,
+        capacity: usize,
+    ) -> Result<Self> {
+        let element_builder = make_variant_to_shredded_variant_arrow_row_builder(
+            element_data_type,
+            cast_options,
+            capacity,
+            NullValue::ArrayElement,
+        )?;
+        Ok(Self {
+            field,
+            list_size,
+            element_builder: Box::new(element_builder),
+            nulls: NullBufferBuilder::new(capacity),
+            cast_options,
+        })
+    }
+
+    fn append_null(&mut self) -> Result<()> {
+        for _ in 0..self.list_size {
+            self.element_builder.append_null()?;
+        }
+        self.nulls.append_null();
+        Ok(())
+    }
+
+    fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
+        match value {
+            Variant::List(list) => {
+                let len = list.len();
+                if len != self.list_size as usize {
+                    return Err(ArrowError::CastError(format!(
+                        "Expected fixed size list of size {}, got size {}",
+                        self.list_size, len
+                    )));
+                }
+                for element in list.iter() {
+                    self.element_builder.append_value(element)?;
+                }
+                self.nulls.append_non_null();
+                Ok(true)
+            }
+            _ if self.cast_options.safe => {
+                self.append_null()?;
+                Ok(false)
+            }
+            _ => Err(ArrowError::CastError(format!(
+                "Failed to extract fixed size list from variant {:?}",
+                value
+            ))),
+        }
+    }
+
+    fn finish(mut self) -> Result<ArrayRef> {
+        let (value, typed_value, nulls) = self.element_builder.finish()?;
+        let element_array =
+            ShreddedVariantFieldArray::from_parts(Some(value), Some(typed_value), nulls);
+        let field = Arc::new(
+            self.field
+                .as_ref()
+                .clone()
+                .with_data_type(element_array.data_type().clone()),
+        );
+        let fixed_size_list_array = FixedSizeListArray::try_new(
+            field,
+            self.list_size,
+            ArrayRef::from(element_array),
+            self.nulls.finish(),
+        )?;
+        Ok(Arc::new(fixed_size_list_array))
+    }
+}
+
 /// Builder for creating VariantArray output (for path extraction without type conversion)
 pub(crate) struct VariantToBinaryVariantArrowRowBuilder {
     metadata: BinaryViewArray,
