Skip to content

[Variant] add variant_to_arrow union builder #10313

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sdf-jkl wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[Variant] add variant_to_arrow union builder #10313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1b341ac
add `variant_to_arrow` union builder
sdf-jkl Jul 9, 2026
3a35750
Merge branch 'main' of https://github.com/apache/arrow-rs into varian…
sdf-jkl Jul 20, 2026
97ddd93
add map support tests
sdf-jkl Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
337 changes: 336 additions & 1 deletion parquet-variant-compute/src/variant_get.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -497,6 +497,7 @@ mod test {
LargeStringArray, ListArray, ListBuilder, ListViewArray, MapBuilder, NullArray,
NullBuilder, StringArray, StringBuilder, StringViewArray, StructArray,
Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
UnionArray,
};
use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow::compute::{CastOptions, cast};
Expand All @@ -505,7 +506,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,
Expand Down Expand Up @@ -5195,4 +5198,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),
],
)
.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_with_map_field() {
// With no Struct field in the union, an object routes to the Map child.
let fields = UnionFields::try_new(
vec![0, 1],
vec![
Field::new("map", map_data_type(DataType::Int64), true),
Field::new("str", DataType::Utf8, true),
],
)
.unwrap();
let json = StringArray::from(vec![r#"{"a": 1, "b": 2}"#, "\"hi\""]);
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 map_builder = MapBuilder::new(None, StringBuilder::new(), Int64Builder::new());
map_builder.keys().append_value("a");
map_builder.values().append_value(1);
map_builder.keys().append_value("b");
map_builder.values().append_value(2);
map_builder.append(true).unwrap();
let expected: ArrayRef = Arc::new(
UnionArray::try_new(
fields,
ScalarBuffer::from(vec![0i8, 1]),
Some(ScalarBuffer::from(vec![0i32, 0])),
vec![
Arc::new(map_builder.finish()),
Arc::new(StringArray::from(vec!["hi"])),
],
)
.unwrap(),
);
assert_eq!(&result, &expected);
}

#[test]
fn get_variant_as_union_prefers_struct_over_map() {
// Both a Struct and a Map field can hold an object; the object routes to Struct because
// it represents the object more exactly (rank 0 vs 1).
let fields = UnionFields::try_new(
vec![0, 1],
vec![
Field::new("map", map_data_type(DataType::Int64), true),
Field::new(
"struct",
DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int64, true)])),
true,
),
],
)
.unwrap();
let json = StringArray::from(vec![r#"{"a": 1}"#]);
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 union = result.as_any().downcast_ref::<UnionArray>().unwrap();
// type_id 1 == the struct child
assert_eq!(union.type_ids(), &[1i8]);
}

#[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}"
);
}
}
Loading
Loading