sdf-jkl commented on code in PR #10313:
URL: https://github.com/apache/arrow-rs/pull/10313#discussion_r3731615398
##########
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:
thanks, addressed here -
https://github.com/apache/arrow-rs/pull/10313/commits/4bac07b29b6e9aa9968623ac3436430bedf335cc
--
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]
