feat: add dense union encoding by HarukiMoriarty · Pull Request #9367 · vortex-data/vortex

HarukiMoriarty · 2026-08-11T19:20:09Z

Rationale for this change

GeoArrow mixed-geometry arrays use dense unions, while the Vortex canonical union representation is sparse. This PR adds a generic external dense physical encoding without moving encoding policy into vortex-array.

What changes are included in this PR?

Add the external vortex-dense-union encoding for logical DType::Union values.
Store row-aligned type IDs and offsets with compact variant children.
Preserve DenseUnion through slice, filter, take, and mask operations.
Canonicalize unsupported operations to sparse UnionArray using dictionary-backed children without copying payload values.
Register DenseUnion for Vortex file serialization and deserialization.
Add behavior, validation, serde, and benchmark coverage.

What APIs are changed? Are there any user-facing changes?

Adds the vortex-dense-union crate and its DenseUnion construction/accessor APIs.
DType::Union and canonical sparse UnionArray remain in vortex-array.
This does not yet add a GeoUnion extension or GeoArrow dense-union conversion.

Performance

Representation	Median executed `take`	Relative time
Sparse `UnionArray`	42.29–42.56 µs	1.00×
DenseUnion	88.91–90.35 µs	2.10–2.12×

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

connortsui20 · 2026-08-11T19:36:50Z

+fn dense_take(bencher: Bencher) {
+    bencher
+        .with_inputs(|| (dense_union(), indices()))
+        .bench_values(|(array, indices)| divan::black_box(array.take(indices).unwrap()));
+}
+
+#[divan::bench]
+fn sparse_take(bencher: Bencher) {
+    bencher
+        .with_inputs(|| (sparse_union(), indices()))
+        .bench_values(|(array, indices)| divan::black_box(array.take(indices).unwrap()));
+}


You need to execute the array here otherwise there is no work that is done

Ah, that's correct, for now it is lazy-take. Fixed.

im interested in what the new benchmark results are?

Already showed in the PR content. DenseUnion is 2.10-2.12x slower.

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

connortsui20 · 2026-08-12T01:45:19Z

It seems like this actually has a similar issue to ListView, where we can filter, take, slice the type IDs and offsets, but then we have no way of garbage collecting the unused data (which kind of defeats the purpose of the Dense encoding)?

And if someone inevitably wants that behavior, they actually have no way to express it? I guess technically we can express it by canonicalizing into a sparse encoding by literally rebuilding the whole thing from scratch, but that is very inefficient.

Maybe its time we add that garbage collection array that we've talked about for almost a year?

connortsui20 · 2026-08-13T14:25:57Z

+        let type_ids = array.type_ids().take(indices.clone())?;
+        let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable());
+        let offset_indices = indices.clone().fill_null(fill_scalar)?;
+        let offsets = array.offsets().take(offset_indices)?;


I think this does not actually work? If we look at https://arrow.apache.org/docs/format/Columnar.html#dense-union it says this:

The respective offsets for each child value array must be in order / increasing.

So if we take on the offsets we can reorder these?

So "not correct" here just means not zero-copyable to Arrow for this encoding. But maybe we don't care? It seems generally useful to have things out of order for more efficient take?

@robert3005 do you have any thoughts?

connortsui20 · 2026-08-13T14:26:04Z

see Tracking Issue: Garbage Collection #9397

connortsui20 · 2026-08-13T18:53:07Z

the AI review comments were verbose so I deleted them (claude cant edit review comments seemingly).

I will ask it to post again, but with more concise explanations

connortsui20

Reviewed 568bf2a. Differential tests against the sparse UnionArray across slice, filter, take, mask, and canonicalize agreed on every case, including nullable variants, null take indices, empty children, and zero-length results, so no correctness bug turned up in the data paths.

Two findings worth blocking on. The canonicalize hot loop does a per-row linear tag scan and allocates len * variants zeroed codes, which is where the 2.1x goes. DenseUnion also cannot be written to a file, because no edition declares vortex.dense_union. The benchmark behind the performance table never runs in CI either.

The Arrow offset-ordering question on take is still open. The differential tests would not have caught it, because they compare against Vortex sparse unions rather than Arrow.

Generated by Claude Code

connortsui20 · 2026-08-13T18:55:20Z

+
+    let mut assign_row = |row: usize| -> VortexResult<()> {
+        let type_id = type_id_values[row];
+        let child_index = variants


tag_to_child_index is a linear scan over up to 256 tags, and array.child(..) re-resolves a slot. Both run per row here, so build a [u32; 256] tag table and materialize the child lengths before the loop. The bounds check only needs child_len, so the child lookup goes away entirely.

Generated by Claude Code

connortsui20 · 2026-08-13T18:55:27Z

+    let type_id_values = type_ids.as_slice::<u8>();
+    let offset_values = offsets.as_slice::<i32>();
+    let valid_rows = type_ids.validity()?.execute_mask(len, ctx)?;
+    let mut codes_by_child = vec![vec![0u32; len]; variants.len()];


This zero-fills len * variants.len() codes, and then PrimitiveArray::from_iter copies each Vec again. That is 112 MB at 1M rows and 28 variants, including variants that no row selects. Use BufferMut<u32>, and give unselected variants a ConstantArray of 0 instead.

Generated by Claude Code

connortsui20 · 2026-08-13T18:55:35Z

+            offset < child_len,
+            "DenseUnion offset {offset} is out of bounds for child {child_index} of length {child_len}"
+        );
+        codes_by_child[child_index][row] = u32::try_from(offset)


