feat: add dense union encoding

.github/workflows/codspeed.yml

@@ -55,7 +55,7 @@ jobs:
           - { shard: 4, name: "Encodings 1", packages: "vortex-alp vortex-bytebool vortex-datetime-parts" }
           - { shard: 5, name: "Encodings 2", packages: "vortex-decimal-byte-parts vortex-fastlanes vortex-fsst", features: "--features _test-harness" }
           - { shard: 6, name: "Encodings 3", packages: "vortex-pco vortex-runend vortex-sequence" }
-          - { shard: 7, name: "Encodings 4", packages: "vortex-sparse vortex-zigzag vortex-zstd" }
+          - { shard: 7, name: "Encodings 4", packages: "vortex-dense-union vortex-sparse vortex-zigzag vortex-zstd" }
           - { shard: 8, name: "Storage formats & row encoding", packages: "vortex-flatbuffers vortex-proto vortex-btrblocks vortex-row" }
           - { shard: 9, name: "Tensor & spatial", packages: "vortex-tensor vortex-spatial" }
     name: "Benchmark with Codspeed (Shard #${{ matrix.shard }})"

Cargo.toml

@@ -62,6 +62,7 @@ members = [
     "encodings/bytebool",
     "encodings/parquet-variant",
     "encodings/onpair",
+    "encodings/dense-union",
     # Benchmarks
     "benchmarks/bench-support",
     "benchmarks/lance-bench",
@@ -306,6 +307,7 @@ vortex-compute = { version = "0.1.0", path = "./vortex-compute", default-feature
 vortex-datafusion = { version = "0.1.0", path = "./vortex-datafusion", default-features = false }
 vortex-datetime-parts = { version = "0.1.0", path = "./encodings/datetime-parts", default-features = false }
 vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte-parts", default-features = false }
+vortex-dense-union = { version = "0.1.0", path = "./encodings/dense-union", default-features = false }
 vortex-edition = { version = "0.1.0", path = "./vortex-edition", default-features = false }
 vortex-error = { version = "0.1.0", path = "./vortex-error", default-features = false }
 vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false }

encodings/dense-union/Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+name = "vortex-dense-union"
+authors = { workspace = true }
+categories = { workspace = true }
+description = "Dense union encoding for Vortex arrays"
+edition = { workspace = true }
+homepage = { workspace = true }
+include = { workspace = true }
+keywords = { workspace = true }
+license = { workspace = true }
+readme = "README.md"
+repository = { workspace = true }
+rust-version = { workspace = true }
+version = { workspace = true }
+
+[dependencies]
+prost = { workspace = true }
+vortex-array = { workspace = true }
+vortex-buffer = { workspace = true }
+vortex-error = { workspace = true }
+vortex-mask = { workspace = true }
+vortex-session = { workspace = true }
+
+[dev-dependencies]
+divan = { workspace = true }
+rstest = { workspace = true }
+vortex-array = { workspace = true, features = ["_test-harness"] }
+
+[lints]
+workspace = true
+
+[[bench]]
+name = "take"
+harness = false

encodings/dense-union/README.md

@@ -0,0 +1,3 @@
+# Vortex Dense Union
+
+An external dense physical encoding for Vortex's logical `DType::Union`.

encodings/dense-union/benches/take.rs

@@ -0,0 +1,118 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(clippy::unwrap_used)]
+#![expect(clippy::cast_possible_truncation)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::RecursiveCanonical;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::UnionArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::FieldNames;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
+use vortex_array::dtype::UnionVariants;
+use vortex_dense_union::DenseUnion;
+use vortex_dense_union::initialize;
+use vortex_session::VortexSession;
+
+const LEN: usize = 65_536;
+const N_VARIANTS: usize = 28;
+const TAKE_LEN: usize = 4_096;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = array_session();
+    initialize(&session);
+    session
+});
+
+fn variants() -> UnionVariants {
+    let names = FieldNames::from_iter((0..N_VARIANTS).map(|index| format!("variant_{index}")));
+    let dtypes = vec![DType::Primitive(PType::I32, Nullability::NonNullable); N_VARIANTS];
+    let type_ids = (1..=N_VARIANTS).map(|type_id| type_id as u8).collect();
+    UnionVariants::try_new(names, dtypes, type_ids).unwrap()
+}
+
+fn selectors() -> (ArrayRef, ArrayRef, Vec<usize>) {
+    let mut child_lengths = vec![0usize; N_VARIANTS];
+    let mut type_ids = Vec::with_capacity(LEN);
+    let mut offsets = Vec::with_capacity(LEN);
+    for row in 0..LEN {
+        let child_index = row % N_VARIANTS;
+        type_ids.push((child_index + 1) as u8);
+        offsets.push(child_lengths[child_index] as i32);
+        child_lengths[child_index] += 1;
+    }
+    (
+        PrimitiveArray::from_iter(type_ids).into_array(),
+        PrimitiveArray::from_iter(offsets).into_array(),
+        child_lengths,
+    )
+}
+
+fn dense_union() -> ArrayRef {
+    let (type_ids, offsets, child_lengths) = selectors();
+    let children = child_lengths
+        .into_iter()
+        .map(|len| PrimitiveArray::from_iter(0..len as i32).into_array())
+        .collect::<Vec<_>>();
+    DenseUnion::try_new(type_ids, offsets, variants(), children)
+        .unwrap()
+        .into_array()
+}
+
+fn sparse_union() -> ArrayRef {
+    let (type_ids, ..) = selectors();
+    let children = (0..N_VARIANTS)
+        .map(|child_index| {
+            PrimitiveArray::from_iter((0..LEN).map(move |row| {
+                if row % N_VARIANTS == child_index {
+                    (row / N_VARIANTS) as i32
+                } else {
+                    0
+                }
+            }))
+            .into_array()
+        })
+        .collect::<Vec<_>>();
+    UnionArray::try_new(type_ids, variants(), children)
+        .unwrap()
+        .into_array()
+}
+
+fn indices() -> ArrayRef {
+    PrimitiveArray::from_iter((0..TAKE_LEN).rev().map(|index| index as u32)).into_array()
+}
+
+fn bench_take(bencher: Bencher, array: ArrayRef, indices: ArrayRef) {
+    bencher
+        .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, indices, ctx)| {
+            array
+                .take((*indices).clone())
+                .unwrap()
+                .execute::<RecursiveCanonical>(ctx)
+        });
+}
+
+#[divan::bench]
+fn dense_take(bencher: Bencher) {
+    bench_take(bencher, dense_union(), indices());
+}
+
+#[divan::bench]
+fn sparse_take(bencher: Bencher) {
+    bench_take(bencher, sparse_union(), indices());
+}

