From 64186fc17dfb5c4cd9d8436db5b271bcfa3fed02 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:13:44 -0400 Subject: [PATCH 1/4] feat: add dense union encoding Signed-off-by: Nemo Yu --- Cargo.lock | 14 + Cargo.toml | 2 + encodings/dense-union/Cargo.toml | 33 ++ encodings/dense-union/README.md | 3 + encodings/dense-union/benches/take.rs | 97 ++++ encodings/dense-union/src/array.rs | 465 ++++++++++++++++++++ encodings/dense-union/src/canonical.rs | 90 ++++ encodings/dense-union/src/compute/filter.rs | 25 ++ encodings/dense-union/src/compute/mask.rs | 25 ++ encodings/dense-union/src/compute/mod.rs | 7 + encodings/dense-union/src/compute/slice.rs | 26 ++ encodings/dense-union/src/compute/take.rs | 31 ++ encodings/dense-union/src/lib.rs | 25 ++ encodings/dense-union/src/rules.rs | 17 + encodings/dense-union/src/tests.rs | 359 +++++++++++++++ vortex-file/Cargo.toml | 1 + vortex-file/src/lib.rs | 1 + 17 files changed, 1221 insertions(+) create mode 100644 encodings/dense-union/Cargo.toml create mode 100644 encodings/dense-union/README.md create mode 100644 encodings/dense-union/benches/take.rs create mode 100644 encodings/dense-union/src/array.rs create mode 100644 encodings/dense-union/src/canonical.rs create mode 100644 encodings/dense-union/src/compute/filter.rs create mode 100644 encodings/dense-union/src/compute/mask.rs create mode 100644 encodings/dense-union/src/compute/mod.rs create mode 100644 encodings/dense-union/src/compute/slice.rs create mode 100644 encodings/dense-union/src/compute/take.rs create mode 100644 encodings/dense-union/src/lib.rs create mode 100644 encodings/dense-union/src/rules.rs create mode 100644 encodings/dense-union/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..4e58b02f838 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10011,6 +10011,19 @@ dependencies = [ "vortex-session", ] +[[package]] +name = "vortex-dense-union" +version = "0.1.0" +dependencies = [ + "codspeed-divan-compat", + "prost 0.14.4", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + [[package]] name = "vortex-duckdb" version = "0.1.0" @@ -10140,6 +10153,7 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-dense-union", "vortex-edition", "vortex-error", "vortex-fastlanes", diff --git a/Cargo.toml b/Cargo.toml index 44e154eadc3..934ec647a96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,7 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", + "encodings/dense-union", # Benchmarks "benchmarks/lance-bench", "benchmarks/compress-bench", @@ -304,6 +305,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 } diff --git a/encodings/dense-union/Cargo.toml b/encodings/dense-union/Cargo.toml new file mode 100644 index 00000000000..061bc6427fb --- /dev/null +++ b/encodings/dense-union/Cargo.toml @@ -0,0 +1,33 @@ +[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 = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +prost = { workspace = true } +vortex-array = { workspace = true } +vortex-error = { workspace = true } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-buffer = { workspace = true } + +[lints] +workspace = true + +[[bench]] +name = "take" +harness = false diff --git a/encodings/dense-union/README.md b/encodings/dense-union/README.md new file mode 100644 index 00000000000..79132dee06c --- /dev/null +++ b/encodings/dense-union/README.md @@ -0,0 +1,3 @@ +# Vortex Dense Union + +An external dense physical encoding for Vortex's logical `DType::Union`. diff --git a/encodings/dense-union/benches/take.rs b/encodings/dense-union/benches/take.rs new file mode 100644 index 00000000000..19a6111cb50 --- /dev/null +++ b/encodings/dense-union/benches/take.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +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; + +const LEN: usize = 65_536; +const N_VARIANTS: usize = 28; +const TAKE_LEN: usize = 4_096; + +fn main() { + divan::main(); +} + +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) { + 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::>(); + 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::>(); + 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() +} + +#[divan::bench] +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())); +} diff --git a/encodings/dense-union/src/array.rs b/encodings/dense-union/src/array.rs new file mode 100644 index 00000000000..92cfacba282 --- /dev/null +++ b/encodings/dense-union/src/array.rs @@ -0,0 +1,465 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EmptyArrayData; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::OperationsVTable; +use vortex_array::TypedArrayRef; +use vortex_array::VTable; +use vortex_array::ValidityVTable; +use vortex_array::array_slots; +use vortex_array::arrays::Primitive; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::UnionVariants; +use vortex_array::require_child; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_array::with_empty_buffers; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::canonical::canonicalize; +use crate::rules::PARENT_RULES; + +const OFFSETS_DTYPE: DType = DType::Primitive(PType::I32, Nullability::NonNullable); + +/// A [`DenseUnion`]-encoded Vortex array. +pub type DenseUnionArray = Array; + +/// 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, +} + +#[derive(Clone, prost::Message)] +struct DenseUnionMetadata { + #[prost(uint64, repeated, tag = "1")] + child_lengths: Vec, +} + +/// 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, +} + +fn make_parts( + type_ids: ArrayRef, + offsets: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, +) -> ArrayParts { + 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 + '_ { + 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. + fn child_by_name_opt(&self, name: impl AsRef) -> Option<&ArrayRef> { + self.child(self.variants().find(name)?) + } + + /// Return a compact child selected by variant name. + fn child_by_name(&self, name: impl AsRef) -> 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> 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, + ) -> DenseUnionArray { + Self::try_new(type_ids, offsets, variants, children) + .vortex_expect("DenseUnion construction failed") + } + + /// Try to construct a dense union array. + /// + /// 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. + pub fn try_new( + type_ids: ArrayRef, + offsets: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, + ) -> VortexResult { + 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 { + 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, + } + } +} + +fn validate_components( + type_ids: &ArrayRef, + offsets: &ArrayRef, + children: &[&ArrayRef], + dtype: &DType, + len: usize, +) -> VortexResult<()> { + let DType::Union(variants, nullability) = dtype else { + return Err(vortex_err!( + "DenseUnion requires a union dtype, got {dtype}" + )); + }; + vortex_ensure_eq!( + children.len(), + variants.len(), + "DenseUnion has {} compact children but expected {}", + children.len(), + variants.len() + ); + vortex_ensure_eq!( + type_ids.dtype(), + &DType::Primitive(PType::U8, *nullability), + "DenseUnion type_ids have incompatible dtype" + ); + vortex_ensure_eq!( + type_ids.len(), + len, + "DenseUnion type_ids length does not match its logical length" + ); + vortex_ensure_eq!( + offsets.dtype(), + &OFFSETS_DTYPE, + "DenseUnion offsets must be non-nullable i32" + ); + vortex_ensure_eq!( + offsets.len(), + len, + "DenseUnion offsets length does not match its logical length" + ); + + for (index, (variant_dtype, child)) in variants.variants().zip(children).enumerate() { + vortex_ensure_eq!( + child.dtype(), + &variant_dtype, + "DenseUnion child {index} has incompatible dtype" + ); + } + + Ok(()) +} + +impl VTable for DenseUnion { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.dense_union"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let DType::Union(variants, _) = dtype else { + return Err(vortex_err!( + "DenseUnion requires a union dtype, got {dtype}" + )); + }; + vortex_ensure_eq!( + slots.len(), + DenseUnionSlots::CHILDREN_OFFSET + variants.len(), + "DenseUnion has an unexpected number of slots" + ); + let type_ids = slots[DenseUnionSlots::TYPE_IDS] + .as_ref() + .ok_or_else(|| vortex_err!("DenseUnion is missing its type_ids slot"))?; + let offsets = slots[DenseUnionSlots::OFFSETS] + .as_ref() + .ok_or_else(|| vortex_err!("DenseUnion is missing its offsets slot"))?; + let children = slots[DenseUnionSlots::CHILDREN_OFFSET..] + .iter() + .enumerate() + .map(|(index, child)| { + child + .as_ref() + .ok_or_else(|| vortex_err!("DenseUnion is missing compact child {index}")) + }) + .collect::>>()?; + + validate_components(type_ids, offsets, &children, dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("DenseUnion buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("DenseUnion buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + let child_lengths = array + .iter_children() + .map(|child| { + u64::try_from(child.len()) + .map_err(|_| vortex_err!("DenseUnion child length does not fit in u64")) + }) + .collect::>>()?; + Ok(Some(DenseUnionMetadata { child_lengths }.encode_to_vec())) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!(buffers.is_empty(), "DenseUnion expects no buffers"); + let DType::Union(variants, nullability) = dtype else { + return Err(vortex_err!( + "DenseUnion requires a union dtype, got {dtype}" + )); + }; + let metadata = DenseUnionMetadata::decode(metadata)?; + vortex_ensure_eq!( + metadata.child_lengths.len(), + variants.len(), + "DenseUnion metadata has an unexpected number of child lengths" + ); + vortex_ensure_eq!( + children.len(), + DenseUnionSlots::CHILDREN_OFFSET + variants.len(), + "DenseUnion has an unexpected number of serialized children" + ); + + let type_ids = children.get( + DenseUnionSlots::TYPE_IDS, + &DType::Primitive(PType::U8, *nullability), + len, + )?; + let offsets = children.get(DenseUnionSlots::OFFSETS, &OFFSETS_DTYPE, len)?; + let compact_children = variants + .variants() + .zip(metadata.child_lengths) + .enumerate() + .map(|(index, (variant_dtype, child_len))| { + let child_len = usize::try_from(child_len) + .map_err(|_| vortex_err!("DenseUnion child length does not fit in usize"))?; + children.get( + DenseUnionSlots::CHILDREN_OFFSET + index, + &variant_dtype, + child_len, + ) + }) + .collect::>>()?; + + Ok(make_parts( + type_ids, + offsets, + variants.clone(), + compact_children, + )) + } + + fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + DenseUnionSlots::TYPE_IDS => "type_ids".to_string(), + DenseUnionSlots::OFFSETS => "offsets".to_string(), + _ => array.variants().names()[idx - DenseUnionSlots::CHILDREN_OFFSET].to_string(), + } + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let array = require_child!( + array, + array.type_ids(), + DenseUnionSlots::TYPE_IDS => Primitive + ); + let array = require_child!( + array, + array.offsets(), + DenseUnionSlots::OFFSETS => Primitive + ); + canonicalize(array, ctx).map(ExecutionResult::done) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for DenseUnion { + fn scalar_at( + array: ArrayView<'_, DenseUnion>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let type_id_scalar = array.type_ids().execute_scalar(index, ctx)?; + let Some(type_id) = type_id_scalar.as_primitive().typed_value::() else { + return Ok(Scalar::null(array.dtype().clone())); + }; + let child_index = array + .variants() + .tag_to_child_index(type_id) + .ok_or_else(|| vortex_err!("DenseUnion contains unknown type ID {type_id}"))?; + let offset = array + .offsets() + .execute_scalar(index, ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| vortex_err!("DenseUnion contains a null offset at row {index}"))?; + let offset = usize::try_from(offset).map_err(|_| { + vortex_err!("DenseUnion contains negative offset {offset} at row {index}") + })?; + let child = array + .child(child_index) + .ok_or_else(|| vortex_err!("DenseUnion is missing compact child {child_index}"))?; + vortex_ensure!( + offset < child.len(), + "DenseUnion offset {offset} is out of bounds for child {child_index} of length {}", + child.len() + ); + + Scalar::union( + array.variants().clone(), + type_id, + child.execute_scalar(offset, ctx)?, + array.dtype().nullability(), + ) + } +} + +impl ValidityVTable for DenseUnion { + fn validity(array: ArrayView<'_, DenseUnion>) -> VortexResult { + array.type_ids().validity() + } +} diff --git a/encodings/dense-union/src/canonical.rs b/encodings/dense-union/src/canonical.rs new file mode 100644 index 00000000000..2b5399f9c2f --- /dev/null +++ b/encodings/dense-union/src/canonical.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Array; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::UnionArray; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; + +use crate::array::DenseUnion; +use crate::array::DenseUnionArrayExt; +use crate::array::DenseUnionArraySlotsExt; + +pub(crate) fn canonicalize( + array: Array, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let variants = array.variants().clone(); + let type_ids = array.type_ids().as_::(); + let offsets = array.offsets().as_::(); + let type_id_values = type_ids.as_slice::(); + let offset_values = offsets.as_slice::(); + let valid_rows = type_ids.validity()?.execute_mask(len, ctx)?; + let mut codes_by_child = vec![vec![0u32; len]; variants.len()]; + + let mut assign_row = |row: usize| -> VortexResult<()> { + let type_id = type_id_values[row]; + let child_index = variants + .tag_to_child_index(type_id) + .ok_or_else(|| vortex_err!("DenseUnion contains unknown type ID {type_id}"))?; + let offset = usize::try_from(offset_values[row]).map_err(|_| { + vortex_err!( + "DenseUnion contains negative offset {} at row {row}", + offset_values[row] + ) + })?; + let child_len = array + .child(child_index) + .ok_or_else(|| vortex_err!("DenseUnion is missing compact child {child_index}"))? + .len(); + vortex_ensure!( + 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) + .map_err(|_| vortex_err!("DenseUnion offset {offset} does not fit in u32"))?; + Ok(()) + }; + + match valid_rows.indices() { + AllOr::All => { + for row in 0..len { + assign_row(row)?; + } + } + AllOr::None => {} + AllOr::Some(rows) => { + for &row in rows { + assign_row(row)?; + } + } + } + + let sparse_children = array + .iter_children() + .zip(codes_by_child) + .map(|(child, codes)| { + let values = if child.is_empty() { + ConstantArray::new(Scalar::default_value(child.dtype()), 1).into_array() + } else { + child.clone() + }; + DictArray::try_new(PrimitiveArray::from_iter(codes).into_array(), values) + .map(IntoArray::into_array) + }) + .collect::>>()?; + + UnionArray::try_new(type_ids.array().clone(), variants, sparse_children) + .map(IntoArray::into_array) +} diff --git a/encodings/dense-union/src/compute/filter.rs b/encodings/dense-union/src/compute/filter.rs new file mode 100644 index 00000000000..60fb8181bf2 --- /dev/null +++ b/encodings/dense-union/src/compute/filter.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::filter::FilterReduce; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::DenseUnion; +use crate::DenseUnionArrayExt; +use crate::DenseUnionArraySlotsExt; + +impl FilterReduce for DenseUnion { + fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { + DenseUnion::try_new( + array.type_ids().filter(mask.clone())?, + array.offsets().filter(mask.clone())?, + array.variants().clone(), + array.iter_children().cloned(), + ) + .map(|array| Some(array.into_array())) + } +} diff --git a/encodings/dense-union/src/compute/mask.rs b/encodings/dense-union/src/compute/mask.rs new file mode 100644 index 00000000000..85f6489460a --- /dev/null +++ b/encodings/dense-union/src/compute/mask.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_error::VortexResult; + +use crate::DenseUnion; +use crate::DenseUnionArrayExt; +use crate::DenseUnionArraySlotsExt; + +impl MaskReduce for DenseUnion { + fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + DenseUnion::try_new( + array.type_ids().clone().mask(mask.clone())?, + array.offsets().clone(), + array.variants().clone(), + array.iter_children().cloned(), + ) + .map(|array| Some(array.into_array())) + } +} diff --git a/encodings/dense-union/src/compute/mod.rs b/encodings/dense-union/src/compute/mod.rs new file mode 100644 index 00000000000..a9deb01f82e --- /dev/null +++ b/encodings/dense-union/src/compute/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod filter; +mod mask; +mod slice; +mod take; diff --git a/encodings/dense-union/src/compute/slice.rs b/encodings/dense-union/src/compute/slice.rs new file mode 100644 index 00000000000..72a1d2c3113 --- /dev/null +++ b/encodings/dense-union/src/compute/slice.rs @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_error::VortexResult; + +use crate::DenseUnion; +use crate::DenseUnionArrayExt; +use crate::DenseUnionArraySlotsExt; + +impl SliceReduce for DenseUnion { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + DenseUnion::try_new( + array.type_ids().slice(range.clone())?, + array.offsets().slice(range)?, + array.variants().clone(), + array.iter_children().cloned(), + ) + .map(|array| Some(array.into_array())) + } +} diff --git a/encodings/dense-union/src/compute/take.rs b/encodings/dense-union/src/compute/take.rs new file mode 100644 index 00000000000..45035c2d615 --- /dev/null +++ b/encodings/dense-union/src/compute/take.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::dict::TakeReduce; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::DenseUnion; +use crate::DenseUnionArrayExt; +use crate::DenseUnionArraySlotsExt; + +impl TakeReduce for DenseUnion { + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + 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)?; + + DenseUnion::try_new( + type_ids, + offsets, + array.variants().clone(), + array.iter_children().cloned(), + ) + .map(|array| Some(array.into_array())) + } +} diff --git a/encodings/dense-union/src/lib.rs b/encodings/dense-union/src/lib.rs new file mode 100644 index 00000000000..4eadcd7a5b3 --- /dev/null +++ b/encodings/dense-union/src/lib.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A dense physical encoding for Vortex union arrays. +//! +//! [`DenseUnionArray`] stores one type ID and one child offset per logical row. Variant children +//! are compact: unlike the canonical sparse union, they do not contain placeholders for rows that +//! select a different variant. The array still has the logical [`DType::Union`] dtype. + +mod array; +mod canonical; +mod compute; +mod rules; + +pub use array::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Register the dense union encoding in a Vortex session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(DenseUnion); +} + +#[cfg(test)] +mod tests; diff --git a/encodings/dense-union/src/rules.rs b/encodings/dense-union/src/rules.rs new file mode 100644 index 00000000000..0afb9ecfa65 --- /dev/null +++ b/encodings/dense-union/src/rules.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::dict::TakeReduceAdaptor; +use vortex_array::arrays::filter::FilterReduceAdaptor; +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; + +use crate::DenseUnion; + +pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&FilterReduceAdaptor(DenseUnion)), + ParentRuleSet::lift(&MaskReduceAdaptor(DenseUnion)), + ParentRuleSet::lift(&SliceReduceAdaptor(DenseUnion)), + ParentRuleSet::lift(&TakeReduceAdaptor(DenseUnion)), +]); diff --git a/encodings/dense-union/src/tests.rs b/encodings/dense-union/src/tests.rs new file mode 100644 index 00000000000..0b9f9cf1f9c --- /dev/null +++ b/encodings/dense-union/src/tests.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::UnionArray; +use vortex_array::arrays::union::UnionArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::UnionVariants; +use vortex_array::scalar::Scalar; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::DenseUnion; +use crate::DenseUnionArray; +use crate::DenseUnionArraySlotsExt; +use crate::initialize; + +fn variants() -> VortexResult { + UnionVariants::try_new( + ["number", "flag"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + ) +} + +fn dense_union() -> VortexResult { + DenseUnion::try_new( + PrimitiveArray::from_iter([5u8, 9, 5]).into_array(), + PrimitiveArray::from_iter([0i32, 0, 1]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 30]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + ) +} + +fn nullable_variants() -> VortexResult { + UnionVariants::try_new( + ["number", "optional"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::Nullable), + ], + vec![5, 9], + ) +} + +fn nullable_dense_union() -> VortexResult { + DenseUnion::try_new( + PrimitiveArray::from_option_iter([Some(5u8), None, Some(9), Some(9)]).into_array(), + PrimitiveArray::from_iter([0i32, 0, 0, 1]).into_array(), + nullable_variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + PrimitiveArray::from_option_iter([None, Some(40i64)]).into_array(), + ], + ) +} + +fn session() -> VortexSession { + let session = vortex_array::array_session(); + initialize(&session); + session +} + +fn assert_rows( + array: &ArrayRef, + expected: Vec, + session: &VortexSession, +) -> VortexResult<()> { + let mut ctx = session.create_execution_ctx(); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(array.execute_scalar(index, &mut ctx)?, expected); + } + Ok(()) +} + +fn assert_same_rows( + left: &ArrayRef, + right: &ArrayRef, + session: &VortexSession, +) -> VortexResult<()> { + assert_eq!(left.dtype(), right.dtype()); + assert_eq!(left.len(), right.len()); + let mut ctx = session.create_execution_ctx(); + for index in 0..left.len() { + assert_eq!( + left.execute_scalar(index, &mut ctx)?, + right.execute_scalar(index, &mut ctx)? + ); + } + Ok(()) +} + +#[test] +fn scalar_at_uses_type_id_and_offset() -> VortexResult<()> { + let session = session(); + let array = dense_union()?.into_array(); + assert_rows( + &array, + vec![ + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)?, + ], + &session, + ) +} + +#[test] +fn outer_and_selected_child_nulls_are_distinct() -> VortexResult<()> { + let session = session(); + let variants = nullable_variants()?; + let array = nullable_dense_union()?.into_array(); + assert_rows( + &array, + vec![ + Scalar::union(variants.clone(), 5, 10i32.into(), Nullability::Nullable)?, + Scalar::null(DType::Union(variants.clone(), Nullability::Nullable)), + Scalar::union( + variants.clone(), + 9, + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + Nullability::Nullable, + )?, + Scalar::union( + variants, + 9, + Scalar::primitive(40i64, Nullability::Nullable), + Nullability::Nullable, + )?, + ], + &session, + ) +} + +#[test] +fn slice_filter_take_and_mask_preserve_dense_encoding() -> VortexResult<()> { + let session = session(); + let array = dense_union()?.into_array(); + + let sliced = array.slice(1..3)?; + let filtered = array.filter(Mask::from_iter([true, false, true]))?; + let taken = array.take(PrimitiveArray::from_iter([2u32, 0, 1]).into_array())?; + let masked = array.mask(BoolArray::from_iter([true, false, true]).into_array())?; + + assert!(sliced.is::()); + assert!(filtered.is::()); + assert!(taken.is::()); + assert!(masked.is::()); + assert_eq!(sliced.as_::().children()[0].len(), 2); + assert_eq!(filtered.as_::().children()[0].len(), 2); + assert_eq!(taken.as_::().children()[0].len(), 2); + assert_eq!(masked.as_::().children()[0].len(), 2); + + assert_rows( + &sliced, + vec![ + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)?, + ], + &session, + )?; + assert_rows( + &filtered, + vec![ + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)?, + ], + &session, + )?; + assert_rows( + &taken, + vec![ + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?, + ], + &session, + )?; + assert_rows( + &masked, + vec![ + Scalar::union(variants()?, 5, 10i32.into(), Nullability::Nullable)?, + Scalar::null(DType::Union(variants()?, Nullability::Nullable)), + Scalar::union(variants()?, 5, 30i32.into(), Nullability::Nullable)?, + ], + &session, + ) +} + +#[test] +fn nullable_take_indices_become_outer_nulls() -> VortexResult<()> { + let session = session(); + let taken = dense_union()? + .into_array() + .take(PrimitiveArray::from_option_iter([Some(2u32), None, Some(0)]).into_array())?; + assert!(taken.is::()); + assert_rows( + &taken, + vec![ + Scalar::union(variants()?, 5, 30i32.into(), Nullability::Nullable)?, + Scalar::null(DType::Union(variants()?, Nullability::Nullable)), + Scalar::union(variants()?, 5, 10i32.into(), Nullability::Nullable)?, + ], + &session, + ) +} + +#[test] +fn canonicalization_uses_sparse_dictionary_children() -> VortexResult<()> { + let session = session(); + let array = dense_union()?.into_array(); + let mut ctx = session.create_execution_ctx(); + let canonical = array.clone().execute::(&mut ctx)?; + + assert!(canonical.children()[0].is::()); + assert!(canonical.children()[1].is::()); + assert_same_rows(&array, &canonical.into_array(), &session) +} + +#[test] +fn canonicalization_handles_unselected_empty_child() -> VortexResult<()> { + let session = session(); + let array = DenseUnion::try_new( + PrimitiveArray::from_iter([5u8, 5]).into_array(), + PrimitiveArray::from_iter([0i32, 1]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32, 20]).into_array(), + BoolArray::from_iter(Vec::::new()).into_array(), + ], + )? + .into_array(); + let mut ctx = session.create_execution_ctx(); + let canonical = array.clone().execute::(&mut ctx)?.into_array(); + + assert_same_rows(&array, &canonical, &session) +} + +#[test] +fn invalid_type_id_and_offsets_return_errors() -> VortexResult<()> { + let session = session(); + let mut ctx = session.create_execution_ctx(); + + let unknown_type_id = DenseUnion::try_new( + PrimitiveArray::from_iter([7u8]).into_array(), + PrimitiveArray::from_iter([0i32]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + )?; + assert!(unknown_type_id.execute_scalar(0, &mut ctx).is_err()); + + let negative_offset = DenseUnion::try_new( + PrimitiveArray::from_iter([5u8]).into_array(), + PrimitiveArray::from_iter([-1i32]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + )?; + assert!(negative_offset.execute_scalar(0, &mut ctx).is_err()); + + let out_of_bounds = DenseUnion::try_new( + PrimitiveArray::from_iter([9u8]).into_array(), + PrimitiveArray::from_iter([1i32]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + )?; + assert!(out_of_bounds.execute_scalar(0, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn validates_structural_components() -> VortexResult<()> { + assert!( + DenseUnion::try_new( + PrimitiveArray::from_iter([5u16]).into_array(), + PrimitiveArray::from_iter([0i32]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + ) + .is_err() + ); + assert!( + DenseUnion::try_new( + PrimitiveArray::from_iter([5u8]).into_array(), + PrimitiveArray::from_iter([0u32]).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ], + ) + .is_err() + ); + assert!( + DenseUnion::try_new( + PrimitiveArray::from_iter([5u8]).into_array(), + PrimitiveArray::from_iter([0i32]).into_array(), + variants()?, + vec![PrimitiveArray::from_iter([10i32]).into_array()], + ) + .is_err() + ); + Ok(()) +} + +#[test] +fn serde_roundtrip() -> VortexResult<()> { + let session = session(); + let array = nullable_dense_union()?.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(concat.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &session, + )?; + + assert!(decoded.is::()); + assert_same_rows(&array, &decoded, &session) +} diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index fc406f7133e..d0084f0d53d 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -39,6 +39,7 @@ vortex-bytebool = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-dense-union = { workspace = true } vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index e57b7ff344b..a9587124337 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -189,6 +189,7 @@ pub fn register_default_encodings(session: &VortexSession) { vortex_alp::initialize(session); vortex_datetime_parts::initialize(session); vortex_decimal_byte_parts::initialize(session); + vortex_dense_union::initialize(session); vortex_fastlanes::initialize(session); vortex_runend::initialize(session); vortex_sequence::initialize(session); From a44ec27637e28bba1876cc843d76f6ea67a36a01 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:30:22 -0400 Subject: [PATCH 2/4] docs: fix dense union dtype link Signed-off-by: Nemo Yu --- encodings/dense-union/src/lib.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/encodings/dense-union/src/lib.rs b/encodings/dense-union/src/lib.rs index 4eadcd7a5b3..5c237632209 100644 --- a/encodings/dense-union/src/lib.rs +++ b/encodings/dense-union/src/lib.rs @@ -5,7 +5,8 @@ //! //! [`DenseUnionArray`] stores one type ID and one child offset per logical row. Variant children //! are compact: unlike the canonical sparse union, they do not contain placeholders for rows that -//! select a different variant. The array still has the logical [`DType::Union`] dtype. +//! select a different variant. The array still has the logical +//! [`DType::Union`](vortex_array::dtype::DType::Union) dtype. mod array; mod canonical; From f8e4f84229ba8513ed9408251dd3830b924cb140 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:41:40 -0400 Subject: [PATCH 3/4] bench: execute dense union take Signed-off-by: Nemo Yu --- encodings/dense-union/benches/take.rs | 33 ++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/encodings/dense-union/benches/take.rs b/encodings/dense-union/benches/take.rs index 19a6111cb50..f239240f24f 100644 --- a/encodings/dense-union/benches/take.rs +++ b/encodings/dense-union/benches/take.rs @@ -4,9 +4,14 @@ #![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; @@ -15,15 +20,24 @@ 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 = 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]; @@ -82,16 +96,23 @@ 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::(ctx) + }); +} + #[divan::bench] fn dense_take(bencher: Bencher) { - bencher - .with_inputs(|| (dense_union(), indices())) - .bench_values(|(array, indices)| divan::black_box(array.take(indices).unwrap())); + bench_take(bencher, dense_union(), indices()); } #[divan::bench] fn sparse_take(bencher: Bencher) { - bencher - .with_inputs(|| (sparse_union(), indices())) - .bench_values(|(array, indices)| divan::black_box(array.take(indices).unwrap())); + bench_take(bencher, sparse_union(), indices()); } From 568bf2a1e4d3e13400d9db5f273391cb40d567e5 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:42:40 -0400 Subject: [PATCH 4/4] fix: package dense union readme Signed-off-by: Nemo Yu --- encodings/dense-union/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/encodings/dense-union/Cargo.toml b/encodings/dense-union/Cargo.toml index 061bc6427fb..2b9a4ca3e28 100644 --- a/encodings/dense-union/Cargo.toml +++ b/encodings/dense-union/Cargo.toml @@ -8,7 +8,7 @@ homepage = { workspace = true } include = { workspace = true } keywords = { workspace = true } license = { workspace = true } -readme = { workspace = true } +readme = "README.md" repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true }