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. ///