diff --git a/vortex-array/src/array/haystack/bytes.rs b/vortex-array/src/array/haystack/bytes.rs new file mode 100644 index 00000000000..cea30ca5e2a --- /dev/null +++ b/vortex-array/src/array/haystack/bytes.rs @@ -0,0 +1,321 @@ +// 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_bail; +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::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; +use crate::IntoArray; +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`. +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 struct BytesHaystack; + +impl BytesHaystack { + /// 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 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 = Self; + type ValidityVTable = HaystackValidity; + + 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 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: 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) + .as_::() + .into_owned(); + let element_views = elements.views(); + let element_buffers = data_buffers(&elements); + + 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..4cba6e294c5 --- /dev/null +++ b/vortex-array/src/array/haystack/integer.rs @@ -0,0 +1,529 @@ +// 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_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::registry::CachedId; + +use super::ArrayHaystack; +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; +use crate::array::Array; +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; +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::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; + +/// Binds `$T` to the integer type of the primitive or decimal array `$values`, `$values` to its +/// 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(), $needles) { + (DType::Primitive(ptype, _), CanonicalView::Primitive(needles)) => { + let ptype = bit_pattern_ptype(*ptype); + let values = $values.as_::().reinterpret_cast(ptype); + let needles = needles.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, _), CanonicalView::Decimal(needles)) => { + let values_type = DecimalType::smallest_decimal_value_type(decimal); + let values = $values.as_::().into_owned(); + let needles = needles.into_owned(); + 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, needles) => vortex_bail!( + "An integer haystack of {dtype} cannot look up needles of {}", + needles.to_array_ref().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 struct SortedHaystack; + +impl SortedHaystack { + /// The slot of the sorted values. + pub 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 = Self; + type ValidityVTable = HaystackValidity; + + 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!( + 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(()) + } + + 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 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: CanonicalView<'_>, + 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 struct BitmapHaystack; + +impl BitmapHaystack { + /// 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 { + 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(slots: &[Option]) -> BitBuffer { + slot(slots, Self::BITMAP).as_::().to_bit_buffer() + } +} + +impl VTable for BitmapHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; + + 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!( + 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(()) + } + + 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()); + 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: CanonicalView<'_>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + let allocator = ctx.allocator().clone(); + let bitmap = Self::bitmap(array.slots()); + 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, + ) + } + ) + } +} + +/// 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 { + 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())) +} + +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..35e6c93eac6 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-array/src/array/haystack/mod.rs @@ -2,73 +2,95 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Lookups of needles among the elements of an array. - -use std::hash::BuildHasher; - -use num_traits::ToPrimitive; -use num_traits::WrappingSub; +//! +//! 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; + +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_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::CanonicalView; 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::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::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 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: 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, CanonicalView<'_>, &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 +98,105 @@ impl ArrayHaystack { array.filter(valid)? }; - Ok(Self { - state: new_state(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 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( - &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)), + 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 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, - ), - } + /// The lookup array, which canonicalizes to the elements of the haystack. + pub fn array(&self) -> &ArrayRef { + &self.array } -} - -/// Primitive integers, or floats by their bit patterns. -/// -/// 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); -impl HaystackState for PrimitiveSet { - fn contains( + /// One bit per needle, set when the needle is an element, and the validity of the needles. + /// + /// The needles are canonical, with the dtype of the lookup array ignoring nullability. + /// + /// # Errors + /// + /// 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)> { - let needles = needles - .clone() - .execute::(ctx)? - .reinterpret_cast(T::PTYPE); - let bits = self.0.contains(needles.as_slice::(), ctx.allocator()); - Ok((bits, needles.validity()?)) + (self.contains)(&self.array, needles, ctx) } } -/// 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); +fn contains_of( + array: &ArrayRef, + needles: CanonicalView<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(BitBuffer, Validity)> { + V::contains(array.as_::(), needles, ctx) +} -impl HaystackState for DecimalSet { - fn contains( - &self, - needles: &ArrayRef, - 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)) +/// The validity that the lookup arrays of Vortex share: every row is valid. +#[derive(Clone, Debug)] +pub struct HaystackValidity; + +impl ValidityVTable for HaystackValidity { + fn validity(array: ArrayView<'_, V>) -> VortexResult { + Ok(Validity::from(array.dtype().nullability())) } } -/// 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..e45043b4781 --- /dev/null +++ b/vortex-array/src/array/haystack/rows.rs @@ -0,0 +1,198 @@ +// 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; +use vortex_error::vortex_ensure_eq; +use vortex_session::registry::CachedId; + +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; +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::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; + +/// 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 struct RowsHaystack; + +impl RowsHaystack { + /// 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"]; +} + +impl VTable for RowsHaystack { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = HaystackValidity; + + 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!( + order.is::(), + "A row haystack holds its order as a primitive array, got {}", + order.encoding_id() + ); + 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 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: 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()), + needles.validity()?, + )); + } + + let elements = slot(array.slots(), Self::ELEMENTS); + let order = slot(array.slots(), Self::ORDER).as_::(); + let order = order.as_slice::(); + // 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( + 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..d288b6d126b --- /dev/null +++ b/vortex-array/src/array/haystack/tests.rs @@ -0,0 +1,324 @@ +// 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_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::ArrayHaystack; +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; +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::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; +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::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; + +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.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_rows_lookup_array_reads_a_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let haystack = build(struct_of(&[3, 1, 3]), &mut ctx)?; + + assert_eq!( + haystack.array().execute_scalar(1, &mut ctx)?, + struct_of(&[3]).execute_scalar(0, &mut ctx)? + ); + 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 = Self; + type ValidityVTable = HaystackValidity; + + 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 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: CanonicalView<'_>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<(BitBuffer, Validity)> { + 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() + .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_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())); + 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..e38f98fa4eb 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -33,7 +33,14 @@ pub use erased::*; mod plugin; pub use plugin::*; mod haystack; -pub(crate) use haystack::ArrayHaystack; +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/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/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 aaaf047c1d9..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 @@ -11,23 +11,27 @@ 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; 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; 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; @@ -35,8 +39,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 +56,8 @@ 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 }) - } + /// The haystack of the non-null elements, given or built on the first lookup. + haystack: OnceLock>, } impl PreparedSetData { @@ -111,10 +96,33 @@ impl PreparedSetData { literal, has_null_element, is_empty, - set: 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 { @@ -131,20 +139,22 @@ impl PreparedSetData { self.0.literal.as_::() } - /// The elements of the list, one row each, null elements included. + /// 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 the haystack. - pub fn elements(&self, ctx: &mut ExecutionCtx) -> VortexResult<&ArrayRef> { - Ok(&self.element_set(ctx)?.elements) - } - - /// 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> { + /// 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(|| { + 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()) } @@ -170,7 +180,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 /// @@ -187,8 +203,15 @@ impl PreparedSetData { } self.check_needle_dtype(needles.dtype())?; - let (bits, needle_validity) = self.element_set(ctx)?.haystack.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`. @@ -206,12 +229,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.element_set(ctx)?.haystack.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/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() 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)?;