diff --git a/Cargo.lock b/Cargo.lock index b0f98336b0b..9529d2a3997 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10615,6 +10615,8 @@ dependencies = [ "mimalloc", "prost 0.14.4", "rstest", + "tokio", + "vortex", "vortex-array", "vortex-arrow", "vortex-buffer", 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-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/vortex-spatial/benches/dense_union_take.rs b/vortex-spatial/benches/dense_union_take.rs new file mode 100644 index 00000000000..f6023d77feb --- /dev/null +++ b/vortex-spatial/benches/dense_union_take.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::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_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; +const TAKE_LEN: usize = 4_096; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +fn variants() -> UnionVariants { + let names = FieldNames::from_iter((0..N_VARIANTS).map(|index| format!("variant_{index}"))); + let dtypes = vec![DType::Primitive(PType::I32, Nullability::NonNullable); N_VARIANTS]; + let type_ids = (1..=N_VARIANTS).map(|type_id| type_id as u8).collect(); + UnionVariants::try_new(names, dtypes, type_ids).unwrap() +} + +fn selectors() -> (ArrayRef, ArrayRef, Vec) { + 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() +} + +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) { + bench_take(bencher, dense_union(), indices()); +} + +#[divan::bench] +fn sparse_take(bencher: Bencher) { + bench_take(bencher, sparse_union(), indices()); +} diff --git a/vortex-spatial/src/dense_union/array.rs b/vortex-spatial/src/dense_union/array.rs new file mode 100644 index 00000000000..552c20306a1 --- /dev/null +++ b/vortex-spatial/src/dense_union/array.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Array; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::EmptyArrayData; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::dtype::DType; +use vortex_array::dtype::UnionVariants; +use vortex_error::VortexResult; + +/// 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)] + 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, +} + +pub(crate) 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) + } +} + +impl> DenseUnionArrayExt for T {} + +/// The dense physical encoding for the logical [`DType::Union`] type. +#[derive(Clone, Debug)] +pub struct DenseUnion; + +impl DenseUnion { + /// Try to construct a dense union array. + /// + /// The logical union's nullability is inherited from `type_ids`; nullable type IDs represent + /// outer union nulls. `type_ids` must be a nullable or non-nullable `u8` array, `offsets` must + /// be a non-nullable `i32` array of the same length, and the compact children must match the + /// variant count and dtypes. Type IDs and offsets are structurally validated, but their + /// individual values are checked only when the array is accessed or converted to its canonical + /// sparse representation. + /// + /// # Errors + /// + /// Returns an error when the selector arrays or compact children do not satisfy these + /// structural invariants. + pub fn try_new( + type_ids: ArrayRef, + offsets: ArrayRef, + variants: UnionVariants, + children: impl IntoIterator, + ) -> VortexResult { + Array::try_from_parts(make_parts(type_ids, offsets, variants, children)) + } +} diff --git a/vortex-spatial/src/dense_union/canonical.rs b/vortex-spatial/src/dense_union/canonical.rs new file mode 100644 index 00000000000..70f133d9869 --- /dev/null +++ b/vortex-spatial/src/dense_union/canonical.rs @@ -0,0 +1,109 @@ +// 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_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; + +use super::array::DenseUnion; +use super::array::DenseUnionArrayExt; +use super::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, +) -> 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 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 = child_indices[usize::from(type_id)] + .ok_or_else(|| vortex_err!("DenseUnion contains unknown type ID {type_id}"))?; + 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 = child_lengths[child_index]; + vortex_ensure!( + (offset as usize) < child_len, + "DenseUnion offset {offset} is out of bounds for child {child_index} of length {child_len}" + ); + let codes = codes_by_child[child_index].get_or_insert_with(|| BufferMut::zeroed(len)); + codes[row] = offset; + 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 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(codes, values).map(IntoArray::into_array) + }) + .collect::>>()?; + + UnionArray::try_new(type_ids.array().clone(), variants, sparse_children) + .map(IntoArray::into_array) +} diff --git a/vortex-spatial/src/dense_union/compute/filter.rs b/vortex-spatial/src/dense_union/compute/filter.rs new file mode 100644 index 00000000000..e313411b181 --- /dev/null +++ b/vortex-spatial/src/dense_union/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::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> { + 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/vortex-spatial/src/dense_union/compute/mask.rs b/vortex-spatial/src/dense_union/compute/mask.rs new file mode 100644 index 00000000000..59e7b3eb33a --- /dev/null +++ b/vortex-spatial/src/dense_union/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::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> { + 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/vortex-spatial/src/dense_union/compute/mod.rs b/vortex-spatial/src/dense_union/compute/mod.rs new file mode 100644 index 00000000000..68f0709a964 --- /dev/null +++ b/vortex-spatial/src/dense_union/compute/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compute kernels for dense unions. + +mod filter; +mod mask; +mod slice; +mod take; diff --git a/vortex-spatial/src/dense_union/compute/slice.rs b/vortex-spatial/src/dense_union/compute/slice.rs new file mode 100644 index 00000000000..2715c78770a --- /dev/null +++ b/vortex-spatial/src/dense_union/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::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> { + 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/vortex-spatial/src/dense_union/compute/take.rs b/vortex-spatial/src/dense_union/compute/take.rs new file mode 100644 index 00000000000..ab317dde667 --- /dev/null +++ b/vortex-spatial/src/dense_union/compute/take.rs @@ -0,0 +1,37 @@ +// 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; +use vortex_array::arrays::dict::TakeReduce; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +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> { + 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/vortex-spatial/src/dense_union/mod.rs b/vortex-spatial/src/dense_union/mod.rs new file mode 100644 index 00000000000..aa9c91edf4c --- /dev/null +++ b/vortex-spatial/src/dense_union/mod.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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 +//! 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; +use vortex_session::VortexSession; + +pub(crate) fn initialize(session: &VortexSession) { + session.arrays().register(DenseUnion); +} + +#[cfg(test)] +mod tests; diff --git a/vortex-spatial/src/dense_union/rules.rs b/vortex-spatial/src/dense_union/rules.rs new file mode 100644 index 00000000000..c08cc74e768 --- /dev/null +++ b/vortex-spatial/src/dense_union/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 super::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/vortex-spatial/src/dense_union/tests.rs b/vortex-spatial/src/dense_union/tests.rs new file mode 100644 index 00000000000..ea9129aad21 --- /dev/null +++ b/vortex-spatial/src/dense_union/tests.rs @@ -0,0 +1,402 @@ +// 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; +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::assert_arrays_eq; +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_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 super::DenseUnion; +use super::DenseUnionArray; +use super::DenseUnionArraySlotsExt; + +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(); + crate::initialize(&session); + session +} + +#[track_caller] +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(()) +} + +#[track_caller] +fn assert_same_rows( + left: &ArrayRef, + right: &ArrayRef, + session: &VortexSession, +) -> VortexResult<()> { + let mut ctx = session.create_execution_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(()) +} + +#[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) +} + +#[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 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(), + ], + )?; + 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(()) +} + +#[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!( + 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(); + 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) +} + +#[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/vortex-spatial/src/dense_union/vtable.rs b/vortex-spatial/src/dense_union/vtable.rs new file mode 100644 index 00000000000..872db7f2951 --- /dev/null +++ b/vortex-spatial/src/dense_union/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 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); + +#[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.st.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-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;