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/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..bd5a8c5196c 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,48 @@ impl Array { } } + /// Construct a `len`-row union in which every row holds `scalar`. + /// + /// 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 + /// + /// 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)), + }; + + 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..5ad0a6f1287 --- /dev/null +++ b/vortex-array/src/arrays/union/compute/take.rs @@ -0,0 +1,52 @@ +// 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; + +/// Gathers the type IDs and every sparse child at `indices`. +/// +/// 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` 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 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) + .map(Some); + } + + 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 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)?; + + 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..0dda06e2766 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. +//! +//! 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.rs b/vortex-array/src/arrays/union/tests/mod.rs similarity index 77% rename from vortex-array/src/arrays/union/tests.rs rename to vortex-array/src/arrays/union/tests/mod.rs index be82654e17a..fc0abec37b8 100644 --- a/vortex-array/src/arrays/union/tests.rs +++ b/vortex-array/src/arrays/union/tests/mod.rs @@ -5,14 +5,18 @@ 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; 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; use crate::arrays::PrimitiveArray; use crate::arrays::Union; use crate::arrays::UnionArray; @@ -27,6 +31,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 +240,68 @@ fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> { Ok(()) } +/// A constant union canonicalizes into a sparse union whose selected child repeats the scalar's +/// 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, + ), + ]; + + 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(()) +} + +#[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 `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); + + 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..0088ace43b2 --- /dev/null +++ b/vortex-array/src/arrays/union/tests/take.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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`, asserting the result reduced back to a union instead of staying a dictionary. +#[track_caller] +fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { + Ok(array + .clone() + .into_array() + .take(indices)? + .as_::() + .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(), + )?; + + 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(), + )?; + + assert_eq!(taken.dtype(), &DType::Union(variants.clone(), nullability)); + + // A nullable index widens the union but never its variants. + 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_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(), + )?; + + 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, + )?, // + ], + ) +} + +/// `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 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 via_take = empty + .take(indices.clone())? + .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(); + + let expected = || vec![Scalar::null(DType::Union(variants.clone(), nullability)); 2]; + + assert_rows(&via_take, expected())?; + assert_rows(&via_reduce, expected()) +} + +#[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(()) +} 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); +}