Skip to content

Preserve nested Variant extension fields in variant_get #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
osipovartem merged 1 commit into from
Sep 24, 2026
Merged

Preserve nested Variant extension fields in variant_get #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 112 additions & 1 deletion parquet-variant-compute/src/variant_get.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down Expand Up @@ -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::<MapArray>().unwrap();
assert!(map.is_null(1));
let lists = map.values().as_any().downcast_ref::<ListArray>().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::<MapArray>().unwrap();
let objects = map.values().as_any().downcast_ref::<StructArray>().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::<MapArray>().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::<Vec<_>>();
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![
Expand Down
105 changes: 95 additions & 10 deletions parquet-variant-compute/src/variant_to_arrow.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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>),
Expand All @@ -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(),
}
}
Expand All @@ -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),
}
}
Expand All @@ -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::<Result<Vec<_>>>()?;
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<VariantToArrowRowBuilder<'a>> {
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,
Expand Down Expand Up @@ -599,7 +673,8 @@ impl<'a> StructVariantToArrowRowBuilder<'a> {
) -> Result<Self> {
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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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))
};

Expand Down Expand Up @@ -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 {
Expand Down
Loading