sdf-jkl commented on code in PR #9663:
URL: https://github.com/apache/arrow-rs/pull/9663#discussion_r3046383765
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -1050,6 +1058,92 @@ where
}
}
+pub(crate) struct VariantToFixedSizeListArrowRowBuilder<'a> {
+ field: FieldRef,
+ list_size: i32,
+ element_builder: Box<VariantToShreddedVariantRowBuilder<'a>>,
Review Comment:
Should reuse the new `ListElementBuilder`. Otherwise, `variant_to_arrow` on
`FixedSizeArray` always returns a shredded array. The same error for other List
builders was fixed here - #9631
```suggestion
element_builder: ListElementBuilder<'a>,
```
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -1050,6 +1058,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,
Review Comment:
Add a shredded param like the List builder above.
```suggestion
capacity: usize,
shredded: bool,
```
Can reuse the rest of the logic from `VariantToListArrowRowBuilder`
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -1050,6 +1058,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) => {
Review Comment:
Can reuse `variant_cast_with_options`
```suggestion
match variant_cast_with_options(value, self.cast_options, Variant::as_list) {
```
##########
parquet-variant-compute/src/variant_get.rs:
##########
@@ -4352,10 +4353,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() {
Review Comment:
This test is wrong. We're trying to follow `arrow-cast` semantics for casts.
This basically is a cast from `List` to a `FixedSizeList(_, 2)`. We should only
return `Err` if `safe` flag is false.
A snippet from `arrow-cast`:
https://github.com/apache/arrow-rs/blob/43d984e0b3584841991e7169cf62a35b7025b691/arrow-cast/src/cast/mod.rs#L629-L630
##########
parquet-variant-compute/src/variant_get.rs:
##########
@@ -4352,10 +4353,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() {
Review Comment:
Once `variant_to_arrow` returns non-shredded arrays the whole test could be
moved to `test_variant_get_list_like_safe_cast` test above.
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
