Jefffrey commented on code in PR #10313:
URL: https://github.com/apache/arrow-rs/pull/10313#discussion_r3721389603
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -661,6 +674,216 @@ 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<UnionChildBuilder<'a>>,
+ type_ids: Vec<i8>,
+ /// Dense mode only
+ offsets: Vec<i32>,
+ 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<Self> {
+ // 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);
Review Comment:
do we need to consider nullability of the field?
##########
parquet-variant-compute/src/variant_get.rs:
##########
@@ -5171,4 +5174,336 @@ 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),
Review Comment:
should the null field be moved to be second then, if this is to prove it
doesnt go to the first field?
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -661,6 +674,216 @@ 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<UnionChildBuilder<'a>>,
+ type_ids: Vec<i8>,
+ /// Dense mode only
+ offsets: Vec<i32>,
+ 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<Self> {
+ // 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<bool> {
+ // `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<usize> {
+ 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<bool> {
+ 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<ArrayRef> {
+ 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<u8> {
+ 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,
Review Comment:
what about if the struct fields doesnt match what fields the object has
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
