This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new e1de892def [Variant] Add `variant_to_arrow` `Map` type support (#10307)
e1de892def is described below
commit e1de892deff45a4aa496b04278f4188d17233314
Author: Konstantin Tarasov <[email protected]>
AuthorDate: Fri Jul 10 10:43:51 2026 -0400
[Variant] Add `variant_to_arrow` `Map` type support (#10307)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Closes #10012.
# Rationale for this change
Check issue.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
- Added `MapVariantToArrowRowBuilder`
- Added unit tests
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
- yes, added unit tests
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
- Users now can cast variant to map
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
parquet-variant-compute/src/variant_get.rs | 166 +++++++++++++++++++++++-
parquet-variant-compute/src/variant_to_arrow.rs | 157 +++++++++++++++++++++-
2 files changed, 316 insertions(+), 7 deletions(-)
diff --git a/parquet-variant-compute/src/variant_get.rs
b/parquet-variant-compute/src/variant_get.rs
index fc01ca8d3b..59407c617c 100644
--- a/parquet-variant-compute/src/variant_get.rs
+++ b/parquet-variant-compute/src/variant_get.rs
@@ -493,10 +493,10 @@ mod test {
Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray,
Date32Array,
Date64Array, Decimal32Array, Decimal64Array, Decimal128Array,
Decimal256Array,
FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array,
Int32Array,
- Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray,
LargeStringArray,
- ListArray, ListViewArray, NullArray, NullBuilder, StringArray,
StringViewArray,
- StructArray, Time32MillisecondArray, Time32SecondArray,
Time64MicrosecondArray,
- Time64NanosecondArray,
+ Int64Array, Int64Builder, LargeBinaryArray, LargeListArray,
LargeListViewArray,
+ LargeStringArray, ListArray, ListBuilder, ListViewArray, MapBuilder,
NullArray,
+ NullBuilder, StringArray, StringBuilder, StringViewArray, StructArray,
+ Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
Time64NanosecondArray,
};
use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow::compute::{CastOptions, cast};
@@ -4433,6 +4433,164 @@ mod test {
assert_eq!(decoded.as_ref(), &expected);
}
+ /// Map data type with `MapBuilder`'s default field names, so results can
be
+ /// compared against arrays built by `MapBuilder`.
+ fn map_data_type(value_type: DataType) -> DataType {
+ DataType::Map(
+ Arc::new(Field::new(
+ "entries",
+ DataType::Struct(Fields::from(vec![
+ Field::new("keys", DataType::Utf8, false),
+ Field::new("values", value_type, true),
+ ])),
+ false,
+ )),
+ false,
+ )
+ }
+
+ fn map_get_options(data_type: &DataType) -> GetOptions<'static> {
+ GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
+ "map",
+ data_type.clone(),
+ true,
+ ))))
+ }
+
+ #[test]
+ fn get_variant_as_map() {
+ let input: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"a": 1, "b": 2}"#),
+ Some(r#"{"c": 3}"#),
+ None,
+ Some("{}"),
+ Some(r#"{"d": null}"#),
+ ]));
+ let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
+
+ let data_type = map_data_type(DataType::Int64);
+ let result = variant_get(&variant_array,
map_get_options(&data_type)).unwrap();
+ assert_eq!(result.data_type(), &data_type);
+
+ let mut expected = MapBuilder::new(None, StringBuilder::new(),
Int64Builder::new());
+ expected.keys().append_value("a");
+ expected.values().append_value(1);
+ expected.keys().append_value("b");
+ expected.values().append_value(2);
+ expected.append(true).unwrap();
+ expected.keys().append_value("c");
+ expected.values().append_value(3);
+ expected.append(true).unwrap();
+ expected.append(false).unwrap(); // null row
+ expected.append(true).unwrap(); // empty object -> empty map
+ expected.keys().append_value("d");
+ expected.values().append_null(); // variant null -> null map value
+ expected.append(true).unwrap();
+ let expected = expected.finish();
+ assert_eq!(result.as_ref(), &expected);
+ }
+
+ #[test]
+ fn get_variant_as_map_of_lists() {
+ let input: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"a": [1, 2], "b": []}"#),
+ Some(r#"{"c": [3]}"#),
+ ]));
+ let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
+
+ let data_type = map_data_type(DataType::List(Arc::new(Field::new(
+ "item",
+ DataType::Int64,
+ true,
+ ))));
+ let result = variant_get(&variant_array,
map_get_options(&data_type)).unwrap();
+ assert_eq!(result.data_type(), &data_type);
+
+ let mut expected = MapBuilder::new(
+ None,
+ StringBuilder::new(),
+ ListBuilder::new(Int64Builder::new()),
+ );
+ expected.keys().append_value("a");
+ expected.values().append_value([Some(1), Some(2)]);
+ expected.keys().append_value("b");
+ expected.values().append_value([]);
+ expected.append(true).unwrap();
+ expected.keys().append_value("c");
+ expected.values().append_value([Some(3)]);
+ expected.append(true).unwrap();
+ let expected = expected.finish();
+ assert_eq!(result.as_ref(), &expected);
+ }
+
+ #[test]
+ fn get_variant_as_map_non_object_rows() {
+ let input: ArrayRef = Arc::new(StringArray::from(vec![
+ Some(r#"{"a": 1}"#),
+ Some("42"), // not an object
+ ]));
+ let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
+ let data_type = map_data_type(DataType::Int64);
+
+ // With safe casting (the default), non-object rows become null
+ let result = variant_get(&variant_array,
map_get_options(&data_type)).unwrap();
+ let mut expected = MapBuilder::new(None, StringBuilder::new(),
Int64Builder::new());
+ expected.keys().append_value("a");
+ expected.values().append_value(1);
+ expected.append(true).unwrap();
+ expected.append(false).unwrap();
+ let expected = expected.finish();
+ assert_eq!(result.as_ref(), &expected);
+
+ // With strict casting, non-object rows are an error
+ let options =
map_get_options(&data_type).with_cast_options(CastOptions {
+ safe: false,
+ format_options: FormatOptions::default(),
+ });
+ let err = variant_get(&variant_array, options).unwrap_err();
+ assert!(
+ err.to_string().contains("Failed to extract object"),
+ "unexpected error: {err}"
+ );
+ }
+
+ #[test]
+ fn get_variant_as_map_invalid_entries() {
+ let input: ArrayRef = Arc::new(StringArray::from(vec![Some(r#"{"a":
1}"#)]));
+ let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
+
+ // Entries type is not a struct
+ let data_type = DataType::Map(
+ Arc::new(Field::new("entries", DataType::Int32, false)),
+ false,
+ );
+ let err = variant_get(&variant_array,
map_get_options(&data_type)).unwrap_err();
+ assert!(
+ err.to_string().contains("Map entries must be Struct"),
+ "unexpected error: {err}"
+ );
+
+ // Entries struct does not have exactly two fields
+ let data_type = DataType::Map(
+ Arc::new(Field::new(
+ "entries",
+ DataType::Struct(Fields::from(vec![Field::new(
+ "keys",
+ DataType::Utf8,
+ false,
+ )])),
+ false,
+ )),
+ false,
+ );
+ let err = variant_get(&variant_array,
map_get_options(&data_type)).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("Map entries must have exactly two fields"),
+ "unexpected error: {err}"
+ );
+ }
+
fn invalid_time_variant_array() -> ArrayRef {
let mut builder = VariantArrayBuilder::new(3);
// 86401000000 is invalid for Time64Microsecond (max is 86400000000)
diff --git a/parquet-variant-compute/src/variant_to_arrow.rs
b/parquet-variant-compute/src/variant_to_arrow.rs
index 9841da555d..2ba274a3c2 100644
--- a/parquet-variant-compute/src/variant_to_arrow.rs
+++ b/parquet-variant-compute/src/variant_to_arrow.rs
@@ -28,9 +28,9 @@ use crate::{VariantArray, VariantValueArrayBuilder};
use arrow::array::{
ArrayRef, ArrowNativeTypeOp, BinaryBuilder, BinaryLikeArrayBuilder,
BinaryViewBuilder,
BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray,
GenericListArray,
- GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, NullArray,
NullBufferBuilder,
- OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder,
StringViewBuilder,
- StructArray,
+ 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};
@@ -48,6 +48,7 @@ pub(crate) enum VariantToArrowRowBuilder<'a> {
Primitive(PrimitiveVariantToArrowRowBuilder<'a>),
Array(ArrayVariantToArrowRowBuilder<'a>),
Struct(StructVariantToArrowRowBuilder<'a>),
+ Map(MapVariantToArrowRowBuilder<'a>),
Encoded(EncodedVariantToArrowRowBuilder<'a>),
BinaryVariant(VariantToBinaryVariantArrowRowBuilder),
@@ -62,6 +63,7 @@ impl<'a> VariantToArrowRowBuilder<'a> {
Primitive(b) => b.append_null(),
Array(b) => b.append_null(),
Struct(b) => b.append_null(),
+ Map(b) => b.append_null(),
Encoded(b) => b.append_null(),
BinaryVariant(b) => b.append_null(),
WithPath(path_builder) => path_builder.append_null(),
@@ -74,6 +76,7 @@ impl<'a> VariantToArrowRowBuilder<'a> {
Primitive(b) => b.append_value(&value),
Array(b) => b.append_value(&value),
Struct(b) => b.append_value(&value),
+ Map(b) => b.append_value(&value),
Encoded(b) => b.append_value(value),
BinaryVariant(b) => b.append_value(value),
WithPath(path_builder) => path_builder.append_value(value),
@@ -86,6 +89,7 @@ impl<'a> VariantToArrowRowBuilder<'a> {
Primitive(b) => b.finish(),
Array(b) => b.finish(),
Struct(b) => b.finish(),
+ Map(b) => b.finish(),
Encoded(b) => b.finish(),
BinaryVariant(b) => b.finish(),
WithPath(path_builder) => path_builder.finish(),
@@ -114,6 +118,15 @@ fn make_typed_variant_to_arrow_row_builder<'a>(
ArrayVariantToArrowRowBuilder::try_new(data_type,
cast_options, capacity, false)?;
Ok(Array(builder))
}
+ DataType::Map(entries_field, ordered) => {
+ let builder = MapVariantToArrowRowBuilder::try_new(
+ entries_field,
+ *ordered,
+ cast_options,
+ capacity,
+ )?;
+ Ok(Map(builder))
+ }
DataType::Dictionary(_, value_type) => {
let builder = EncodedVariantToArrowRowBuilder::try_new(
data_type,
@@ -648,6 +661,144 @@ impl<'a> StructVariantToArrowRowBuilder<'a> {
}
}
+/// Builder for converting variant objects into Arrow `MapArray`s.
+///
+/// Each variant object field becomes one map entry: the field name is the map
key and the field
+/// value is converted to the requested value type. Variant objects store
their fields in
+/// lexicographic key order, so maps with string keys come out sorted by key.
+pub(crate) struct MapVariantToArrowRowBuilder<'a> {
+ entries_field: &'a FieldRef,
+ key_field: &'a FieldRef,
+ value_field: &'a FieldRef,
+ ordered: bool,
+ key_builder: Box<VariantToArrowRowBuilder<'a>>,
+ value_builder: Box<VariantToArrowRowBuilder<'a>>,
+ offsets: Vec<i32>,
+ current_offset: i32,
+ nulls: NullBufferBuilder,
+ cast_options: &'a CastOptions<'a>,
+}
+
+impl<'a> MapVariantToArrowRowBuilder<'a> {
+ fn try_new(
+ entries_field: &'a FieldRef,
+ ordered: bool,
+ cast_options: &'a CastOptions,
+ capacity: usize,
+ ) -> Result<Self> {
+ let DataType::Struct(entry_fields) = entries_field.data_type() else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Map entries must be Struct, got {:?}",
+ entries_field.data_type()
+ )));
+ };
+ let (key_field, value_field) = match entry_fields.as_ref() {
+ [key_field, value_field] => (key_field, value_field),
+ fields => {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Map entries must have exactly two fields (key, value),
got {}",
+ fields.len()
+ )));
+ }
+ };
+ let key_builder = Box::new(make_typed_variant_to_arrow_row_builder(
+ key_field.data_type(),
+ cast_options,
+ capacity,
+ )?);
+ let value_builder = Box::new(make_typed_variant_to_arrow_row_builder(
+ value_field.data_type(),
+ cast_options,
+ capacity,
+ )?);
+ if capacity >= isize::MAX as usize {
+ return Err(ArrowError::ComputeError(
+ "Capacity exceeds isize::MAX when reserving map
offsets".to_string(),
+ ));
+ }
+ let mut offsets = Vec::with_capacity(capacity + 1);
+ offsets.push(0);
+ Ok(Self {
+ entries_field,
+ key_field,
+ value_field,
+ ordered,
+ key_builder,
+ value_builder,
+ offsets,
+ current_offset: 0,
+ nulls: NullBufferBuilder::new(capacity),
+ cast_options,
+ })
+ }
+
+ fn append_null(&mut self) -> Result<()> {
+ self.offsets.push(self.current_offset);
+ self.nulls.append_null();
+ Ok(())
+ }
+
+ fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
+ match variant_cast_with_options(value, self.cast_options,
Variant::as_object) {
+ Ok(Some(obj)) => {
+ for (key, field_value) in obj.iter() {
+ // Map keys cannot be null, so a failed key conversion is
an error even when
+ // cast_options.safe is true.
+ if !self.key_builder.append_value(Variant::from(key))? {
+ return Err(ArrowError::CastError(format!(
+ "Failed to cast map key {key:?} to {:?}",
+ self.key_field.data_type()
+ )));
+ }
+ self.value_builder.append_value(field_value)?;
+ self.current_offset = self.current_offset.add_checked(1)?;
+ }
+ self.offsets.push(self.current_offset);
+ self.nulls.append_non_null();
+ Ok(true)
+ }
+ Ok(None) => {
+ self.append_null()?;
+ Ok(false)
+ }
+ Err(_) => Err(ArrowError::CastError(format!(
+ "Failed to extract object from variant {value:?}"
+ ))),
+ }
+ }
+
+ fn finish(mut self) -> Result<ArrayRef> {
+ let keys = self.key_builder.finish()?;
+ let values = self.value_builder.finish()?;
+
+ let entry_fields = Fields::from(vec![
+ self.key_field
+ .as_ref()
+ .clone()
+ .with_data_type(keys.data_type().clone()),
+ self.value_field
+ .as_ref()
+ .clone()
+ .with_data_type(values.data_type().clone()),
+ ]);
+ let entries = StructArray::try_new(entry_fields.clone(), vec![keys,
values], None)?;
+ let entries_field = Arc::new(
+ self.entries_field
+ .as_ref()
+ .clone()
+ .with_data_type(DataType::Struct(entry_fields)),
+ );
+ let map_array = MapArray::try_new(
+ entries_field,
+ OffsetBuffer::new(ScalarBuffer::from(self.offsets)),
+ entries,
+ self.nulls.finish(),
+ self.ordered,
+ )?;
+ Ok(Arc::new(map_array))
+ }
+}
+
impl<'a> ArrayVariantToArrowRowBuilder<'a> {
/// Creates a new list builder for the given data type.
///
