From a217552ab341a487077aec9d51fb0b0d8fb1e0c2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 19:46:49 +0000 Subject: [PATCH 1/7] feat: implement take for canonical sparse Union arrays Registers TakeReduce for Union and wires the union branch of the dictionary execution path, removing both TODOs. Take gathers the type IDs with the original indices so a null index becomes an outer union null, and gathers every sparse child with the null indices filled in so each child keeps its declared variant dtype. Every child has to be visited because sparse children are row-aligned with the union, so take costs O(variants * indices). Also adds UnionArray::constant, which take needs for an empty source and which lets constant Union arrays canonicalize into a sparse union whose unselected children hold their variant's default value. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/Cargo.toml | 4 + vortex-array/benches/take_union.rs | 105 ++++++++++ .../src/arrays/constant/vtable/canonical.rs | 9 +- vortex-array/src/arrays/dict/execute.rs | 18 +- vortex-array/src/arrays/union/array.rs | 49 +++++ vortex-array/src/arrays/union/compute/mod.rs | 1 + .../src/arrays/union/compute/rules.rs | 5 +- vortex-array/src/arrays/union/compute/take.rs | 53 +++++ vortex-array/src/arrays/union/mod.rs | 3 + .../arrays/union/{tests.rs => tests/mod.rs} | 34 +++ vortex-array/src/arrays/union/tests/take.rs | 194 ++++++++++++++++++ 11 files changed, 459 insertions(+), 16 deletions(-) create mode 100644 vortex-array/benches/take_union.rs create mode 100644 vortex-array/src/arrays/union/compute/take.rs rename vortex-array/src/arrays/union/{tests.rs => tests/mod.rs} (89%) create mode 100644 vortex-array/src/arrays/union/tests/take.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 22a8768af5d..e31f2e41aa3 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -243,6 +243,10 @@ harness = false name = "take_struct" harness = false +[[bench]] +name = "take_union" +harness = false + [[bench]] name = "take_fsl" harness = false diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs new file mode 100644 index 00000000000..d82c5faa9cd --- /dev/null +++ b/vortex-array/benches/take_union.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Take on a canonical sparse union. +//! +//! A sparse union keeps every child row-aligned with the union, so take gathers all of them and +//! costs `O(variants * indices)`. The variant count is the axis worth measuring, so these +//! benchmarks pin the array and index counts and sweep it. The nullable-indices case additionally +//! pays for the fill-null pass that keeps the children off the union's outer nullability. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::UnionArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::UnionVariants; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const ARRAY_SIZE: usize = 100_000; +const TAKE_SIZE: usize = 1_000; +const VARIANT_COUNTS: [usize; 4] = [2, 4, 8, 16]; + +/// A sparse union of `variant_count` `i64` variants whose type IDs cycle through every variant. +fn union_array(variant_count: usize, rng: &mut StdRng) -> ArrayRef { + let names: FieldNames = (0..variant_count).map(|i| format!("v{i}")).collect(); + let dtypes = vec![DType::Primitive(PType::I64, Nullability::NonNullable); variant_count]; + let variants = UnionVariants::new(names, dtypes).unwrap(); + + let type_ids = PrimitiveArray::from_iter( + (0..ARRAY_SIZE).map(|i| u8::try_from(i % variant_count).unwrap()), + ); + let children = (0..variant_count) + .map(|_| { + (0..ARRAY_SIZE) + .map(|_| rng.random::()) + .collect::>() + .into_array() + }) + .collect::>(); + + UnionArray::new(type_ids.into_array(), variants, children).into_array() +} + +#[divan::bench(args = VARIANT_COUNTS)] +fn take_union(bencher: Bencher, variant_count: usize) { + let mut rng = StdRng::seed_from_u64(0); + let array = union_array(variant_count, &mut rng); + + let indices = (0..TAKE_SIZE) + .map(|_| rng.random_range(0..ARRAY_SIZE) as u64) + .collect::>() + .into_array(); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + }); +} + +#[divan::bench(args = VARIANT_COUNTS)] +fn take_union_nullable_indices(bencher: Bencher, variant_count: usize) { + let mut rng = StdRng::seed_from_u64(0); + let array = union_array(variant_count, &mut rng); + + // Every tenth index is null, which is what turns a gathered row into an outer union null. + let indices = PrimitiveArray::from_option_iter( + (0..TAKE_SIZE).map(|i| (i % 10 != 0).then(|| rng.random_range(0..ARRAY_SIZE) as u64)), + ) + .into_array(); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + }); +} diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 6929f30f815..a668831e224 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -24,6 +24,7 @@ use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; +use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::arrays::varbinview::BinaryView; @@ -175,13 +176,7 @@ pub(crate) fn constant_canonicalize( StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) }) } - DType::Union(..) => { - todo!( - "TODO(connor)[Union]: canonicalize constant Union arrays in a focused follow-up \ - after defining placeholder values for every inactive sparse child, including \ - nested Struct and Union variants" - ) - } + DType::Union(..) => Canonical::Union(UnionArray::constant(scalar, array.len())?), DType::Variant(_) => Canonical::Variant(VariantArray::try_new( array.array().clone().into_array(), None, diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index a4afb7edf20..098ce932e04 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -28,6 +28,8 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::StructArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; @@ -56,12 +58,7 @@ pub(crate) fn take_canonical( Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx)) } CanonicalView::Struct(a) => Canonical::Struct(take_struct(a, codes)), - CanonicalView::Union(_) => { - todo!( - "TODO(connor)[Union]: implement dictionary execution after Union take supports \ - nullable indices and outer null propagation" - ) - } + CanonicalView::Union(a) => Canonical::Union(take_union(a, codes)), CanonicalView::Extension(a) => Canonical::Extension(take_extension(a, codes, ctx)), CanonicalView::Variant(a) => { let indices = codes.array().clone(); @@ -178,6 +175,15 @@ fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) -> .into_owned() } +fn take_union(array: ArrayView<'_, Union>, codes: ArrayView<'_, Primitive>) -> UnionArray { + let codes_ref = codes.array(); + ::take(array, codes_ref) + .vortex_expect("take union array") + .vortex_expect("take union should not return None") + .as_::() + .into_owned() +} + fn take_extension( array: ArrayView<'_, Extension>, codes: ArrayView<'_, Primitive>, diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs index abd9aead436..0d032a1f1e2 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -14,12 +14,15 @@ use crate::array::ArrayParts; use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; use crate::array_slots; +use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Union; +use crate::arrays::union::union_type_ids_dtype; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::UnionVariants; +use crate::scalar::Scalar; /// Slot layout of a canonical sparse union array. #[array_slots(Union)] @@ -174,6 +177,52 @@ impl Array { } } + /// Construct a `len`-row union in which every row holds `scalar`. + /// + /// A sparse union keeps every child row-aligned even though at most one of them is selected, + /// so the unselected children are filled with their variant's default value: a null for a + /// nullable variant and a zero for a non-nullable one. An outer null `scalar` therefore + /// produces null type IDs and a placeholder for every child. + /// + /// # Errors + /// + /// Returns an error if `scalar` does not have a union dtype. + pub fn constant(scalar: &Scalar, len: usize) -> VortexResult { + let union = scalar + .as_union_opt() + .ok_or_else(|| vortex_err!("Expected a union scalar, got {}", scalar.dtype()))?; + let variants = union.variants().clone(); + let nullability = union.nullability(); + + let type_ids = match union.type_id() { + Some(type_id) => Scalar::primitive(type_id, nullability), + None => Scalar::null(union_type_ids_dtype(nullability)), + }; + + // The selected variant carries the scalar's value. Every other child is a placeholder that + // exists only to keep the sparse layout row-aligned. + let selected = union.child_index().zip(union.child()); + + let children = variants + .variants() + .enumerate() + .map(|(index, dtype)| { + let value = match &selected { + Some((selected_index, child)) if *selected_index == index => child.clone(), + _ => Scalar::default_value(&dtype), + }; + + ConstantArray::new(value, len).into_array() + }) + .collect::>(); + + Self::try_new( + ConstantArray::new(type_ids, len).into_array(), + variants, + children, + ) + } + /// Create an empty array for a union dtype. pub(crate) fn empty(variants: UnionVariants, nullability: Nullability) -> Self { let type_ids = PrimitiveArray::empty::(nullability).into_array(); diff --git a/vortex-array/src/arrays/union/compute/mod.rs b/vortex-array/src/arrays/union/compute/mod.rs index 965f5c41e12..b4143ad873f 100644 --- a/vortex-array/src/arrays/union/compute/mod.rs +++ b/vortex-array/src/arrays/union/compute/mod.rs @@ -5,3 +5,4 @@ pub(crate) mod rules; mod mask; mod slice; +mod take; diff --git a/vortex-array/src/arrays/union/compute/rules.rs b/vortex-array/src/arrays/union/compute/rules.rs index 92a558fc9a9..921d66a24d3 100644 --- a/vortex-array/src/arrays/union/compute/rules.rs +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use crate::arrays::Union; +use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; use crate::scalar_fn::fns::mask::MaskReduceAdaptor; @@ -9,7 +10,5 @@ use crate::scalar_fn::fns::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Union)), ParentRuleSet::lift(&SliceReduceAdaptor(Union)), + ParentRuleSet::lift(&TakeReduceAdaptor(Union)), ]); - -// TODO(connor)[Union]: Register TakeReduce only once nullable indices can introduce outer Union -// nulls while preserving the sparse children's dtypes and row alignment. diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs new file mode 100644 index 00000000000..56be3889a33 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; +use crate::builtins::ArrayBuiltins; +use crate::scalar::Scalar; + +/// Structural take for [`UnionArray`]: gathers the type IDs and every sparse child at `indices`. +/// +/// A sparse union keeps every child row-aligned with the union, so a gather must visit all of them +/// even though at most one is active per row. Take therefore costs `O(variants * indices)`. +/// Reducing that cost requires the dense union encoding, not a different sparse gather. +/// +/// The type IDs carry the union's validity, so gathering them with the original `indices` is what +/// turns a null index into an outer union null. The children are gathered with the null indices +/// filled in, which keeps each child's dtype exactly as the variant schema declares it. +impl TakeReduce for Union { + fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { + // An empty union has no row for a child to point at, so the only legal indices are all + // null and every output row is an outer union null. + if array.is_empty() { + return UnionArray::constant(&Scalar::null(array.dtype().as_nullable()), indices.len()) + .map(UnionArray::into_array) + .map(Some); + } + + let type_ids = array.type_ids().take(indices.clone())?; + + // Nullability is stripped so that the children keep their declared variant dtypes. The + // type IDs already record which rows are null. + let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable()); + let child_indices = indices.clone().fill_null(fill_scalar)?; + + let children: Vec = array + .iter_children() + .map(|child| child.take(child_indices.clone())) + .try_collect()?; + + UnionArray::try_new(type_ids, array.variants().clone(), children) + .map(UnionArray::into_array) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs index 33a906d5cae..c59a33c391c 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -10,6 +10,9 @@ //! //! Type ID values are not validated during construction. Accessing a non-null row whose type ID //! is not declared by the union variants will panic. +//! +//! Slice, mask, and take are structural: each rewrites the type IDs and every child, so each costs +//! `O(variants)` child operations. use crate::dtype::DType; use crate::dtype::Nullability; diff --git a/vortex-array/src/arrays/union/tests.rs b/vortex-array/src/arrays/union/tests/mod.rs similarity index 89% rename from vortex-array/src/arrays/union/tests.rs rename to vortex-array/src/arrays/union/tests/mod.rs index be82654e17a..75aec48a656 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -9,10 +9,12 @@ use vortex_mask::Mask; use vortex_session::registry::ReadContext; use crate::ArrayContext; +use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Union; use crate::arrays::UnionArray; @@ -27,6 +29,8 @@ use crate::serde::SerializeOptions; use crate::serde::SerializedArray; use crate::validity::Validity; +mod take; + fn variants() -> VortexResult { UnionVariants::try_new( ["number", "flag"].into(), @@ -234,6 +238,36 @@ fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> { Ok(()) } +#[test] +fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Union(variants()?, Nullability::Nullable); + + for scalar in [ + Scalar::union(variants()?, 9, true.into(), Nullability::Nullable)?, + Scalar::null(dtype.clone()), + ] { + let canonical = ConstantArray::new(scalar.clone(), 3) + .into_array() + .execute::(&mut ctx)? + .into_union(); + + assert_eq!(canonical.dtype(), &dtype); + + // The unselected variant is only a placeholder, so it keeps its declared dtype. + assert_eq!( + canonical.child_by_name("number")?.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + + for index in 0..canonical.len() { + assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); + } + } + + Ok(()) +} + #[test] fn serde_roundtrip() -> VortexResult<()> { let session = array_session(); diff --git a/vortex-array/src/arrays/union/tests/take.rs b/vortex-array/src/arrays/union/tests/take.rs new file mode 100644 index 00000000000..5b6854fae48 --- /dev/null +++ b/vortex-array/src/arrays/union/tests/take.rs @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Take coverage for [`UnionArray`]: variant selection, outer nulls, and the empty source. + +use rstest::rstest; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::nullable_union_array; +use super::nullable_variants; +use super::union_array; +use super::variants; +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::arrays::Union; +use crate::arrays::UnionArray; +use crate::arrays::dict::TakeReduce; +use crate::arrays::union::UnionArrayExt; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; + +/// Take `indices` from `array`, asserting that the result reduced back to a union rather than +/// staying a lazy dictionary. +#[track_caller] +fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { + Ok(array + .clone() + .into_array() + .take(indices)? + .as_::() + .into_owned()) +} + +#[test] +fn take_reorders_and_repeats_variant_selection() -> VortexResult<()> { + let taken = take( + &union_array()?, + PrimitiveArray::from_iter([2u64, 1, 0, 1]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + taken.dtype(), + &DType::Union(variants()?, Nullability::NonNullable) + ); + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)? + ); + assert_eq!( + taken.execute_scalar(1, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)? + ); + assert_eq!( + taken.execute_scalar(2, &mut ctx)?, + Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)? + ); + assert_eq!( + taken.execute_scalar(3, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)? + ); + + Ok(()) +} + +#[test] +fn null_indices_become_outer_nulls_and_leave_children_alone() -> VortexResult<()> { + let taken = take( + &union_array()?, + PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + taken.dtype(), + &DType::Union(variants()?, Nullability::Nullable) + ); + + // A nullable index widens the union, but never the variants: the type IDs own the outer + // nullability and the sparse children keep the dtypes the schema declares. + assert_eq!( + taken.child_by_name("number")?.dtype(), + &DType::Primitive(PType::I32, Nullability::NonNullable) + ); + assert_eq!( + taken.child_by_name("flag")?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::union(variants()?, 9, true.into(), Nullability::Nullable)? + ); + assert_eq!( + taken.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::Union(variants()?, Nullability::Nullable)) + ); + assert_eq!( + taken.execute_scalar(2, &mut ctx)?, + Scalar::union(variants()?, 5, 10i32.into(), Nullability::Nullable)? + ); + + Ok(()) +} + +#[test] +fn take_keeps_outer_and_inner_nulls_distinct() -> VortexResult<()> { + let variants = nullable_variants()?; + + // Row 1 is an outer null and row 2 is a present union selecting a null child. + let taken = take( + &nullable_union_array()?, + PrimitiveArray::from_iter([1u64, 2, 3]).into_array(), + )?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::null(DType::Union(variants.clone(), Nullability::Nullable)) + ); + assert_eq!( + taken.execute_scalar(1, &mut ctx)?, + Scalar::union( + variants.clone(), + 9, + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + Nullability::Nullable, + )? + ); + assert_eq!( + taken.execute_scalar(2, &mut ctx)?, + Scalar::union( + variants, + 9, + Scalar::primitive(40i64, Nullability::Nullable), + Nullability::Nullable, + )? + ); + + Ok(()) +} + +/// An empty union has no row for its sparse children to point at, so an all-null gather must +/// synthesize placeholders instead. `take` short-circuits an empty source into a constant before +/// the union ever sees it, so cover that path and the union's own gather together. +#[test] +fn take_from_empty_union_is_all_null() -> VortexResult<()> { + let empty = UnionArray::empty(variants()?, Nullability::Nullable).into_array(); + let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let via_take = empty + .take(indices.clone())? + .execute::(&mut ctx)? + .into_union(); + let via_reduce = ::take(empty.as_::(), &indices)? + .ok_or_else(|| vortex_err!("Union take must never decline"))? + .as_::() + .into_owned(); + + for taken in [via_take, via_reduce] { + assert_eq!(taken.len(), 2); + assert_eq!( + taken.dtype(), + &DType::Union(variants()?, Nullability::Nullable) + ); + + for index in 0..taken.len() { + assert!(taken.execute_scalar(index, &mut ctx)?.is_null()); + } + } + + Ok(()) +} + +#[rstest] +#[case::non_nullable(union_array())] +#[case::nullable(nullable_union_array())] +fn take_conformance(#[case] array: VortexResult) -> VortexResult<()> { + test_take_conformance( + &array?.into_array(), + &mut array_session().create_execution_ctx(), + ); + + Ok(()) +} From fc4468da6edd69cdbafa6bcdcb9b7a80eedaba52 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 19:46:49 +0000 Subject: [PATCH 2/7] fix: address review on Union take Corrects the union module doc, which claimed mask rewrites every child. Mask only rewrites the type IDs, because outer nulls live there. Removes the UncompressedSizeInBytes todo for constant Union arrays. Its comment named constant Union canonicalization as the blocker, and that is now in place, so the arm joins the group that canonicalizes and recurses. Notes at the fill-null site that the lazy node is executed once per child, and why TakeReduce cannot materialize it. Widens the constant Union test to cover the non-nullable and nullable-present cases alongside the outer null, and asserts the placeholder value rather than only its dtype. Shrinks the take benchmark to 256 indices over 2, 4, and 8 variants so the widest case stays inside the sub-millisecond budget for microbenchmarks. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/benches/take_union.rs | 8 +- .../fns/uncompressed_size_in_bytes/mod.rs | 7 +- vortex-array/src/arrays/union/compute/take.rs | 5 ++ vortex-array/src/arrays/union/mod.rs | 4 +- vortex-array/src/arrays/union/tests/mod.rs | 83 +++++++++++++------ 5 files changed, 73 insertions(+), 34 deletions(-) diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs index d82c5faa9cd..328ffd26488 100644 --- a/vortex-array/benches/take_union.rs +++ b/vortex-array/benches/take_union.rs @@ -39,8 +39,12 @@ fn main() { static SESSION: LazyLock = LazyLock::new(array_session); const ARRAY_SIZE: usize = 100_000; -const TAKE_SIZE: usize = 1_000; -const VARIANT_COUNTS: [usize; 4] = [2, 4, 8, 16]; + +/// The index count is held low so that the widest variant case stays inside the sub-millisecond +/// budget for microbenchmarks. Cost is linear in it, so the variant sweep still reads the same. +const TAKE_SIZE: usize = 256; + +const VARIANT_COUNTS: [usize; 3] = [2, 4, 8]; /// A sparse union of `variant_count` `i64` variants whose type IDs cycle through every variant. fn union_array(variant_count: usize, rng: &mut StdRng) -> ArrayRef { diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index ee9426c7919..a10eb816580 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -242,16 +242,11 @@ pub(crate) fn constant_uncompressed_size_in_bytes( | DType::Map(..) | DType::FixedSizeList(..) | DType::Struct(..) + | DType::Union(..) | DType::Extension(_) => { let canonical = array.array().clone().execute::(ctx)?; return canonical_uncompressed_size_in_bytes(&canonical, ctx); } - DType::Union(..) => { - todo!( - "TODO(connor)[Union]: support constant Union size accounting after constant Union \ - canonicalization defines inactive sparse-child placeholders" - ) - } DType::Variant(_) => { vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays") } diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs index 56be3889a33..4565837a900 100644 --- a/vortex-array/src/arrays/union/compute/take.rs +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -38,6 +38,11 @@ impl TakeReduce for Union { // Nullability is stripped so that the children keep their declared variant dtypes. The // type IDs already record which rows are null. + // + // This stays a lazy node that every child then executes for itself, so the fill runs once + // per variant. `TakeReduce` has no `ExecutionCtx` to materialize it with, and the cost is + // proportional to the indices rather than to the data, so it is left alone. Non-nullable + // indices skip it entirely because `FillNull::simplify` returns its input unchanged. let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable()); let child_indices = indices.clone().fill_null(fill_scalar)?; diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs index c59a33c391c..1ff3a10d7f2 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -11,8 +11,8 @@ //! Type ID values are not validated during construction. Accessing a non-null row whose type ID //! is not declared by the union variants will panic. //! -//! Slice, mask, and take are structural: each rewrites the type IDs and every child, so each costs -//! `O(variants)` child operations. +//! Mask only rewrites the type IDs, because outer nulls live there. Slice and take have to rewrite +//! every child to keep them row-aligned, so each costs `O(variants)` child operations. use crate::dtype::DType; use crate::dtype::Nullability; diff --git a/vortex-array/src/arrays/union/tests/mod.rs b/vortex-array/src/arrays/union/tests/mod.rs index 75aec48a656..6d9187b2a96 100644 --- a/vortex-array/src/arrays/union/tests/mod.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use rstest::rstest; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::ReadContext; @@ -12,6 +14,7 @@ use crate::ArrayContext; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; @@ -238,36 +241,68 @@ fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> { Ok(()) } -#[test] -fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> { +/// A constant union canonicalizes into a sparse union whose selected child repeats the scalar's +/// value and whose unselected children hold placeholders. The cases cover a non-nullable union, a +/// nullable union that is present, and an outer null. +#[rstest] +#[case::non_nullable( + Scalar::union(variants().vortex_expect("valid Union variants"), 9, true.into(), Nullability::NonNullable) + .vortex_expect("valid Union scalar"), + DType::Primitive(PType::I32, Nullability::NonNullable), +)] +#[case::nullable_present( + Scalar::union(nullable_variants().vortex_expect("valid Union variants"), 5, 7i32.into(), Nullability::Nullable) + .vortex_expect("valid Union scalar"), + DType::Primitive(PType::I64, Nullability::Nullable), +)] +#[case::outer_null( + Scalar::null(DType::Union(variants().vortex_expect("valid Union variants"), Nullability::Nullable)), + DType::Primitive(PType::I32, Nullability::NonNullable), +)] +fn constant_union_canonicalizes_to_sparse_union( + #[case] scalar: Scalar, + #[case] placeholder_dtype: DType, +) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Union(variants()?, Nullability::Nullable); - - for scalar in [ - Scalar::union(variants()?, 9, true.into(), Nullability::Nullable)?, - Scalar::null(dtype.clone()), - ] { - let canonical = ConstantArray::new(scalar.clone(), 3) - .into_array() - .execute::(&mut ctx)? - .into_union(); - - assert_eq!(canonical.dtype(), &dtype); - - // The unselected variant is only a placeholder, so it keeps its declared dtype. - assert_eq!( - canonical.child_by_name("number")?.dtype(), - &DType::Primitive(PType::I32, Nullability::NonNullable) - ); - - for index in 0..canonical.len() { - assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); - } + + let canonical = ConstantArray::new(scalar.clone(), 3) + .into_array() + .execute::(&mut ctx)? + .into_union(); + + assert_eq!(canonical.dtype(), scalar.dtype()); + + // An unselected child is only a placeholder, so it keeps the dtype the schema declares. That + // is a zero for a non-nullable variant and a null for a nullable one. + let unselected = canonical + .iter_children() + .find(|child| child.dtype() == &placeholder_dtype) + .ok_or_else(|| vortex_err!("No child with dtype {placeholder_dtype}"))?; + assert_eq!( + unselected.execute_scalar(0, &mut ctx)?, + Scalar::default_value(&placeholder_dtype) + ); + + for index in 0..canonical.len() { + assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); } Ok(()) } +#[test] +fn constant_union_reports_uncompressed_size() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?; + + // Four `i32` rows plus four `bool` placeholder bits, rounded up to a byte. + let size = uncompressed_size_in_bytes(&ConstantArray::new(scalar, 4).into_array(), &mut ctx)?; + + assert_eq!(size, 4 * 4 + 4 + 1); + + Ok(()) +} + #[test] fn serde_roundtrip() -> VortexResult<()> { let session = array_session(); From 40326d0ffe18dfc9731ba361ff4d3a56a0a20d8b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 20:22:41 +0000 Subject: [PATCH 3/7] style: trim Union take comments and tests Cuts the comments down to what the code does not already say. The take doc no longer restates the impl header, the fill-null note drops the contract it duplicated from that doc, and the constant-union placeholder comment is gone because the doc above it covers the same ground. Collapses the repeated per-row scalar assertions in the take tests behind an `assert_rows` helper, and replaces the `rstest` cases on the constant-union test with a table in the body so the scalars are built with `?` instead of `vortex_expect` inside an attribute. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/benches/take_union.rs | 13 +- vortex-array/src/arrays/union/array.rs | 8 +- vortex-array/src/arrays/union/compute/take.rs | 25 +-- vortex-array/src/arrays/union/mod.rs | 4 +- vortex-array/src/arrays/union/tests/mod.rs | 83 +++++---- vortex-array/src/arrays/union/tests/take.rs | 159 ++++++++---------- 6 files changed, 127 insertions(+), 165 deletions(-) diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs index 328ffd26488..b81f382c539 100644 --- a/vortex-array/benches/take_union.rs +++ b/vortex-array/benches/take_union.rs @@ -1,12 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Take on a canonical sparse union. +//! Take on a canonical sparse union, swept over the variant count. //! -//! A sparse union keeps every child row-aligned with the union, so take gathers all of them and -//! costs `O(variants * indices)`. The variant count is the axis worth measuring, so these -//! benchmarks pin the array and index counts and sweep it. The nullable-indices case additionally -//! pays for the fill-null pass that keeps the children off the union's outer nullability. +//! Take gathers every sparse child, so cost is linear in the variant count. The nullable-indices +//! case also pays for the fill-null pass. #![expect(clippy::unwrap_used)] @@ -40,8 +38,7 @@ static SESSION: LazyLock = LazyLock::new(array_session); const ARRAY_SIZE: usize = 100_000; -/// The index count is held low so that the widest variant case stays inside the sub-millisecond -/// budget for microbenchmarks. Cost is linear in it, so the variant sweep still reads the same. +/// Held low so the widest variant case stays inside the sub-millisecond microbenchmark budget. const TAKE_SIZE: usize = 256; const VARIANT_COUNTS: [usize; 3] = [2, 4, 8]; @@ -92,7 +89,7 @@ fn take_union_nullable_indices(bencher: Bencher, variant_count: usize) { let mut rng = StdRng::seed_from_u64(0); let array = union_array(variant_count, &mut rng); - // Every tenth index is null, which is what turns a gathered row into an outer union null. + // Every tenth index is null, which produces an outer union null. let indices = PrimitiveArray::from_option_iter( (0..TAKE_SIZE).map(|i| (i % 10 != 0).then(|| rng.random_range(0..ARRAY_SIZE) as u64)), ) diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs index 0d032a1f1e2..bd5a8c5196c 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -179,10 +179,8 @@ impl Array { /// Construct a `len`-row union in which every row holds `scalar`. /// - /// A sparse union keeps every child row-aligned even though at most one of them is selected, - /// so the unselected children are filled with their variant's default value: a null for a - /// nullable variant and a zero for a non-nullable one. An outer null `scalar` therefore - /// produces null type IDs and a placeholder for every child. + /// Unselected children are filled with their variant's default value, a null for a nullable + /// variant and a zero for a non-nullable one. An outer null `scalar` selects no child at all. /// /// # Errors /// @@ -199,8 +197,6 @@ impl Array { None => Scalar::null(union_type_ids_dtype(nullability)), }; - // The selected variant carries the scalar's value. Every other child is a placeholder that - // exists only to keep the sparse layout row-aligned. let selected = union.child_index().zip(union.child()); let children = variants diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs index 4565837a900..250d549d1d3 100644 --- a/vortex-array/src/arrays/union/compute/take.rs +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -15,19 +15,17 @@ use crate::arrays::union::UnionArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::scalar::Scalar; -/// Structural take for [`UnionArray`]: gathers the type IDs and every sparse child at `indices`. +/// Gathers the type IDs and every sparse child at `indices`. /// -/// A sparse union keeps every child row-aligned with the union, so a gather must visit all of them -/// even though at most one is active per row. Take therefore costs `O(variants * indices)`. -/// Reducing that cost requires the dense union encoding, not a different sparse gather. +/// Sparse children are row-aligned with the union, so a gather must visit all of them. Take costs +/// `O(variants * indices)`, which only the dense encoding fixes. /// -/// The type IDs carry the union's validity, so gathering them with the original `indices` is what -/// turns a null index into an outer union null. The children are gathered with the null indices -/// filled in, which keeps each child's dtype exactly as the variant schema declares it. +/// The type IDs carry the union's validity, so gathering them with the original `indices` turns a +/// null index into an outer union null. The children are gathered with the nulls filled in, which +/// keeps their declared variant dtypes. impl TakeReduce for Union { fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult> { - // An empty union has no row for a child to point at, so the only legal indices are all - // null and every output row is an outer union null. + // An empty union has no row to point at, so the indices must be all null. if array.is_empty() { return UnionArray::constant(&Scalar::null(array.dtype().as_nullable()), indices.len()) .map(UnionArray::into_array) @@ -36,13 +34,8 @@ impl TakeReduce for Union { let type_ids = array.type_ids().take(indices.clone())?; - // Nullability is stripped so that the children keep their declared variant dtypes. The - // type IDs already record which rows are null. - // - // This stays a lazy node that every child then executes for itself, so the fill runs once - // per variant. `TakeReduce` has no `ExecutionCtx` to materialize it with, and the cost is - // proportional to the indices rather than to the data, so it is left alone. Non-nullable - // indices skip it entirely because `FillNull::simplify` returns its input unchanged. + // This stays a lazy node, so the fill runs once per child. `TakeReduce` has no + // `ExecutionCtx` to materialize it with, and the cost is per index rather than per element. let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable()); let child_indices = indices.clone().fill_null(fill_scalar)?; diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs index 1ff3a10d7f2..0dda06e2766 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -11,8 +11,8 @@ //! Type ID values are not validated during construction. Accessing a non-null row whose type ID //! is not declared by the union variants will panic. //! -//! Mask only rewrites the type IDs, because outer nulls live there. Slice and take have to rewrite -//! every child to keep them row-aligned, so each costs `O(variants)` child operations. +//! Mask rewrites only the type IDs. Slice and take rewrite every child to keep them row-aligned, +//! so each costs `O(variants)` child operations. use crate::dtype::DType; use crate::dtype::Nullability; diff --git a/vortex-array/src/arrays/union/tests/mod.rs b/vortex-array/src/arrays/union/tests/mod.rs index 6d9187b2a96..af1c4459550 100644 --- a/vortex-array/src/arrays/union/tests/mod.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use rstest::rstest; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -242,49 +241,49 @@ fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> { } /// A constant union canonicalizes into a sparse union whose selected child repeats the scalar's -/// value and whose unselected children hold placeholders. The cases cover a non-nullable union, a -/// nullable union that is present, and an outer null. -#[rstest] -#[case::non_nullable( - Scalar::union(variants().vortex_expect("valid Union variants"), 9, true.into(), Nullability::NonNullable) - .vortex_expect("valid Union scalar"), - DType::Primitive(PType::I32, Nullability::NonNullable), -)] -#[case::nullable_present( - Scalar::union(nullable_variants().vortex_expect("valid Union variants"), 5, 7i32.into(), Nullability::Nullable) - .vortex_expect("valid Union scalar"), - DType::Primitive(PType::I64, Nullability::Nullable), -)] -#[case::outer_null( - Scalar::null(DType::Union(variants().vortex_expect("valid Union variants"), Nullability::Nullable)), - DType::Primitive(PType::I32, Nullability::NonNullable), -)] -fn constant_union_canonicalizes_to_sparse_union( - #[case] scalar: Scalar, - #[case] placeholder_dtype: DType, -) -> VortexResult<()> { +/// value and whose unselected children hold placeholders. +#[test] +fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); + let i32_variant = DType::Primitive(PType::I32, Nullability::NonNullable); + let i64_variant = DType::Primitive(PType::I64, Nullability::Nullable); + + // Each case pairs a constant scalar with the dtype of a variant it does not select. + let cases = [ + ( + Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?, + i32_variant.clone(), + ), + ( + Scalar::union(nullable_variants()?, 5, 7i32.into(), Nullability::Nullable)?, + i64_variant, + ), + ( + Scalar::null(DType::Union(variants()?, Nullability::Nullable)), + i32_variant, + ), + ]; - let canonical = ConstantArray::new(scalar.clone(), 3) - .into_array() - .execute::(&mut ctx)? - .into_union(); - - assert_eq!(canonical.dtype(), scalar.dtype()); - - // An unselected child is only a placeholder, so it keeps the dtype the schema declares. That - // is a zero for a non-nullable variant and a null for a nullable one. - let unselected = canonical - .iter_children() - .find(|child| child.dtype() == &placeholder_dtype) - .ok_or_else(|| vortex_err!("No child with dtype {placeholder_dtype}"))?; - assert_eq!( - unselected.execute_scalar(0, &mut ctx)?, - Scalar::default_value(&placeholder_dtype) - ); - - for index in 0..canonical.len() { - assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); + for (scalar, unselected_dtype) in cases { + let canonical = ConstantArray::new(scalar.clone(), 3) + .into_array() + .execute::(&mut ctx)? + .into_union(); + + assert_eq!(canonical.dtype(), scalar.dtype()); + + let unselected = canonical + .iter_children() + .find(|child| child.dtype() == &unselected_dtype) + .ok_or_else(|| vortex_err!("No child with dtype {unselected_dtype}"))?; + assert_eq!( + unselected.execute_scalar(0, &mut ctx)?, + Scalar::default_value(&unselected_dtype) + ); + + for index in 0..canonical.len() { + assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar); + } } Ok(()) diff --git a/vortex-array/src/arrays/union/tests/take.rs b/vortex-array/src/arrays/union/tests/take.rs index 5b6854fae48..0088ace43b2 100644 --- a/vortex-array/src/arrays/union/tests/take.rs +++ b/vortex-array/src/arrays/union/tests/take.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Take coverage for [`UnionArray`]: variant selection, outer nulls, and the empty source. - use rstest::rstest; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -27,8 +25,7 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::scalar::Scalar; -/// Take `indices` from `array`, asserting that the result reduced back to a union rather than -/// staying a lazy dictionary. +/// Take `indices`, asserting the result reduced back to a union instead of staying a dictionary. #[track_caller] fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { Ok(array @@ -39,53 +36,54 @@ fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { .into_owned()) } +/// Assert that `array` holds exactly `expected`, row for row. +#[track_caller] +fn assert_rows(array: &UnionArray, expected: Vec) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.len(), expected.len()); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(array.execute_scalar(index, &mut ctx)?, expected); + } + + Ok(()) +} + #[test] fn take_reorders_and_repeats_variant_selection() -> VortexResult<()> { + let variants = variants()?; + let nullability = Nullability::NonNullable; + let taken = take( &union_array()?, PrimitiveArray::from_iter([2u64, 1, 0, 1]).into_array(), )?; - let mut ctx = array_session().create_execution_ctx(); - - assert_eq!( - taken.dtype(), - &DType::Union(variants()?, Nullability::NonNullable) - ); - assert_eq!( - taken.execute_scalar(0, &mut ctx)?, - Scalar::union(variants()?, 5, 30i32.into(), Nullability::NonNullable)? - ); - assert_eq!( - taken.execute_scalar(1, &mut ctx)?, - Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)? - ); - assert_eq!( - taken.execute_scalar(2, &mut ctx)?, - Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)? - ); - assert_eq!( - taken.execute_scalar(3, &mut ctx)?, - Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)? - ); - Ok(()) + assert_eq!(taken.dtype(), &DType::Union(variants.clone(), nullability)); + assert_rows( + &taken, + vec![ + Scalar::union(variants.clone(), 5, 30i32.into(), nullability)?, // + Scalar::union(variants.clone(), 9, true.into(), nullability)?, // + Scalar::union(variants.clone(), 5, 10i32.into(), nullability)?, // + Scalar::union(variants, 9, true.into(), nullability)?, // + ], + ) } #[test] fn null_indices_become_outer_nulls_and_leave_children_alone() -> VortexResult<()> { + let variants = variants()?; + let nullability = Nullability::Nullable; + let taken = take( &union_array()?, PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array(), )?; - let mut ctx = array_session().create_execution_ctx(); - assert_eq!( - taken.dtype(), - &DType::Union(variants()?, Nullability::Nullable) - ); + assert_eq!(taken.dtype(), &DType::Union(variants.clone(), nullability)); - // A nullable index widens the union, but never the variants: the type IDs own the outer - // nullability and the sparse children keep the dtypes the schema declares. + // A nullable index widens the union but never its variants. assert_eq!( taken.child_by_name("number")?.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) @@ -95,90 +93,69 @@ fn null_indices_become_outer_nulls_and_leave_children_alone() -> VortexResult<() &DType::Bool(Nullability::NonNullable) ); - assert_eq!( - taken.execute_scalar(0, &mut ctx)?, - Scalar::union(variants()?, 9, true.into(), Nullability::Nullable)? - ); - assert_eq!( - taken.execute_scalar(1, &mut ctx)?, - Scalar::null(DType::Union(variants()?, Nullability::Nullable)) - ); - assert_eq!( - taken.execute_scalar(2, &mut ctx)?, - Scalar::union(variants()?, 5, 10i32.into(), Nullability::Nullable)? - ); - - Ok(()) + assert_rows( + &taken, + vec![ + Scalar::union(variants.clone(), 9, true.into(), nullability)?, // + Scalar::null(DType::Union(variants.clone(), nullability)), // + Scalar::union(variants, 5, 10i32.into(), nullability)?, // + ], + ) } #[test] fn take_keeps_outer_and_inner_nulls_distinct() -> VortexResult<()> { let variants = nullable_variants()?; + let nullability = Nullability::Nullable; // Row 1 is an outer null and row 2 is a present union selecting a null child. let taken = take( &nullable_union_array()?, PrimitiveArray::from_iter([1u64, 2, 3]).into_array(), )?; - let mut ctx = array_session().create_execution_ctx(); - assert_eq!( - taken.execute_scalar(0, &mut ctx)?, - Scalar::null(DType::Union(variants.clone(), Nullability::Nullable)) - ); - assert_eq!( - taken.execute_scalar(1, &mut ctx)?, - Scalar::union( - variants.clone(), - 9, - Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), - Nullability::Nullable, - )? - ); - assert_eq!( - taken.execute_scalar(2, &mut ctx)?, - Scalar::union( - variants, - 9, - Scalar::primitive(40i64, Nullability::Nullable), - Nullability::Nullable, - )? - ); - - Ok(()) + assert_rows( + &taken, + vec![ + Scalar::null(DType::Union(variants.clone(), nullability)), // + Scalar::union( + variants.clone(), + 9, + Scalar::null(DType::Primitive(PType::I64, nullability)), + nullability, + )?, // + Scalar::union( + variants, + 9, + Scalar::primitive(40i64, nullability), + nullability, + )?, // + ], + ) } -/// An empty union has no row for its sparse children to point at, so an all-null gather must -/// synthesize placeholders instead. `take` short-circuits an empty source into a constant before -/// the union ever sees it, so cover that path and the union's own gather together. +/// `take` short-circuits an empty source into a constant before the union sees it, so this covers +/// that path and the union's own gather. #[test] fn take_from_empty_union_is_all_null() -> VortexResult<()> { - let empty = UnionArray::empty(variants()?, Nullability::Nullable).into_array(); + let variants = variants()?; + let nullability = Nullability::Nullable; + let empty = UnionArray::empty(variants.clone(), nullability).into_array(); let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); let via_take = empty .take(indices.clone())? - .execute::(&mut ctx)? + .execute::(&mut array_session().create_execution_ctx())? .into_union(); let via_reduce = ::take(empty.as_::(), &indices)? .ok_or_else(|| vortex_err!("Union take must never decline"))? .as_::() .into_owned(); - for taken in [via_take, via_reduce] { - assert_eq!(taken.len(), 2); - assert_eq!( - taken.dtype(), - &DType::Union(variants()?, Nullability::Nullable) - ); + let expected = || vec![Scalar::null(DType::Union(variants.clone(), nullability)); 2]; - for index in 0..taken.len() { - assert!(taken.execute_scalar(index, &mut ctx)?.is_null()); - } - } - - Ok(()) + assert_rows(&via_take, expected())?; + assert_rows(&via_reduce, expected()) } #[rstest] From b0f8c63c0f948f576a57cc14abf0f82605f1ec38 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 20:33:48 +0000 Subject: [PATCH 4/7] style: deduplicate the take benchmark and fix two comments Pulls the shared bencher plumbing out of the two take benchmarks into `bench_take`, so each one only builds its indices. Corrects the uncompressed-size assertion comment, which listed the `i32` rows and the `bool` placeholder bits but not the type IDs, and reorders the expression to match. Restates the fill-null cost note in terms of the indices rather than "per element", which read as if it were per output row. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/benches/take_union.rs | 30 +++++++++---------- vortex-array/src/arrays/union/compute/take.rs | 3 +- vortex-array/src/arrays/union/tests/mod.rs | 4 +-- 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs index b81f382c539..af4c10660b1 100644 --- a/vortex-array/benches/take_union.rs +++ b/vortex-array/benches/take_union.rs @@ -64,6 +64,18 @@ fn union_array(variant_count: usize, rng: &mut StdRng) -> ArrayRef { UnionArray::new(type_ids.into_array(), variants, children).into_array() } +/// Take `indices` and execute the result, so that the child gathers actually run. +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(args = VARIANT_COUNTS)] fn take_union(bencher: Bencher, variant_count: usize) { let mut rng = StdRng::seed_from_u64(0); @@ -74,14 +86,7 @@ fn take_union(bencher: Bencher, variant_count: usize) { .collect::>() .into_array(); - bencher - .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) - .bench_refs(|(array, indices, ctx)| { - array - .take((*indices).clone()) - .unwrap() - .execute::(ctx) - }); + bench_take(bencher, array, indices); } #[divan::bench(args = VARIANT_COUNTS)] @@ -95,12 +100,5 @@ fn take_union_nullable_indices(bencher: Bencher, variant_count: usize) { ) .into_array(); - bencher - .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) - .bench_refs(|(array, indices, ctx)| { - array - .take((*indices).clone()) - .unwrap() - .execute::(ctx) - }); + bench_take(bencher, array, indices); } diff --git a/vortex-array/src/arrays/union/compute/take.rs b/vortex-array/src/arrays/union/compute/take.rs index 250d549d1d3..5ad0a6f1287 100644 --- a/vortex-array/src/arrays/union/compute/take.rs +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -35,7 +35,8 @@ impl TakeReduce for Union { let type_ids = array.type_ids().take(indices.clone())?; // This stays a lazy node, so the fill runs once per child. `TakeReduce` has no - // `ExecutionCtx` to materialize it with, and the cost is per index rather than per element. + // `ExecutionCtx` to materialize it with, and the cost scales with the indices, not the + // data behind them. let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable()); let child_indices = indices.clone().fill_null(fill_scalar)?; diff --git a/vortex-array/src/arrays/union/tests/mod.rs b/vortex-array/src/arrays/union/tests/mod.rs index af1c4459550..fc0abec37b8 100644 --- a/vortex-array/src/arrays/union/tests/mod.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -294,10 +294,10 @@ fn constant_union_reports_uncompressed_size() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?; - // Four `i32` rows plus four `bool` placeholder bits, rounded up to a byte. + // Four `u8` type IDs, four `i32` rows, and four `bool` placeholder bits rounded up to a byte. let size = uncompressed_size_in_bytes(&ConstantArray::new(scalar, 4).into_array(), &mut ctx)?; - assert_eq!(size, 4 * 4 + 4 + 1); + assert_eq!(size, 4 + 4 * 4 + 1); Ok(()) } From d4099e856917b3f7d4a0afaf19118fd70d4cd864 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 21:01:39 +0000 Subject: [PATCH 5/7] bench: give the union take benchmark heterogeneous variants The pool was `variant_count` identical `i64` children, which understated take because every child gathered at the same cheap rate. It now runs `i64`, `list`, and `utf8`, so widening the union adds a child that gathers differently. The mix costs roughly 2.5x the primitive-only version per child, so the sweep drops from 2, 4, and 8 variants to 1, 2, and 3, and the index count drops to 128. The single-variant case is the baseline that isolates the type IDs gather from the child gathers. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/benches/take_union.rs | 90 +++++++++++++++++++++++------- 1 file changed, 70 insertions(+), 20 deletions(-) diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs index af4c10660b1..ef94d47a893 100644 --- a/vortex-array/benches/take_union.rs +++ b/vortex-array/benches/take_union.rs @@ -3,8 +3,10 @@ //! Take on a canonical sparse union, swept over the variant count. //! -//! Take gathers every sparse child, so cost is linear in the variant count. The nullable-indices -//! case also pays for the fill-null pass. +//! Take gathers every sparse child, so a wider union costs more. The children mix encodings on +//! purpose, because each one carries its own gather cost and a union of identical primitives +//! understates the total. At this index count the per-child setup dominates the gather itself, +//! and that setup is exactly what the variant count multiplies. #![expect(clippy::unwrap_used)] @@ -19,13 +21,12 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ListViewArray; 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::arrays::VarBinViewArray; use vortex_array::dtype::UnionVariants; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -39,27 +40,76 @@ static SESSION: LazyLock = LazyLock::new(array_session); const ARRAY_SIZE: usize = 100_000; /// Held low so the widest variant case stays inside the sub-millisecond microbenchmark budget. -const TAKE_SIZE: usize = 256; +const TAKE_SIZE: usize = 128; + +/// One entry per pool variant, starting from a single-variant baseline that isolates the type IDs +/// gather from the child gathers. +const VARIANT_COUNTS: [usize; 3] = [1, 2, 3]; + +/// A `List` of `ARRAY_SIZE` short lists over one shared element buffer. +fn list_child(rng: &mut StdRng) -> ArrayRef { + const MAX_LIST_LEN: i32 = 8; + + let sizes: Buffer = (0..ARRAY_SIZE) + .map(|_| rng.random_range(0..MAX_LIST_LEN)) + .collect(); + let offsets: Buffer = sizes + .iter() + .scan(0i32, |offset, size| { + let start = *offset; + *offset += size; + Some(start) + }) + .collect(); + + let total = offsets.last().unwrap() + sizes.last().unwrap(); + let elements = (0..total) + .map(|_| rng.random::()) + .collect::>() + .into_array(); + + ListViewArray::new( + elements, + offsets.into_array(), + sizes.into_array(), + Validity::NonNullable, + ) + .into_array() +} -const VARIANT_COUNTS: [usize; 3] = [2, 4, 8]; +/// The `index`th variant of the pool, as a name and an `ARRAY_SIZE`-row child. +/// +/// The pool runs `i64`, `list`, and `utf8`, so widening the union adds a child that gathers +/// differently rather than one more primitive gather. +fn variant_child(index: usize, rng: &mut StdRng) -> (String, ArrayRef) { + let child = match index { + 0 => (0..ARRAY_SIZE) + .map(|_| rng.random::()) + .collect::>() + .into_array(), + 1 => list_child(rng), + _ => VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|i| format!("value-{i}"))) + .into_array(), + }; + + (format!("v{index}"), child) +} -/// A sparse union of `variant_count` `i64` variants whose type IDs cycle through every variant. +/// A sparse union over the first `variant_count` pool entries, with type IDs cycling through them. fn union_array(variant_count: usize, rng: &mut StdRng) -> ArrayRef { - let names: FieldNames = (0..variant_count).map(|i| format!("v{i}")).collect(); - let dtypes = vec![DType::Primitive(PType::I64, Nullability::NonNullable); variant_count]; - let variants = UnionVariants::new(names, dtypes).unwrap(); + let (names, children): (Vec, Vec) = (0..variant_count) + .map(|index| variant_child(index, rng)) + .unzip(); + + let variants = UnionVariants::new( + names.into_iter().collect(), + children.iter().map(|child| child.dtype().clone()).collect(), + ) + .unwrap(); let type_ids = PrimitiveArray::from_iter( (0..ARRAY_SIZE).map(|i| u8::try_from(i % variant_count).unwrap()), ); - let children = (0..variant_count) - .map(|_| { - (0..ARRAY_SIZE) - .map(|_| rng.random::()) - .collect::>() - .into_array() - }) - .collect::>(); UnionArray::new(type_ids.into_array(), variants, children).into_array() } From 8bbec19e20d091d00aec876f108aacff2e143599 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 21:08:40 +0000 Subject: [PATCH 6/7] bench: take a union of integers, strings, and lists Replaces the variant-count sweep over a synthetic pool with one realistic schema. The sweep measured a contrived axis, and its children were weak: the strings were all under the 12-byte `VarBinView` inline limit, so the gather never touched a data buffer, and the lists were built around a `MAX_LIST_LEN` constant that read as a fixed size when `FixedSizeList` is the type for that. The union is now `i64`, `utf8`, and `list`. The strings mix inline and out-of-line lengths, and the list lengths vary so the gather has to rebuild offsets and sizes. Two benchmarks remain, one per index nullability, down from six. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/benches/take_union.rs | 94 ++++++++++++++---------------- 1 file changed, 45 insertions(+), 49 deletions(-) diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs index ef94d47a893..0f1f4b51b3a 100644 --- a/vortex-array/benches/take_union.rs +++ b/vortex-array/benches/take_union.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Take on a canonical sparse union, swept over the variant count. +//! Take on a canonical sparse union of integers, strings, and lists. //! -//! Take gathers every sparse child, so a wider union costs more. The children mix encodings on -//! purpose, because each one carries its own gather cost and a union of identical primitives -//! understates the total. At this index count the per-child setup dominates the gather itself, -//! and that setup is exactly what the variant count multiplies. +//! Take gathers every sparse child, not only the selected one, so the cost is the sum over all +//! three variants. The encodings differ on purpose: a primitive gather is close to a memcpy, a +//! `VarBinView` gather copies views and leaves the data buffers alone, and a `ListView` gather +//! rebuilds offsets and sizes over a shared element buffer. Nullable indices add a fill-null pass +//! per child on top of that. #![expect(clippy::unwrap_used)] @@ -39,20 +40,36 @@ static SESSION: LazyLock = LazyLock::new(array_session); const ARRAY_SIZE: usize = 100_000; -/// Held low so the widest variant case stays inside the sub-millisecond microbenchmark budget. +/// Held low so that both cases stay inside the sub-millisecond microbenchmark budget. Per-child +/// setup dominates the gather at this size, and that setup is paid once per variant. const TAKE_SIZE: usize = 128; -/// One entry per pool variant, starting from a single-variant baseline that isolates the type IDs -/// gather from the child gathers. -const VARIANT_COUNTS: [usize; 3] = [1, 2, 3]; +fn integer_child(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random::()) + .collect::>() + .into_array() +} -/// A `List` of `ARRAY_SIZE` short lists over one shared element buffer. -fn list_child(rng: &mut StdRng) -> ArrayRef { - const MAX_LIST_LEN: i32 = 8; +/// A `utf8` child mixing strings short enough to inline into a `VarBinView` view with longer ones +/// that have to point into the data buffer. +/// +/// A `VarBinView` view inlines up to 12 bytes. A child of uniformly short strings would gather +/// without ever reading the data buffer, which is not what a real string column costs. +fn string_child() -> ArrayRef { + VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|i| match i % 4 { + 0 => format!("s{i}"), + _ => format!("a considerably longer string value, number {i}"), + })) + .into_array() +} - let sizes: Buffer = (0..ARRAY_SIZE) - .map(|_| rng.random_range(0..MAX_LIST_LEN)) - .collect(); +/// A `list` child of variable-length lists over one shared element buffer. +/// +/// The lengths vary on purpose. Equal-length lists are what `FixedSizeList` is for, and they would +/// not exercise the offsets and sizes that a `ListView` gather has to rebuild. +fn list_child(rng: &mut StdRng) -> ArrayRef { + let sizes: Buffer = (0..ARRAY_SIZE).map(|_| rng.random_range(0..16)).collect(); let offsets: Buffer = sizes .iter() .scan(0i32, |offset, size| { @@ -64,8 +81,8 @@ fn list_child(rng: &mut StdRng) -> ArrayRef { let total = offsets.last().unwrap() + sizes.last().unwrap(); let elements = (0..total) - .map(|_| rng.random::()) - .collect::>() + .map(|_| rng.random::()) + .collect::>() .into_array(); ListViewArray::new( @@ -77,38 +94,17 @@ fn list_child(rng: &mut StdRng) -> ArrayRef { .into_array() } -/// The `index`th variant of the pool, as a name and an `ARRAY_SIZE`-row child. -/// -/// The pool runs `i64`, `list`, and `utf8`, so widening the union adds a child that gathers -/// differently rather than one more primitive gather. -fn variant_child(index: usize, rng: &mut StdRng) -> (String, ArrayRef) { - let child = match index { - 0 => (0..ARRAY_SIZE) - .map(|_| rng.random::()) - .collect::>() - .into_array(), - 1 => list_child(rng), - _ => VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|i| format!("value-{i}"))) - .into_array(), - }; - - (format!("v{index}"), child) -} - -/// A sparse union over the first `variant_count` pool entries, with type IDs cycling through them. -fn union_array(variant_count: usize, rng: &mut StdRng) -> ArrayRef { - let (names, children): (Vec, Vec) = (0..variant_count) - .map(|index| variant_child(index, rng)) - .unzip(); - +/// A sparse union whose rows cycle through an integer, a string, and a list. +fn union_array(rng: &mut StdRng) -> ArrayRef { + let children = vec![integer_child(rng), string_child(), list_child(rng)]; let variants = UnionVariants::new( - names.into_iter().collect(), + ["ints", "strings", "lists"].into(), children.iter().map(|child| child.dtype().clone()).collect(), ) .unwrap(); let type_ids = PrimitiveArray::from_iter( - (0..ARRAY_SIZE).map(|i| u8::try_from(i % variant_count).unwrap()), + (0..ARRAY_SIZE).map(|i| u8::try_from(i % children.len()).unwrap()), ); UnionArray::new(type_ids.into_array(), variants, children).into_array() @@ -126,10 +122,10 @@ fn bench_take(bencher: Bencher, array: ArrayRef, indices: ArrayRef) { }); } -#[divan::bench(args = VARIANT_COUNTS)] -fn take_union(bencher: Bencher, variant_count: usize) { +#[divan::bench] +fn take_union(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); - let array = union_array(variant_count, &mut rng); + let array = union_array(&mut rng); let indices = (0..TAKE_SIZE) .map(|_| rng.random_range(0..ARRAY_SIZE) as u64) @@ -139,10 +135,10 @@ fn take_union(bencher: Bencher, variant_count: usize) { bench_take(bencher, array, indices); } -#[divan::bench(args = VARIANT_COUNTS)] -fn take_union_nullable_indices(bencher: Bencher, variant_count: usize) { +#[divan::bench] +fn take_union_nullable_indices(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); - let array = union_array(variant_count, &mut rng); + let array = union_array(&mut rng); // Every tenth index is null, which produces an outer union null. let indices = PrimitiveArray::from_option_iter( From dfa124c42ff200f226b5a99b8b8ba1d012e714ee Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 6 Aug 2026 21:53:42 +0000 Subject: [PATCH 7/7] bench: compare dense and compressed union children Moves the union take benchmark to the `vortex` crate, which already depends on both `vortex-array` and `vortex-sparse`, and adds a compressed shape alongside the dense one. Rows are skewed 98% integers to 1% strings and 1% lists, which is the distribution the sparse layout exists for. The `dense_children` cases materialize every child at the union's length. The `compressed_children` cases keep the dominant integer child canonical and store each rare child as a `SparseArray`, which is what a compressor produces. Compressed children turn out to be 3.3x slower to take, not faster. The cost is cache-missing binary searches in `Patches::take`: with 1000 patches against 128 indices the ratio stays above `PREFER_MAP_WHEN_PATCHES_OVER_INDICES_LESS_THAN`, so every call searches. That is a `Patches` cost rather than a union one, and union take multiplies it by the variant count. The compressed cases run over the sub-millisecond budget the other microbenchmarks hold to. The gap only appears once a child's patch indices outgrow the cache, so a size inside the budget measures nothing. Signed-off-by: Connor Tsui Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_014VBqADUGaR81PoijKkvoUs --- vortex-array/Cargo.toml | 4 - vortex-array/benches/take_union.rs | 150 ------------------ vortex/Cargo.toml | 5 + vortex/benches/take_union.rs | 246 +++++++++++++++++++++++++++++ 4 files changed, 251 insertions(+), 154 deletions(-) delete mode 100644 vortex-array/benches/take_union.rs create mode 100644 vortex/benches/take_union.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index e31f2e41aa3..22a8768af5d 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -243,10 +243,6 @@ harness = false name = "take_struct" harness = false -[[bench]] -name = "take_union" -harness = false - [[bench]] name = "take_fsl" harness = false diff --git a/vortex-array/benches/take_union.rs b/vortex-array/benches/take_union.rs deleted file mode 100644 index 0f1f4b51b3a..00000000000 --- a/vortex-array/benches/take_union.rs +++ /dev/null @@ -1,150 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Take on a canonical sparse union of integers, strings, and lists. -//! -//! Take gathers every sparse child, not only the selected one, so the cost is the sum over all -//! three variants. The encodings differ on purpose: a primitive gather is close to a memcpy, a -//! `VarBinView` gather copies views and leaves the data buffers alone, and a `ListView` gather -//! rebuilds offsets and sizes over a shared element buffer. Nullable indices add a fill-null pass -//! per child on top of that. - -#![expect(clippy::unwrap_used)] - -use std::sync::LazyLock; - -use divan::Bencher; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -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::ListViewArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::UnionArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::dtype::UnionVariants; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_session::VortexSession; - -fn main() { - LazyLock::force(&SESSION); - divan::main(); -} - -static SESSION: LazyLock = LazyLock::new(array_session); - -const ARRAY_SIZE: usize = 100_000; - -/// Held low so that both cases stay inside the sub-millisecond microbenchmark budget. Per-child -/// setup dominates the gather at this size, and that setup is paid once per variant. -const TAKE_SIZE: usize = 128; - -fn integer_child(rng: &mut StdRng) -> ArrayRef { - (0..ARRAY_SIZE) - .map(|_| rng.random::()) - .collect::>() - .into_array() -} - -/// A `utf8` child mixing strings short enough to inline into a `VarBinView` view with longer ones -/// that have to point into the data buffer. -/// -/// A `VarBinView` view inlines up to 12 bytes. A child of uniformly short strings would gather -/// without ever reading the data buffer, which is not what a real string column costs. -fn string_child() -> ArrayRef { - VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|i| match i % 4 { - 0 => format!("s{i}"), - _ => format!("a considerably longer string value, number {i}"), - })) - .into_array() -} - -/// A `list` child of variable-length lists over one shared element buffer. -/// -/// The lengths vary on purpose. Equal-length lists are what `FixedSizeList` is for, and they would -/// not exercise the offsets and sizes that a `ListView` gather has to rebuild. -fn list_child(rng: &mut StdRng) -> ArrayRef { - let sizes: Buffer = (0..ARRAY_SIZE).map(|_| rng.random_range(0..16)).collect(); - let offsets: Buffer = sizes - .iter() - .scan(0i32, |offset, size| { - let start = *offset; - *offset += size; - Some(start) - }) - .collect(); - - let total = offsets.last().unwrap() + sizes.last().unwrap(); - let elements = (0..total) - .map(|_| rng.random::()) - .collect::>() - .into_array(); - - ListViewArray::new( - elements, - offsets.into_array(), - sizes.into_array(), - Validity::NonNullable, - ) - .into_array() -} - -/// A sparse union whose rows cycle through an integer, a string, and a list. -fn union_array(rng: &mut StdRng) -> ArrayRef { - let children = vec![integer_child(rng), string_child(), list_child(rng)]; - let variants = UnionVariants::new( - ["ints", "strings", "lists"].into(), - children.iter().map(|child| child.dtype().clone()).collect(), - ) - .unwrap(); - - let type_ids = PrimitiveArray::from_iter( - (0..ARRAY_SIZE).map(|i| u8::try_from(i % children.len()).unwrap()), - ); - - UnionArray::new(type_ids.into_array(), variants, children).into_array() -} - -/// Take `indices` and execute the result, so that the child gathers actually run. -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 take_union(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let array = union_array(&mut rng); - - let indices = (0..TAKE_SIZE) - .map(|_| rng.random_range(0..ARRAY_SIZE) as u64) - .collect::>() - .into_array(); - - bench_take(bencher, array, indices); -} - -#[divan::bench] -fn take_union_nullable_indices(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let array = union_array(&mut rng); - - // Every tenth index is null, which produces an outer union null. - let indices = PrimitiveArray::from_option_iter( - (0..TAKE_SIZE).map(|i| (i % 10 != 0).then(|| rng.random_range(0..ARRAY_SIZE) as u64)), - ) - .into_array(); - - bench_take(bencher, array, indices); -} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 5b26944c4ed..5aeba6c17bb 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -105,6 +105,11 @@ name = "common_encoding_tree_throughput" harness = false test = false +[[bench]] +name = "take_union" +harness = false +test = false + [[bench]] name = "pipeline" harness = false diff --git a/vortex/benches/take_union.rs b/vortex/benches/take_union.rs new file mode 100644 index 00000000000..a7072254275 --- /dev/null +++ b/vortex/benches/take_union.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Take on a canonical sparse union of integers, strings, and lists. +//! +//! Take gathers every sparse child, not only the one each row selects, so the cost is the sum over +//! all three variants. That is the price of the sparse layout, and it is also what makes the layout +//! affordable: Vortex does not require the children of a canonical array to be canonical, so the +//! inactive slots of a child can be encoded away. +//! +//! Rows are skewed: 98% are integers and the strings and lists are 1% each. That skew is the case +//! the layout is designed for, and it is what makes a rare variant worth encoding sparsely. +//! +//! The `dense_children` cases hold every child materialized at the union's length, which is what a +//! writer produces before compression. The `compressed_children` cases keep the dominant integer +//! child canonical and store each rare child as a `SparseArray` holding only the rows its variant +//! selects, which is what a compressor produces. +//! +//! The `compressed_children` cases run over the sub-millisecond budget the other microbenchmarks +//! hold to. What they measure is cache-missing binary searches in `Patches::take`, which only +//! appears once a child's patch indices outgrow the cache. Shrinking the array to fit the budget +//! closes the gap between the two shapes to noise and measures nothing. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::IntoArray; +use vortex::array::RecursiveCanonical; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::UnionArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::dtype::UnionVariants; +use vortex::array::scalar::Scalar; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::encodings::sparse::Sparse; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(VortexSession::default); + +const ARRAY_SIZE: usize = 100_000; + +const TAKE_SIZE: usize = 128; + +const VARIANT_NAMES: [&str; 3] = ["ints", "strings", "lists"]; + +/// The variant that almost every row selects. +const DOMINANT_VARIANT: usize = 0; + +/// The variant that `row` selects. Integers dominate, and strings and lists are rare. +fn type_id_for(row: usize) -> u8 { + match row % 100 { + 98 => 1, + 99 => 2, + _ => 0, + } +} + +fn integer_child(len: usize, rng: &mut StdRng) -> ArrayRef { + (0..len) + .map(|_| rng.random::()) + .collect::>() + .into_array() +} + +/// A `utf8` child mixing strings short enough to inline into a `VarBinView` view with longer ones +/// that have to point into the data buffer. +/// +/// A `VarBinView` view inlines up to 12 bytes. A child of uniformly short strings would gather +/// without ever reading the data buffer, which is not what a real string column costs. +fn string_child(len: usize) -> ArrayRef { + VarBinViewArray::from_iter_str((0..len).map(|i| match i % 4 { + 0 => format!("s{i}"), + _ => format!("a considerably longer string value, number {i}"), + })) + .into_array() +} + +/// A `list` child of variable-length lists over one shared element buffer. +/// +/// The lengths vary on purpose. Equal-length lists are what `FixedSizeList` is for, and they would +/// not exercise the offsets and sizes that a `ListView` gather has to rebuild. +fn list_child(len: usize, rng: &mut StdRng) -> ArrayRef { + let sizes: Buffer = (0..len).map(|_| rng.random_range(0..16)).collect(); + let offsets: Buffer = sizes + .iter() + .scan(0i32, |offset, size| { + let start = *offset; + *offset += size; + Some(start) + }) + .collect(); + + let total = offsets.last().unwrap() + sizes.last().unwrap(); + let elements = (0..total) + .map(|_| rng.random::()) + .collect::>() + .into_array(); + + ListViewArray::new( + elements, + offsets.into_array(), + sizes.into_array(), + Validity::NonNullable, + ) + .into_array() +} + +fn variant_child(variant: usize, len: usize, rng: &mut StdRng) -> ArrayRef { + match variant { + 0 => integer_child(len, rng), + 1 => string_child(len), + _ => list_child(len, rng), + } +} + +/// The rows whose type ID selects `variant`. +fn variant_indices(variant: usize) -> Buffer { + (0..ARRAY_SIZE) + .filter(|&row| usize::from(type_id_for(row)) == variant) + .map(|row| row as u64) + .collect() +} + +fn union_array(children: Vec) -> ArrayRef { + let variants = UnionVariants::new( + VARIANT_NAMES.into(), + children.iter().map(|child| child.dtype().clone()).collect(), + ) + .unwrap(); + + let type_ids = PrimitiveArray::from_iter((0..ARRAY_SIZE).map(type_id_for)); + + UnionArray::new(type_ids.into_array(), variants, children).into_array() +} + +/// A union whose children are each materialized at the union's full length. +fn dense_children_union(rng: &mut StdRng) -> ArrayRef { + union_array( + (0..VARIANT_NAMES.len()) + .map(|variant| variant_child(variant, ARRAY_SIZE, rng)) + .collect(), + ) +} + +/// A union whose rare children store only the rows their variant selects, leaving the inactive +/// slots to the sparse fill value. +/// +/// The dominant integer child stays canonical. At 98% density a sparse child would store a patch +/// index for nearly every row and buy nothing. +fn compressed_children_union(rng: &mut StdRng) -> ArrayRef { + union_array( + (0..VARIANT_NAMES.len()) + .map(|variant| { + if variant == DOMINANT_VARIANT { + return variant_child(variant, ARRAY_SIZE, rng); + } + + let indices = variant_indices(variant); + let values = variant_child(variant, indices.len(), rng); + let fill = Scalar::default_value(values.dtype()); + + Sparse::try_new(indices.into_array(), values, ARRAY_SIZE, fill) + .unwrap() + .into_array() + }) + .collect(), + ) +} + +fn random_indices(rng: &mut StdRng) -> ArrayRef { + (0..TAKE_SIZE) + .map(|_| rng.random_range(0..ARRAY_SIZE) as u64) + .collect::>() + .into_array() +} + +/// Random indices where every tenth one is null, which produces an outer union null. +fn random_nullable_indices(rng: &mut StdRng) -> ArrayRef { + PrimitiveArray::from_option_iter( + (0..TAKE_SIZE).map(|i| (i % 10 != 0).then(|| rng.random_range(0..ARRAY_SIZE) as u64)), + ) + .into_array() +} + +/// Take `indices` and execute the result, so that the child gathers actually run. +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 take_union_dense_children(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let array = dense_children_union(&mut rng); + let indices = random_indices(&mut rng); + + bench_take(bencher, array, indices); +} + +#[divan::bench] +fn take_union_compressed_children(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let array = compressed_children_union(&mut rng); + let indices = random_indices(&mut rng); + + bench_take(bencher, array, indices); +} + +#[divan::bench] +fn take_union_dense_children_nullable_indices(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let array = dense_children_union(&mut rng); + let indices = random_nullable_indices(&mut rng); + + bench_take(bencher, array, indices); +} + +#[divan::bench] +fn take_union_compressed_children_nullable_indices(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let array = compressed_children_union(&mut rng); + let indices = random_nullable_indices(&mut rng); + + bench_take(bencher, array, indices); +}