From 7a727c7f49afb3b3291ef538c61338ec6408df42 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 14 Oct 2026 18:19:31 +0100 Subject: [PATCH 1/4] refactor(array): make the list_contains haystack a lookup array An `ArrayHaystack` now holds its set as an array whose encoding is built for lookups and canonicalizes to the set's elements, so the lookup structures are ordinary arrays that kernels can read. Vortex builds four in-memory lookup encodings, none of which is ever serialized: - `SortedHaystack`: integers, floats by bit pattern and decimals by unscaled value, sorted without duplicates, found by binary search. - `BitmapHaystack`: such integers over a dense span, as a smallest value and a bitmap. - `BytesHaystack`: UTF-8 or binary elements, with their hash tables kept as the array's in-memory data. - `RowsHaystack`: any other elements, recursively canonical, with the indices of the distinct elements in sorted order. An encoding outside Vortex becomes a lookup array by implementing the public `HaystackVTable`, and `ArrayHaystack::new` wraps its array. `PreparedSetData::try_with_haystack` prepares a set that looks up needles through such a haystack, and `PreparedSetData::haystack` gives a needle kernel the lookup array. A set without its own haystack still builds one lazily on the first lookup and shares it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Robert Kruszewski --- vortex-array/src/array/haystack/bytes.rs | 296 ++++++++++ vortex-array/src/array/haystack/integer.rs | 490 +++++++++++++++ vortex-array/src/array/haystack/mod.rs | 557 +++++------------- vortex-array/src/array/haystack/rows.rs | 174 ++++++ vortex-array/src/array/haystack/tests.rs | 267 +++++++++ vortex-array/src/array/mod.rs | 3 +- .../fns/list_contains/prepared/data.rs | 100 ++-- 7 files changed, 1442 insertions(+), 445 deletions(-) create mode 100644 vortex-array/src/array/haystack/bytes.rs create mode 100644 vortex-array/src/array/haystack/integer.rs create mode 100644 vortex-array/src/array/haystack/rows.rs create mode 100644 vortex-array/src/array/haystack/tests.rs diff --git a/vortex-array/src/array/haystack/bytes.rs b/vortex-array/src/array/haystack/bytes.rs new file mode 100644 index 00000000000..4de939b3903 --- /dev/null +++ b/vortex-array/src/array/haystack/bytes.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A lookup array of UTF-8 or binary values, found through hash tables. + +use std::fmt::Debug; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::BuildHasher; +use std::hash::Hasher; +use std::sync::Arc; + +use smallvec::smallvec; +use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::collect_bits; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_utils::aliases::hash_map::HashTable; +use vortex_utils::aliases::hash_map::HashTableEntry; +use vortex_utils::aliases::hash_map::RandomState; + +use super::HaystackOps; +use super::HaystackVTable; +use super::haystack_vtable_common; +use super::slot; +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::DType; +use crate::validity::Validity; + +/// Builds the bytes haystack of the non-null `members`. +pub(super) fn try_new( + members: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let elements = members.execute::(ctx)?; + let tables = BytesTables::new(&elements); + let elements = elements.into_array(); + + Array::try_from_parts(ArrayParts::new( + BytesHaystack, + elements.dtype().clone(), + elements.len(), + tables, + smallvec![Some(elements)], + )) +} + +/// A lookup array of UTF-8 or binary values, found by their views where the view decides. +/// +/// Its one child holds the elements, which may repeat. Its hash tables are its in-memory data. +/// +/// A value of at most 12 bytes is inlined whole in its view, zero-padded as the compare kernel +/// requires, so a short needle is an element exactly when its whole view is the view of a short +/// element: no data buffer is read and no bytes are hashed. A long needle is found through a +/// table of the long elements, hashed by their bytes. The first 8 bytes of its view, the length +/// and a 4-byte prefix, are compared before the bytes after the prefix. +/// +/// A one-bit filter over the heads of the elements was measured in front of these lookups and +/// removed: it helped short needles by about a tenth and cost long needles that share a prefix, +/// whose heads collide, up to a quarter. +#[derive(Clone, Debug)] +pub(crate) struct BytesHaystack; + +impl BytesHaystack { + const ELEMENTS: usize = 0; +} + +/// The hash tables of a [`BytesHaystack`], over the views of its elements. +#[derive(Clone)] +pub(crate) struct BytesTables(Arc); + +struct Tables { + hasher: RandomState, + /// The distinct whole views of the elements of at most 12 bytes. + short: HashTable, + /// The indices of the distinct long elements, hashed by their bytes. + long: HashTable, +} + +impl BytesTables { + fn new(elements: &VarBinViewArray) -> Self { + let hasher = RandomState::default(); + let views = elements.views(); + let buffers = data_buffers(elements); + + // Sized up front, so that no insert grows a table and hashes every element again. + let short_len = views.iter().filter(|view| view.is_inlined()).count(); + let mut short = HashTable::with_capacity(short_len); + let mut long = HashTable::with_capacity(views.len() - short_len); + + for (idx, view) in views.iter().enumerate() { + if view.is_inlined() { + let whole = view.as_u128(); + if let HashTableEntry::Vacant(vacant) = short.entry( + hasher.hash_one(whole), + |&other| other == whole, + |&other| hasher.hash_one(other), + ) { + vacant.insert(whole); + } + continue; + } + + let value = view.bytes(&buffers); + let bytes = |other: u32| views[other as usize].bytes(&buffers); + if let HashTableEntry::Vacant(vacant) = long.entry( + hasher.hash_one(value), + |&other| bytes(other) == value, + |&other| hasher.hash_one(bytes(other)), + ) { + vacant.insert( + u32::try_from(idx).vortex_expect("a list holds fewer than 2^32 elements"), + ); + } + } + + Self(Arc::new(Tables { + hasher, + short, + long, + })) + } + + /// Whether `view`, which points into `buffers`, is the view of an element. + #[inline] + fn contains_view( + &self, + view: &BinaryView, + buffers: &[&[u8]], + element_views: &[BinaryView], + element_buffers: &[&[u8]], + ) -> bool { + let tables = &self.0; + if view.is_inlined() { + let whole = view.as_u128(); + return tables + .short + .find(tables.hasher.hash_one(whole), |&other| other == whole) + .is_some(); + } + + // Once the heads are equal, the bytes can differ only after the 4-byte prefix. + let head = view_head(view); + let value = view.bytes(buffers); + tables + .long + .find(tables.hasher.hash_one(value), |&idx| { + let element = &element_views[idx as usize]; + view_head(element) == head && element.bytes(element_buffers)[4..] == value[4..] + }) + .is_some() + } +} + +impl Debug for BytesTables { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BytesTables") + .field("short", &self.0.short.len()) + .field("long", &self.0.long.len()) + .finish() + } +} + +impl Display for BytesTables { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "short: {}, long: {}", + self.0.short.len(), + self.0.long.len() + ) + } +} + +// The tables follow from the elements, which the array compares and hashes as its child. +impl ArrayEq for BytesTables { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl ArrayHash for BytesTables { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl VTable for BytesHaystack { + type TypedArrayData = BytesTables; + type OperationsVTable = HaystackOps; + type ValidityVTable = HaystackOps; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.haystack.bytes"); + *ID + } + + fn validate( + &self, + _data: &BytesTables, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), 1, "A bytes haystack has one child"); + let elements = slots[Self::ELEMENTS] + .as_ref() + .ok_or_else(|| vortex_err!("A bytes haystack needs its elements"))?; + vortex_ensure!( + elements.is::(), + "A bytes haystack indexes the views of a VarBinView array" + ); + vortex_ensure_eq!( + elements.dtype(), + dtype, + "A bytes haystack has its elements' dtype" + ); + vortex_ensure_eq!( + elements.len(), + len, + "A bytes haystack has a row per element" + ); + Ok(()) + } + + haystack_vtable_common!("BytesHaystack"); + + fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { + "elements".to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + slot(array.slots(), Self::ELEMENTS).clone(), + )) + } +} + +impl HaystackVTable for BytesHaystack { + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + let elements = slot(array.slots(), Self::ELEMENTS) + .clone() + .execute::(ctx)?; + let element_views = elements.views(); + let element_buffers = data_buffers(&elements); + + let needles = needles.clone().execute::(ctx)?; + let buffers = data_buffers(&needles); + let tables = array.data(); + let bits = collect_bits( + needles.views(), + |view: BinaryView| { + tables.contains_view(&view, &buffers, element_views, &element_buffers) + }, + ctx.allocator(), + ); + Ok((bits, needles.validity()?)) + } +} + +/// The head of a view: the `u32` length of its value and the first 4 bytes of the value, +/// zero-padded for a value shorter than 4 bytes. +#[inline] +#[expect( + clippy::cast_possible_truncation, + reason = "the head is the low 8 bytes" +)] +fn view_head(view: &BinaryView) -> u64 { + view.as_u128() as u64 +} + +/// The host slices of an array's data buffers, indexed by a view's buffer index. +fn data_buffers(array: &VarBinViewArray) -> Vec<&[u8]> { + (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).as_slice()) + .collect() +} diff --git a/vortex-array/src/array/haystack/integer.rs b/vortex-array/src/array/haystack/integer.rs new file mode 100644 index 00000000000..e55d9fa91d2 --- /dev/null +++ b/vortex-array/src/array/haystack/integer.rs @@ -0,0 +1,490 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lookup arrays of integers: primitives by their bit patterns, and decimals by their unscaled +//! values. +//! +//! A float is an element exactly when the compare kernel would call it equal to an element, +//! which is when their bit patterns match. That distinguishes `-0.0` from `0.0` and one NaN +//! payload from another, so floats are looked up by their bits, as integers. A decimal is looked +//! up by its unscaled value, as the narrowest type that holds every value of its precision. A +//! needle has the precision of the elements, so a valid needle converts to that type whatever its +//! own storage width, and a needle stored at that type converts without a copy. + +use num_traits::ToPrimitive; +use num_traits::WrappingSub; +use smallvec::smallvec; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; +use vortex_compute::lane_kernels::collect_bits; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use super::ArrayHaystack; +use super::HaystackOps; +use super::HaystackVTable; +use super::haystack_vtable_common; +use super::slot; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::BoolArray; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::decimal::converted_buffer; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::dtype::DType; +use crate::dtype::DecimalType; +use crate::dtype::PType; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::match_each_integer_ptype; +use crate::validity::Validity; + +/// A set whose span of values needs at most this many bits per element is looked up through a +/// bitmap over the span, bounding the bitmap to a few words per element. +const BITMAP_BITS_PER_ELEMENT: u128 = 64; +/// A span this narrow is looked up through a bitmap whatever the size of the set. +const BITMAP_MIN_BITS: u128 = 1 << 12; + +/// Binds `$T` to the integer type of the primitive or decimal array `$values`, `$values` to its +/// integers, and `$needles` to the integers of the needles of its dtype, then evaluates `$body` +/// to the membership bits. Gives the bits and the validity of the needles. +macro_rules! with_integers { + ($values:expr, $needles:expr, $ctx:expr, | $T:ident, $vals:ident, $nds:ident | $body:expr) => {{ + match $values.dtype() { + DType::Primitive(ptype, _) => { + let ptype = bit_pattern_ptype(*ptype); + let values = $values + .clone() + .execute::($ctx)? + .reinterpret_cast(ptype); + let needles = $needles + .clone() + .execute::($ctx)? + .reinterpret_cast(ptype); + let validity = needles.validity()?; + + match_each_integer_ptype!(ptype, |$T| { + let $vals: &[$T] = values.as_slice::<$T>(); + let $nds: &[$T] = needles.as_slice::<$T>(); + Ok(($body, validity)) + }) + } + DType::Decimal(decimal, _) => { + let values_type = DecimalType::smallest_decimal_value_type(decimal); + let values = $values.clone().execute::($ctx)?; + let needles = $needles.clone().execute::($ctx)?; + let validity = needles.validity()?; + let valid = validity.execute_mask(needles.len(), $ctx)?; + + match_each_decimal_value_type!(values_type, |$T| { + let values = converted_buffer::<$T>(&values, &Mask::new_true(values.len()))?; + let needles = converted_buffer::<$T>(&needles, &valid)?; + let $vals: &[$T] = &values; + let $nds: &[$T] = &needles; + Ok(($body, validity)) + }) + } + dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), + } + }}; +} + +/// Builds the bitmap haystack of the non-null `members` when their span is dense, and the +/// sorted haystack otherwise. +pub(super) fn try_new(members: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match members.dtype().clone() { + DType::Primitive(ptype, _) => { + let bits_ptype = bit_pattern_ptype(ptype); + let members = members + .execute::(ctx)? + .reinterpret_cast(bits_ptype); + + match_each_integer_ptype!(bits_ptype, |T| { + let values = members.into_buffer::(); + let primitive = |values: Buffer| { + PrimitiveArray::new(values, Validity::NonNullable) + .reinterpret_cast(ptype) + .into_array() + }; + + Ok(match integer_bitmap(&values, ctx.allocator()) { + Some((min, bitmap)) => ArrayHaystack::new(BitmapHaystack::new( + primitive(Buffer::from(vec![min])), + bitmap, + )), + None => { + ArrayHaystack::new(SortedHaystack::new(primitive(sorted_values(values)))) + } + }) + }) + } + DType::Decimal(decimal, _) => { + let values_type = DecimalType::smallest_decimal_value_type(&decimal); + let members = members.execute::(ctx)?; + let all_valid = Mask::new_true(members.len()); + + match_each_decimal_value_type!(values_type, |T| { + let values = converted_buffer::(&members, &all_valid)?; + let decimals = |values: Buffer| { + DecimalArray::new(values, decimal, Validity::NonNullable).into_array() + }; + + Ok(match integer_bitmap(&values, ctx.allocator()) { + Some((min, bitmap)) => ArrayHaystack::new(BitmapHaystack::new( + decimals(Buffer::from(vec![min])), + bitmap, + )), + None => { + ArrayHaystack::new(SortedHaystack::new(decimals(sorted_values(values)))) + } + }) + }) + } + dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), + } +} + +/// A lookup array of integers sorted without duplicates, found by binary search. +/// +/// Its one child holds the values, a primitive array sorted by bit pattern or a decimal array +/// sorted by unscaled value. A hash set and, for a handful of elements, a linear scan both lost +/// to the binary search at every set size measured by the `list_contains_set` benchmark, up to +/// 16 384 elements. +#[derive(Clone, Debug)] +pub(crate) struct SortedHaystack; + +impl SortedHaystack { + const VALUES: usize = 0; + + fn new(values: ArrayRef) -> Array { + let dtype = values.dtype().clone(); + let len = values.len(); + Array::try_from_parts(ArrayParts::new( + Self, + dtype, + len, + EmptyArrayData, + smallvec![Some(values)], + )) + .vortex_expect("A sorted haystack is valid by construction") + } +} + +impl VTable for SortedHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = HaystackOps; + type ValidityVTable = HaystackOps; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.haystack.sorted"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), 1, "A sorted haystack has one child"); + let values = slots[Self::VALUES] + .as_ref() + .ok_or_else(|| vortex_error::vortex_err!("A sorted haystack needs its values"))?; + vortex_ensure_eq!( + values.dtype(), + dtype, + "A sorted haystack has its values' dtype" + ); + vortex_ensure_eq!(values.len(), len, "A sorted haystack has a row per value"); + Ok(()) + } + + haystack_vtable_common!("SortedHaystack"); + + fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { + "values".to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + slot(array.slots(), Self::VALUES).clone(), + )) + } +} + +impl HaystackVTable for SortedHaystack { + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + let allocator = ctx.allocator().clone(); + with_integers!( + slot(array.slots(), Self::VALUES), + needles, + ctx, + |T, sorted, needles| collect_bits( + needles, + |needle| sorted.binary_search(&needle).is_ok(), + &allocator + ) + ) + } +} + +/// A lookup array of integers spanning a dense range: one bit per integer of the span. +/// +/// Its children are the smallest integer, a primitive or decimal array of one row, and a bitmap +/// over the span above it. Its rows are the integers of the set bits, in increasing order of +/// offset from the smallest. +#[derive(Clone, Debug)] +pub(crate) struct BitmapHaystack; + +impl BitmapHaystack { + const MIN: usize = 0; + const BITMAP: usize = 1; + const NAMES: [&'static str; 2] = ["min", "bitmap"]; + + fn new(min: ArrayRef, bitmap: BitBuffer) -> Array { + let dtype = min.dtype().clone(); + let len = bitmap.true_count(); + let bitmap = BoolArray::new(bitmap, Validity::NonNullable).into_array(); + Array::try_from_parts(ArrayParts::new( + Self, + dtype, + len, + EmptyArrayData, + smallvec![Some(min), Some(bitmap)], + )) + .vortex_expect("A bitmap haystack is valid by construction") + } + + fn bitmap(array: &[Option], ctx: &mut ExecutionCtx) -> VortexResult { + Ok(slot(array, Self::BITMAP) + .clone() + .execute::(ctx)? + .to_bit_buffer()) + } +} + +impl VTable for BitmapHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = HaystackOps; + type ValidityVTable = HaystackOps; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.haystack.bitmap"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + _len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), 2, "A bitmap haystack has two children"); + let (Some(min), Some(bitmap)) = (&slots[Self::MIN], &slots[Self::BITMAP]) else { + vortex_bail!("A bitmap haystack needs its smallest value and its bitmap"); + }; + vortex_ensure_eq!( + min.dtype(), + dtype, + "A bitmap haystack has its values' dtype" + ); + vortex_ensure_eq!(min.len(), 1, "A bitmap haystack has one smallest value"); + vortex_ensure!( + matches!(bitmap.dtype(), DType::Bool(_)), + "A bitmap haystack needs a boolean bitmap, got {}", + bitmap.dtype() + ); + Ok(()) + } + + haystack_vtable_common!("BitmapHaystack"); + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + Self::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let bitmap = Self::bitmap(array.slots(), ctx)?; + let min = slot(array.slots(), Self::MIN).clone(); + + let values = match min.dtype() { + DType::Primitive(ptype, _) => { + let ptype = *ptype; + let bits_ptype = bit_pattern_ptype(ptype); + let min = min + .execute::(ctx)? + .reinterpret_cast(bits_ptype); + match_each_integer_ptype!(bits_ptype, |T| { + let values = bitmap_values(min.as_slice::()[0], &bitmap); + PrimitiveArray::new(values, Validity::NonNullable) + .reinterpret_cast(ptype) + .into_array() + }) + } + DType::Decimal(decimal, _) => { + let decimal = *decimal; + let values_type = DecimalType::smallest_decimal_value_type(&decimal); + let min = min.execute::(ctx)?; + match_each_decimal_value_type!(values_type, |T| { + let min = converted_buffer::(&min, &Mask::new_true(1))?[0]; + DecimalArray::new(bitmap_values(min, &bitmap), decimal, Validity::NonNullable) + .into_array() + }) + } + dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), + }; + + Ok(ExecutionResult::done(values)) + } +} + +impl HaystackVTable for BitmapHaystack { + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + let allocator = ctx.allocator().clone(); + let bitmap = Self::bitmap(array.slots(), ctx)?; + with_integers!( + slot(array.slots(), Self::MIN), + needles, + ctx, + |T, min, needles| { + let min = min[0]; + collect_bits( + needles, + // A needle below the smallest element wraps past the bitmap, so one comparison + // checks both bounds. + |needle| { + needle + .offset_from(min) + .is_some_and(|offset| offset < bitmap.len() && bitmap.value(offset)) + }, + &allocator, + ) + } + ) + } +} + +/// The integer type with a float's bit pattern, or the type itself for an integer. +fn bit_pattern_ptype(ptype: PType) -> PType { + match ptype { + PType::F16 => PType::U16, + PType::F32 => PType::U32, + PType::F64 => PType::U64, + _ => ptype, + } +} + +/// An unsigned modular distance, rejecting offsets too wide to address a bitmap. +/// Keeping the arithmetic at the physical width also handles signed ranges spanning zero. +trait SetInteger: Copy + Ord + Send + Sync + 'static { + fn offset_from(self, min: Self) -> Option; + + /// The integer `offset` above `self`, for an offset within a bitmap's span. + fn add_offset(self, offset: usize) -> Self; +} + +macro_rules! impl_set_integer { + ($($signed:ty => $unsigned:ty),* $(,)?) => { + $(impl SetInteger for $signed { + fn offset_from(self, min: Self) -> Option { + usize::try_from((self as $unsigned).wrapping_sub(min as $unsigned)).ok() + } + + fn add_offset(self, offset: usize) -> Self { + let offset = <$unsigned>::try_from(offset) + .vortex_expect("A bitmap offset fits the integer type"); + (self as $unsigned).wrapping_add(offset) as $signed + } + })* + }; +} + +impl_set_integer!( + u8 => u8, u16 => u16, u32 => u32, u64 => u64, + i8 => u8, i16 => u16, i32 => u32, i64 => u64, i128 => u128, +); + +impl SetInteger for i256 { + fn offset_from(self, min: Self) -> Option { + self.wrapping_sub(&min).to_usize() + } + + fn add_offset(self, offset: usize) -> Self { + let offset = u64::try_from(offset).vortex_expect("A bitmap offset fits in 64 bits"); + i256::wrapping_add(&self, i256::from(offset)) + } +} + +/// The smallest integer and a bitmap over the span of `values` above it, when the span is dense. +fn integer_bitmap( + values: &[T], + allocator: &BufferAllocatorRef, +) -> Option<(T, BitBuffer)> { + let min = *values.iter().min()?; + let max = *values.iter().max()?; + let span = max.offset_from(min)?; + if span == usize::MAX + || span as u128 >= (values.len() as u128 * BITMAP_BITS_PER_ELEMENT).max(BITMAP_MIN_BITS) + { + return None; + } + let mut bitmap = BitBufferMut::from_buffer( + BufferMut::zeroed_in((span + 1).div_ceil(8), allocator.clone()), + 0, + span + 1, + ); + for &value in values { + bitmap.set( + value + .offset_from(min) + .vortex_expect("value within bitmap span"), + ); + } + Some((min, bitmap.freeze())) +} + +/// The integers of the set bits of `bitmap`, above `min`. +fn bitmap_values(min: T, bitmap: &BitBuffer) -> Buffer { + bitmap + .set_indices() + .map(|offset| min.add_offset(offset)) + .collect() +} + +fn sorted_values(values: Buffer) -> Buffer { + if values.is_sorted_by(|a, b| a < b) { + return values; + } + let mut sorted = values.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + Buffer::from(sorted) +} diff --git a/vortex-array/src/array/haystack/mod.rs b/vortex-array/src/array/haystack/mod.rs index f1e04f9c686..7dc919400ef 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-array/src/array/haystack/mod.rs @@ -2,73 +2,89 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Lookups of needles among the elements of an array. +//! +//! An [`ArrayHaystack`] holds the non-null elements of a set as a lookup array: an array whose +//! encoding is built to find needles, and which canonicalizes to the elements it holds. Vortex +//! builds a sorted, a bitmap, a hashed or a row haystack from the elements. An encoding outside +//! Vortex becomes a lookup array by implementing [`HaystackVTable`]. + +mod bytes; +mod integer; +mod rows; +#[cfg(test)] +mod tests; -use std::hash::BuildHasher; - -use num_traits::ToPrimitive; -use num_traits::WrappingSub; use vortex_buffer::BitBuffer; -use vortex_buffer::BitBufferMut; -use vortex_buffer::Buffer; -use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; -use vortex_compute::lane_kernels::collect_bits; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_mask::Mask; -use vortex_utils::aliases::hash_map::HashTable; -use vortex_utils::aliases::hash_map::HashTableEntry; -use vortex_utils::aliases::hash_map::RandomState; use crate::ArrayRef; +use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -use crate::RecursiveCanonical; -use crate::arrays::DecimalArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::decimal::converted_buffer; -use crate::arrays::primitive::PrimitiveArrayExt; -use crate::arrays::varbinview::BinaryView; +use crate::array::Array; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; use crate::dtype::DType; -use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::match_each_integer_ptype; -use crate::scalar_fn::fns::binary::build_row_comparator; +use crate::scalar::Scalar; use crate::validity::Validity; -/// A set whose span of values needs at most this many bits per element is looked up through a -/// bitmap over the span, bounding the bitmap to a few words per element. -const BITMAP_BITS_PER_ELEMENT: u128 = 64; -/// A span this narrow is looked up through a bitmap whatever the size of the set. -const BITMAP_MIN_BITS: u128 = 1 << 12; +/// An encoding whose arrays can be the lookup array of an [`ArrayHaystack`]. +/// +/// Every row of a lookup array is an element of the set, so a lookup array holds no null rows. It +/// may hold an element more than once. +pub trait HaystackVTable: VTable { + /// One bit per needle, set when the needle is an element of `array`, and the validity of the + /// needles. + /// + /// The needles have the dtype of `array`, ignoring nullability. The bit of a null needle has + /// no effect. + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)>; +} + +/// The [`HaystackVTable::contains`] of the encoding of a lookup array, with the array erased. +type ContainsFn = + fn(&ArrayRef, &ArrayRef, &mut ExecutionCtx) -> VortexResult<(BitBuffer, Validity)>; -/// Looks up needles among the elements of an array, for many batches of needles. +/// Looks up needles among the elements of a lookup array, for many batches of needles. /// /// The haystack counterpart of [`RepeatedArrayProbe`](crate::RepeatedArrayProbe): where a probe /// reads the rows of an array by index, a haystack answers whether each row of a batch of /// needles is an element of the array. A haystack pays off only over repeated lookups, so it /// has no one-off form. It is built eagerly. A holder that must build it lazily, or share it /// between threads, keeps it in a shared cell. -/// -/// Integers, and floats by their bit patterns, are looked up in an [`IntegerSet`] of their own -/// type, and decimals in an [`IntegerSet`] of their unscaled values. UTF-8 and binary values are -/// found through a hash table, and nested values through sorted row indices with the same -/// comparator as equality. No haystack constructs per-element expressions or materializes -/// scalars in its loop. -pub(crate) struct ArrayHaystack { - state: Box, +#[derive(Clone)] +pub struct ArrayHaystack { + array: ArrayRef, + contains: ContainsFn, } impl ArrayHaystack { + /// A haystack over a lookup array of your own encoding. + pub fn new(array: Array) -> Self { + Self { + array: array.into_array(), + contains: contains_of::, + } + } + /// Builds the haystack of the elements of `array`. /// - /// A null element never equals a needle, so the haystack leaves null elements out. - pub(crate) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + /// A null element never equals a needle, so the haystack leaves null elements out. Integers, + /// and floats by their bit patterns, go into a bitmap over their span when the span is dense, + /// and into a sorted array otherwise, as do decimals by their unscaled values. UTF-8 and + /// binary values go into hash tables, and nested values into sorted row indices with the same + /// comparator as equality. + /// + /// # Errors + /// + /// Fails when the elements cannot be executed. + pub fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let valid = array.validity()?.execute_mask(array.len(), ctx)?; let members = if valid.all_true() { array.clone() @@ -76,404 +92,125 @@ impl ArrayHaystack { array.filter(valid)? }; - Ok(Self { - state: new_state(members, ctx)?, - }) + match members.dtype() { + DType::Primitive(..) | DType::Decimal(..) => integer::try_new(members, ctx), + DType::Utf8(_) | DType::Binary(_) => Ok(Self::new(bytes::try_new(members, ctx)?)), + _ => Ok(Self::new(rows::try_new(members, ctx)?)), + } } - /// One membership bit per needle, and the validity of the needles, which have the dtype of - /// the elements. - pub(crate) fn contains( - &self, - needles: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)> { - self.state.contains(needles, ctx) + /// The lookup array, which canonicalizes to the elements of the haystack. + pub fn array(&self) -> &ArrayRef { + &self.array } -} -/// The state that an [`ArrayHaystack`] keeps for the canonical encoding of its elements. -trait HaystackState: Send + Sync { - /// One membership bit per needle, and the validity of the needles. - fn contains( + /// One bit per needle, set when the needle is an element, and the validity of the needles. + /// + /// The needles have the dtype of the lookup array, ignoring nullability. + /// + /// # Errors + /// + /// Fails when the needles cannot be executed. + pub fn contains( &self, needles: &ArrayRef, ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)>; -} - -/// Builds the state of the non-null `elements`. -fn new_state(elements: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - Ok(match elements.dtype() { - DType::Primitive(ptype, _) => { - let ptype = bit_pattern_ptype(*ptype); - let elements = elements - .execute::(ctx)? - .reinterpret_cast(ptype); - match_each_integer_ptype!(ptype, |T| { - let set = IntegerSet::::new(elements.into_buffer(), ctx.allocator()); - Box::new(PrimitiveSet(set)) as Box - }) - } - DType::Decimal(decimal, _) => { - let values_type = DecimalType::smallest_decimal_value_type(decimal); - let elements = elements.execute::(ctx)?; - let all_valid = Mask::new_true(elements.len()); - match_each_decimal_value_type!(values_type, |T| { - let values = converted_buffer::(&elements, &all_valid)?; - Box::new(DecimalSet(IntegerSet::new(values, ctx.allocator()))) - as Box - }) - } - DType::Utf8(_) | DType::Binary(_) => { - Box::new(BytesSet::new(elements.execute::(ctx)?)) - } - _ => Box::new(RowSet::try_new(elements, ctx)?), - }) -} - -/// The distinct values of a set of integers of one type. -enum IntegerSet { - /// Values spanning a dense range: one bit per value of the span above `min`. - Bitmap { min: T, bitmap: BitBuffer }, - /// Values sorted without duplicates. - Sorted(Buffer), -} - -impl IntegerSet { - /// A bitmap over the values' span when the span is dense, and a sorted slice otherwise. - /// - /// A hash set and, for a handful of elements, a linear scan both lost to the binary search at - /// every set size measured by the `list_contains_set` benchmark, up to 16 384 elements. - fn new(values: Buffer, allocator: &BufferAllocatorRef) -> Self { - match integer_bitmap(&values, allocator) { - Some((min, bitmap)) => Self::Bitmap { min, bitmap }, - None => Self::Sorted(sorted_values(values)), - } + ) -> VortexResult<(BitBuffer, Validity)> { + (self.contains)(&self.array, needles, ctx) } +} - /// One bit per needle, set when the needle is an element. - fn contains(&self, needles: &[T], allocator: &BufferAllocatorRef) -> BitBuffer { - match self { - Self::Bitmap { min, bitmap } => collect_bits( - needles, - // A needle below the smallest element wraps past the bitmap, so one comparison - // checks both bounds. - |needle| { - needle - .offset_from(*min) - .is_some_and(|offset| offset < bitmap.len() && bitmap.value(offset)) - }, - allocator, - ), - Self::Sorted(sorted) => collect_bits( - needles, - |needle| sorted.binary_search(&needle).is_ok(), - allocator, - ), - } - } +fn contains_of( + array: &ArrayRef, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(BitBuffer, Validity)> { + V::contains(array.as_::(), needles, ctx) } -/// Primitive integers, or floats by their bit patterns. +/// The operations and validity that the lookup arrays of Vortex share. /// -/// A float is a member exactly when the compare kernel would call it equal to an element, which -/// is when their bit patterns match — distinguishing `-0.0` from `0.0` and one NaN payload from -/// another — so floats are looked up by their bits, as integers. -struct PrimitiveSet(IntegerSet); +/// Every row is valid. A row is read from the canonical form, since a lookup array exists for +/// lookups, not for reads. +#[derive(Clone, Debug)] +pub(crate) struct HaystackOps; -impl HaystackState for PrimitiveSet { - fn contains( - &self, - needles: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)> { - let needles = needles - .clone() - .execute::(ctx)? - .reinterpret_cast(T::PTYPE); - let bits = self.0.contains(needles.as_slice::(), ctx.allocator()); - Ok((bits, needles.validity()?)) +impl ValidityVTable for HaystackOps { + fn validity(array: ArrayView<'_, V>) -> VortexResult { + Ok(Validity::from(array.dtype().nullability())) } } -/// The unscaled values of decimals, as the narrowest type that holds every value of their -/// precision. -/// -/// A needle has the precision of the elements, so a valid needle converts to that type whatever -/// its own storage width, and a needle stored at that type converts without a copy. -struct DecimalSet(IntegerSet); +impl OperationsVTable for HaystackOps { + type ProbeState = (); -impl HaystackState for DecimalSet { - fn contains( - &self, - needles: &ArrayRef, + fn scalar_at( + array: ArrayView<'_, V>, + index: usize, ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)> { - let needles = needles.clone().execute::(ctx)?; - let validity = needles.validity()?; - let valid = validity.execute_mask(needles.len(), ctx)?; - let values = converted_buffer::(&needles, &valid)?; - Ok((self.0.contains(&values, ctx.allocator()), validity)) + ) -> VortexResult { + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .execute_scalar(index, ctx) } } -/// UTF-8 or binary elements, found by their views where the view decides. -/// -/// A value of at most 12 bytes is inlined whole in its view, zero-padded as the compare kernel -/// requires, so a short needle is an element exactly when its whole view is the view of a short -/// element: no data buffer is read and no bytes are hashed. A long needle is found through a -/// table of the long elements, hashed by their bytes. The first 8 bytes of its view, the length -/// and a 4-byte prefix, are compared before the bytes after the prefix. -/// -/// A one-bit filter over the heads of the elements was measured in front of these lookups and -/// removed: it helped short needles by about a tenth and cost long needles that share a prefix, -/// whose heads collide, up to a quarter. -struct BytesSet { - hasher: RandomState, - /// The distinct whole views of the elements of at most 12 bytes. - short: HashTable, - /// The elements, for the bytes of the long ones. - elements: VarBinViewArray, - /// The indices of the distinct long elements, hashed by their bytes. - long: HashTable, +/// The child in `slot` of a lookup array of Vortex, which validation found present. +fn slot(slots: &[Option], slot: usize) -> &ArrayRef { + slots[slot] + .as_ref() + .unwrap_or_else(|| vortex_error::vortex_panic!("Haystack slot {slot} must be present")) } -impl BytesSet { - fn new(elements: VarBinViewArray) -> Self { - let hasher = RandomState::default(); - let views = elements.views(); - let buffers = data_buffers(&elements); - - // Sized up front, so that no insert grows a table and hashes every element again. - let short_len = views.iter().filter(|view| view.is_inlined()).count(); - let mut short = HashTable::with_capacity(short_len); - let mut long = HashTable::with_capacity(views.len() - short_len); - - for (idx, view) in views.iter().enumerate() { - if view.is_inlined() { - let whole = view.as_u128(); - if let HashTableEntry::Vacant(vacant) = short.entry( - hasher.hash_one(whole), - |&other| other == whole, - |&other| hasher.hash_one(other), - ) { - vacant.insert(whole); - } - continue; - } - - let value = view.bytes(&buffers); - let bytes = |other: u32| views[other as usize].bytes(&buffers); - if let HashTableEntry::Vacant(vacant) = long.entry( - hasher.hash_one(value), - |&other| bytes(other) == value, - |&other| hasher.hash_one(bytes(other)), - ) { - vacant.insert( - u32::try_from(idx).vortex_expect("a list holds fewer than 2^32 elements"), - ); - } +/// The buffer and serialization methods that every lookup array of Vortex shares: it keeps its +/// data in children and in memory, and it is never serialized. +macro_rules! haystack_vtable_common { + ($name:literal) => { + fn nbuffers(_array: $crate::array::ArrayView<'_, Self>) -> usize { + 0 } - Self { - hasher, - short, - elements, - long, + fn buffer( + _array: $crate::array::ArrayView<'_, Self>, + _idx: usize, + ) -> $crate::buffer::BufferHandle { + ::vortex_error::vortex_panic!("A haystack has no buffers") } - } - /// Whether `view`, which points into `buffers`, is the view of an element. - #[inline] - fn contains_view( - &self, - view: &BinaryView, - buffers: &[&[u8]], - element_views: &[BinaryView], - element_buffers: &[&[u8]], - ) -> bool { - if view.is_inlined() { - let whole = view.as_u128(); - return self - .short - .find(self.hasher.hash_one(whole), |&other| other == whole) - .is_some(); + fn buffer_name(_array: $crate::array::ArrayView<'_, Self>, _idx: usize) -> Option { + None } - // Once the heads are equal, the bytes can differ only after the 4-byte prefix. - let head = view_head(view); - let value = view.bytes(buffers); - self.long - .find(self.hasher.hash_one(value), |&idx| { - let element = &element_views[idx as usize]; - view_head(element) == head && element.bytes(element_buffers)[4..] == value[4..] - }) - .is_some() - } -} - -impl HaystackState for BytesSet { - fn contains( - &self, - needles: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)> { - let element_views = self.elements.views(); - let element_buffers = data_buffers(&self.elements); - let array = needles.clone().execute::(ctx)?; - let buffers = data_buffers(&array); - let bits = collect_bits( - array.views(), - |view: BinaryView| self.contains_view(&view, &buffers, element_views, &element_buffers), - ctx.allocator(), - ); - Ok((bits, array.validity()?)) - } -} - -/// The head of a view: the `u32` length of its value and the first 4 bytes of the value, -/// zero-padded for a value shorter than 4 bytes. -#[inline] -#[expect( - clippy::cast_possible_truncation, - reason = "the head is the low 8 bytes" -)] -fn view_head(view: &BinaryView) -> u64 { - view.as_u128() as u64 -} - -/// Recursively canonical elements, indexed in sorted order with duplicates removed. -struct RowSet { - elements: ArrayRef, - indices: Vec, -} - -impl RowSet { - fn try_new(elements: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let elements = elements.execute::(ctx)?.0.into_array(); - let mut indices: Vec = (0..elements.len()).collect(); - if !indices.is_empty() { - let compare = build_row_comparator(&elements, &elements, ctx)?; - if !indices.is_sorted_by(|&lhs, &rhs| compare(lhs, rhs).is_le()) { - indices.sort_unstable_by(|&lhs, &rhs| compare(lhs, rhs)); - } - indices.dedup_by(|lhs, rhs| compare(*lhs, *rhs).is_eq()); + fn with_buffers( + &self, + array: $crate::array::ArrayView<'_, Self>, + buffers: &[$crate::buffer::BufferHandle], + ) -> ::vortex_error::VortexResult<$crate::array::ArrayParts> { + $crate::array::with_empty_buffers(self, array, buffers) } - Ok(Self { elements, indices }) - } -} -impl HaystackState for RowSet { - fn contains( - &self, - needles: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult<(BitBuffer, Validity)> { - if self.indices.is_empty() { - return Ok(( - BitBuffer::full_in(false, needles.len(), ctx.allocator().clone()), - needles.validity()?, - )); + fn serialize( + _array: $crate::array::ArrayView<'_, Self>, + _session: &::vortex_session::VortexSession, + ) -> ::vortex_error::VortexResult>> { + ::vortex_error::vortex_bail!("{} is never serialized", $name) } - let needles = needles - .clone() - .execute::(ctx)? - .0 - .into_array(); - // The comparator materializes child buffers and validity once, including decimal - // widening and nested offsets. Each search then reads those buffers directly. - let compare = build_row_comparator(&self.elements, &needles, ctx)?; - let bits = BitBuffer::collect_bool_in( - needles.len(), - |row| { - self.indices - .binary_search_by(|&element| compare(element, row)) - .is_ok() - }, - ctx.allocator().clone(), - ); - Ok((bits, needles.validity()?)) - } -} - -/// The host slices of an array's data buffers, indexed by a view's buffer index. -fn data_buffers(array: &VarBinViewArray) -> Vec<&[u8]> { - (0..array.data_buffers().len()) - .map(|idx| array.buffer(idx).as_slice()) - .collect() -} - -/// The integer type with a float's bit pattern, or the type itself for an integer. -fn bit_pattern_ptype(ptype: PType) -> PType { - match ptype { - PType::F16 => PType::U16, - PType::F32 => PType::U32, - PType::F64 => PType::U64, - _ => ptype, - } -} -/// An unsigned modular distance, rejecting offsets too wide to address a bitmap. -/// Keeping the subtraction at the physical width also handles signed ranges spanning zero. -trait SetInteger: Copy + Ord + Send + Sync + 'static { - fn offset_from(self, min: Self) -> Option; -} - -macro_rules! impl_set_integer { - ($($signed:ty => $unsigned:ty),* $(,)?) => { - $(impl SetInteger for $signed { - fn offset_from(self, min: Self) -> Option { - usize::try_from((self as $unsigned).wrapping_sub(min as $unsigned)).ok() - } - })* + fn deserialize( + &self, + _dtype: &$crate::dtype::DType, + _len: usize, + _metadata: &[u8], + _buffers: &[$crate::buffer::BufferHandle], + _children: &dyn $crate::serde::ArrayChildren, + _session: &::vortex_session::VortexSession, + ) -> ::vortex_error::VortexResult<$crate::array::ArrayParts> { + ::vortex_error::vortex_bail!("{} is never serialized", $name) + } }; } -impl_set_integer!( - u8 => u8, u16 => u16, u32 => u32, u64 => u64, - i8 => u8, i16 => u16, i32 => u32, i64 => u64, i128 => u128, -); - -impl SetInteger for i256 { - fn offset_from(self, min: Self) -> Option { - self.wrapping_sub(&min).to_usize() - } -} - -fn integer_bitmap( - values: &[T], - allocator: &BufferAllocatorRef, -) -> Option<(T, BitBuffer)> { - let min = *values.iter().min()?; - let max = *values.iter().max()?; - let span = max.offset_from(min)?; - if span == usize::MAX - || span as u128 >= (values.len() as u128 * BITMAP_BITS_PER_ELEMENT).max(BITMAP_MIN_BITS) - { - return None; - } - let mut bitmap = BitBufferMut::from_buffer( - BufferMut::zeroed_in((span + 1).div_ceil(8), allocator.clone()), - 0, - span + 1, - ); - for &value in values { - bitmap.set( - value - .offset_from(min) - .vortex_expect("value within bitmap span"), - ); - } - Some((min, bitmap.freeze())) -} - -fn sorted_values(values: Buffer) -> Buffer { - if values.is_sorted_by(|a, b| a < b) { - return values; - } - let mut sorted = values.to_vec(); - sorted.sort_unstable(); - sorted.dedup(); - Buffer::from(sorted) -} +use haystack_vtable_common; diff --git a/vortex-array/src/array/haystack/rows.rs b/vortex-array/src/array/haystack/rows.rs new file mode 100644 index 00000000000..bdb25bf2b02 --- /dev/null +++ b/vortex-array/src/array/haystack/rows.rs @@ -0,0 +1,174 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A lookup array of any values, found by binary search over their sorted rows. + +use smallvec::smallvec; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::registry::CachedId; + +use super::HaystackOps; +use super::HaystackVTable; +use super::haystack_vtable_common; +use super::slot; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar_fn::fns::binary::build_row_comparator; +use crate::validity::Validity; + +/// Builds the row haystack of the non-null `members`. +pub(super) fn try_new( + members: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let elements = members.execute::(ctx)?.0.into_array(); + let mut order: Vec = (0..elements.len() as u64).collect(); + if !order.is_empty() { + let compare = build_row_comparator(&elements, &elements, ctx)?; + let compare = |lhs: &u64, rhs: &u64| compare(row(*lhs), row(*rhs)); + if !order.is_sorted_by(|lhs, rhs| compare(lhs, rhs).is_le()) { + order.sort_unstable_by(compare); + } + order.dedup_by(|lhs, rhs| compare(lhs, rhs).is_eq()); + } + let order = PrimitiveArray::new(Buffer::from(order), Validity::NonNullable).into_array(); + + Array::try_from_parts(ArrayParts::new( + RowsHaystack, + elements.dtype().clone(), + order.len(), + EmptyArrayData, + smallvec![Some(elements), Some(order)], + )) +} + +/// A lookup array of any values, found by binary search with the same comparator as equality. +/// +/// Its children are the elements, recursively canonical, and the indices of the distinct +/// elements in sorted order. Its rows are the distinct elements in that order. The comparator +/// materializes child buffers and validity once per batch of needles, including decimal widening +/// and nested offsets. Each search then reads those buffers directly. +#[derive(Clone, Debug)] +pub(crate) struct RowsHaystack; + +impl RowsHaystack { + const ELEMENTS: usize = 0; + const ORDER: usize = 1; + const NAMES: [&'static str; 2] = ["elements", "order"]; +} + +impl VTable for RowsHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = HaystackOps; + type ValidityVTable = HaystackOps; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.haystack.rows"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), 2, "A row haystack has two children"); + let (Some(elements), Some(order)) = (&slots[Self::ELEMENTS], &slots[Self::ORDER]) else { + vortex_bail!("A row haystack needs its elements and their order"); + }; + vortex_ensure_eq!( + elements.dtype(), + dtype, + "A row haystack has its elements' dtype" + ); + vortex_ensure_eq!( + order.dtype(), + &DType::Primitive(PType::U64, Nullability::NonNullable), + "A row haystack orders its elements by u64 indices" + ); + vortex_ensure_eq!( + order.len(), + len, + "A row haystack has a row per distinct element" + ); + Ok(()) + } + + haystack_vtable_common!("RowsHaystack"); + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + Self::NAMES[idx].to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let elements = slot(array.slots(), Self::ELEMENTS); + let order = slot(array.slots(), Self::ORDER).clone(); + Ok(ExecutionResult::done(elements.take(order)?)) + } +} + +impl HaystackVTable for RowsHaystack { + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + if array.is_empty() { + return Ok(( + BitBuffer::full_in(false, needles.len(), ctx.allocator().clone()), + needles.validity()?, + )); + } + + let elements = slot(array.slots(), Self::ELEMENTS); + let order = slot(array.slots(), Self::ORDER) + .clone() + .execute::(ctx)?; + let order = order.as_slice::(); + let needles = needles + .clone() + .execute::(ctx)? + .0 + .into_array(); + + let compare = build_row_comparator(elements, &needles, ctx)?; + let bits = BitBuffer::collect_bool_in( + needles.len(), + |needle| { + order + .binary_search_by(|&element| compare(row(element), needle)) + .is_ok() + }, + ctx.allocator().clone(), + ); + Ok((bits, needles.validity()?)) + } +} + +/// A row index of the elements, which fit in memory and so in `usize`. +#[expect( + clippy::cast_possible_truncation, + reason = "the elements of a haystack fit in memory" +)] +fn row(index: u64) -> usize { + index as usize +} diff --git a/vortex-array/src/array/haystack/tests.rs b/vortex-array/src/array/haystack/tests.rs new file mode 100644 index 00000000000..ef12cfe96a1 --- /dev/null +++ b/vortex-array/src/array/haystack/tests.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use smallvec::smallvec; +use vortex_buffer::BitBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::ArrayHaystack; +use super::HaystackOps; +use super::HaystackVTable; +use super::bytes::BytesHaystack; +use super::haystack_vtable_common; +use super::integer::BitmapHaystack; +use super::integer::SortedHaystack; +use super::rows::RowsHaystack; +use super::slot; +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::VTable; +use crate::arrays::BoolArray; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::list_contains::ListContainsOptions; +use crate::scalar_fn::fns::list_contains::PreparedSetData; +use crate::validity::Validity; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// A struct array with one `i32` field of `values`. +fn struct_of(values: &[i32]) -> ArrayRef { + StructArray::try_new( + FieldNames::from(["a"]), + vec![PrimitiveArray::from_iter(values.iter().copied()).into_array()], + values.len(), + Validity::NonNullable, + ) + .unwrap() + .into_array() +} + +fn build(elements: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + ArrayHaystack::try_new(&elements, ctx) +} + +/// Each kind of element goes into its lookup array, which canonicalizes to the distinct non-null +/// elements in the order of the lookup. +#[rstest] +#[case::dense_integers( + PrimitiveArray::from_option_iter([Some(5i64), Some(-3), None, Some(5)]).into_array(), + BitmapHaystack.id(), + PrimitiveArray::from_iter([-3i64, 5]).into_array(), +)] +#[case::sparse_integers( + PrimitiveArray::from_iter([1i64 << 40, 7, -(1i64 << 40), 7]).into_array(), + SortedHaystack.id(), + PrimitiveArray::from_iter([-(1i64 << 40), 7, 1i64 << 40]).into_array(), +)] +#[case::floats_by_bits( + PrimitiveArray::from_iter([1.5f64, -0.0, 0.0]).into_array(), + SortedHaystack.id(), + PrimitiveArray::from_iter([0.0f64, 1.5, -0.0]).into_array(), +)] +#[case::dense_decimals( + DecimalArray::from_option_iter::([Some(12), None, Some(10)], DecimalDType::new(5, 1)) + .into_array(), + BitmapHaystack.id(), + DecimalArray::from_iter::([10, 12], DecimalDType::new(5, 1)).into_array(), +)] +#[case::rows( + struct_of(&[3, 1, 3]), + RowsHaystack.id(), + struct_of(&[1, 3]), +)] +fn test_lookup_array_of_each_kind( + #[case] elements: ArrayRef, + #[case] encoding: ArrayId, + #[case] canonical: ArrayRef, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = build(elements, &mut ctx)?; + + assert_eq!(haystack.array().encoding_id(), encoding); + assert_arrays_eq!(haystack.array().clone(), canonical, &mut ctx); + Ok(()) +} + +#[test] +fn test_bytes_lookup_array_keeps_the_non_null_elements() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let elements = + VarBinViewArray::from_iter_nullable_str([Some("a"), None, Some("a long string value")]) + .into_array(); + let haystack = build(elements, &mut ctx)?; + + assert_eq!(haystack.array().encoding_id(), BytesHaystack.id()); + assert_arrays_eq!( + haystack.array().clone(), + VarBinViewArray::from_iter_nullable_str([Some("a"), Some("a long string value")]), + &mut ctx + ); + + let needles = VarBinViewArray::from_iter_nullable_str([ + Some("a long string value"), + Some("a long string valuf"), + Some("a"), + None, + ]) + .into_array(); + let (bits, validity) = haystack.contains(&needles, &mut ctx)?; + assert_eq!(bits, BitBuffer::from_iter([true, false, true, false])); + assert_eq!(validity.execute_mask(4, &mut ctx)?.true_count(), 3); + Ok(()) +} + +#[test] +fn test_lookup_array_reads_a_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = build(buffer![9u8, 2, 4].into_array(), &mut ctx)?; + + assert_eq!( + haystack.array().execute_scalar(1, &mut ctx)?, + Scalar::from(4u8) + ); + Ok(()) +} + +/// A lookup array of your own encoding: `i32` values found by a linear scan. +#[derive(Clone, Debug)] +struct LinearHaystack; + +impl LinearHaystack { + fn new(values: ArrayRef) -> Array { + Array::try_from_parts(ArrayParts::new( + Self, + values.dtype().clone(), + values.len(), + EmptyArrayData, + smallvec![Some(values)], + )) + .unwrap() + } +} + +impl VTable for LinearHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = HaystackOps; + type ValidityVTable = HaystackOps; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("test.haystack.linear"); + *ID + } + + fn validate( + &self, + _data: &EmptyArrayData, + _dtype: &DType, + _len: usize, + _slots: &[Option], + ) -> VortexResult<()> { + Ok(()) + } + + haystack_vtable_common!("LinearHaystack"); + + fn slot_name(_array: ArrayView<'_, Self>, _idx: usize) -> String { + "values".to_string() + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(slot(array.slots(), 0).clone())) + } +} + +impl HaystackVTable for LinearHaystack { + fn contains( + array: ArrayView<'_, Self>, + needles: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + let values = slot(array.slots(), 0) + .clone() + .execute::(ctx)?; + let needles = needles.clone().execute::(ctx)?; + let bits = needles + .as_slice::() + .iter() + .map(|needle| values.as_slice::().contains(needle)) + .collect(); + Ok((bits, needles.validity()?)) + } +} + +fn i32_list(values: &[i32]) -> Scalar { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + Scalar::list( + std::sync::Arc::new(dtype), + values.iter().map(|&value| Scalar::from(value)).collect(), + Nullability::NonNullable, + ) +} + +#[test] +fn test_prepared_set_looks_up_through_your_own_haystack() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // The lookup array holds 7 only, so the set finds 7 and not 3, although the list holds both. + let haystack = ArrayHaystack::new(LinearHaystack::new(buffer![7i32].into_array())); + let set = PreparedSetData::try_with_haystack(i32_list(&[3, 7]), haystack)?; + + assert_eq!( + set.haystack(&mut ctx)?.array().encoding_id(), + LinearHaystack.id() + ); + assert_arrays_eq!( + set.contains( + &buffer![7i32, 3, 1].into_array(), + &ListContainsOptions::default(), + &mut ctx + )?, + BoolArray::from_iter([true, false, false]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_your_own_haystack_must_hold_the_element_dtype() { + let haystack = ArrayHaystack::new(LinearHaystack::new(buffer![7i64].into_array())); + assert!(PreparedSetData::try_with_haystack(i32_list(&[7]), haystack).is_err()); +} + +#[test] +fn test_a_built_haystack_canonicalizes() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = build(buffer![1u16, 1, 3].into_array(), &mut ctx)?; + let canonical = haystack.array().clone().execute::(&mut ctx)?; + assert_arrays_eq!( + canonical.into_array(), + buffer![1u16, 3].into_array(), + &mut ctx + ); + Ok(()) +} diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index fae796e1714..2967f9c39df 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -33,7 +33,8 @@ pub use erased::*; mod plugin; pub use plugin::*; mod haystack; -pub(crate) use haystack::ArrayHaystack; +pub use haystack::ArrayHaystack; +pub use haystack::HaystackVTable; mod probe; pub use probe::*; diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs index aaaf047c1d9..f601a7ae3ed 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs @@ -35,8 +35,10 @@ use crate::validity::Validity; /// The options of a [`PreparedSet`]: the list that every row holds, and its set. /// /// The set is shared by every clone, so every batch a [`PreparedSet`] is applied to, and every -/// slice or filter of the array it gives, looks up needles in the one set. The haystack needs -/// an [`ExecutionCtx`] to materialize the elements, so it is built on the first lookup. +/// slice or filter of the array it gives, looks up needles in the one set. The set looks up +/// needles through an [`ArrayHaystack`]: one of your own, given to [`Self::try_with_haystack`], +/// or else one that [`ArrayHaystack::try_new`] builds on the first lookup, since it needs an +/// [`ExecutionCtx`] to materialize the elements. /// /// [`PreparedSet`]: crate::scalar_fn::fns::list_contains::PreparedSet #[derive(Clone)] @@ -50,29 +52,11 @@ struct SharedSet { has_null_element: bool, /// Whether the list holds no element at all, counting null elements. is_empty: bool, - /// The elements, and the haystack of the non-null ones, built on the first lookup. - set: OnceLock>, -} - -/// The elements of the list, and their haystack. -struct ElementSet { - /// All elements of the list, null elements included, with the element dtype of the list. - elements: ArrayRef, - haystack: ArrayHaystack, -} - -impl ElementSet { - fn try_new(list: &Scalar, ctx: &mut ExecutionCtx) -> VortexResult { - // The canonical form of a constant list holds its elements once, whatever its length. - let elements = ConstantArray::new(list.clone(), 1) - .into_array() - .execute::(ctx)? - .elements() - .clone(); - let haystack = ArrayHaystack::try_new(&elements, ctx)?; - - Ok(Self { elements, haystack }) - } + /// All elements of the list, null elements included, with the element dtype of the list, + /// built on first use. + elements: OnceLock>, + /// The haystack of the non-null elements, given or built on the first lookup. + haystack: OnceLock>, } impl PreparedSetData { @@ -111,10 +95,34 @@ impl PreparedSetData { literal, has_null_element, is_empty, - set: OnceLock::new(), + elements: OnceLock::new(), + haystack: OnceLock::new(), }))) } + /// Prepares the non-null list scalar `list` as a set that looks up needles through + /// `haystack`, a haystack of your own over the elements of the list. + /// + /// # Errors + /// + /// Fails when `list` is not a list, or is null, or when the lookup array of `haystack` does + /// not have the element dtype of the list. + pub fn try_with_haystack(list: Scalar, haystack: ArrayHaystack) -> VortexResult { + let set = Self::try_new(list)?; + vortex_ensure!( + set.element_dtype() + .eq_ignore_nullability(haystack.array().dtype()), + "A haystack of {} cannot look up elements of {}", + haystack.array().dtype(), + set.element_dtype() + ); + + if set.0.haystack.set(Ok(haystack)).is_err() { + vortex_panic!("A new prepared set has no haystack yet"); + } + Ok(set) + } + /// Whether `other` shares this set rather than preparing its own. #[cfg(test)] pub(super) fn shares_set_with(&self, other: &Self) -> bool { @@ -135,16 +143,40 @@ impl PreparedSetData { /// /// # Errors /// - /// Fails when the elements cannot be executed into the haystack. + /// Fails when the list cannot be executed into its elements. pub fn elements(&self, ctx: &mut ExecutionCtx) -> VortexResult<&ArrayRef> { - Ok(&self.element_set(ctx)?.elements) + self.0 + .elements + .get_or_init(|| { + // The canonical form of a constant list holds its elements once, whatever its + // length. + ConstantArray::new(self.list().clone(), 1) + .into_array() + .execute::(ctx) + .map(|list| list.elements().clone()) + .map_err(Arc::new) + }) + .as_ref() + .map_err(|err| Arc::clone(err).into()) } - /// The elements and their haystack, built on the first call and shared after it. - fn element_set<'a>(&'a self, ctx: &mut ExecutionCtx) -> VortexResult<&'a ElementSet> { + /// The haystack that looks up needles among the elements, built on the first call unless + /// one was given, and shared after it. + /// + /// A kernel of a needle encoding can read the lookup array of the haystack, for example to + /// look up its own values in a lookup array of its own encoding. + /// + /// # Errors + /// + /// Fails when the elements cannot be executed into a haystack. + pub fn haystack(&self, ctx: &mut ExecutionCtx) -> VortexResult<&ArrayHaystack> { self.0 - .set - .get_or_init(|| ElementSet::try_new(self.list(), ctx).map_err(Arc::new)) + .haystack + .get_or_init(|| { + self.elements(ctx) + .and_then(|elements| ArrayHaystack::try_new(elements, ctx)) + .map_err(Arc::new) + }) .as_ref() .map_err(|err| Arc::clone(err).into()) } @@ -187,7 +219,7 @@ impl PreparedSetData { } self.check_needle_dtype(needles.dtype())?; - let (bits, needle_validity) = self.element_set(ctx)?.haystack.contains(needles, ctx)?; + let (bits, needle_validity) = self.haystack(ctx)?.contains(needles, ctx)?; self.result_from_bits(bits, needle_validity, needles.dtype(), options) } @@ -207,7 +239,7 @@ impl PreparedSetData { ) -> VortexResult { self.check_needle_dtype(needle.dtype())?; let needles = ConstantArray::new(needle.clone(), 1).into_array(); - let (bits, needle_validity) = self.element_set(ctx)?.haystack.contains(&needles, ctx)?; + let (bits, needle_validity) = self.haystack(ctx)?.contains(&needles, ctx)?; self.result_from_bits(bits, needle_validity, needle.dtype(), options)? .execute_scalar(0, ctx) } From 4c6567b07c1ce834686ab30ae9b582ab7ceaa651 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 14 Oct 2026 23:54:59 +0000 Subject: [PATCH 2/4] refactor(array): let needle encodings drive list_contains lookups A haystack no longer decodes its needles. `HaystackVTable::contains` takes canonical needles, and `PreparedSetData::contains` returns a lazy `ListContains` over needles that are not canonical. Execution steps the needles one decoding step at a time, so a `ListContainsElementKernel` of each encoding they pass through can take the lookup before they reach the haystack. The lookup arrays read their own children by downcasting, which validation now guarantees, and read a row through their children rather than executing to canonical: a bitmap selects its nth set bit, and a row haystack indexes its elements through its order. `PreparedSetData` keeps only the haystack. The elements it cached were read only to build the haystack, whose lookup array canonicalizes to the non-null elements. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TC7eug8jjTe1e1pbDreBo9 --- vortex-array/src/array/haystack/bytes.rs | 40 +++- vortex-array/src/array/haystack/integer.rs | 178 +++++++++++------- vortex-array/src/array/haystack/mod.rs | 49 ++--- vortex-array/src/array/haystack/rows.rs | 46 +++-- vortex-array/src/array/haystack/tests.rs | 66 +++++-- .../src/scalar_fn/fns/list_contains/mod.rs | 6 +- .../fns/list_contains/prepared/data.rs | 81 ++++---- .../fns/list_contains/prepared/tests.rs | 6 +- 8 files changed, 297 insertions(+), 175 deletions(-) diff --git a/vortex-array/src/array/haystack/bytes.rs b/vortex-array/src/array/haystack/bytes.rs index 4de939b3903..8ae0084557f 100644 --- a/vortex-array/src/array/haystack/bytes.rs +++ b/vortex-array/src/array/haystack/bytes.rs @@ -15,6 +15,7 @@ use vortex_buffer::BitBuffer; use vortex_compute::lane_kernels::collect_bits; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -23,13 +24,14 @@ use vortex_utils::aliases::hash_map::HashTable; use vortex_utils::aliases::hash_map::HashTableEntry; use vortex_utils::aliases::hash_map::RandomState; -use super::HaystackOps; use super::HaystackVTable; +use super::HaystackValidity; use super::haystack_vtable_common; use super::slot; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; +use crate::CanonicalView; use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -38,10 +40,13 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; +use crate::array::OperationsVTable; use crate::array::VTable; +use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; +use crate::scalar::Scalar; use crate::validity::Validity; /// Builds the bytes haystack of the non-null `members`. @@ -202,8 +207,8 @@ impl ArrayHash for BytesTables { impl VTable for BytesHaystack { type TypedArrayData = BytesTables; - type OperationsVTable = HaystackOps; - type ValidityVTable = HaystackOps; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; fn id(&self) -> ArrayId { static ID: CachedId = CachedId::new("vortex.haystack.bytes"); @@ -222,7 +227,7 @@ impl VTable for BytesHaystack { .as_ref() .ok_or_else(|| vortex_err!("A bytes haystack needs its elements"))?; vortex_ensure!( - elements.is::(), + elements.is::(), "A bytes haystack indexes the views of a VarBinView array" ); vortex_ensure_eq!( @@ -251,19 +256,38 @@ impl VTable for BytesHaystack { } } +impl OperationsVTable for BytesHaystack { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + slot(array.slots(), Self::ELEMENTS).execute_scalar(index, ctx) + } +} + impl HaystackVTable for BytesHaystack { fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { + let CanonicalView::VarBinView(needles) = needles else { + vortex_bail!( + "A bytes haystack of {} cannot look up needles of {}", + array.dtype(), + needles.to_array_ref().dtype() + ); + }; + let needles = needles.into_owned(); let elements = slot(array.slots(), Self::ELEMENTS) - .clone() - .execute::(ctx)?; + .as_::() + .into_owned(); let element_views = elements.views(); let element_buffers = data_buffers(&elements); - let needles = needles.clone().execute::(ctx)?; let buffers = data_buffers(&needles); let tables = array.data(); let bits = collect_bits( diff --git a/vortex-array/src/array/haystack/integer.rs b/vortex-array/src/array/haystack/integer.rs index e55d9fa91d2..878826a8d23 100644 --- a/vortex-array/src/array/haystack/integer.rs +++ b/vortex-array/src/array/haystack/integer.rs @@ -25,15 +25,17 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::registry::CachedId; use super::ArrayHaystack; -use super::HaystackOps; use super::HaystackVTable; +use super::HaystackValidity; use super::haystack_vtable_common; use super::slot; use crate::ArrayRef; +use crate::CanonicalView; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; @@ -42,9 +44,13 @@ use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; use crate::array::VTable; +use crate::arrays::Bool; use crate::arrays::BoolArray; +use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::decimal::converted_buffer; @@ -55,6 +61,7 @@ use crate::dtype::PType; use crate::dtype::i256; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; +use crate::scalar::Scalar; use crate::validity::Validity; /// A set whose span of values needs at most this many bits per element is looked up through a @@ -64,21 +71,15 @@ const BITMAP_BITS_PER_ELEMENT: u128 = 64; const BITMAP_MIN_BITS: u128 = 1 << 12; /// Binds `$T` to the integer type of the primitive or decimal array `$values`, `$values` to its -/// integers, and `$needles` to the integers of the needles of its dtype, then evaluates `$body` -/// to the membership bits. Gives the bits and the validity of the needles. +/// integers, and `$needles` to the integers of the canonical needles of its dtype, then evaluates +/// `$body` to the membership bits. Gives the bits and the validity of the needles. macro_rules! with_integers { ($values:expr, $needles:expr, $ctx:expr, | $T:ident, $vals:ident, $nds:ident | $body:expr) => {{ - match $values.dtype() { - DType::Primitive(ptype, _) => { + match ($values.dtype(), $needles) { + (DType::Primitive(ptype, _), CanonicalView::Primitive(needles)) => { let ptype = bit_pattern_ptype(*ptype); - let values = $values - .clone() - .execute::($ctx)? - .reinterpret_cast(ptype); - let needles = $needles - .clone() - .execute::($ctx)? - .reinterpret_cast(ptype); + let values = $values.as_::().reinterpret_cast(ptype); + let needles = needles.reinterpret_cast(ptype); let validity = needles.validity()?; match_each_integer_ptype!(ptype, |$T| { @@ -87,10 +88,10 @@ macro_rules! with_integers { Ok(($body, validity)) }) } - DType::Decimal(decimal, _) => { + (DType::Decimal(decimal, _), CanonicalView::Decimal(needles)) => { let values_type = DecimalType::smallest_decimal_value_type(decimal); - let values = $values.clone().execute::($ctx)?; - let needles = $needles.clone().execute::($ctx)?; + let values = $values.as_::().into_owned(); + let needles = needles.into_owned(); let validity = needles.validity()?; let valid = validity.execute_mask(needles.len(), $ctx)?; @@ -102,7 +103,10 @@ macro_rules! with_integers { Ok(($body, validity)) }) } - dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), + (dtype, needles) => vortex_bail!( + "An integer haystack of {dtype} cannot look up needles of {}", + needles.to_array_ref().dtype() + ), } }}; } @@ -190,8 +194,8 @@ impl SortedHaystack { impl VTable for SortedHaystack { type TypedArrayData = EmptyArrayData; - type OperationsVTable = HaystackOps; - type ValidityVTable = HaystackOps; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; fn id(&self) -> ArrayId { static ID: CachedId = CachedId::new("vortex.haystack.sorted"); @@ -214,6 +218,11 @@ impl VTable for SortedHaystack { dtype, "A sorted haystack has its values' dtype" ); + vortex_ensure!( + is_integer_array(values), + "A sorted haystack holds a primitive or decimal array, got {}", + values.encoding_id() + ); vortex_ensure_eq!(values.len(), len, "A sorted haystack has a row per value"); Ok(()) } @@ -231,10 +240,22 @@ impl VTable for SortedHaystack { } } +impl OperationsVTable for SortedHaystack { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + slot(array.slots(), Self::VALUES).execute_scalar(index, ctx) + } +} + impl HaystackVTable for SortedHaystack { fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { let allocator = ctx.allocator().clone(); @@ -278,18 +299,15 @@ impl BitmapHaystack { .vortex_expect("A bitmap haystack is valid by construction") } - fn bitmap(array: &[Option], ctx: &mut ExecutionCtx) -> VortexResult { - Ok(slot(array, Self::BITMAP) - .clone() - .execute::(ctx)? - .to_bit_buffer()) + fn bitmap(slots: &[Option]) -> BitBuffer { + slot(slots, Self::BITMAP).as_::().to_bit_buffer() } } impl VTable for BitmapHaystack { type TypedArrayData = EmptyArrayData; - type OperationsVTable = HaystackOps; - type ValidityVTable = HaystackOps; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; fn id(&self) -> ArrayId { static ID: CachedId = CachedId::new("vortex.haystack.bitmap"); @@ -314,9 +332,14 @@ impl VTable for BitmapHaystack { ); vortex_ensure_eq!(min.len(), 1, "A bitmap haystack has one smallest value"); vortex_ensure!( - matches!(bitmap.dtype(), DType::Bool(_)), - "A bitmap haystack needs a boolean bitmap, got {}", - bitmap.dtype() + is_integer_array(min), + "A bitmap haystack holds its smallest value as a primitive or decimal array, got {}", + min.encoding_id() + ); + vortex_ensure!( + bitmap.is::(), + "A bitmap haystack needs a bool array bitmap, got {}", + bitmap.encoding_id() ); Ok(()) } @@ -327,49 +350,36 @@ impl VTable for BitmapHaystack { Self::NAMES[idx].to_string() } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let bitmap = Self::bitmap(array.slots(), ctx)?; - let min = slot(array.slots(), Self::MIN).clone(); - - let values = match min.dtype() { - DType::Primitive(ptype, _) => { - let ptype = *ptype; - let bits_ptype = bit_pattern_ptype(ptype); - let min = min - .execute::(ctx)? - .reinterpret_cast(bits_ptype); - match_each_integer_ptype!(bits_ptype, |T| { - let values = bitmap_values(min.as_slice::()[0], &bitmap); - PrimitiveArray::new(values, Validity::NonNullable) - .reinterpret_cast(ptype) - .into_array() - }) - } - DType::Decimal(decimal, _) => { - let decimal = *decimal; - let values_type = DecimalType::smallest_decimal_value_type(&decimal); - let min = min.execute::(ctx)?; - match_each_decimal_value_type!(values_type, |T| { - let min = converted_buffer::(&min, &Mask::new_true(1))?[0]; - DecimalArray::new(bitmap_values(min, &bitmap), decimal, Validity::NonNullable) - .into_array() - }) - } - dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), - }; - + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let bitmap = Self::bitmap(array.slots()); + let values = offset_values(slot(array.slots(), Self::MIN), bitmap.set_indices())?; Ok(ExecutionResult::done(values)) } } +impl OperationsVTable for BitmapHaystack { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let offset = Self::bitmap(array.slots()) + .select(index) + .ok_or_else(|| vortex_err!("A bitmap haystack has no row {index}"))?; + offset_values(slot(array.slots(), Self::MIN), [offset].into_iter())?.execute_scalar(0, ctx) + } +} + impl HaystackVTable for BitmapHaystack { fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { let allocator = ctx.allocator().clone(); - let bitmap = Self::bitmap(array.slots(), ctx)?; + let bitmap = Self::bitmap(array.slots()); with_integers!( slot(array.slots(), Self::MIN), needles, @@ -392,6 +402,40 @@ impl HaystackVTable for BitmapHaystack { } } +/// Whether `array` holds integers in a form the integer haystacks read without executing it. +fn is_integer_array(array: &ArrayRef) -> bool { + array.is::() || array.is::() +} + +/// The integers `offsets` above the one integer of `min`, a primitive or decimal array. +fn offset_values(min: &ArrayRef, offsets: impl Iterator) -> VortexResult { + Ok(match min.dtype() { + DType::Primitive(ptype, _) => { + let ptype = *ptype; + let bits_ptype = bit_pattern_ptype(ptype); + let min = min.as_::().reinterpret_cast(bits_ptype); + match_each_integer_ptype!(bits_ptype, |T| { + let min = min.as_slice::()[0]; + let values: Buffer = offsets.map(|offset| min.add_offset(offset)).collect(); + PrimitiveArray::new(values, Validity::NonNullable) + .reinterpret_cast(ptype) + .into_array() + }) + } + DType::Decimal(decimal, _) => { + let decimal = *decimal; + let values_type = DecimalType::smallest_decimal_value_type(&decimal); + let min = min.as_::().into_owned(); + match_each_decimal_value_type!(values_type, |T| { + let min = converted_buffer::(&min, &Mask::new_true(1))?[0]; + let values: Buffer = offsets.map(|offset| min.add_offset(offset)).collect(); + DecimalArray::new(values, decimal, Validity::NonNullable).into_array() + }) + } + dtype => vortex_bail!("An integer haystack holds primitives or decimals, not {dtype}"), + }) +} + /// The integer type with a float's bit pattern, or the type itself for an integer. fn bit_pattern_ptype(ptype: PType) -> PType { match ptype { @@ -471,14 +515,6 @@ fn integer_bitmap( Some((min, bitmap.freeze())) } -/// The integers of the set bits of `bitmap`, above `min`. -fn bitmap_values(min: T, bitmap: &BitBuffer) -> Buffer { - bitmap - .set_indices() - .map(|offset| min.add_offset(offset)) - .collect() -} - fn sorted_values(values: Buffer) -> Buffer { if values.is_sorted_by(|a, b| a < b) { return values; diff --git a/vortex-array/src/array/haystack/mod.rs b/vortex-array/src/array/haystack/mod.rs index 7dc919400ef..8cb330a43ff 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-array/src/array/haystack/mod.rs @@ -18,16 +18,14 @@ use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use crate::ArrayRef; -use crate::Canonical; +use crate::CanonicalView; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::Array; use crate::array::ArrayView; -use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::dtype::DType; -use crate::scalar::Scalar; use crate::validity::Validity; /// An encoding whose arrays can be the lookup array of an [`ArrayHaystack`]. @@ -38,18 +36,21 @@ pub trait HaystackVTable: VTable { /// One bit per needle, set when the needle is an element of `array`, and the validity of the /// needles. /// - /// The needles have the dtype of `array`, ignoring nullability. The bit of a null needle has - /// no effect. + /// The needles are canonical, with the dtype of `array` ignoring nullability. They reach the + /// haystack only once their own encoding has declined the lookup at every step of its + /// decoding, through a + /// [`ListContainsElementKernel`](crate::scalar_fn::fns::list_contains::ListContainsElementKernel). + /// The bit of a null needle has no effect. fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)>; } /// The [`HaystackVTable::contains`] of the encoding of a lookup array, with the array erased. type ContainsFn = - fn(&ArrayRef, &ArrayRef, &mut ExecutionCtx) -> VortexResult<(BitBuffer, Validity)>; + fn(&ArrayRef, CanonicalView<'_>, &mut ExecutionCtx) -> VortexResult<(BitBuffer, Validity)>; /// Looks up needles among the elements of a lookup array, for many batches of needles. /// @@ -106,14 +107,14 @@ impl ArrayHaystack { /// One bit per needle, set when the needle is an element, and the validity of the needles. /// - /// The needles have the dtype of the lookup array, ignoring nullability. + /// The needles are canonical, with the dtype of the lookup array ignoring nullability. /// /// # Errors /// - /// Fails when the needles cannot be executed. + /// Fails when the needles do not have the dtype of the lookup array. pub fn contains( &self, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { (self.contains)(&self.array, needles, ctx) @@ -122,42 +123,22 @@ impl ArrayHaystack { fn contains_of( array: &ArrayRef, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { V::contains(array.as_::(), needles, ctx) } -/// The operations and validity that the lookup arrays of Vortex share. -/// -/// Every row is valid. A row is read from the canonical form, since a lookup array exists for -/// lookups, not for reads. +/// The validity that the lookup arrays of Vortex share: every row is valid. #[derive(Clone, Debug)] -pub(crate) struct HaystackOps; +pub(crate) struct HaystackValidity; -impl ValidityVTable for HaystackOps { +impl ValidityVTable for HaystackValidity { fn validity(array: ArrayView<'_, V>) -> VortexResult { Ok(Validity::from(array.dtype().nullability())) } } -impl OperationsVTable for HaystackOps { - type ProbeState = (); - - fn scalar_at( - array: ArrayView<'_, V>, - index: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .execute_scalar(index, ctx) - } -} - /// The child in `slot` of a lookup array of Vortex, which validation found present. fn slot(slots: &[Option], slot: usize) -> &ArrayRef { slots[slot] diff --git a/vortex-array/src/array/haystack/rows.rs b/vortex-array/src/array/haystack/rows.rs index bdb25bf2b02..189cb891bf5 100644 --- a/vortex-array/src/array/haystack/rows.rs +++ b/vortex-array/src/array/haystack/rows.rs @@ -8,14 +8,16 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_session::registry::CachedId; -use super::HaystackOps; use super::HaystackVTable; +use super::HaystackValidity; use super::haystack_vtable_common; use super::slot; use crate::ArrayRef; +use crate::CanonicalView; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; @@ -25,11 +27,14 @@ use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; use crate::array::VTable; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::scalar::Scalar; use crate::scalar_fn::fns::binary::build_row_comparator; use crate::validity::Validity; @@ -76,8 +81,8 @@ impl RowsHaystack { impl VTable for RowsHaystack { type TypedArrayData = EmptyArrayData; - type OperationsVTable = HaystackOps; - type ValidityVTable = HaystackOps; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; fn id(&self) -> ArrayId { static ID: CachedId = CachedId::new("vortex.haystack.rows"); @@ -105,6 +110,11 @@ impl VTable for RowsHaystack { &DType::Primitive(PType::U64, Nullability::NonNullable), "A row haystack orders its elements by u64 indices" ); + vortex_ensure!( + order.is::(), + "A row haystack holds its order as a primitive array, got {}", + order.encoding_id() + ); vortex_ensure_eq!( order.len(), len, @@ -126,12 +136,28 @@ impl VTable for RowsHaystack { } } +impl OperationsVTable for RowsHaystack { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let element = slot(array.slots(), Self::ORDER) + .as_::() + .as_slice::()[index]; + slot(array.slots(), Self::ELEMENTS).execute_scalar(row(element), ctx) + } +} + impl HaystackVTable for RowsHaystack { fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, + needles: CanonicalView<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { + let needles = needles.to_array_ref(); if array.is_empty() { return Ok(( BitBuffer::full_in(false, needles.len(), ctx.allocator().clone()), @@ -140,15 +166,11 @@ impl HaystackVTable for RowsHaystack { } let elements = slot(array.slots(), Self::ELEMENTS); - let order = slot(array.slots(), Self::ORDER) - .clone() - .execute::(ctx)?; + let order = slot(array.slots(), Self::ORDER).as_::(); let order = order.as_slice::(); - let needles = needles - .clone() - .execute::(ctx)? - .0 - .into_array(); + // The needles are canonical, but the row comparator also reads the children of nested + // needles, such as the fields of a struct, as canonical arrays. + let needles = needles.execute::(ctx)?.0.into_array(); let compare = build_row_comparator(elements, &needles, ctx)?; let bits = BitBuffer::collect_bool_in( diff --git a/vortex-array/src/array/haystack/tests.rs b/vortex-array/src/array/haystack/tests.rs index ef12cfe96a1..df2d3068dc5 100644 --- a/vortex-array/src/array/haystack/tests.rs +++ b/vortex-array/src/array/haystack/tests.rs @@ -8,20 +8,23 @@ use smallvec::smallvec; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use super::ArrayHaystack; -use super::HaystackOps; use super::HaystackVTable; +use super::HaystackValidity; use super::bytes::BytesHaystack; use super::haystack_vtable_common; use super::integer::BitmapHaystack; use super::integer::SortedHaystack; use super::rows::RowsHaystack; use super::slot; +use crate::AnyCanonical; use crate::ArrayRef; use crate::Canonical; +use crate::CanonicalView; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; @@ -31,9 +34,11 @@ use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; use crate::array::VTable; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; @@ -130,20 +135,43 @@ fn test_bytes_lookup_array_keeps_the_non_null_elements() -> VortexResult<()> { None, ]) .into_array(); - let (bits, validity) = haystack.contains(&needles, &mut ctx)?; + let (bits, validity) = haystack.contains(needles.as_::(), &mut ctx)?; assert_eq!(bits, BitBuffer::from_iter([true, false, true, false])); assert_eq!(validity.execute_mask(4, &mut ctx)?.true_count(), 3); Ok(()) } +/// Each kind of lookup array reads a row in the order of its lookup. +#[rstest] +#[case::bitmap(buffer![9u8, 2, 4].into_array(), Scalar::from(4u8))] +#[case::sorted(buffer![1u64 << 40, 7, 1 << 41].into_array(), Scalar::from(1u64 << 40))] +#[case::decimal_bitmap( + DecimalArray::from_iter::([12, 10, 11], DecimalDType::new(5, 1)).into_array(), + Scalar::decimal(11i32.into(), DecimalDType::new(5, 1), Nullability::NonNullable), +)] +#[case::bytes( + VarBinViewArray::from_iter_str(["b", "a long string value"]).into_array(), + Scalar::from("a long string value"), +)] +fn test_lookup_array_reads_a_row( + #[case] elements: ArrayRef, + #[case] expected: Scalar, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = build(elements, &mut ctx)?; + + assert_eq!(haystack.array().execute_scalar(1, &mut ctx)?, expected); + Ok(()) +} + #[test] -fn test_lookup_array_reads_a_row() -> VortexResult<()> { +fn test_rows_lookup_array_reads_a_row() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let haystack = build(buffer![9u8, 2, 4].into_array(), &mut ctx)?; + let haystack = build(struct_of(&[3, 1, 3]), &mut ctx)?; assert_eq!( haystack.array().execute_scalar(1, &mut ctx)?, - Scalar::from(4u8) + struct_of(&[3]).execute_scalar(0, &mut ctx)? ); Ok(()) } @@ -167,8 +195,8 @@ impl LinearHaystack { impl VTable for LinearHaystack { type TypedArrayData = EmptyArrayData; - type OperationsVTable = HaystackOps; - type ValidityVTable = HaystackOps; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; fn id(&self) -> ArrayId { static ID: CachedId = CachedId::new("test.haystack.linear"); @@ -196,16 +224,28 @@ impl VTable for LinearHaystack { } } +impl OperationsVTable for LinearHaystack { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, Self>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + slot(array.slots(), 0).execute_scalar(index, ctx) + } +} + impl HaystackVTable for LinearHaystack { fn contains( array: ArrayView<'_, Self>, - needles: &ArrayRef, - ctx: &mut ExecutionCtx, + needles: CanonicalView<'_>, + _ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { - let values = slot(array.slots(), 0) - .clone() - .execute::(ctx)?; - let needles = needles.clone().execute::(ctx)?; + let CanonicalView::Primitive(needles) = needles else { + vortex_bail!("A linear haystack looks up primitive needles"); + }; + let values = slot(array.slots(), 0).as_::(); let bits = needles .as_slice::() .iter() diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 5d92baac6f6..7747a23edc7 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -354,9 +354,11 @@ fn compute_list_contains( .into_array()); } - // The reduce step has already replaced a constant list with a prepared set. + // The reduce step has already replaced a constant list with a prepared set. The needles take + // one step of their decoding at a time, so that a kernel of each encoding they pass through + // can take the lookup before they are canonical. if let Some(set) = PreparedSet::set_of(array) { - return set.contains(value, options, ctx); + return set.contains(&value.clone().execute::(ctx)?, options, ctx); } let nullability = options.result_nullability(array.dtype(), value.dtype()); diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs index f601a7ae3ed..3637664c62e 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs @@ -11,13 +11,18 @@ use std::sync::OnceLock; use vortex_buffer::BitBuffer; use vortex_error::SharedVortexResult; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use super::PreparedSet; +use crate::AnyCanonical; use crate::ArrayHaystack; use crate::ArrayRef; +use crate::Canonical; +use crate::CanonicalView; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; @@ -28,6 +33,7 @@ use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListContainsOptions; use crate::scalar_fn::fns::literal::Literal; use crate::validity::Validity; @@ -52,9 +58,6 @@ struct SharedSet { has_null_element: bool, /// Whether the list holds no element at all, counting null elements. is_empty: bool, - /// All elements of the list, null elements included, with the element dtype of the list, - /// built on first use. - elements: OnceLock>, /// The haystack of the non-null elements, given or built on the first lookup. haystack: OnceLock>, } @@ -95,7 +98,6 @@ impl PreparedSetData { literal, has_null_element, is_empty, - elements: OnceLock::new(), haystack: OnceLock::new(), }))) } @@ -139,27 +141,6 @@ impl PreparedSetData { self.0.literal.as_::() } - /// The elements of the list, one row each, null elements included. - /// - /// # Errors - /// - /// Fails when the list cannot be executed into its elements. - pub fn elements(&self, ctx: &mut ExecutionCtx) -> VortexResult<&ArrayRef> { - self.0 - .elements - .get_or_init(|| { - // The canonical form of a constant list holds its elements once, whatever its - // length. - ConstantArray::new(self.list().clone(), 1) - .into_array() - .execute::(ctx) - .map(|list| list.elements().clone()) - .map_err(Arc::new) - }) - .as_ref() - .map_err(|err| Arc::clone(err).into()) - } - /// The haystack that looks up needles among the elements, built on the first call unless /// one was given, and shared after it. /// @@ -173,8 +154,12 @@ impl PreparedSetData { self.0 .haystack .get_or_init(|| { - self.elements(ctx) - .and_then(|elements| ArrayHaystack::try_new(elements, ctx)) + // The canonical form of a constant list holds its elements once, whatever its + // length. + ConstantArray::new(self.list().clone(), 1) + .into_array() + .execute::(ctx) + .and_then(|list| ArrayHaystack::try_new(list.elements(), ctx)) .map_err(Arc::new) }) .as_ref() @@ -202,7 +187,13 @@ impl PreparedSetData { /// semantics, which gives `false` for every needle. Under SQL null semantics, a list that holds /// a null element gives `null` for a needle that matches no element. /// - /// A constant needle is looked up once, and gives a constant result. + /// A constant needle is looked up once, and gives a constant result. Canonical needles are + /// looked up in the haystack. Any other needles are not decoded here: the result is a lazy + /// [`ListContains`] over them, so that their encoding chooses how to look them up, through a + /// [`ListContainsElementKernel`] at any step of its decoding. A kernel therefore passes the + /// parts of its array, such as the values of a dictionary, and never the array itself. + /// + /// [`ListContainsElementKernel`]: crate::scalar_fn::fns::list_contains::ListContainsElementKernel /// /// # Errors /// @@ -219,8 +210,15 @@ impl PreparedSetData { } self.check_needle_dtype(needles.dtype())?; - let (bits, needle_validity) = self.haystack(ctx)?.contains(needles, ctx)?; - self.result_from_bits(bits, needle_validity, needles.dtype(), options) + match needles.as_opt::() { + Some(canonical) => self.lookup(canonical, needles.dtype(), options, ctx), + None => Ok(ListContains::try_new_opts( + PreparedSet::new_array(self.clone(), needles.len()), + needles.clone(), + *options, + )? + .into_array()), + } } /// Whether the constant `needle` is an element of the list, under `options`. @@ -238,12 +236,29 @@ impl PreparedSetData { ctx: &mut ExecutionCtx, ) -> VortexResult { self.check_needle_dtype(needle.dtype())?; - let needles = ConstantArray::new(needle.clone(), 1).into_array(); - let (bits, needle_validity) = self.haystack(ctx)?.contains(&needles, ctx)?; - self.result_from_bits(bits, needle_validity, needle.dtype(), options)? + let needles = ConstantArray::new(needle.clone(), 1) + .into_array() + .execute::(ctx)? + .into_array(); + let needles = needles + .as_opt::() + .vortex_expect("A canonical array matches AnyCanonical"); + self.lookup(needles, needle.dtype(), options, ctx)? .execute_scalar(0, ctx) } + /// Looks up canonical needles in the haystack. + fn lookup( + &self, + needles: CanonicalView<'_>, + needle_dtype: &DType, + options: &ListContainsOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let (bits, needle_validity) = self.haystack(ctx)?.contains(needles, ctx)?; + self.result_from_bits(bits, needle_validity, needle_dtype, options) + } + /// Makes the result of [`Self::contains`] from one membership bit per needle. /// /// A kernel that finds the matches of its needles without this haystack, for example from diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs index fe3c6b46da6..d6067ab028a 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs @@ -11,6 +11,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use super::PreparedSet; +use crate::AnyCanonical; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -142,8 +143,9 @@ fn test_row_set_returns_membership_bits( let options = ListContainsOptions { sql_null_semantics }; let set = prepare(&list)?; let result = set.contains(&needles, &options, &mut ctx)?; - // The old fallback returned a lazy OR tree here. - assert!(result.is::()); + // Canonical needles are looked up at once. Any others stay a lazy lookup, for their encoding + // to take. + assert_eq!(result.is::(), needles.is::()); let expected = (0..needles.len()) .map(|idx| { let value = needles.execute_scalar(idx, &mut ctx)?; From 0e51ee9d4808de43a45bbaea883a20ce3a1f3979 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 15 Oct 2026 09:04:26 +0000 Subject: [PATCH 3/4] feat(array): make the lookup encodings and PreparedSet::new_array public A needle kernel in another crate can now recognize the lookup array of a prepared set, `SortedHaystack`, `BitmapHaystack`, `BytesHaystack` or `RowsHaystack`, and read its children through their public slot indices, so that it can look up its own encoded needles against that encoding. `PreparedSet::new_array` is public, so that a set made by `PreparedSetData::try_with_haystack` can be the list of a `ListContains` and reach execution. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TC7eug8jjTe1e1pbDreBo9 --- vortex-array/src/array/haystack/bytes.rs | 7 ++++--- vortex-array/src/array/haystack/integer.rs | 13 ++++++++----- vortex-array/src/array/haystack/mod.rs | 7 ++++++- vortex-array/src/array/haystack/rows.rs | 8 +++++--- vortex-array/src/array/haystack/tests.rs | 17 +++++++++++++++++ vortex-array/src/array/mod.rs | 6 ++++++ .../fns/list_contains/prepared/function.rs | 7 ++++++- 7 files changed, 52 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/array/haystack/bytes.rs b/vortex-array/src/array/haystack/bytes.rs index 8ae0084557f..cea30ca5e2a 100644 --- a/vortex-array/src/array/haystack/bytes.rs +++ b/vortex-array/src/array/haystack/bytes.rs @@ -81,15 +81,16 @@ pub(super) fn try_new( /// removed: it helped short needles by about a tenth and cost long needles that share a prefix, /// whose heads collide, up to a quarter. #[derive(Clone, Debug)] -pub(crate) struct BytesHaystack; +pub struct BytesHaystack; impl BytesHaystack { - const ELEMENTS: usize = 0; + /// The slot of the elements, a `VarBinView` array. + pub const ELEMENTS: usize = 0; } /// The hash tables of a [`BytesHaystack`], over the views of its elements. #[derive(Clone)] -pub(crate) struct BytesTables(Arc); +pub struct BytesTables(Arc); struct Tables { hasher: RandomState, diff --git a/vortex-array/src/array/haystack/integer.rs b/vortex-array/src/array/haystack/integer.rs index 878826a8d23..4cba6e294c5 100644 --- a/vortex-array/src/array/haystack/integer.rs +++ b/vortex-array/src/array/haystack/integer.rs @@ -173,10 +173,11 @@ pub(super) fn try_new(members: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// to the binary search at every set size measured by the `list_contains_set` benchmark, up to /// 16 384 elements. #[derive(Clone, Debug)] -pub(crate) struct SortedHaystack; +pub struct SortedHaystack; impl SortedHaystack { - const VALUES: usize = 0; + /// The slot of the sorted values. + pub const VALUES: usize = 0; fn new(values: ArrayRef) -> Array { let dtype = values.dtype().clone(); @@ -278,11 +279,13 @@ impl HaystackVTable for SortedHaystack { /// over the span above it. Its rows are the integers of the set bits, in increasing order of /// offset from the smallest. #[derive(Clone, Debug)] -pub(crate) struct BitmapHaystack; +pub struct BitmapHaystack; impl BitmapHaystack { - const MIN: usize = 0; - const BITMAP: usize = 1; + /// The slot of the smallest value, an array of one row. + pub const MIN: usize = 0; + /// The slot of the bitmap over the span above the smallest value. + pub const BITMAP: usize = 1; const NAMES: [&'static str; 2] = ["min", "bitmap"]; fn new(min: ArrayRef, bitmap: BitBuffer) -> Array { diff --git a/vortex-array/src/array/haystack/mod.rs b/vortex-array/src/array/haystack/mod.rs index 8cb330a43ff..35e6c93eac6 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-array/src/array/haystack/mod.rs @@ -14,6 +14,11 @@ mod rows; #[cfg(test)] mod tests; +pub use bytes::BytesHaystack; +pub use bytes::BytesTables; +pub use integer::BitmapHaystack; +pub use integer::SortedHaystack; +pub use rows::RowsHaystack; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -131,7 +136,7 @@ fn contains_of( /// The validity that the lookup arrays of Vortex share: every row is valid. #[derive(Clone, Debug)] -pub(crate) struct HaystackValidity; +pub struct HaystackValidity; impl ValidityVTable for HaystackValidity { fn validity(array: ArrayView<'_, V>) -> VortexResult { diff --git a/vortex-array/src/array/haystack/rows.rs b/vortex-array/src/array/haystack/rows.rs index 189cb891bf5..e45043b4781 100644 --- a/vortex-array/src/array/haystack/rows.rs +++ b/vortex-array/src/array/haystack/rows.rs @@ -71,11 +71,13 @@ pub(super) fn try_new( /// materializes child buffers and validity once per batch of needles, including decimal widening /// and nested offsets. Each search then reads those buffers directly. #[derive(Clone, Debug)] -pub(crate) struct RowsHaystack; +pub struct RowsHaystack; impl RowsHaystack { - const ELEMENTS: usize = 0; - const ORDER: usize = 1; + /// The slot of the elements, recursively canonical. + pub const ELEMENTS: usize = 0; + /// The slot of the `u64` indices of the distinct elements in sorted order. + pub const ORDER: usize = 1; const NAMES: [&'static str; 2] = ["elements", "order"]; } diff --git a/vortex-array/src/array/haystack/tests.rs b/vortex-array/src/array/haystack/tests.rs index df2d3068dc5..d288b6d126b 100644 --- a/vortex-array/src/array/haystack/tests.rs +++ b/vortex-array/src/array/haystack/tests.rs @@ -49,7 +49,9 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; use crate::scalar::Scalar; +use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListContainsOptions; +use crate::scalar_fn::fns::list_contains::PreparedSet; use crate::scalar_fn::fns::list_contains::PreparedSetData; use crate::validity::Validity; @@ -287,6 +289,21 @@ fn test_prepared_set_looks_up_through_your_own_haystack() -> VortexResult<()> { Ok(()) } +#[test] +fn test_list_contains_executes_through_your_own_haystack() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = ArrayHaystack::new(LinearHaystack::new(buffer![7i32].into_array())); + let set = PreparedSetData::try_with_haystack(i32_list(&[3, 7]), haystack)?; + let result = ListContains::try_new( + PreparedSet::new_array(set, 3), + buffer![7i32, 3, 1].into_array(), + )? + .into_array(); + + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + Ok(()) +} + #[test] fn test_your_own_haystack_must_hold_the_element_dtype() { let haystack = ArrayHaystack::new(LinearHaystack::new(buffer![7i64].into_array())); diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 2967f9c39df..e38f98fa4eb 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -34,7 +34,13 @@ mod plugin; pub use plugin::*; mod haystack; pub use haystack::ArrayHaystack; +pub use haystack::BitmapHaystack; +pub use haystack::BytesHaystack; +pub use haystack::BytesTables; pub use haystack::HaystackVTable; +pub use haystack::HaystackValidity; +pub use haystack::RowsHaystack; +pub use haystack::SortedHaystack; mod probe; pub use probe::*; diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs index eb18ed35372..7f109ef4904 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs @@ -68,7 +68,12 @@ impl PreparedSet { } /// An array of `len` rows that share the prepared set `set`. - pub(crate) fn new_array(set: PreparedSetData, len: usize) -> ArrayRef { + /// + /// A set made by [`PreparedSetData::try_with_haystack`] reaches execution this way, as the list + /// of a [`ListContains`]. + /// + /// [`ListContains`]: crate::scalar_fn::fns::list_contains::ListContains + pub fn new_array(set: PreparedSetData, len: usize) -> ArrayRef { ScalarFnArray::try_new_with_len(Self.bind(set), [], len) .vortex_expect("A prepared set has no children, and its dtype is the list dtype") .into_array() From 50383937d765d53b18e323c84fc5e66c8b93b725 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 25 Oct 2026 12:17:22 +0000 Subject: [PATCH 4/4] refactor(array): build the prepared set's haystack from the list's elements The haystack was built by wrapping the list in a one-row constant array, executing it to a list view and taking that view's elements. The list view's elements come from `list_scalar_elements`, so the haystack now calls it through `ListScalar::elements_array` and skips the constant array, the execution and the list view. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TC7eug8jjTe1e1pbDreBo9 --- vortex-array/src/arrays/constant/mod.rs | 1 + vortex-array/src/arrays/constant/vtable/canonical.rs | 2 +- vortex-array/src/scalar/typed_view/list.rs | 11 +++++++++++ .../src/scalar_fn/fns/list_contains/prepared/data.rs | 11 ++--------- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/constant/mod.rs b/vortex-array/src/arrays/constant/mod.rs index 36df554d4c3..dbe4403fea2 100644 --- a/vortex-array/src/arrays/constant/mod.rs +++ b/vortex-array/src/arrays/constant/mod.rs @@ -9,6 +9,7 @@ pub use arbitrary::ArbitraryConstantArray; mod array; pub use array::ConstantData; pub use vtable::ConstantArray; +pub(crate) use vtable::canonical::list_scalar_elements; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index bbea85c4077..81b43bd7b8e 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -311,7 +311,7 @@ fn constant_canonical_list_array( } /// The elements of a list scalar as an array, one row per element; empty for a null list. -fn list_scalar_elements(list: &ListScalar, allocator: &BufferAllocatorRef) -> ArrayRef { +pub(crate) fn list_scalar_elements(list: &ListScalar, allocator: &BufferAllocatorRef) -> ArrayRef { let element_dtype = list.element_dtype(); let Some(elements) = list.element_values() else { return Canonical::empty(element_dtype).into_array(); diff --git a/vortex-array/src/scalar/typed_view/list.rs b/vortex-array/src/scalar/typed_view/list.rs index 01e91800432..8b1c9fe5709 100644 --- a/vortex-array/src/scalar/typed_view/list.rs +++ b/vortex-array/src/scalar/typed_view/list.rs @@ -9,12 +9,15 @@ use std::hash::Hash; use std::sync::Arc; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use crate::ArrayRef; +use crate::arrays::constant::list_scalar_elements; use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -168,6 +171,14 @@ impl<'a> ListScalar<'a> { self.elements } + /// Returns the elements of the list as an array, one row per element, null elements + /// included. + /// + /// Returns an empty array if the list is null. + pub(crate) fn elements_array(&self, allocator: &BufferAllocatorRef) -> ArrayRef { + list_scalar_elements(self, allocator) + } + /// Returns all elements in the list as a vector of scalars. /// /// Returns None if the list is null. diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs index 3637664c62e..18afa6e436d 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs @@ -27,8 +27,6 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; -use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArraySlotsExt; use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnRef; @@ -154,13 +152,8 @@ impl PreparedSetData { self.0 .haystack .get_or_init(|| { - // The canonical form of a constant list holds its elements once, whatever its - // length. - ConstantArray::new(self.list().clone(), 1) - .into_array() - .execute::(ctx) - .and_then(|list| ArrayHaystack::try_new(list.elements(), ctx)) - .map_err(Arc::new) + let elements = self.list().as_list().elements_array(ctx.allocator()); + ArrayHaystack::try_new(&elements, ctx).map_err(Arc::new) }) .as_ref() .map_err(|err| Arc::clone(err).into())