From 1b341acbecadf3b3bc4495a6989c3404ae6485cf Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Thu, 9 Jul 2026 13:06:09 -0400 Subject: [PATCH] add `variant_to_arrow` union builder Casting Variant to dense or sparse Union dispatches each value to the union field that most exactly represents its runtime type (lossless widening allowed, declaration order breaks ties). Null rows land in a Null-typed child if declared, otherwise the first child, since unions have no top-level null buffer. Co-Authored-By: Claude Fable 5 --- parquet-variant-compute/src/variant_get.rs | 282 +++++++++++++++++- .../src/variant_to_arrow.rs | 228 +++++++++++++- 2 files changed, 503 insertions(+), 7 deletions(-) diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index fc01ca8d3b20..809dd43ce631 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, NullArray, NullBuilder, + StringArray, StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray, + Time64MicrosecondArray, Time64NanosecondArray, UnionArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::{CastOptions, cast}; @@ -505,7 +505,9 @@ mod test { use arrow::util::display::FormatOptions; use arrow_schema::ArrowError; use arrow_schema::DataType::{Boolean, Float32, Float64, Int8}; - use arrow_schema::{DataType, Field, FieldRef, Fields, IntervalUnit, TimeUnit}; + use arrow_schema::{ + DataType, Field, FieldRef, Fields, IntervalUnit, TimeUnit, UnionFields, UnionMode, + }; use chrono::DateTime; use parquet_variant::{ EMPTY_VARIANT_METADATA_BYTES, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16, @@ -5037,4 +5039,274 @@ mod test { .with_precision_and_scale(20, 3) .unwrap() ); + + fn union_get_options(fields: &UnionFields, mode: UnionMode) -> GetOptions<'static> { + let field = Field::new("union", DataType::Union(fields.clone(), mode), true); + GetOptions::new().with_as_type(Some(FieldRef::from(field))) + } + + fn int_str_bool_union_fields() -> UnionFields { + UnionFields::try_new( + vec![0, 1, 2], + vec![ + Field::new("int", DataType::Int64, true), + Field::new("str", DataType::Utf8, true), + Field::new("bool", DataType::Boolean, true), + ], + ) + .unwrap() + } + + /// int8, string, bool, array-level null, `Variant::Null`, double (no matching field), int64 + fn mixed_variant_array() -> ArrayRef { + let mut builder = VariantArrayBuilder::new(7); + builder.append_variant(Variant::Int8(1)); + builder.append_variant(Variant::from("hello")); + builder.append_variant(Variant::from(true)); + builder.append_null(); + builder.append_variant(Variant::Null); + builder.append_variant(Variant::Double(2.5)); + builder.append_variant(Variant::Int64(5_000_000_000)); + ArrayRef::from(builder.build()) + } + + #[test] + fn get_variant_as_dense_union() { + let fields = int_str_bool_union_fields(); + let array = mixed_variant_array(); + let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap(); + + // nulls, `Variant::Null`, and the unmatched Double all land as nulls in the first child + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0i8, 1, 2, 0, 0, 0, 0]), + Some(ScalarBuffer::from(vec![0i32, 0, 0, 1, 2, 3, 4])), + vec![ + Arc::new(Int64Array::from(vec![ + Some(1), + None, + None, + None, + Some(5_000_000_000), + ])), + Arc::new(StringArray::from(vec!["hello"])), + Arc::new(BooleanArray::from(vec![true])), + ], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + } + + #[test] + fn get_variant_as_sparse_union() { + let fields = int_str_bool_union_fields(); + let array = mixed_variant_array(); + let result = variant_get(&array, union_get_options(&fields, UnionMode::Sparse)).unwrap(); + + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0i8, 1, 2, 0, 0, 0, 0]), + None, + vec![ + Arc::new(Int64Array::from(vec![ + Some(1), + None, + None, + None, + None, + None, + Some(5_000_000_000), + ])), + Arc::new(StringArray::from(vec![ + None, + Some("hello"), + None, + None, + None, + None, + None, + ])), + Arc::new(BooleanArray::from(vec![ + None, + None, + Some(true), + None, + None, + None, + None, + ])), + ], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + } + + #[test] + fn get_variant_as_union_prefers_most_exact_field() { + // Int8 picks the later-declared Int32 over Int64: exactness wins over declaration order + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("big", DataType::Int64, true), + Field::new("small", DataType::Int32, true), + ], + ) + .unwrap(); + let mut builder = VariantArrayBuilder::new(3); + builder.append_variant(Variant::Int8(1)); + builder.append_variant(Variant::Int32(2)); + builder.append_variant(Variant::Int64(3)); + let array = ArrayRef::from(builder.build()); + + let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap(); + + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields, + ScalarBuffer::from(vec![1i8, 1, 0]), + Some(ScalarBuffer::from(vec![0i32, 1, 0])), + vec![ + Arc::new(Int64Array::from(vec![3])), + Arc::new(Int32Array::from(vec![1, 2])), + ], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + } + + #[test] + fn get_variant_as_union_with_null_field() { + // nulls and unmatched values land in the Null-typed field instead of the first one + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("null", DataType::Null, true), + Field::new("int", DataType::Int64, true), + ], + ) + .unwrap(); + let mut builder = VariantArrayBuilder::new(4); + builder.append_variant(Variant::Int8(1)); + builder.append_null(); + builder.append_variant(Variant::Null); + builder.append_variant(Variant::from("no matching field")); + let array = ArrayRef::from(builder.build()); + + let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap(); + + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields, + ScalarBuffer::from(vec![1i8, 0, 0, 0]), + Some(ScalarBuffer::from(vec![0i32, 0, 1, 2])), + vec![ + Arc::new(NullArray::new(3)), + Arc::new(Int64Array::from(vec![1])), + ], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + } + + #[test] + fn get_variant_as_union_of_nested_types() { + let fields = UnionFields::try_new( + vec![0, 1, 2], + vec![ + Field::new( + "struct", + DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int64, true)])), + true, + ), + Field::new( + "list", + DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), + true, + ), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(); + let json = StringArray::from(vec![r#"{"a": 1}"#, "[1, 2, 3]", "\"s\""]); + let array = ArrayRef::from(json_to_variant(&(Arc::new(json) as ArrayRef)).unwrap()); + + let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap(); + + let mut list_builder = ListBuilder::new(Int64Builder::new()); + list_builder.append_value([Some(1), Some(2), Some(3)]); + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0i8, 1, 2]), + Some(ScalarBuffer::from(vec![0i32, 0, 0])), + vec![ + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int64, true)), + Arc::new(Int64Array::from(vec![1])) as ArrayRef, + )])), + Arc::new(list_builder.finish()), + Arc::new(StringArray::from(vec!["s"])), + ], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + } + + #[test] + fn get_variant_as_union_no_matching_field() { + let fields = + UnionFields::try_new(vec![0], vec![Field::new("str", DataType::Utf8, true)]).unwrap(); + let mut builder = VariantArrayBuilder::new(2); + builder.append_variant(Variant::from("kept")); + builder.append_variant(Variant::Int8(1)); + let array = ArrayRef::from(builder.build()); + + // Safe mode: the Int8 row becomes a null in the first (only) child. + let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap(); + let expected: ArrayRef = Arc::new( + UnionArray::try_new( + fields.clone(), + ScalarBuffer::from(vec![0i8, 0]), + Some(ScalarBuffer::from(vec![0i32, 1])), + vec![Arc::new(StringArray::from(vec![Some("kept"), None]))], + ) + .unwrap(), + ); + assert_eq!(&result, &expected); + + // Strict mode: the same row is a cast error. + let options = union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions { + safe: false, + ..Default::default() + }); + let err = variant_get(&array, options).unwrap_err(); + assert!( + err.to_string().contains("no field can represent it"), + "unexpected error: {err}" + ); + } + + #[test] + fn get_variant_as_union_empty_fields_errors() { + let mut builder = VariantArrayBuilder::new(1); + builder.append_variant(Variant::Int8(1)); + let array = ArrayRef::from(builder.build()); + + let err = variant_get( + &array, + union_get_options(&UnionFields::empty(), UnionMode::Dense), + ) + .unwrap_err(); + assert!( + err.to_string().contains("at least one union field"), + "unexpected error: {err}" + ); + } } diff --git a/parquet-variant-compute/src/variant_to_arrow.rs b/parquet-variant-compute/src/variant_to_arrow.rs index 9841da555da0..12b6e51fc780 100644 --- a/parquet-variant-compute/src/variant_to_arrow.rs +++ b/parquet-variant-compute/src/variant_to_arrow.rs @@ -30,13 +30,13 @@ use arrow::array::{ BooleanBuilder, FixedSizeBinaryBuilder, FixedSizeListArray, GenericListArray, GenericListViewArray, LargeBinaryBuilder, LargeStringBuilder, NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveBuilder, StringBuilder, StringLikeArrayBuilder, StringViewBuilder, - StructArray, + StructArray, UnionArray, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::{CastOptions, DecimalCast, cast_with_options}; use arrow::datatypes::{self, DataType, DecimalType}; use arrow::error::{ArrowError, Result}; -use arrow_schema::{FieldRef, Fields, TimeUnit}; +use arrow_schema::{FieldRef, Fields, TimeUnit, UnionFields, UnionMode}; use parquet_variant::{Variant, VariantPath}; use std::sync::Arc; @@ -48,6 +48,7 @@ pub(crate) enum VariantToArrowRowBuilder<'a> { Primitive(PrimitiveVariantToArrowRowBuilder<'a>), Array(ArrayVariantToArrowRowBuilder<'a>), Struct(StructVariantToArrowRowBuilder<'a>), + Union(UnionVariantToArrowRowBuilder<'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(), + Union(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), + Union(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(), + Union(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::Union(union_fields, mode) => { + let builder = UnionVariantToArrowRowBuilder::try_new( + union_fields, + *mode, + cast_options, + capacity, + )?; + Ok(Union(builder)) + } DataType::Dictionary(_, value_type) => { let builder = EncodedVariantToArrowRowBuilder::try_new( data_type, @@ -648,6 +661,217 @@ impl<'a> StructVariantToArrowRowBuilder<'a> { } } +/// Builder for converting variant values into a [`UnionArray`]. +/// +/// Each value is dispatched to the union field that most exactly represents its runtime type +/// (see [`union_child_rank`]), with ties broken by declaration order. Unions have no top-level +/// null buffer, so null rows -- and, in safe mode, values no field can represent -- become a +/// null in the [`DataType::Null`] child if the union declares one, otherwise in the first child. +pub(crate) struct UnionVariantToArrowRowBuilder<'a> { + fields: &'a UnionFields, + mode: UnionMode, + children: Vec>, + type_ids: Vec, + /// Dense mode only + offsets: Vec, + null_child: usize, + cast_options: &'a CastOptions<'a>, +} + +struct UnionChildBuilder<'a> { + type_id: i8, + builder: VariantToArrowRowBuilder<'a>, + len: i32, +} + +impl<'a> UnionVariantToArrowRowBuilder<'a> { + fn try_new( + fields: &'a UnionFields, + mode: UnionMode, + cast_options: &'a CastOptions<'a>, + capacity: usize, + ) -> Result { + // null rows need a child to land in + if fields.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "Casting Variant to a union requires at least one union field".to_string(), + )); + } + let mut children = Vec::with_capacity(fields.len()); + for (type_id, field) in fields.iter() { + children.push(UnionChildBuilder { + type_id, + builder: make_typed_variant_to_arrow_row_builder( + field.data_type(), + cast_options, + capacity, + )?, + len: 0, + }); + } + let null_child = fields + .iter() + .position(|(_, field)| field.data_type() == &DataType::Null) + .unwrap_or(0); + let offsets = match mode { + UnionMode::Dense => Vec::with_capacity(capacity), + UnionMode::Sparse => Vec::new(), + }; + Ok(Self { + fields, + mode, + children, + type_ids: Vec::with_capacity(capacity), + offsets, + null_child, + cast_options, + }) + } + + fn append_null(&mut self) -> Result<()> { + self.append_to_child(self.null_child, None)?; + Ok(()) + } + + fn append_value(&mut self, value: &Variant<'_, '_>) -> Result { + // `Variant::Null` becomes null even in strict mode, like in the other builders + if matches!(value, Variant::Null) { + self.append_null()?; + return Ok(false); + } + match self.select_child(value) { + Some(index) => self.append_to_child(index, Some(value)), + None if self.cast_options.safe => { + self.append_null()?; + Ok(false) + } + None => Err(ArrowError::CastError(format!( + "Failed to cast variant {value:?} to union: no field can represent it" + ))), + } + } + + fn select_child(&self, value: &Variant<'_, '_>) -> Option { + let mut best: Option<(u8, usize)> = None; + for (index, (_, field)) in self.fields.iter().enumerate() { + let Some(rank) = union_child_rank(value, field.data_type()) else { + continue; + }; + if best.is_none_or(|(best_rank, _)| rank < best_rank) { + best = Some((rank, index)); + } + } + best.map(|(_, index)| index) + } + + fn append_to_child(&mut self, index: usize, value: Option<&Variant<'_, '_>>) -> Result { + self.type_ids.push(self.children[index].type_id); + match self.mode { + UnionMode::Dense => { + let child = &mut self.children[index]; + self.offsets.push(child.len); + child.len = child.len.add_checked(1)?; + match value { + Some(value) => child.builder.append_value(value.clone()), + None => { + child.builder.append_null()?; + Ok(false) + } + } + } + UnionMode::Sparse => { + let mut appended = false; + for (child_index, child) in self.children.iter_mut().enumerate() { + match value { + Some(value) if child_index == index => { + appended = child.builder.append_value(value.clone())?; + } + _ => child.builder.append_null()?, + } + } + Ok(appended) + } + } + } + + fn finish(self) -> Result { + let mut type_ids = Vec::with_capacity(self.children.len()); + let mut fields = Vec::with_capacity(self.children.len()); + let mut arrays = Vec::with_capacity(self.children.len()); + for (child, (_, field)) in self.children.into_iter().zip(self.fields.iter()) { + let array = child.builder.finish()?; + type_ids.push(child.type_id); + fields.push( + field + .as_ref() + .clone() + .with_data_type(array.data_type().clone()), + ); + arrays.push(array); + } + let fields = UnionFields::try_new(type_ids, fields)?; + let offsets = (self.mode == UnionMode::Dense).then(|| ScalarBuffer::from(self.offsets)); + let array = + UnionArray::try_new(fields, ScalarBuffer::from(self.type_ids), offsets, arrays)?; + Ok(Arc::new(array)) + } +} + +/// Ranks how exactly a union child of type `data_type` can represent a variant value's runtime +/// type: 0 is the value's natural Arrow type, higher ranks are lossless widenings, and `None` +/// means the child cannot represent the value losslessly. Every pair admitted here must be +/// convertible by the corresponding row builder. +fn union_child_rank(value: &Variant<'_, '_>, data_type: &DataType) -> Option { + use DataType::*; + let rank = match (value, data_type) { + (Variant::BooleanTrue | Variant::BooleanFalse, Boolean) => 0, + (Variant::Int8(_), Int8) => 0, + (Variant::Int8(_), Int16) => 1, + (Variant::Int8(_), Int32) => 2, + (Variant::Int8(_), Int64) => 3, + (Variant::Int16(_), Int16) => 0, + (Variant::Int16(_), Int32) => 1, + (Variant::Int16(_), Int64) => 2, + (Variant::Int32(_), Int32) => 0, + (Variant::Int32(_), Int64) => 1, + (Variant::Int64(_), Int64) => 0, + (Variant::Float(_), Float32) => 0, + (Variant::Float(_), Float64) => 1, + (Variant::Double(_), Float64) => 0, + (Variant::Decimal4(_), Decimal32(..)) => 0, + (Variant::Decimal4(_), Decimal64(..)) => 1, + (Variant::Decimal4(_), Decimal128(..)) => 2, + (Variant::Decimal4(_), Decimal256(..)) => 3, + (Variant::Decimal8(_), Decimal64(..)) => 0, + (Variant::Decimal8(_), Decimal128(..)) => 1, + (Variant::Decimal8(_), Decimal256(..)) => 2, + (Variant::Decimal16(_), Decimal128(..)) => 0, + (Variant::Decimal16(_), Decimal256(..)) => 1, + (Variant::Date(_), Date32) => 0, + (Variant::Date(_), Date64) => 1, + (Variant::TimestampMicros(_), Timestamp(TimeUnit::Microsecond, Some(_))) => 0, + (Variant::TimestampMicros(_), Timestamp(TimeUnit::Nanosecond, Some(_))) => 1, + (Variant::TimestampNanos(_), Timestamp(TimeUnit::Nanosecond, Some(_))) => 0, + (Variant::TimestampNtzMicros(_), Timestamp(TimeUnit::Microsecond, None)) => 0, + (Variant::TimestampNtzMicros(_), Timestamp(TimeUnit::Nanosecond, None)) => 1, + (Variant::TimestampNtzNanos(_), Timestamp(TimeUnit::Nanosecond, None)) => 0, + (Variant::Time(_), Time64(TimeUnit::Microsecond)) => 0, + (Variant::Time(_), Time64(TimeUnit::Nanosecond)) => 1, + (Variant::String(_) | Variant::ShortString(_), Utf8 | LargeUtf8 | Utf8View) => 0, + (Variant::Binary(_), Binary | LargeBinary | BinaryView) => 0, + (Variant::Uuid(_), FixedSizeBinary(16)) => 0, + (Variant::Object(_), Struct(_)) => 0, + // unreachable until casting Variant to Map lands (#10012) + (Variant::Object(_), Map(..)) => 1, + (Variant::List(_), List(_)) => 0, + (Variant::List(_), LargeList(_)) => 1, + (Variant::List(_), ListView(_)) => 2, + (Variant::List(_), LargeListView(_)) => 3, + _ => return None, + }; + Some(rank) +} + impl<'a> ArrayVariantToArrowRowBuilder<'a> { /// Creates a new list builder for the given data type. ///