offset came from an i32 through usize::try_from, so it always fits in u32 and this error branch is dead. Convert once from offset_values[row] instead.

Generated by Claude Code

connortsui20 · 2026-08-13T18:55:42Z

+use crate::array::DenseUnionArrayExt;
+use crate::array::DenseUnionArraySlotsExt;
+
+pub(crate) fn canonicalize(


This file needs a doc. Three things are not derivable from the code: the children are DictArrays over the original compact children so payload values are never copied, unselected rows keep code 0 because they are unreachable through that type ID, and an empty child becomes a ConstantArray because DictArray requires non-empty values. An # Errors section covers the unknown type ID and the two offset cases.

Generated by Claude Code

connortsui20 · 2026-08-13T18:55:48Z

+    vortex_ensure_eq!(
+        type_ids.dtype(),
+        &DType::Primitive(PType::U8, *nullability),
+        "DenseUnion type_ids have incompatible dtype"


vortex_ensure_eq! appends the compared values only when there is no custom message, so this reports neither the actual nor the expected dtype. Same at 229, 234, 239, 246, 278, 352, and 357. Line 214 formats both values and is the model to follow.

Generated by Claude Code

connortsui20 · 2026-08-13T18:56:31Z

+use crate::DenseUnionArrayExt;
+use crate::DenseUnionArraySlotsExt;
+
+impl TakeReduce for DenseUnion {


A one-row take retains every child in full. That is the right trade-off for O(selectors) operations, but document it here the way Union::take documents its own cost model.

Generated by Claude Code

connortsui20 · 2026-08-13T18:56:37Z

+    Ok(())
+}
+
+fn assert_same_rows(


assert_arrays_eq! instead of the per-index loop. Both helpers also want #[track_caller].

Generated by Claude Code

connortsui20 · 2026-08-13T18:56:43Z

+
+#[test]
+fn invalid_type_id_and_offsets_return_errors() -> VortexResult<()> {
+    let session = session();


Three near-identical bodies in one test, so the first failure hides the other two. #[rstest] cases named for what each one rejects fixes that, and asserting on the error stops them passing for an unrelated failure. A zero-length case is also missing.

Generated by Claude Code

connortsui20 · 2026-08-13T18:56:50Z

This benchmark never runs. vortex-dense-union is absent from the package matrix in .github/workflows/codspeed.yml, so nothing watches the 2.1x gap.

Generated by Claude Code

connortsui20 · 2026-08-13T18:56:58Z

    vortex_alp::initialize(session);
    vortex_datetime_parts::initialize(session);
    vortex_decimal_byte_parts::initialize(session);
+    vortex_dense_union::initialize(session);


This makes DenseUnion readable but not writable, because writes are gated on the enabled editions and no edition declares vortex.dense_union. A write fails with normalize forbids encoding (vortex.dense_union) (vortex-array/src/normalize.rs:46). fastlanes.delta and vortex.onpair sit in the same position, so this can be deliberate, but the PR description claims serialization.

Generated by Claude Code

connortsui20 · 2026-08-13T19:01:59Z

+pub(crate) fn canonicalize(
+    array: Array<DenseUnion>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {


in general I think that this whole function slightly less efficient than ideal? But at the same time what we really want is to just wrap everything in SparseArray and effectively do no work. But we can't do that since SparseArray is not in here...

cc @robert3005

connortsui20 · 2026-08-13T19:03:47Z

+    Ok(())
+}
+
+impl VTable for DenseUnion {


could you split this vtable out into its own file vtable.rs? There are other examples in the codebase of how we modularize our encodings into different files

feat: add dense union encoding

64186fc

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

HarukiMoriarty requested a review from connortsui20 August 11, 2026 19:20

HarukiMoriarty added the changelog/feature A new feature label Aug 11, 2026

docs: fix dense union dtype link

a44ec27

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

connortsui20 reviewed Aug 11, 2026

View reviewed changes

HarukiMoriarty added 2 commits August 11, 2026 15:41

bench: execute dense union take

f8e4f84

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

fix: package dense union readme

568bf2a

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

HarukiMoriarty requested a review from connortsui20 August 11, 2026 19:43

HarukiMoriarty mentioned this pull request Aug 11, 2026

feat(vortex-spatial): support GeoArrow Geometry unions #9368

Open

connortsui20 requested review from joseph-isaacs and robert3005 August 12, 2026 01:50

connortsui20 reviewed Aug 13, 2026

View reviewed changes

This comment was marked as outdated.

Sign in to view

connortsui20 reviewed Aug 13, 2026

View reviewed changes

Uh oh!

Conversation

HarukiMoriarty commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

What APIs are changed? Are there any user-facing changes?

Performance

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

connortsui20 commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

connortsui20 commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

connortsui20 commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

connortsui20 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

HarukiMoriarty commented Aug 11, 2026 •

edited

Loading

connortsui20 commented Aug 12, 2026 •

edited

Loading

connortsui20 commented Aug 13, 2026 •

edited

Loading

connortsui20 commented Aug 13, 2026 •

edited

Loading