From 9a13499fdde1747fb1ff3f2844d1b5e0eae6df73 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:13:44 -0400 Subject: [PATCH 1/7] 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 b0f98336b0b..157f316e1c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10021,6 +10021,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" @@ -10151,6 +10164,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 65452f18300..629255524b6 100644 --- a/Cargo.toml +++ b/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 } 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 72a47d88f50..92767259977 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -40,6 +40,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 25760f0890e..11c113ab837 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 ee4547dfec25edd8a1fe51fc409f6fb5781b93e4 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:30:22 -0400 Subject: [PATCH 2/7] 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 4e1e86bf28c1767af289c84c31f15c9d02dab0a6 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:41:40 -0400 Subject: [PATCH 3/7] 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 dc7f161126b77f610507d171424706ec73733f38 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 11 Aug 2026 15:42:40 -0400 Subject: [PATCH 4/7] 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 } From 7ad10277d8c35a032a1844768f5e0081486321f4 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 14 Aug 2026 10:41:36 -0400 Subject: [PATCH 5/7] fix: address dense union review feedback Signed-off-by: Nemo Yu --- .github/workflows/codspeed.yml | 2 +- Cargo.lock | 2 + encodings/dense-union/Cargo.toml | 3 +- encodings/dense-union/src/array.rs | 315 +-------------------- encodings/dense-union/src/canonical.rs | 45 ++- encodings/dense-union/src/compute/take.rs | 6 + encodings/dense-union/src/lib.rs | 5 + encodings/dense-union/src/tests.rs | 154 +++++----- encodings/dense-union/src/vtable.rs | 324 ++++++++++++++++++++++ vortex-layout/src/layouts/file_stats.rs | 2 +- vortex/Cargo.toml | 1 + vortex/src/editions/tests.rs | 49 ++++ vortex/src/editions/unstable/v2026_06.rs | 5 +- vortex/src/lib.rs | 5 + 14 files changed, 530 insertions(+), 388 deletions(-) create mode 100644 encodings/dense-union/src/vtable.rs diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 13393072b8c..39de3eadbf8 100644 --- a/.github/workflows/codspeed.yml +++ b/.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 }})" diff --git a/Cargo.lock b/Cargo.lock index 157f316e1c0..fdce018d351 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9541,6 +9541,7 @@ dependencies = [ "vortex-cloud", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-dense-union", "vortex-edition", "vortex-error", "vortex-fastlanes", @@ -10027,6 +10028,7 @@ version = "0.1.0" dependencies = [ "codspeed-divan-compat", "prost 0.14.4", + "rstest", "vortex-array", "vortex-buffer", "vortex-error", diff --git a/encodings/dense-union/Cargo.toml b/encodings/dense-union/Cargo.toml index 2b9a4ca3e28..0865f200353 100644 --- a/encodings/dense-union/Cargo.toml +++ b/encodings/dense-union/Cargo.toml @@ -16,14 +16,15 @@ 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"] } -vortex-buffer = { workspace = true } [lints] workspace = true diff --git a/encodings/dense-union/src/array.rs b/encodings/dense-union/src/array.rs index 92cfacba282..967c8ac9635 100644 --- a/encodings/dense-union/src/array.rs +++ b/encodings/dense-union/src/array.rs @@ -1,45 +1,18 @@ // 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; @@ -58,12 +31,6 @@ pub struct DenseUnionSlots { 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. @@ -76,7 +43,7 @@ pub struct DenseUnionDataParts { pub children: Vec, } -fn make_parts( +pub(crate) fn make_parts( type_ids: ArrayRef, offsets: ArrayRef, variants: UnionVariants, @@ -125,7 +92,7 @@ pub trait DenseUnionArrayExt: DenseUnionArraySlotsExt { self.child(self.variants().tag_to_child_index(type_id)?) } - /// Return a compact child selected by variant name. + /// Return a compact child selected by variant name, if present. fn child_by_name_opt(&self, name: impl AsRef) -> Option<&ArrayRef> { self.child(self.variants().find(name)?) } @@ -166,8 +133,17 @@ impl DenseUnion { /// 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. + /// 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, @@ -198,268 +174,3 @@ impl DenseUnionArrayOwnedExt for Array { } } } - -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 index 2b5399f9c2f..f6f285461b2 100644 --- a/encodings/dense-union/src/canonical.rs +++ b/encodings/dense-union/src/canonical.rs @@ -11,6 +11,8 @@ use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::UnionArray; use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -20,6 +22,16 @@ use crate::array::DenseUnion; use crate::array::DenseUnionArrayExt; use crate::array::DenseUnionArraySlotsExt; +/// Converts a dense union to its canonical sparse representation. +/// +/// Each child is a dictionary over the original compact child, so values are not copied. Codes for +/// other variants stay zero because their type IDs make them unreachable. Unused variants use a +/// constant zero code array, and empty children use a one-value constant because dictionaries +/// require non-empty values. +/// +/// # Errors +/// +/// Returns an error for unknown type IDs, invalid offsets, or failed array construction. pub(crate) fn canonicalize( array: Array, ctx: &mut ExecutionCtx, @@ -31,29 +43,31 @@ pub(crate) fn canonicalize( 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 child_lengths = array.iter_children().map(ArrayRef::len).collect::>(); + + let mut child_indices = [None; 256]; + for (child_index, type_id) in variants.type_ids().iter().copied().enumerate() { + child_indices[usize::from(type_id)] = Some(child_index); + } + let mut codes_by_child: Vec>> = vec![None; 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) + let child_index = child_indices[usize::from(type_id)] .ok_or_else(|| vortex_err!("DenseUnion contains unknown type ID {type_id}"))?; - let offset = usize::try_from(offset_values[row]).map_err(|_| { + let offset = u32::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(); + let child_len = child_lengths[child_index]; vortex_ensure!( - offset < child_len, + (offset as usize) < 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"))?; + let codes = codes_by_child[child_index].get_or_insert_with(|| BufferMut::zeroed(len)); + codes[row] = offset; Ok(()) }; @@ -75,13 +89,18 @@ pub(crate) fn canonicalize( .iter_children() .zip(codes_by_child) .map(|(child, codes)| { + let codes = match codes { + Some(codes) => { + PrimitiveArray::new(codes.freeze(), Validity::NonNullable).into_array() + } + None => ConstantArray::new(0u32, len).into_array(), + }; 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) + DictArray::try_new(codes, values).map(IntoArray::into_array) }) .collect::>>()?; diff --git a/encodings/dense-union/src/compute/take.rs b/encodings/dense-union/src/compute/take.rs index 45035c2d615..e176c8c982b 100644 --- a/encodings/dense-union/src/compute/take.rs +++ b/encodings/dense-union/src/compute/take.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Selector-only take for dense unions. +//! +//! For performance, take gathers only the row selectors and retains every compact child in full. +//! This is O(selected rows) but may reorder per-child offsets; Arrow dense unions instead require +//! those offsets to increase. + use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; diff --git a/encodings/dense-union/src/lib.rs b/encodings/dense-union/src/lib.rs index 5c237632209..18f19ccd723 100644 --- a/encodings/dense-union/src/lib.rs +++ b/encodings/dense-union/src/lib.rs @@ -7,11 +7,16 @@ //! 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`](vortex_array::dtype::DType::Union) dtype. +//! +//! Vortex does not require offsets for each child to increase. Selector-only operations can retain +//! the original compact children and reorder their offsets, so this encoding is not necessarily a +//! directly exportable Arrow dense-union layout without compaction and offset rebasing. mod array; mod canonical; mod compute; mod rules; +mod vtable; pub use array::*; use vortex_array::session::ArraySessionExt; diff --git a/encodings/dense-union/src/tests.rs b/encodings/dense-union/src/tests.rs index 0b9f9cf1f9c..ad07828b641 100644 --- a/encodings/dense-union/src/tests.rs +++ b/encodings/dense-union/src/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -11,6 +12,7 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::UnionArray; use vortex_array::arrays::union::UnionArraySlotsExt; +use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -82,6 +84,7 @@ fn session() -> VortexSession { session } +#[track_caller] fn assert_rows( array: &ArrayRef, expected: Vec, @@ -94,19 +97,20 @@ fn assert_rows( Ok(()) } +#[track_caller] 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)? - ); + assert_eq!(left.dtype(), right.dtype()); + let left = left.clone().execute::(&mut ctx)?; + let right = right.clone().execute::(&mut ctx)?; + assert_arrays_eq!(left.type_ids(), right.type_ids(), &mut ctx); + assert_eq!(left.children().len(), right.children().len()); + for (left, right) in left.children().iter().zip(right.children().iter()) { + assert_arrays_eq!(left, right, &mut ctx); } Ok(()) } @@ -257,84 +261,96 @@ fn canonicalization_handles_unselected_empty_child() -> VortexResult<()> { assert_same_rows(&array, &canonical, &session) } -#[test] -fn invalid_type_id_and_offsets_return_errors() -> VortexResult<()> { +#[rstest] +#[case::unknown_type_id(7, 0, "unknown type ID 7")] +#[case::negative_offset(5, -1, "negative offset -1")] +#[case::out_of_bounds_offset(9, 1, "out of bounds for child 1 of length 1")] +fn invalid_type_id_and_offsets_return_errors( + #[case] type_id: u8, + #[case] offset: i32, + #[case] expected: &str, +) -> 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(), + let array = DenseUnion::try_new( + PrimitiveArray::from_iter([type_id]).into_array(), + PrimitiveArray::from_iter([offset]).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()); + let Err(error) = array.execute_scalar(0, &mut ctx) else { + panic!("DenseUnion must reject {expected}"); + }; + assert!( + error.to_string().contains(expected), + "error should mention {expected:?}, got: {error}" + ); 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() - ); +#[derive(Clone, Copy)] +enum InvalidStructure { + TypeIdsDType, + OffsetsDType, + ChildCount, +} + +#[rstest] +#[case::type_ids_dtype(InvalidStructure::TypeIdsDType, "type_ids have dtype u16")] +#[case::offsets_dtype(InvalidStructure::OffsetsDType, "offsets have dtype u32")] +#[case::child_count(InvalidStructure::ChildCount, "3 slots, expected 4")] +fn validates_structural_components( + #[case] invalid: InvalidStructure, + #[case] expected: &str, +) -> VortexResult<()> { + let type_ids = match invalid { + InvalidStructure::TypeIdsDType => PrimitiveArray::from_iter([5u16]).into_array(), + _ => PrimitiveArray::from_iter([5u8]).into_array(), + }; + let offsets = match invalid { + InvalidStructure::OffsetsDType => PrimitiveArray::from_iter([0u32]).into_array(), + _ => PrimitiveArray::from_iter([0i32]).into_array(), + }; + let mut children = vec![ + PrimitiveArray::from_iter([10i32]).into_array(), + BoolArray::from_iter([true]).into_array(), + ]; + if matches!(invalid, InvalidStructure::ChildCount) { + children.pop(); + } + + let Err(error) = DenseUnion::try_new(type_ids, offsets, variants()?, children) else { + panic!("DenseUnion must reject {expected}"); + }; 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() + error.to_string().contains(expected), + "error should mention {expected:?}, got: {error}" ); Ok(()) } +#[test] +fn canonicalization_handles_zero_length_array() -> VortexResult<()> { + let session = session(); + let array = DenseUnion::try_new( + PrimitiveArray::from_iter(Vec::::new()).into_array(), + PrimitiveArray::from_iter(Vec::::new()).into_array(), + variants()?, + vec![ + PrimitiveArray::from_iter(Vec::::new()).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 serde_roundtrip() -> VortexResult<()> { let session = session(); diff --git a/encodings/dense-union/src/vtable.rs b/encodings/dense-union/src/vtable.rs new file mode 100644 index 00000000000..f4332c2a50b --- /dev/null +++ b/encodings/dense-union/src/vtable.rs @@ -0,0 +1,324 @@ +// 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::ArrayView; +use vortex_array::EmptyArrayData; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::OperationsVTable; +use vortex_array::VTable; +use vortex_array::ValidityVTable; +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::VortexResult; +use vortex_error::vortex_bail; +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::array::DenseUnion; +use crate::array::DenseUnionArrayExt; +use crate::array::DenseUnionArraySlotsExt; +use crate::array::DenseUnionSlots; +use crate::array::make_parts; +use crate::canonical::canonicalize; +use crate::rules::PARENT_RULES; + +const OFFSETS_DTYPE: DType = DType::Primitive(PType::I32, Nullability::NonNullable); + +#[derive(Clone, prost::Message)] +struct DenseUnionMetadata { + /// The length of each compact child in variant order, used to size children during decoding. + #[prost(uint64, repeated, tag = "1")] + child_lengths: Vec, +} + +fn union_dtype(dtype: &DType) -> VortexResult<(&UnionVariants, Nullability)> { + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("DenseUnion requires a union dtype, got {dtype}"); + }; + Ok((variants, *nullability)) +} + +fn validate_components( + type_ids: &ArrayRef, + offsets: &ArrayRef, + children: &[&ArrayRef], + dtype: &DType, + len: usize, +) -> VortexResult<()> { + let (variants, nullability) = union_dtype(dtype)?; + vortex_ensure_eq!( + children.len(), + variants.len(), + "DenseUnion has {} compact children but expected {}", + children.len(), + variants.len() + ); + let expected_type_ids_dtype = DType::Primitive(PType::U8, nullability); + vortex_ensure_eq!( + type_ids.dtype(), + &expected_type_ids_dtype, + "DenseUnion type_ids have dtype {}, expected {}", + type_ids.dtype(), + expected_type_ids_dtype + ); + vortex_ensure_eq!( + type_ids.len(), + len, + "DenseUnion type_ids have length {}, expected {len}", + type_ids.len() + ); + vortex_ensure_eq!( + offsets.dtype(), + &OFFSETS_DTYPE, + "DenseUnion offsets have dtype {}, expected {OFFSETS_DTYPE}", + offsets.dtype() + ); + vortex_ensure_eq!( + offsets.len(), + len, + "DenseUnion offsets have length {}, expected {len}", + offsets.len() + ); + + for (index, (variant_dtype, child)) in variants.variants().zip(children).enumerate() { + vortex_ensure_eq!( + child.dtype(), + &variant_dtype, + "DenseUnion child {index} has dtype {}, expected {variant_dtype}", + child.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 (variants, _) = union_dtype(dtype)?; + let expected_slots = DenseUnionSlots::CHILDREN_OFFSET + variants.len(); + vortex_ensure_eq!( + slots.len(), + expected_slots, + "DenseUnion has {} slots, expected {expected_slots}", + slots.len() + ); + 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 (variants, nullability) = union_dtype(dtype)?; + let metadata = DenseUnionMetadata::decode(metadata)?; + vortex_ensure_eq!( + metadata.child_lengths.len(), + variants.len(), + "DenseUnion metadata has {} child lengths, expected {}", + metadata.child_lengths.len(), + variants.len() + ); + let expected_children = DenseUnionSlots::CHILDREN_OFFSET + variants.len(); + vortex_ensure_eq!( + children.len(), + expected_children, + "DenseUnion has {} serialized children, expected {expected_children}", + children.len() + ); + + 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/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index dd0c4d6c13c..96fe1c03720 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -206,7 +206,7 @@ fn stat_was_truncated( } fn supports_file_stats(dtype: &DType) -> bool { - !matches!(dtype, DType::Variant(_)) + !matches!(dtype, DType::Union(..) | DType::Variant(_)) } fn is_varlen_dtype(dtype: &DType) -> bool { diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 57392ed627d..efd779e0eed 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -32,6 +32,7 @@ vortex-bytebool = { workspace = true } vortex-cloud = { workspace = true, optional = 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/src/editions/tests.rs b/vortex/src/editions/tests.rs index d418a8aa412..8897bb60176 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -7,17 +7,25 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::array_session; +#[cfg(feature = "unstable_encodings")] +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; +#[cfg(feature = "unstable_encodings")] +use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +#[cfg(feature = "unstable_encodings")] +use vortex_array::dtype::UnionVariants; use vortex_array::field_path; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::ByteBufferMut; +#[cfg(feature = "unstable_encodings")] +use vortex_dense_union::DenseUnion; use vortex_edition::ComponentKind; use vortex_edition::Edition; use vortex_edition::EditionDeclaration; @@ -137,6 +145,7 @@ fn encodings_in_editions_unions_families() { assert!(both.len() > core_only.len()); assert!(both.iter().any(|id| id.as_str() == "fastlanes.delta")); + assert!(both.iter().any(|id| id.as_str() == "vortex.dense_union")); assert!(both.iter().any(|id| id.as_str() == "vortex.onpair")); assert!(core_only.iter().all(|id| both.contains(id))); } @@ -283,6 +292,46 @@ fn sequential_integers() -> PrimitiveArray { PrimitiveArray::from_iter(0..65_536i32) } +#[cfg(feature = "unstable_encodings")] +#[tokio::test] +async fn default_unstable_edition_writes_dense_union() -> VortexResult<()> { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + let variants = UnionVariants::try_new( + FieldNames::from(["number", "flag"]), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![5, 9], + )?; + let array = 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(), + ], + )? + .into_array(); + let dtype = array.dtype().clone(); + let mut buffer = ByteBufferMut::empty(); + + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut buffer, array.to_array_stream()) + .await?; + + let file = session.open_options().open_buffer(buffer)?; + assert_eq!(file.dtype(), &dtype); + let round_tripped = file.scan()?.into_array_stream()?.read_all().await?; + assert!(round_tripped.is::()); + Ok(()) +} + const WRITER_TEST_EDITION: EditionId = EditionId::new("writer-test", 2026, 7, 0); static WRITER_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { diff --git a/vortex/src/editions/unstable/v2026_06.rs b/vortex/src/editions/unstable/v2026_06.rs index 86872f488a2..5686c72905a 100644 --- a/vortex/src/editions/unstable/v2026_06.rs +++ b/vortex/src/editions/unstable/v2026_06.rs @@ -17,5 +17,8 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { id: UNSTABLE_2026_06_0, min_vortex_version: None, }, - added: &[EditionMember::array(&"vortex.onpair")], + added: &[ + EditionMember::array(&"vortex.dense_union"), + EditionMember::array(&"vortex.onpair"), + ], }; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 289500e2543..b2cf23a0a93 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -255,6 +255,11 @@ pub mod encodings { pub use vortex_decimal_byte_parts::*; } + /// Dense physical encoding for union arrays. + pub mod dense_union { + pub use vortex_dense_union::*; + } + /// FastLanes integer encodings: bit-packing, delta, frame-of-reference, and RLE. pub mod fastlanes { pub use vortex_fastlanes::*; From 95a0804a7f21ca431e5d8e6aa01e3f7d50dee787 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 14 Aug 2026 12:10:26 -0400 Subject: [PATCH 6/7] refactor: move dense union into vortex-spatial Signed-off-by: Nemo Yu --- .github/workflows/codspeed.yml | 2 +- Cargo.lock | 18 +----- Cargo.toml | 2 - encodings/dense-union/Cargo.toml | 34 ----------- encodings/dense-union/README.md | 3 - vortex-file/Cargo.toml | 1 - vortex-file/src/lib.rs | 1 - vortex-spatial/Cargo.toml | 6 ++ .../benches/dense_union_take.rs | 11 +--- .../src/dense_union}/array.rs | 58 +------------------ .../src/dense_union}/canonical.rs | 6 +- .../src/dense_union}/compute/filter.rs | 6 +- .../src/dense_union}/compute/mask.rs | 6 +- .../src/dense_union}/compute/mod.rs | 2 + .../src/dense_union}/compute/slice.rs | 6 +- .../src/dense_union}/compute/take.rs | 6 +- .../src/dense_union/mod.rs | 5 +- .../src/dense_union}/rules.rs | 2 +- .../src/dense_union}/tests.rs | 37 ++++++++++-- .../src/dense_union}/vtable.rs | 16 ++--- vortex-spatial/src/editions.rs | 17 +++++- vortex-spatial/src/lib.rs | 3 + vortex-spatial/src/test_harness.rs | 1 + vortex/Cargo.toml | 1 - vortex/src/editions/tests.rs | 49 ---------------- vortex/src/editions/unstable/v2026_06.rs | 5 +- vortex/src/lib.rs | 5 -- 27 files changed, 95 insertions(+), 214 deletions(-) delete mode 100644 encodings/dense-union/Cargo.toml delete mode 100644 encodings/dense-union/README.md rename encodings/dense-union/benches/take.rs => vortex-spatial/benches/dense_union_take.rs (93%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/array.rs (67%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/canonical.rs (97%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/compute/filter.rs (84%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/compute/mask.rs (84%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/compute/mod.rs (78%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/compute/slice.rs (84%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/compute/take.rs (90%) rename encodings/dense-union/src/lib.rs => vortex-spatial/src/dense_union/mod.rs (85%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/rules.rs (96%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/tests.rs (91%) rename {encodings/dense-union/src => vortex-spatial/src/dense_union}/vtable.rs (97%) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 39de3eadbf8..13393072b8c 100644 --- a/.github/workflows/codspeed.yml +++ b/.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-dense-union vortex-sparse vortex-zigzag vortex-zstd" } + - { shard: 7, name: "Encodings 4", packages: "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 }})" diff --git a/Cargo.lock b/Cargo.lock index fdce018d351..9529d2a3997 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9541,7 +9541,6 @@ dependencies = [ "vortex-cloud", "vortex-datetime-parts", "vortex-decimal-byte-parts", - "vortex-dense-union", "vortex-edition", "vortex-error", "vortex-fastlanes", @@ -10022,20 +10021,6 @@ dependencies = [ "vortex-session", ] -[[package]] -name = "vortex-dense-union" -version = "0.1.0" -dependencies = [ - "codspeed-divan-compat", - "prost 0.14.4", - "rstest", - "vortex-array", - "vortex-buffer", - "vortex-error", - "vortex-mask", - "vortex-session", -] - [[package]] name = "vortex-duckdb" version = "0.1.0" @@ -10166,7 +10151,6 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", - "vortex-dense-union", "vortex-edition", "vortex-error", "vortex-fastlanes", @@ -10631,6 +10615,8 @@ dependencies = [ "mimalloc", "prost 0.14.4", "rstest", + "tokio", + "vortex", "vortex-array", "vortex-arrow", "vortex-buffer", diff --git a/Cargo.toml b/Cargo.toml index 629255524b6..65452f18300 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,7 +62,6 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", - "encodings/dense-union", # Benchmarks "benchmarks/bench-support", "benchmarks/lance-bench", @@ -307,7 +306,6 @@ 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 deleted file mode 100644 index 0865f200353..00000000000 --- a/encodings/dense-union/Cargo.toml +++ /dev/null @@ -1,34 +0,0 @@ -[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 diff --git a/encodings/dense-union/README.md b/encodings/dense-union/README.md deleted file mode 100644 index 79132dee06c..00000000000 --- a/encodings/dense-union/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Vortex Dense Union - -An external dense physical encoding for Vortex's logical `DType::Union`. diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 92767259977..72a47d88f50 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -40,7 +40,6 @@ 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 11c113ab837..25760f0890e 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -189,7 +189,6 @@ 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); diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 3be2b2d9d66..508aedb88d1 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -39,6 +39,8 @@ _test-harness = [] divan = { workspace = true } mimalloc = { workspace = true } rstest = { workspace = true } +tokio = { workspace = true, features = ["full"] } +vortex = { workspace = true, features = ["files", "tokio"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-layout = { workspace = true } vortex-spatial = { path = ".", features = ["_test-harness"] } @@ -79,5 +81,9 @@ harness = false name = "convex_hull" harness = false +[[bench]] +name = "dense_union_take" +harness = false + [lints] workspace = true diff --git a/encodings/dense-union/benches/take.rs b/vortex-spatial/benches/dense_union_take.rs similarity index 93% rename from encodings/dense-union/benches/take.rs rename to vortex-spatial/benches/dense_union_take.rs index f239240f24f..f6023d77feb 100644 --- a/encodings/dense-union/benches/take.rs +++ b/vortex-spatial/benches/dense_union_take.rs @@ -11,7 +11,6 @@ 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; @@ -19,9 +18,9 @@ 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; +use vortex_spatial::test_harness::DenseUnion; +use vortex_spatial::test_harness::spatial_session; const LEN: usize = 65_536; const N_VARIANTS: usize = 28; @@ -32,11 +31,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(spatial_session); fn variants() -> UnionVariants { let names = FieldNames::from_iter((0..N_VARIANTS).map(|index| format!("variant_{index}"))); diff --git a/encodings/dense-union/src/array.rs b/vortex-spatial/src/dense_union/array.rs similarity index 67% rename from encodings/dense-union/src/array.rs rename to vortex-spatial/src/dense_union/array.rs index 967c8ac9635..2c5b4915851 100644 --- a/encodings/dense-union/src/array.rs +++ b/vortex-spatial/src/dense_union/array.rs @@ -12,13 +12,13 @@ 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. +/// A Vortex array encoded as a [`DenseUnion`]. pub type DenseUnionArray = Array; /// Slot layout of a dense union array. #[array_slots(DenseUnion)] +#[allow(dead_code)] pub struct DenseUnionSlots { /// The row-aligned type IDs selecting union variants. #[slot(0)] @@ -31,18 +31,6 @@ pub struct DenseUnionSlots { pub children: 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, -} - pub(crate) fn make_parts( type_ids: ArrayRef, offsets: ArrayRef, @@ -86,27 +74,6 @@ pub trait DenseUnionArrayExt: DenseUnionArraySlotsExt { 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) -> 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 {} @@ -153,24 +120,3 @@ impl DenseUnion { 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, - } - } -} diff --git a/encodings/dense-union/src/canonical.rs b/vortex-spatial/src/dense_union/canonical.rs similarity index 97% rename from encodings/dense-union/src/canonical.rs rename to vortex-spatial/src/dense_union/canonical.rs index f6f285461b2..70f133d9869 100644 --- a/encodings/dense-union/src/canonical.rs +++ b/vortex-spatial/src/dense_union/canonical.rs @@ -18,9 +18,9 @@ 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; +use super::array::DenseUnion; +use super::array::DenseUnionArrayExt; +use super::array::DenseUnionArraySlotsExt; /// Converts a dense union to its canonical sparse representation. /// diff --git a/encodings/dense-union/src/compute/filter.rs b/vortex-spatial/src/dense_union/compute/filter.rs similarity index 84% rename from encodings/dense-union/src/compute/filter.rs rename to vortex-spatial/src/dense_union/compute/filter.rs index 60fb8181bf2..e313411b181 100644 --- a/encodings/dense-union/src/compute/filter.rs +++ b/vortex-spatial/src/dense_union/compute/filter.rs @@ -8,9 +8,9 @@ use vortex_array::arrays::filter::FilterReduce; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::DenseUnion; -use crate::DenseUnionArrayExt; -use crate::DenseUnionArraySlotsExt; +use crate::dense_union::DenseUnion; +use crate::dense_union::DenseUnionArrayExt; +use crate::dense_union::DenseUnionArraySlotsExt; impl FilterReduce for DenseUnion { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { diff --git a/encodings/dense-union/src/compute/mask.rs b/vortex-spatial/src/dense_union/compute/mask.rs similarity index 84% rename from encodings/dense-union/src/compute/mask.rs rename to vortex-spatial/src/dense_union/compute/mask.rs index 85f6489460a..59e7b3eb33a 100644 --- a/encodings/dense-union/src/compute/mask.rs +++ b/vortex-spatial/src/dense_union/compute/mask.rs @@ -8,9 +8,9 @@ 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; +use crate::dense_union::DenseUnion; +use crate::dense_union::DenseUnionArrayExt; +use crate::dense_union::DenseUnionArraySlotsExt; impl MaskReduce for DenseUnion { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { diff --git a/encodings/dense-union/src/compute/mod.rs b/vortex-spatial/src/dense_union/compute/mod.rs similarity index 78% rename from encodings/dense-union/src/compute/mod.rs rename to vortex-spatial/src/dense_union/compute/mod.rs index a9deb01f82e..68f0709a964 100644 --- a/encodings/dense-union/src/compute/mod.rs +++ b/vortex-spatial/src/dense_union/compute/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Compute kernels for dense unions. + mod filter; mod mask; mod slice; diff --git a/encodings/dense-union/src/compute/slice.rs b/vortex-spatial/src/dense_union/compute/slice.rs similarity index 84% rename from encodings/dense-union/src/compute/slice.rs rename to vortex-spatial/src/dense_union/compute/slice.rs index 72a1d2c3113..2715c78770a 100644 --- a/encodings/dense-union/src/compute/slice.rs +++ b/vortex-spatial/src/dense_union/compute/slice.rs @@ -9,9 +9,9 @@ use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; -use crate::DenseUnion; -use crate::DenseUnionArrayExt; -use crate::DenseUnionArraySlotsExt; +use crate::dense_union::DenseUnion; +use crate::dense_union::DenseUnionArrayExt; +use crate::dense_union::DenseUnionArraySlotsExt; impl SliceReduce for DenseUnion { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/dense-union/src/compute/take.rs b/vortex-spatial/src/dense_union/compute/take.rs similarity index 90% rename from encodings/dense-union/src/compute/take.rs rename to vortex-spatial/src/dense_union/compute/take.rs index e176c8c982b..ab317dde667 100644 --- a/encodings/dense-union/src/compute/take.rs +++ b/vortex-spatial/src/dense_union/compute/take.rs @@ -15,9 +15,9 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; -use crate::DenseUnion; -use crate::DenseUnionArrayExt; -use crate::DenseUnionArraySlotsExt; +use crate::dense_union::DenseUnion; +use crate::dense_union::DenseUnionArrayExt; +use crate::dense_union::DenseUnionArraySlotsExt; impl TakeReduce for DenseUnion { fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { diff --git a/encodings/dense-union/src/lib.rs b/vortex-spatial/src/dense_union/mod.rs similarity index 85% rename from encodings/dense-union/src/lib.rs rename to vortex-spatial/src/dense_union/mod.rs index 18f19ccd723..aa9c91edf4c 100644 --- a/encodings/dense-union/src/lib.rs +++ b/vortex-spatial/src/dense_union/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A dense physical encoding for Vortex union arrays. +//! A dense physical encoding for spatial 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 @@ -22,8 +22,7 @@ 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) { +pub(crate) fn initialize(session: &VortexSession) { session.arrays().register(DenseUnion); } diff --git a/encodings/dense-union/src/rules.rs b/vortex-spatial/src/dense_union/rules.rs similarity index 96% rename from encodings/dense-union/src/rules.rs rename to vortex-spatial/src/dense_union/rules.rs index 0afb9ecfa65..c08cc74e768 100644 --- a/encodings/dense-union/src/rules.rs +++ b/vortex-spatial/src/dense_union/rules.rs @@ -7,7 +7,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; -use crate::DenseUnion; +use super::DenseUnion; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(DenseUnion)), diff --git a/encodings/dense-union/src/tests.rs b/vortex-spatial/src/dense_union/tests.rs similarity index 91% rename from encodings/dense-union/src/tests.rs rename to vortex-spatial/src/dense_union/tests.rs index ad07828b641..ea9129aad21 100644 --- a/encodings/dense-union/src/tests.rs +++ b/vortex-spatial/src/dense_union/tests.rs @@ -1,7 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use rstest::rstest; +use vortex::VortexSessionDefault; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -21,16 +27,16 @@ use vortex_array::dtype::UnionVariants; use vortex_array::scalar::Scalar; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; +use vortex_array::stream::ArrayStreamExt; 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; +use super::DenseUnion; +use super::DenseUnionArray; +use super::DenseUnionArraySlotsExt; fn variants() -> VortexResult { UnionVariants::try_new( @@ -80,7 +86,7 @@ fn nullable_dense_union() -> VortexResult { fn session() -> VortexSession { let session = vortex_array::array_session(); - initialize(&session); + crate::initialize(&session); session } @@ -373,3 +379,24 @@ fn serde_roundtrip() -> VortexResult<()> { assert!(decoded.is::()); assert_same_rows(&array, &decoded, &session) } + +#[tokio::test] +async fn file_roundtrip_preserves_dense_union() -> VortexResult<()> { + let session = VortexSession::default(); + crate::initialize(&session); + let array = nullable_dense_union()?.into_array(); + let dtype = array.dtype().clone(); + let mut buffer = ByteBufferMut::empty(); + + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut buffer, array.clone().to_array_stream()) + .await?; + + let file = session.open_options().open_buffer(buffer)?; + assert_eq!(file.dtype(), &dtype); + let round_tripped = file.scan()?.into_array_stream()?.read_all().await?; + assert!(round_tripped.is::()); + assert_same_rows(&array, &round_tripped, &session) +} diff --git a/encodings/dense-union/src/vtable.rs b/vortex-spatial/src/dense_union/vtable.rs similarity index 97% rename from encodings/dense-union/src/vtable.rs rename to vortex-spatial/src/dense_union/vtable.rs index f4332c2a50b..872db7f2951 100644 --- a/encodings/dense-union/src/vtable.rs +++ b/vortex-spatial/src/dense_union/vtable.rs @@ -33,13 +33,13 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::array::DenseUnion; -use crate::array::DenseUnionArrayExt; -use crate::array::DenseUnionArraySlotsExt; -use crate::array::DenseUnionSlots; -use crate::array::make_parts; -use crate::canonical::canonicalize; -use crate::rules::PARENT_RULES; +use super::array::DenseUnion; +use super::array::DenseUnionArrayExt; +use super::array::DenseUnionArraySlotsExt; +use super::array::DenseUnionSlots; +use super::array::make_parts; +use super::canonical::canonicalize; +use super::rules::PARENT_RULES; const OFFSETS_DTYPE: DType = DType::Primitive(PType::I32, Nullability::NonNullable); @@ -117,7 +117,7 @@ impl VTable for DenseUnion { type ValidityVTable = Self; fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.dense_union"); + static ID: CachedId = CachedId::new("vortex.st.dense_union"); *ID } diff --git a/vortex-spatial/src/editions.rs b/vortex-spatial/src/editions.rs index 6856e05bb04..2b22447cb82 100644 --- a/vortex-spatial/src/editions.rs +++ b/vortex-spatial/src/editions.rs @@ -26,7 +26,10 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { id: SPATIAL_2026_08, min_vortex_version: None, }, - added: &[EditionMember::aggregate(&"vortex.st.aabb")], + added: &[ + EditionMember::array(&"vortex.st.dense_union"), + EditionMember::aggregate(&"vortex.st.aabb"), + ], }; #[cfg(test)] @@ -55,4 +58,16 @@ mod tests { "spatial session permits {enabled:?}" ); } + + #[test] + fn initialize_permits_dense_union() { + let session = crate::test_harness::spatial_session(); + let enabled = session.enabled_component_ids(ComponentKind::Array); + assert!( + enabled + .iter() + .any(|id| id.as_str() == "vortex.st.dense_union"), + "spatial session permits {enabled:?}" + ); + } } diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 919652e01c4..65e71270422 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -35,6 +35,7 @@ use crate::scalar_fn::length::SpatialLength; use crate::scalar_fn::make_line::SpatialMakeLine; pub mod aggregate_fn; +mod dense_union; pub mod editions; pub mod extension; pub mod prune; @@ -46,6 +47,8 @@ mod tests; /// Set up a session with support for spatial extension types, encodings and layouts. pub fn initialize(session: &VortexSession) { + dense_union::initialize(session); + // Register the spatial extension types. session.dtypes().register(WellKnownBinary); session.arrow().register_exporter(Arc::new(WellKnownBinary)); diff --git a/vortex-spatial/src/test_harness.rs b/vortex-spatial/src/test_harness.rs index 7b471bdf2c4..af9afb6b23a 100644 --- a/vortex-spatial/src/test_harness.rs +++ b/vortex-spatial/src/test_harness.rs @@ -20,6 +20,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; +pub use crate::dense_union::DenseUnion; use crate::extension::LineString; use crate::extension::MultiLineString; use crate::extension::MultiPoint; diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index efd779e0eed..57392ed627d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -32,7 +32,6 @@ vortex-bytebool = { workspace = true } vortex-cloud = { workspace = true, optional = 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/src/editions/tests.rs b/vortex/src/editions/tests.rs index 8897bb60176..d418a8aa412 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -7,25 +7,17 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::array_session; -#[cfg(feature = "unstable_encodings")] -use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; -#[cfg(feature = "unstable_encodings")] -use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -#[cfg(feature = "unstable_encodings")] -use vortex_array::dtype::UnionVariants; use vortex_array::field_path; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::ByteBufferMut; -#[cfg(feature = "unstable_encodings")] -use vortex_dense_union::DenseUnion; use vortex_edition::ComponentKind; use vortex_edition::Edition; use vortex_edition::EditionDeclaration; @@ -145,7 +137,6 @@ fn encodings_in_editions_unions_families() { assert!(both.len() > core_only.len()); assert!(both.iter().any(|id| id.as_str() == "fastlanes.delta")); - assert!(both.iter().any(|id| id.as_str() == "vortex.dense_union")); assert!(both.iter().any(|id| id.as_str() == "vortex.onpair")); assert!(core_only.iter().all(|id| both.contains(id))); } @@ -292,46 +283,6 @@ fn sequential_integers() -> PrimitiveArray { PrimitiveArray::from_iter(0..65_536i32) } -#[cfg(feature = "unstable_encodings")] -#[tokio::test] -async fn default_unstable_edition_writes_dense_union() -> VortexResult<()> { - use crate::VortexSessionDefault; - - let session = VortexSession::default(); - let variants = UnionVariants::try_new( - FieldNames::from(["number", "flag"]), - vec![ - DType::Primitive(PType::I32, Nullability::NonNullable), - DType::Bool(Nullability::NonNullable), - ], - vec![5, 9], - )?; - let array = 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(), - ], - )? - .into_array(); - let dtype = array.dtype().clone(); - let mut buffer = ByteBufferMut::empty(); - - session - .write_options() - .with_strategy(Arc::new(FlatLayoutStrategy::default())) - .write(&mut buffer, array.to_array_stream()) - .await?; - - let file = session.open_options().open_buffer(buffer)?; - assert_eq!(file.dtype(), &dtype); - let round_tripped = file.scan()?.into_array_stream()?.read_all().await?; - assert!(round_tripped.is::()); - Ok(()) -} - const WRITER_TEST_EDITION: EditionId = EditionId::new("writer-test", 2026, 7, 0); static WRITER_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { diff --git a/vortex/src/editions/unstable/v2026_06.rs b/vortex/src/editions/unstable/v2026_06.rs index 5686c72905a..86872f488a2 100644 --- a/vortex/src/editions/unstable/v2026_06.rs +++ b/vortex/src/editions/unstable/v2026_06.rs @@ -17,8 +17,5 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { id: UNSTABLE_2026_06_0, min_vortex_version: None, }, - added: &[ - EditionMember::array(&"vortex.dense_union"), - EditionMember::array(&"vortex.onpair"), - ], + added: &[EditionMember::array(&"vortex.onpair")], }; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index b2cf23a0a93..289500e2543 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -255,11 +255,6 @@ pub mod encodings { pub use vortex_decimal_byte_parts::*; } - /// Dense physical encoding for union arrays. - pub mod dense_union { - pub use vortex_dense_union::*; - } - /// FastLanes integer encodings: bit-packing, delta, frame-of-reference, and RLE. pub mod fastlanes { pub use vortex_fastlanes::*; From 7459934dae98ee484f6e5892580fc28f122c69bf Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 14 Aug 2026 13:27:54 -0400 Subject: [PATCH 7/7] fix: remove unused dense union constructor Signed-off-by: Nemo Yu --- vortex-spatial/src/dense_union/array.rs | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/vortex-spatial/src/dense_union/array.rs b/vortex-spatial/src/dense_union/array.rs index 2c5b4915851..552c20306a1 100644 --- a/vortex-spatial/src/dense_union/array.rs +++ b/vortex-spatial/src/dense_union/array.rs @@ -10,7 +10,6 @@ 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; /// A Vortex array encoded as a [`DenseUnion`]. @@ -83,21 +82,6 @@ impl> DenseUnionArrayExt for T {} 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. /// /// The logical union's nullability is inherited from `type_ids`; nullable type IDs represent