encodings/dense-union/src/array.rs

@@ -0,0 +1,176 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::Array;
+use vortex_array::ArrayParts;
+use vortex_array::ArrayRef;
+use vortex_array::ArraySlots;
+use vortex_array::EmptyArrayData;
+use vortex_array::TypedArrayRef;
+use vortex_array::array_slots;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::UnionVariants;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+/// A [`DenseUnion`]-encoded Vortex array.
+pub type DenseUnionArray = Array<DenseUnion>;
+
+/// Slot layout of a dense union array.
+#[array_slots(DenseUnion)]
+pub struct DenseUnionSlots {
+    /// The row-aligned type IDs selecting union variants.
+    #[slot(0)]
+    pub type_ids: ArrayRef,
+    /// The row-aligned offsets into the selected compact child.
+    #[slot(1)]
+    pub offsets: ArrayRef,
+    /// The compact children in variant order.
+    #[slot(2..)]
+    pub children: Vec<ArrayRef>,
+}
+
+/// Concrete parts of a [`DenseUnionArray`].
+pub struct DenseUnionDataParts {
+    /// The union variant schema.
+    pub variants: UnionVariants,
+    /// The row-aligned type IDs.
+    pub type_ids: ArrayRef,
+    /// The row-aligned compact-child offsets.
+    pub offsets: ArrayRef,
+    /// The compact children in variant order.
+    pub children: Vec<ArrayRef>,
+}
+
+pub(crate) fn make_parts(
+    type_ids: ArrayRef,
+    offsets: ArrayRef,
+    variants: UnionVariants,
+    children: impl IntoIterator<Item = ArrayRef>,
+) -> ArrayParts<DenseUnion> {
+    let len = type_ids.len();
+    let nullability = type_ids.dtype().nullability();
+    let children = children.into_iter();
+    let (lower, _) = children.size_hint();
+    let mut slots = ArraySlots::with_capacity(DenseUnionSlots::CHILDREN_OFFSET + lower);
+    slots.push(Some(type_ids));
+    slots.push(Some(offsets));
+    slots.extend(children.map(Some));
+
+    ArrayParts::new(
+        DenseUnion,
+        DType::Union(variants, nullability),
+        len,
+        EmptyArrayData,
+    )
+    .with_slots(slots)
+}
+
+/// Accessors for a dense union array.
+pub trait DenseUnionArrayExt: DenseUnionArraySlotsExt {
+    /// Return the union's variant schema.
+    fn variants(&self) -> &UnionVariants {
+        match self.as_ref().dtype() {
+            DType::Union(variants, _) => variants,
+            _ => unreachable!("DenseUnionArrayExt requires a union dtype"),
+        }
+    }
+
+    /// Iterate over compact children in variant order.
+    fn iter_children(&self) -> impl ExactSizeIterator<Item = &ArrayRef> + '_ {
+        self.children().iter()
+    }
+
+    /// Return a compact child by variant index.
+    fn child(&self, index: usize) -> Option<&ArrayRef> {
+        self.children().get(index)
+    }
+
+    /// Return a compact child selected by a data-level type ID.
+    fn child_by_type_id(&self, type_id: u8) -> Option<&ArrayRef> {
+        self.child(self.variants().tag_to_child_index(type_id)?)
+    }
+
+    /// Return a compact child selected by variant name, if present.
+    fn child_by_name_opt(&self, name: impl AsRef<str>) -> Option<&ArrayRef> {
+        self.child(self.variants().find(name)?)
+    }
+
+    /// Return a compact child selected by variant name.
+    fn child_by_name(&self, name: impl AsRef<str>) -> VortexResult<&ArrayRef> {
+        let name = name.as_ref();
+        self.child_by_name_opt(name).ok_or_else(|| {
+            vortex_err!(
+                "Variant {name} not found in dense union array with names {:?}",
+                self.variants().names()
+            )
+        })
+    }
+}
+
+impl<T: TypedArrayRef<DenseUnion>> DenseUnionArrayExt for T {}
+
+/// The dense physical encoding for the logical [`DType::Union`] type.
+#[derive(Clone, Debug)]
+pub struct DenseUnion;
+
+impl DenseUnion {
+    /// Construct a dense union array.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the components do not satisfy the invariants documented by [`Self::try_new`].
+    pub fn new(
+        type_ids: ArrayRef,
+        offsets: ArrayRef,
+        variants: UnionVariants,
+        children: impl IntoIterator<Item = ArrayRef>,
+    ) -> DenseUnionArray {
+        Self::try_new(type_ids, offsets, variants, children)
+            .vortex_expect("DenseUnion construction failed")
+    }
+
+    /// Try to construct a dense union array.
+    ///
+    /// The logical union's nullability is inherited from `type_ids`; nullable type IDs represent
+    /// outer union nulls. `type_ids` must be a nullable or non-nullable `u8` array, `offsets` must
+    /// be a non-nullable `i32` array of the same length, and the compact children must match the
+    /// variant count and dtypes. Type IDs and offsets are structurally validated, but their
+    /// individual values are checked only when the array is accessed or converted to its canonical
+    /// sparse representation.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error when the selector arrays or compact children do not satisfy these
+    /// structural invariants.
+    pub fn try_new(
+        type_ids: ArrayRef,
+        offsets: ArrayRef,
+        variants: UnionVariants,
+        children: impl IntoIterator<Item = ArrayRef>,
+    ) -> VortexResult<DenseUnionArray> {
+        Array::try_from_parts(make_parts(type_ids, offsets, variants, children))
+    }
+}
+
+/// Owned accessors for a dense union array.
+pub trait DenseUnionArrayOwnedExt {
+    /// Deconstruct this array into its type IDs, offsets, schema, and compact children.
+    fn into_data_parts(self) -> DenseUnionDataParts;
+}
+
+impl DenseUnionArrayOwnedExt for Array<DenseUnion> {
+    fn into_data_parts(self) -> DenseUnionDataParts {
+        let variants = self.variants().clone();
+        let type_ids = self.type_ids().clone();
+        let offsets = self.offsets().clone();
+        let children = self.iter_children().cloned().collect();
+        DenseUnionDataParts {
+            variants,
+            type_ids,
+            offsets,
+            children,
+        }
+    }
+}