diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index bd28f5c8b7b1..0e34a703bc8b 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -499,7 +499,7 @@ mod test { Date64Array, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, Int64Builder, LargeBinaryArray, LargeListArray, LargeListViewArray, - LargeStringArray, ListArray, ListBuilder, ListViewArray, MapBuilder, NullArray, + LargeStringArray, ListArray, ListBuilder, ListViewArray, MapArray, MapBuilder, NullArray, NullBuilder, StringArray, StringBuilder, StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, }; @@ -4502,6 +4502,117 @@ mod test { assert_eq!(result.as_ref(), &expected); } + #[test] + fn get_variant_as_map_of_variant_lists() { + let input: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a":[1,null],"b":[]}"#), + None, + ])); + let variants = ArrayRef::from(json_to_variant(&input).unwrap()); + let item = VariantArrayBuilder::new(0).build().field("item"); + let data_type = map_data_type(DataType::List(Arc::new(item))); + + let result = variant_get(&variants, map_get_options(&data_type)).unwrap(); + assert_eq!(result.data_type(), &data_type); + let map = result.as_any().downcast_ref::().unwrap(); + assert!(map.is_null(1)); + let lists = map.values().as_any().downcast_ref::().unwrap(); + assert_eq!(lists.value_length(0), 2); + assert_eq!(lists.value_length(1), 0); + let values = lists.value(0); + let nested = VariantArray::try_new(&values).unwrap(); + let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("1"), Some("null")])); + let expected = json_to_variant(&expected).unwrap(); + assert_eq!(nested.try_value(0).unwrap(), expected.try_value(0).unwrap()); + assert_eq!(nested.try_value(1).unwrap(), expected.try_value(1).unwrap()); + } + + #[test] + fn get_variant_as_map_of_variant_structs() { + let input: ArrayRef = Arc::new(StringArray::from(vec![Some( + r#"{"a":{"v":1},"b":{"v":null},"c":{}}"#, + )])); + let variants = ArrayRef::from(json_to_variant(&input).unwrap()); + let value = VariantArrayBuilder::new(0) + .build() + .field("v") + .with_nullable(true); + let data_type = map_data_type(DataType::Struct(Fields::from(vec![value]))); + + let result = variant_get(&variants, map_get_options(&data_type)).unwrap(); + assert_eq!(result.data_type(), &data_type); + let map = result.as_any().downcast_ref::().unwrap(); + let objects = map.values().as_any().downcast_ref::().unwrap(); + let nested = VariantArray::try_new(objects.column(0)).unwrap(); + let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("1"), Some("null")])); + let expected = json_to_variant(&expected).unwrap(); + assert_eq!(nested.try_value(0).unwrap(), expected.try_value(0).unwrap()); + assert_eq!(nested.try_value(1).unwrap(), expected.try_value(1).unwrap()); + assert!(nested.is_null(2)); + } + + #[test] + fn get_variant_as_map_preserves_nested_variant_storage() { + let input: ArrayRef = Arc::new(StringArray::from(vec![Some(r#"{"a":1}"#)])); + let variants = ArrayRef::from(json_to_variant(&input).unwrap()); + let value = VariantArrayBuilder::new(0) + .build() + .field("values") + .with_data_type(DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + ]))); + let data_type = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + value, + ])), + false, + )), + false, + ); + + let result = variant_get(&variants, map_get_options(&data_type)).unwrap(); + assert_eq!(result.data_type(), &data_type); + let map = result.as_any().downcast_ref::().unwrap(); + let nested = VariantArray::try_new(map.values()).unwrap(); + let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("1")])); + let expected = json_to_variant(&expected).unwrap(); + assert_eq!(nested.try_value(0).unwrap(), expected.try_value(0).unwrap()); + } + + #[test] + fn get_variant_as_map_rejects_nested_shredded_variant() { + let input: ArrayRef = Arc::new(StringArray::from(vec![Some(r#"{"a":1}"#)])); + let variants = ArrayRef::from(json_to_variant(&input).unwrap()); + let value = VariantArrayBuilder::new(0).build().field("values"); + let DataType::Struct(fields) = value.data_type() else { + unreachable!("Variant storage is a struct"); + }; + let mut fields = fields + .iter() + .map(|field| field.as_ref().clone()) + .collect::>(); + fields.push(Field::new("typed_value", DataType::Int64, true)); + let value = value.with_data_type(DataType::Struct(Fields::from(fields))); + let data_type = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + value, + ])), + false, + )), + false, + ); + + let error = variant_get(&variants, map_get_options(&data_type)).unwrap_err(); + assert!(error.to_string().contains("shredded nested Variant output")); + } + #[test] fn get_variant_as_map_non_object_rows() { let input: ArrayRef = Arc::new(StringArray::from(vec![ diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs index 66666ee138e8..cf00e08d7234 100644 --- a/parquet-variant-compute/src/variant_to_arrow.rs +++ b/parquet-variant-compute/src/variant_to_arrow.rs @@ -24,18 +24,19 @@ use crate::type_conversion::{ variant_to_unscaled_decimal, }; use crate::variant_array::ShreddedVariantFieldArray; -use crate::{VariantArray, VariantValueArrayBuilder}; +use crate::{VariantArray, VariantArrayBuilder, VariantType, VariantValueArrayBuilder}; use arrow::array::{ - ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder, BinaryViewBuilder, + Array, ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder, BinaryViewBuilder, BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray, GenericListArray, GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, MapArray, NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder, StringViewBuilder, StructArray, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow::compute::{CastOptions, DecimalCast, cast_with_options}; +use arrow::compute::{CastOptions, DecimalCast, cast, cast_with_options}; use arrow::datatypes::{self, DataType, DecimalType}; use arrow::error::{ArrowError, Result}; +use arrow_schema::extension::ExtensionType; use arrow_schema::{FieldRef, Fields, TimeUnit}; use parquet_variant::{Variant, VariantPath}; use std::sync::Arc; @@ -51,6 +52,7 @@ pub(crate) enum VariantToArrowRowBuilder<'a> { Map(MapVariantToArrowRowBuilder<'a>), Encoded(EncodedVariantToArrowRowBuilder<'a>), BinaryVariant(VariantToBinaryVariantArrowRowBuilder), + NativeVariant(VariantArrayBuilder, Fields), // Path extraction wrapper - contains a boxed enum for any of the above WithPath(VariantPathRowBuilder<'a>), @@ -66,6 +68,10 @@ impl<'a> VariantToArrowRowBuilder<'a> { Map(b) => b.append_null(), Encoded(b) => b.append_null(), BinaryVariant(b) => b.append_null(), + NativeVariant(b, _) => { + b.append_null(); + Ok(()) + } WithPath(path_builder) => path_builder.append_null(), } } @@ -79,6 +85,10 @@ impl<'a> VariantToArrowRowBuilder<'a> { Map(b) => b.append_value(&value), Encoded(b) => b.append_value(value), BinaryVariant(b) => b.append_value(value), + NativeVariant(b, _) => { + b.append_variant(value); + Ok(true) + } WithPath(path_builder) => path_builder.append_value(value), } } @@ -92,11 +102,75 @@ impl<'a> VariantToArrowRowBuilder<'a> { Map(b) => b.finish(), Encoded(b) => b.finish(), BinaryVariant(b) => b.finish(), + NativeVariant(b, fields) => { + let built = b.build(); + let inner = built.inner(); + if inner.data_type() == &DataType::Struct(fields.clone()) { + return Ok(ArrayRef::from(built)); + } + let columns = fields + .iter() + .map(|field| { + let column = inner.column_by_name(field.name()).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "Nested Variant is missing the '{}' field", + field.name() + )) + })?; + cast(column.as_ref(), field.data_type()) + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + fields, + columns, + inner.nulls().cloned(), + )?)) + } WithPath(path_builder) => path_builder.finish(), } } } +fn make_typed_variant_to_arrow_field_builder<'a>( + field: &FieldRef, + data_type: &'a DataType, + cast_options: &'a CastOptions, + capacity: usize, +) -> Result> { + if field.extension_type_name() == Some(VariantType::NAME) { + let DataType::Struct(fields) = field.data_type() else { + return Err(ArrowError::InvalidArgumentError( + "Nested Variant extension requires Struct storage".to_string(), + )); + }; + if fields.iter().any(|field| field.name() == "typed_value") { + return Err(ArrowError::NotYetImplemented( + "variant_get with shredded nested Variant output is not yet supported".to_string(), + )); + } + if fields.len() != 2 + || !["metadata", "value"].iter().all(|name| { + fields.iter().any(|field| { + field.name() == *name + && matches!( + field.data_type(), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView + ) + }) + }) + { + return Err(ArrowError::NotYetImplemented( + "Nested Variant output requires binary metadata and value fields".to_string(), + )); + } + return Ok(VariantToArrowRowBuilder::NativeVariant( + VariantArrayBuilder::new(capacity), + fields.clone(), + )); + } + make_typed_variant_to_arrow_row_builder(data_type, cast_options, capacity) +} + fn make_typed_variant_to_arrow_row_builder<'a>( data_type: &'a DataType, cast_options: &'a CastOptions, @@ -599,7 +673,8 @@ impl<'a> StructVariantToArrowRowBuilder<'a> { ) -> Result { let mut field_builders = Vec::with_capacity(fields.len()); for field in fields.iter() { - field_builders.push(make_typed_variant_to_arrow_row_builder( + field_builders.push(make_typed_variant_to_arrow_field_builder( + field, field.data_type(), cast_options, capacity, @@ -701,12 +776,14 @@ impl<'a> MapVariantToArrowRowBuilder<'a> { ))); } }; - let key_builder = Box::new(make_typed_variant_to_arrow_row_builder( + let key_builder = Box::new(make_typed_variant_to_arrow_field_builder( + key_field, key_field.data_type(), cast_options, capacity, )?); - let value_builder = Box::new(make_typed_variant_to_arrow_row_builder( + let value_builder = Box::new(make_typed_variant_to_arrow_field_builder( + value_field, value_field.data_type(), cast_options, capacity, @@ -1200,8 +1277,12 @@ where )?; ListElementBuilder::Shredded(Box::new(builder)) } else { - let builder = - make_typed_variant_to_arrow_row_builder(element_data_type, cast_options, capacity)?; + let builder = make_typed_variant_to_arrow_field_builder( + &field, + element_data_type, + cast_options, + capacity, + )?; ListElementBuilder::Typed(Box::new(builder)) }; @@ -1305,8 +1386,12 @@ impl<'a> VariantToFixedSizeListArrowRowBuilder<'a> { )?; ListElementBuilder::Shredded(Box::new(builder)) } else { - let builder = - make_typed_variant_to_arrow_row_builder(element_data_type, cast_options, capacity)?; + let builder = make_typed_variant_to_arrow_field_builder( + &field, + element_data_type, + cast_options, + capacity, + )?; ListElementBuilder::Typed(Box::new(builder)) }; Ok(Self {