diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs new file mode 100644 index 00000000000..e40cb19b854 --- /dev/null +++ b/encodings/fsst/src/compute/list_contains.rs @@ -0,0 +1,422 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; + +use crate::FSST; +use crate::FSSTArrayExt; + +/// Answers `IN` over compressed strings without decompressing any of them. +/// +/// FSST compression under a fixed symbol table is deterministic and lossless, so +/// two values are equal exactly when their codes are — the same property +/// `compare_fsst_constant` relies on to answer `Eq` in compressed space. +/// Compressing the list once therefore turns membership into a `Binary` +/// membership test over the codes, which is a few bytes per row against the +/// decompressed string the generic path would have to materialize. +/// Rows per list element below which this kernel declines. +/// +/// It compresses the whole list once per batch and saves decompressing every +/// row, so both sides scale and which one wins is a ratio rather than a size. +/// Measured on an 8192-row batch: a column of short words — the least +/// compressible shape, and so the one with the least decompression to save — +/// breaks even near 256 elements, while a column of URLs stays ahead to around +/// 768. The lower crossover is the one to take: declining early costs a missed +/// speedup, accepting late costs a regression. +const MIN_ROWS_PER_LIST_ELEMENT: usize = 32; + +impl ListContainsElementKernel for FSST { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // The list is compressed once and probed by every row, so it has to be + // the same list for every row. + let Some(list_constant) = list.as_opt::() else { + return Ok(None); + }; + let list_scalar = list_constant.scalar().as_list(); + let Some(elements) = list_scalar.element_values() else { + return Ok(None); + }; + + // Executing as a parent kernel bypasses `ListContains`'s own execution, + // and with it the check that the list's elements and the needle are the + // same type. Compressing both sides to `Binary` codes would answer a + // mismatch that the definition rejects, so decline it here and let the + // generic path raise the error. + if !list_scalar + .element_dtype() + .eq_ignore_nullability(element.dtype()) + { + return Ok(None); + } + + if elements.len().saturating_mul(MIN_ROWS_PER_LIST_ELEMENT) > element.len() { + return Ok(None); + } + + // A null element has no compressed form and sends the whole list back to + // the generic path, so find that out before compressing any of it. + if elements.iter().any(Option::is_none) { + return Ok(None); + } + + let compressor = element.compressor(); + let code_dtype = DType::Binary(Nullability::NonNullable); + let mut codes = Vec::with_capacity(elements.len()); + for value in elements { + let bytes = match value { + Some(ScalarValue::Utf8(value)) => value.as_bytes(), + Some(ScalarValue::Binary(value)) => value.as_slice(), + // A null element has no compressed form. The generic path + // answers those, so hand the whole list back to it. + _ => return Ok(None), + }; + codes.push(Scalar::binary( + ByteBuffer::from(compressor.compress(bytes)), + Nullability::NonNullable, + )); + } + + let code_list = Scalar::list( + Arc::new(code_dtype), + codes, + list_constant.scalar().dtype().nullability(), + ); + let needles = element.codes().into_array(); + let len = needles.len(); + + ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(code_list, len).into_array(), needles], + )? + .execute(ctx) + .map(Some) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::bool::BoolArrayExt; + use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::fns::list_contains::ListContains; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use super::MIN_ROWS_PER_LIST_ELEMENT; + use crate::fsst_compress; + use crate::fsst_train_compressor; + + /// The result row by row, with `None` for an invalid slot. + fn bool_answers( + result: &vortex_array::Array, + ctx: &mut vortex_array::ExecutionCtx, + ) -> Vec> { + let bits = result.bit_buffer_view(); + let validity = result + .validity() + .expect("validity") + .execute_mask(bits.len(), ctx) + .expect("validity mask"); + (0..bits.len()) + .map(|i| validity.value(i).then(|| bits.value(i))) + .collect() + } + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + /// The values the tests below compress. Deliberately includes an empty + /// string, a value sharing a long prefix with another, and nulls. + const VALUES: [Option<&str>; 8] = [ + Some("hello"), + None, + Some("world"), + Some(""), + Some("a value long enough to need a buffer 1"), + Some("a value long enough to need a buffer 2"), + Some("hello"), + None, + ]; + + /// Rows in the test column. + /// + /// The kernel declines below [`MIN_ROWS_PER_LIST_ELEMENT`] rows per list + /// element, so the longest list the tests use needs a column this size + /// before the kernel will engage at all. `VALUES` is cycled to fill it, so + /// the answers repeat with the same period. + const ROWS: usize = 8 * 32 * 2; + + fn answers( + list: Vec<&str>, + needle_nullability: Nullability, + ) -> VortexResult>> { + answers_over(&VALUES, list, needle_nullability) + } + + /// Answers over `column` cycled to [`ROWS`] rows, asserting that the + /// compressed and uncompressed paths agree. + fn answers_over( + column: &[Option<&str>], + list: Vec<&str>, + needle_nullability: Nullability, + ) -> VortexResult>> { + assert!( + list.len() * MIN_ROWS_PER_LIST_ELEMENT <= ROWS, + "a {}-element list needs {} rows for the kernel to engage, column has {ROWS}", + list.len(), + list.len() * MIN_ROWS_PER_LIST_ELEMENT, + ); + let mut ctx = SESSION.create_execution_ctx(); + let plain = VarBinArray::from_iter( + (0..ROWS).map(|i| column[i % column.len()]), + DType::Utf8(needle_nullability), + ) + .into_array(); + let compressor = fsst_train_compressor(&plain, &mut ctx)?; + let compressed = fsst_compress(&plain, &compressor, &mut ctx)?.into_array(); + + let list = Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + list.into_iter() + .map(|value| Scalar::utf8(value, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ); + + let mut results = Vec::new(); + for needles in [plain, compressed] { + let len = needles.len(); + let result = ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(list.clone(), len).into_array(), needles], + )? + .execute::(&mut ctx)?; + results.push(bool_answers(&result, &mut ctx)); + } + + // The compressed answer must equal the uncompressed one; the tests then + // only have to assert the latter. + assert_eq!(results[0], results[1], "FSST answer diverged from VarBin"); + + // The column is `VALUES` cycled, so the answers are too. Checking that + // and returning one period keeps the expectations readable while the + // column stays long enough to reach the kernel. + let answers = results.remove(1); + let period = &answers[..column.len()]; + for (chunk, at) in answers + .chunks(column.len()) + .zip((0..).step_by(column.len())) + { + assert_eq!(chunk, period, "answers stopped repeating at row {at}"); + } + Ok(period.to_vec()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn probes_compressed_codes() -> VortexResult<()> { + // Long enough to take the set probe rather than OR-of-equalities. + assert_eq!( + answers( + vec!["hello", "world", "absent", "also absent", "and another"], + Nullability::Nullable + )?, + vec![ + Some(true), + Some(false), + Some(true), + Some(false), + Some(false), + Some(false), + Some(true), + Some(false) + ] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn empty_string_is_a_code_like_any_other() -> VortexResult<()> { + assert_eq!( + answers(vec!["", "a", "b", "c"], Nullability::Nullable)?, + vec![ + Some(false), + Some(false), + Some(false), + Some(true), + Some(false), + Some(false), + Some(false), + Some(false) + ] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn a_shared_prefix_does_not_collide() -> VortexResult<()> { + // The two long values differ only in their final byte, so a code + // comparison that stopped at a prefix would answer both alike. + assert_eq!( + answers( + vec![ + "a value long enough to need a buffer 1", + "q", + "qq", + "qqq", + "qqqq" + ], + Nullability::Nullable + )?, + vec![ + Some(false), + Some(false), + Some(false), + Some(false), + Some(true), + Some(false), + Some(false), + Some(false) + ] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn short_list_below_the_probe_threshold() -> VortexResult<()> { + assert_eq!( + answers(vec!["hello", ""], Nullability::Nullable)?, + vec![ + Some(true), + Some(false), + Some(false), + Some(true), + Some(false), + Some(false), + Some(true), + Some(false) + ] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn empty_list_matches_nothing() -> VortexResult<()> { + assert_eq!( + answers(vec![], Nullability::Nullable)?, + vec![Some(false); VALUES.len()] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn non_nullable_needles() -> VortexResult<()> { + assert_eq!( + answers_over( + &["a", "b", "c", "d"].map(Some), + vec!["a", "c", "e", "f"], + Nullability::NonNullable + )?, + vec![Some(true), Some(false), Some(true), Some(false)] + ); + Ok(()) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn mismatched_element_type_is_rejected() -> VortexResult<()> { + // A `List` against `Utf8` needles is not a valid `list_contains`: + // the generic path rejects it. The FSST kernel must not answer it. + let mut ctx = SESSION.create_execution_ctx(); + let plain = + VarBinArray::from_iter(["a", "b"].map(Some), DType::Utf8(Nullability::NonNullable)) + .into_array(); + let compressor = fsst_train_compressor(&plain, &mut ctx)?; + let compressed = fsst_compress(&plain, &compressor, &mut ctx)?.into_array(); + + let binary_list = Scalar::list( + Arc::new(DType::Binary(Nullability::NonNullable)), + vec![ + Scalar::binary( + vortex_buffer::ByteBuffer::copy_from(b"a"), + Nullability::NonNullable, + ), + Scalar::binary( + vortex_buffer::ByteBuffer::copy_from(b"c"), + Nullability::NonNullable, + ), + ], + Nullability::NonNullable, + ); + + let run = |needles: vortex_array::ArrayRef| { + let len = needles.len(); + ListContains + .try_new_array( + len, + EmptyOptions, + [ + ConstantArray::new(binary_list.clone(), len).into_array(), + needles, + ], + ) + .and_then(|a| a.execute::(&mut SESSION.create_execution_ctx())) + }; + + let varbin = run(plain); + let fsst = run(compressed); + assert!( + varbin.is_err(), + "generic path should reject the type mismatch" + ); + assert!( + fsst.is_err(), + "FSST kernel answered a type mismatch the generic path rejects" + ); + Ok(()) + } +} diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index d447b4b4958..f7cc0582e27 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -6,6 +6,7 @@ mod cast; mod compare; mod filter; mod like; +mod list_contains; use vortex_array::ArrayRef; use vortex_array::ArrayView; diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 30f25006195..47806fd0eac 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -16,6 +16,8 @@ use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; use vortex_session::VortexSession; use crate::FSST; @@ -24,6 +26,11 @@ pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); kernels.register_execute_parent_kernel(Cast.id(), FSST, CastExecuteAdaptor(FSST)); kernels.register_execute_parent_kernel(Binary.id(), FSST, CompareExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel( + ListContains.id(), + FSST, + ListContainsElementExecuteAdaptor(FSST), + ); kernels.register_execute_parent_kernel(Filter.id(), FSST, FilterExecuteAdaptor(FSST)); kernels.register_execute_parent_kernel(Dict.id(), FSST, TakeExecuteAdaptor(FSST)); kernels.register_execute_parent_kernel(Like.id(), FSST, LikeExecuteAdaptor(FSST)); diff --git a/vortex-array/src/arrays/varbinview/mod.rs b/vortex-array/src/arrays/varbinview/mod.rs index a6b8c9db46c..7d08123a971 100644 --- a/vortex-array/src/arrays/varbinview/mod.rs +++ b/vortex-array/src/arrays/varbinview/mod.rs @@ -22,8 +22,10 @@ pub mod build_views; mod view; pub use view::BinaryView; +mod views_side; pub use view::Inlined; pub use view::Ref; +pub(crate) use views_side::ViewsSide; #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/varbinview/views_side.rs b/vortex-array/src/arrays/varbinview/views_side.rs new file mode 100644 index 00000000000..8444981790e --- /dev/null +++ b/vortex-array/src/arrays/varbinview/views_side.rs @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A [`VarBinViewArray`] with its buffers resolved once, for callers that walk +//! every row. + +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; + +/// A resolved view over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices of +/// every data buffer, supporting cheap per-lane byte access. +pub(crate) struct ViewsSide<'a> { + views: &'a [BinaryView], + buffers: Vec<&'a [u8]>, +} + +impl<'a> ViewsSide<'a> { + pub(crate) fn new(array: &'a VarBinViewArray) -> Self { + Self { + views: array.views(), + buffers: (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).as_slice()) + .collect(), + } + } + + pub(crate) fn len(&self) -> usize { + self.views.len() + } + + /// The view structs themselves, for a caller walking every row in order. + pub(crate) fn views(&self) -> &'a [BinaryView] { + self.views + } + + /// The view at `index` without a bounds check. + /// + /// # Safety + /// + /// `index` must be strictly less than `self.len()`. + #[inline] + pub(crate) unsafe fn view_unchecked(&self, index: usize) -> &'a BinaryView { + // SAFETY: caller guarantees index < self.views.len(). + unsafe { self.views.get_unchecked(index) } + } + + /// The full bytes of `view`, which must belong to this side. + #[inline] + pub(crate) fn view_bytes(&self, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &self.buffers[view.buffer_index as usize][view.as_range()] + } + } +} diff --git a/vortex-array/src/scalar/typed_view/list.rs b/vortex-array/src/scalar/typed_view/list.rs index 15de731089a..39e5365a75a 100644 --- a/vortex-array/src/scalar/typed_view/list.rs +++ b/vortex-array/src/scalar/typed_view/list.rs @@ -160,6 +160,17 @@ impl<'a> ListScalar<'a> { }) } + /// Returns the element values, borrowed from the scalar itself. + /// + /// Returns None if the list is null. Unlike [`Self::elements`] this clones + /// nothing, so it is the form to reach for when a caller only reads the + /// values; the cost of materializing a `Vec` is proportional to the + /// length of the list and is paid again on drop. + #[inline] + pub fn element_values(&self) -> Option<&'a [Option]> { + self.elements + } + /// 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/binary/compare/bytes.rs b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs index 70c56411d08..fea24f20cb2 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs @@ -22,6 +22,7 @@ use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::ViewsSide; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -79,50 +80,6 @@ fn constant_bytes(scalar: &Scalar) -> VortexResult> { value.ok_or_else(|| vortex_err!("null constant handled by execute_compare")) } -/// A resolved view over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices of -/// every data buffer, supporting cheap per-lane byte access. -struct ViewsSide<'a> { - views: &'a [BinaryView], - buffers: Vec<&'a [u8]>, -} - -impl<'a> ViewsSide<'a> { - fn new(array: &'a VarBinViewArray) -> Self { - Self { - views: array.views(), - buffers: (0..array.data_buffers().len()) - .map(|idx| array.buffer(idx).as_slice()) - .collect(), - } - } - - fn len(&self) -> usize { - self.views.len() - } - - /// The view at `index` without a bounds check. - /// - /// # Safety - /// - /// `index` must be strictly less than `self.len()`. - #[inline] - unsafe fn view_unchecked(&self, index: usize) -> &'a BinaryView { - // SAFETY: caller guarantees index < self.views.len(). - unsafe { self.views.get_unchecked(index) } - } - - /// The full bytes of `view`, which must belong to this side. - #[inline] - fn view_bytes(&self, view: &'a BinaryView) -> &'a [u8] { - if view.is_inlined() { - view.as_inlined().value() - } else { - let view = view.as_view(); - &self.buffers[view.buffer_index as usize][view.as_range()] - } - } -} - /// The leading 8 bytes of a view: the `u32` length plus the first 4 bytes of the value /// (zero-padded for values shorter than 4 bytes). #[inline] 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 edf530cb75f..8beab939e81 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -3,18 +3,23 @@ mod kernel; +use std::hash::Hash; use std::ops::BitOr; use arrow_buffer::bit_iterator::BitIndexIterator; pub use kernel::*; +use num_traits::AsPrimitive; use num_traits::Zero; use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use vortex_utils::aliases::hash_set::HashSet; use vortex_utils::iter::ReduceBalancedIterExt; use crate::ArrayRef; @@ -23,20 +28,30 @@ use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBin; +use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; +use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::primitive::NativeValue; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::scalar_fn::ScalarFnFactoryExt; +use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbinview::ViewsSide; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::match_each_integer_ptype; +use crate::match_each_native_ptype; use crate::match_each_unsigned_integer_ptype; use crate::scalar::ListScalar; use crate::scalar::Scalar; +use crate::scalar::ScalarValue; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -111,10 +126,13 @@ impl ScalarFnVTable for ListContains { let list_array = args.get(0)?; let value_array = args.get(1)?; - if let Some(list_scalar) = list_array.as_constant() - && let Some(value_scalar) = value_array.as_constant() + // The needle is tested first: for an `IN` list it is an array, so this + // fails on a downcast, whereas `list_array.as_constant()` deep-clones + // every element of the list before the chain can reject it. + if let Some(value_scalar) = value_array.as_constant() + && let Some(list_constant) = list_array.as_opt::() { - let result = compute_contains_scalar(&list_scalar, &value_scalar)?; + let result = compute_contains_scalar(list_constant.scalar(), &value_scalar)?; return Ok(ConstantArray::new(result, args.row_count()).into_array()); } @@ -176,21 +194,62 @@ fn compute_list_contains( if let Some(value_scalar) = value.as_constant() { list_contains_scalar(array, &value_scalar, nullability, ctx) - } else if let Some(list_scalar) = array.as_constant() { - constant_list_scalar_contains(&list_scalar.as_list(), value, nullability) + } else if let Some(list_constant) = array.as_opt::() { + constant_list_scalar_contains(&list_constant.scalar().as_list(), value, nullability, ctx) } else { todo!("unsupported list contains with list and element as arrays") } } +/// List length past which membership is answered by probing a set instead of by +/// OR-ing one equality per element. +/// +/// The equality form costs one full-length comparison per element, so it grows +/// with the list, while the probe is built once per batch and then answers each +/// row in constant time. Where they cross depends on how expensive one +/// comparison is, which is a property of the column rather than of its type. +/// Measured on an 8192-row batch: `i64` needles run the equality form at 3.46, +/// 7.29 and 11.42us for one, two and three elements against a probe at 17.25us +/// for four, crossing just above four; `Utf8` needles run it at 10.92, 26.04 and +/// 38.96us against a probe at 37.08us, crossing just below three. Four sits +/// between them, within about 13% of the equality form at its worst point and +/// ahead of it everywhere after. +const HASH_PROBE_MIN_ELEMENTS: usize = 4; + +/// Slots per element to size the probe set with. +/// +/// Sizing it at exactly the element count leaves the table at its maximum load +/// factor, where the resulting probe chains cost 16-57% more than this across +/// four to eight thousand elements. Four slots an element buys a further 7-9% +/// over two for the mid-range sizes, for at most a few tens of kilobytes. +const PROBE_SET_HEADROOM: usize = 4; + /// There is a constant list scalar (haystack) being compared to an array of needles. fn constant_list_scalar_contains( - list_scalar: &ListScalar, + list_scalar: &ListScalar<'_>, values: &ArrayRef, nullability: Nullability, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let elements = list_scalar.elements().vortex_expect("non null"); + // Borrowed rather than materialized as `Vec`: this runs once per + // batch, and cloning every element of the list back out of the scalar costs + // more than the set built from them. + let element_values = list_scalar.element_values().vortex_expect("non null"); + + // A null element is not a key on any probe path, so the whole list goes back + // to the equality form. Establishing that here rather than inside the probe + // is what stops a nullable list from canonicalizing the needles, and keying + // as much of the list as precedes the null, only to discard both and run the + // equality form anyway — which would make such a list dearer than it was + // before the probe existed. + if element_values.len() >= HASH_PROBE_MIN_ELEMENTS + && element_values.iter().all(Option::is_some) + && let Some(probed) = hash_probe_contains(element_values, values, nullability, ctx)? + { + return Ok(probed); + } + let elements = list_scalar.elements().vortex_expect("non null"); let len = values.len(); let false_scalar = Scalar::bool(false, nullability); @@ -215,6 +274,171 @@ fn constant_list_scalar_contains( Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array())) } +/// Answers membership by building a set from the list once and probing it in a +/// single pass over the needles. +/// +/// Covers primitive, `Utf8` and `Binary` needles. Returns `None` for anything +/// else, and for a list holding a null element, leaving the caller on the +/// OR-of-equalities form, which is the definition of the operation. +/// +/// Validity follows the equality form exactly: it fills a null comparison with +/// `false`, so a null needle is `false` rather than null, and the result carries +/// no invalid slots. +fn hash_probe_contains( + elements: &[Option], + values: &ArrayRef, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let len = values.len(); + let ptype = match values.dtype() { + DType::Primitive(ptype, _) => *ptype, + // Strings and byte strings key on their bytes; see `bytes_probe_contains`. + DType::Utf8(_) | DType::Binary(_) => { + return bytes_probe_contains(elements, values, nullability, ctx); + } + // An extension value *is* its storage value — `ScalarValue` has no + // extension variant — and `Operator::Eq` compares extensions through + // their storage, so the list's values key a storage probe unchanged. + // Both sides were checked against each other before dispatch, so they + // unwrap to the same storage type; one whose storage this path does not + // cover declines below and takes the equality form, as it would have. + DType::Extension(_) => { + let storage = values.clone().execute::(ctx)?; + return hash_probe_contains(elements, storage.storage_array(), nullability, ctx); + } + _ => return Ok(None), + }; + + let needles = values.clone().execute::(ctx)?; + let validity = needles.validity()?.execute_mask(len, ctx)?; + + let bits = match_each_native_ptype!(ptype, |T| { + let Some(set) = primitive_key_set::(elements) else { + return Ok(None); + }; + let slice = needles.as_slice::(); + probe_rows(len, &validity, |idx| set.contains(&NativeValue(slice[idx]))) + }); + + Ok(Some(BoolArray::new(bits, nullability.into()).into_array())) +} + +/// Keys the list on its primitive values, or `None` if any element is not a +/// non-null primitive of `T`. +/// +/// [`NativeValue`] is the key rather than the bare value because its equality is +/// the one the kernel answers with: `NativePType::is_eq` compares floats by +/// their bits, so `NaN` matches itself and `-0.0` does not match `0.0`, and a +/// set keyed on the value would disagree with `Operator::Eq` on both. +/// +/// A null element is not a key, and the equality form maps it to `false` +/// through its own null fill, so such a list goes back to that form rather than +/// being answered with a key missing. +fn primitive_key_set( + elements: &[Option], +) -> Option>> +where + NativeValue: Hash + Eq, +{ + let mut set = HashSet::with_capacity(elements.len() * PROBE_SET_HEADROOM); + for element in elements { + let ScalarValue::Primitive(pvalue) = element.as_ref()? else { + return None; + }; + // The list's element dtype was checked against the needle dtype before + // dispatch, so this holds; `cast` would otherwise convert a value of + // some other width and key it under a number it never equals. + if !pvalue.is_instance_of(&T::PTYPE) { + return None; + } + set.insert(NativeValue(pvalue.cast::().ok()?)); + } + Some(set) +} + +/// Membership for `Utf8`/`Binary` needles, keyed on the element bytes. +/// +/// The keys borrow from `elements`, which outlives the set, so building it +/// copies no string data. +fn bytes_probe_contains( + elements: &[Option], + values: &ArrayRef, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let mut set: HashSet<&[u8]> = HashSet::with_capacity(elements.len() * PROBE_SET_HEADROOM); + for element in elements { + // A null element is not a key; the whole list goes back to the + // equality form rather than being answered with a key missing. + let bytes = match element { + Some(ScalarValue::Utf8(value)) => value.as_bytes(), + Some(ScalarValue::Binary(value)) => value.as_slice(), + _ => return Ok(None), + }; + set.insert(bytes); + } + + let len = values.len(); + + // `VarBin` already holds what the probe needs — a byte buffer and a run of + // offsets — but it is not canonical, so executing it would first build a + // `VarBinView`: sixteen bytes of view per row, to reach bytes that are + // already contiguous. For short values that view costs more to construct + // than the whole comparison it serves, and FSST's codes arrive in exactly + // this shape. + let bits = if let Some(varbin) = values.as_opt::() { + let validity = varbin.varbin_validity().execute_mask(len, ctx)?; + let bytes = varbin.bytes().as_slice(); + let offsets = varbin.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |O| { + // One offset per row plus a final end, so adjacent pairs are the + // rows in order. + let offsets = offsets.as_slice::(); + probe_rows(len, &validity, |idx| { + let (start, end): (usize, usize) = (offsets[idx].as_(), offsets[idx + 1].as_()); + set.contains(&bytes[start..end]) + }) + }) + } else { + let needles = values.clone().execute::(ctx)?; + let validity = needles.validity()?.execute_mask(len, ctx)?; + // Resolved once: reaching a row's bytes through the array re-derives the + // views slice and re-checks the buffer index on every row. + let side = ViewsSide::new(&needles); + let views = side.views(); + probe_rows(len, &validity, |idx| { + set.contains(side.view_bytes(&views[idx])) + }) + }; + + Ok(Some(BoolArray::new(bits, nullability.into()).into_array())) +} + +/// One pass over the needles, setting the bit for each one the set holds. +/// +/// An invalid needle is `false`, which is what the OR-of-equalities form's null +/// fill produces, so the result carries no invalid slots. `hit` is asked about a +/// row only when that row is valid. +fn probe_rows(len: usize, validity: &Mask, hit: impl Fn(usize) -> bool) -> BitBuffer { + match validity { + Mask::AllTrue(_) => (0..len).map(hit).collect(), + Mask::AllFalse(_) => BitBuffer::new_unset(len), + Mask::Values(valid) => { + // Walking the valid rows a word at a time skips runs of nulls + // wholesale, where testing validity per row pays a branch for every + // one of them. + let mut bits = BitBufferMut::new_unset(len); + valid.bit_buffer().for_each_set_index(|idx| { + if hit(idx) { + bits.set(idx); + } + }); + bits.freeze() + } + } +} + /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. fn list_contains_scalar( array: &ArrayRef, @@ -400,16 +624,20 @@ mod tests { use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; + use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::ListContains; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ExtensionArray; use crate::arrays::ListArray; use crate::arrays::VarBinArray; + use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; @@ -426,7 +654,10 @@ mod tests { use crate::expr::or; use crate::expr::root; use crate::expr::stats::Stat; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::fns::list_contains::BoolArray; use crate::scalar_fn::fns::list_contains::ConstantArray; use crate::scalar_fn::fns::list_contains::ListViewArray; @@ -599,23 +830,27 @@ mod tests { Nullability::NonNullable, ); + // The falsifier describes the list as intervals: the scope lies wholly + // outside the list's range, or wholly inside a gap between two adjacent + // list values. For a list this dense the gaps cover every element, so + // it proves exactly what a term per element would. assert_eq!( expr.falsify(&scope, &STATS_SESSION)?, - Some(and( - and( - or( - lt(stat(col("a"), Stat::Max), lit(1i32)), + Some(or( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), + ), + or( + and( gt(stat(col("a"), Stat::Min), lit(1i32)), - ), - or( lt(stat(col("a"), Stat::Max), lit(2i32)), + ), + and( gt(stat(col("a"), Stat::Min), lit(2i32)), - ) + lt(stat(col("a"), Stat::Max), lit(3i32)), + ), ), - or( - lt(stat(col("a"), Stat::Max), lit(3i32)), - gt(stat(col("a"), Stat::Min), lit(3i32)), - ) )) ); Ok(()) @@ -932,4 +1167,136 @@ mod tests { let expected_zero = BoolArray::from_iter([true, false, false, false]); assert_arrays_eq!(result_zero, expected_zero, &mut ctx); } + + // Extension needles are answered through their storage; these came with + // that support and belong beside the path that now provides it. + /// Membership over a timestamp column, whose storage is the `i64` beneath it. + fn timestamps_in(list_values: &[i64], unit: TimeUnit) -> Vec { + let mut ctx = array_session().create_execution_ctx(); + let ext_dtype = Timestamp::new(unit, Nullability::NonNullable).erased(); + let needles = ExtensionArray::new( + ext_dtype.clone(), + buffer![10i64, 20, 30, 40, 50].into_array(), + ) + .into_array(); + + let element_dtype = DType::Extension(ext_dtype); + let list = Scalar::list( + Arc::new(element_dtype.clone()), + list_values + .iter() + .map(|v| { + Scalar::primitive(*v, Nullability::NonNullable) + .cast(&element_dtype) + .expect("i64 into the timestamp it stores") + }) + .collect(), + Nullability::NonNullable, + ); + + let len = needles.len(); + let result = ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(list, len).into_array(), needles], + ) + .expect("build") + .execute::(&mut ctx) + .expect("execute"); + result.bool_vec(&mut ctx) + } + + #[test] + fn a_timestamp_list_is_answered_through_its_storage() { + // Long enough to take the set probe rather than OR-of-equalities. + assert_eq!( + timestamps_in(&[10, 30, 50, 70], TimeUnit::Milliseconds), + vec![true, false, true, false, true] + ); + // Below the probe threshold, so the equality form answers it; both must + // agree. + assert_eq!( + timestamps_in(&[20, 40], TimeUnit::Milliseconds), + vec![false, true, false, true, false] + ); + } + + #[test] + fn a_nullable_list_holding_a_null_falls_back_instead_of_panicking() { + // The dtype check ignores nullability, so a nullable extension list can + // reach this kernel against a non-nullable column. A null element then + // has no storage value to unwrap. + let mut ctx = array_session().create_execution_ctx(); + let nullable = DType::Extension( + Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(), + ); + let needles = ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + buffer![10i64, 20, 30, 40, 50].into_array(), + ) + .into_array(); + + let mut elements: Vec = [10i64, 30, 50, 70] + .iter() + .map(|v| { + Scalar::primitive(*v, Nullability::NonNullable) + .cast(&nullable) + .expect("i64 into the timestamp it stores") + }) + .collect(); + elements.push(Scalar::null(nullable.clone())); + let list = Scalar::list(Arc::new(nullable), elements, Nullability::Nullable); + + let len = needles.len(); + let result = ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(list, len).into_array(), needles], + ) + .expect("build") + .execute::(&mut ctx) + .expect("execute"); + assert_eq!( + result.bool_vec(&mut ctx), + vec![true, false, true, false, true] + ); + } + + #[test] + fn a_different_time_unit_is_not_matched_against_raw_storage() { + // The list is in seconds and the column in milliseconds. Their storage + // integers would compare equal, but the values do not. + let mut ctx = array_session().create_execution_ctx(); + let needles = ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + buffer![10i64, 20, 30].into_array(), + ) + .into_array(); + let seconds = + DType::Extension(Timestamp::new(TimeUnit::Seconds, Nullability::NonNullable).erased()); + let list = Scalar::list( + Arc::new(seconds.clone()), + vec![ + Scalar::primitive(10i64, Nullability::NonNullable) + .cast(&seconds) + .expect("i64 into the timestamp it stores"), + ], + Nullability::NonNullable, + ); + let len = needles.len(); + let result = ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(list, len).into_array(), needles], + ) + .expect("build") + .execute::(&mut ctx); + assert!( + result.is_err(), + "a seconds list must not match a milliseconds column" + ); + } } diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index ea656f24546..27ef85a8bb2 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cmp::Ordering; use std::sync::Arc; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; @@ -25,6 +27,7 @@ use crate::expr::lt_eq; use crate::expr::or; use crate::expr::or_collect; use crate::expr::stats::Stat; +use crate::scalar::Scalar; use crate::scalar::StringLike; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -401,6 +404,30 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { } } +/// Ceiling on how many interior-gap terms the `list_contains` falsifier emits +/// alongside the outer range test. +/// +/// A gap per adjacent pair proves exactly what a term per element would, so a +/// list shorter than this takes them all. A longer list takes only the boundaries +/// that divide it evenly, because each additional term is evaluated against +/// every zone and buys progressively less — the gaps of a wide scattered list +/// are all too narrow for a zone to fit inside however many are emitted. +/// +/// Unlike the per-element form this stays a single top-level `OR`, so it never +/// hands the optimizer the wide `AND` whose pairwise `between` search made +/// deriving the predicate quadratic in the list length. +/// +/// Boundaries are taken by position in the sorted list rather than by width: +/// comparing widths needs arithmetic the element type may not have (strings, +/// decimals across scales), while a position needs only the ordering every +/// element type already provides. +/// +/// Odd on purpose: the boundaries divide the sorted list evenly, so an odd +/// count puts one of them at the median — which is where the single wide gap +/// of a list split into a low group and a high group falls. An even count +/// straddles it and prunes measurably less on exactly that shape. +const FALSIFY_MAX_GAP_TERMS: usize = 31; + fn list_contains_falsify( expr: &Expression, ctx: &StatsRewriteCtx<'_>, @@ -408,12 +435,12 @@ fn list_contains_falsify( let list = expr.child(0); let needle = expr.child(1); - let Some(list_scalar) = literal_stat(list, Stat::Min) else { - return Ok(None); - }; - let elements = list_scalar + // Read the literal directly rather than through `literal_stat`, whose + // `Stat::Min` arm would clone the whole list into a fresh `Literal` only for + // it to be unwrapped again here. + let elements = list .as_opt::() - .and_then(|literal| literal.as_list_opt()) + .and_then(|scalar| scalar.as_list_opt()) .and_then(|list| list.elements()); let Some(elements) = elements else { return Ok(None); @@ -422,6 +449,13 @@ fn list_contains_falsify( return Ok(P::EMIT_UNGUARDED_REWRITES.then(|| lit(true))); } + // Before any work on the list: describing it as intervals means sorting it, + // and a rule that cannot guard this needle would sort it only to throw the + // result away. + let Some(guard) = non_nan_guard::

(ctx, [needle])? else { + return Ok(None); + }; + let Some(value_max) = max(needle, ctx) else { return Ok(None); }; @@ -429,16 +463,64 @@ fn list_contains_falsify( return Ok(None); }; - let value_predicate = and_collect(elements.iter().map(|value| { - or( - lt(value_max.clone(), lit(value.clone())), - gt(value_min.clone(), lit(value.clone())), - ) - })); - value_predicate - .map(|value_predicate| with_non_nan_guards::

(ctx, [needle], value_predicate)) - .transpose() - .map(Option::flatten) + // A null element cannot make `list_contains` true — the equality form maps + // it to `false` through its own null fill — so it is not a value a zone has + // to be checked against. Dropping nulls rather than abandoning the whole + // list is what lets `x IN (1, 2, NULL, 3)` prune at all. + let mut sorted: Vec<&Scalar> = elements.iter().filter(|e| !e.is_null()).collect(); + + // Sorting is what lets the whole list be described by a handful of + // intervals instead of one term per element, and the intervals are only + // sound if that order is the one the statistics comparison uses. Whether + // values order at all is a property of their dtype, which every element of + // a list shares, so one incomparable pair means none of them order and the + // list proves nothing. Floats do order here: they compare totally, so `NaN` + // and the two signed zeros have positions rather than being incomparable. + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal)); + if sorted + .windows(2) + .any(|pair| pair[0].partial_cmp(pair[1]).is_none()) + { + return Ok(None); + } + + // Equal values have nothing between them, so they cannot bound a gap. + // Collapsing them before the boundaries are chosen is what stops a list of + // mostly duplicates from spending every boundary inside one repeated value: + // `[0; 32] + [1000]` has exactly one gap, and sampling the sorted list + // directly would place all 31 boundaries inside the run of zeroes and miss + // it. + sorted.dedup_by(|a, b| (**a).partial_cmp(&**b) == Some(Ordering::Equal)); + + // Every element was null, so nothing orders and nothing is provable. + let (Some(lowest), Some(highest)) = (sorted.first(), sorted.last()) else { + return Ok(None); + }; + + // A scope is proven free of the list when it lies wholly outside the list's + // range, or wholly inside one of the gaps between two adjacent list values — + // nothing orders between an adjacent pair, so a row strictly between them + // cannot be in the list. The gaps are what make a clustered list prunable: + // a list split into a low group and a high group leaves most of the column + // *inside* its range, where the outer bound alone proves nothing. + let mut terms = Vec::with_capacity((sorted.len() + 1).min(FALSIFY_MAX_GAP_TERMS + 2)); + terms.push(lt(value_max.clone(), lit((*lowest).clone()))); + terms.push(gt(value_min.clone(), lit((*highest).clone()))); + + let gaps = (sorted.len() - 1).min(FALSIFY_MAX_GAP_TERMS); + for gap in 1..=gaps { + let idx = gap * sorted.len() / (gaps + 1); + let (below, above) = (sorted[idx - 1], sorted[idx]); + terms.push(and( + gt(value_min.clone(), lit(below.clone())), + lt(value_max.clone(), lit(above.clone())), + )); + } + + // The two outer-range terms are pushed unconditionally, so there is always + // something to reduce. + let value_predicate = or_collect(terms).vortex_expect("terms is never empty"); + Ok(Some(guard.apply(value_predicate))) } #[derive(Debug)] @@ -627,11 +709,36 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option .then(|| stat_fn(expr.clone(), aggregate_fn)) } -fn with_non_nan_guards<'a, P: NonNanProof>( +/// The NaN guard a rewrite carries, resolved independently of the value +/// predicate it will guard. +enum NanGuard { + /// The value predicate has to be conjoined with this check. + Check(Expression), + /// No possible NaN-bearing expression remains, so the value predicate is + /// already guarded and stands alone. + Unguarded, +} + +impl NanGuard { + fn apply(self, value_predicate: Expression) -> Expression { + match self { + Self::Check(check) => and(check, value_predicate), + Self::Unguarded => value_predicate, + } + } +} + +/// Resolves the guard for `exprs`, or `None` when this proof cannot guard them +/// and the rule must therefore emit nothing. +/// +/// Worth resolving before the value predicate is built rather than after: two +/// rules are registered per function and differ only in this proof, so for any +/// given needle one of them commonly emits nothing. Asking first is what stops +/// it from building a predicate it will discard. +fn non_nan_guard<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, exprs: impl IntoIterator, - value_predicate: Expression, -) -> VortexResult> { +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { match P::check(ctx, expr)? { @@ -640,18 +747,24 @@ fn with_non_nan_guards<'a, P: NonNanProof>( NanCheck::Unavailable => return Ok(None), } } - let nan_predicate = and_collect(nan_checks); - - Ok(match nan_predicate { - Some(nan_check) => Some(and(nan_check, value_predicate)), - // No possible NaN-bearing expression remains, so the value predicate is - // already guarded. Only one registered rule emits this unguarded - // rewrite so non-float comparisons are not duplicated. - None if P::EMIT_UNGUARDED_REWRITES => Some(value_predicate), + + Ok(match and_collect(nan_checks) { + Some(check) => Some(NanGuard::Check(check)), + // Only one registered rule emits an unguarded rewrite, so a non-float + // comparison is not rewritten twice. + None if P::EMIT_UNGUARDED_REWRITES => Some(NanGuard::Unguarded), None => None, }) } +fn with_non_nan_guards<'a, P: NonNanProof>( + ctx: &StatsRewriteCtx<'_>, + exprs: impl IntoIterator, + value_predicate: Expression, +) -> VortexResult> { + Ok(non_nan_guard::

(ctx, exprs)?.map(|guard| guard.apply(value_predicate))) +} + fn literal_stat(expr: &Expression, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { @@ -909,23 +1022,27 @@ mod tests { ); let expr = list_contains(lit(list), col("a")); + // The falsifier describes the list as intervals: the scope lies wholly + // outside the list's range, or wholly inside a gap between two adjacent + // list values. For a list this dense the gaps cover every element, so + // it proves exactly what a term per element would. assert_eq!( falsify(&expr)?, - Some(and( - and( - or( - lt(stat(col("a"), Stat::Max), lit(1i32)), + Some(or( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), + ), + or( + and( gt(stat(col("a"), Stat::Min), lit(1i32)), - ), - or( lt(stat(col("a"), Stat::Max), lit(2i32)), + ), + and( gt(stat(col("a"), Stat::Min), lit(2i32)), + lt(stat(col("a"), Stat::Max), lit(3i32)), ), ), - or( - lt(stat(col("a"), Stat::Max), lit(3i32)), - gt(stat(col("a"), Stat::Min), lit(3i32)), - ), )) ); Ok(()) diff --git a/vortex-array/tests/in_list_differential.rs b/vortex-array/tests/in_list_differential.rs new file mode 100644 index 00000000000..ead00bfcec4 --- /dev/null +++ b/vortex-array/tests/in_list_differential.rs @@ -0,0 +1,355 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Differential oracle for `list_contains` over a constant list. +//! +//! Every answer below was recorded by running this test against the commit +//! before the set probe was introduced, so the table pins the previous +//! implementation's results rather than the current one's. A change that makes +//! `list_contains` faster leaves it untouched; a change that makes it answer +//! differently fails here. +//! +//! The matrix is built around what the probe covers and what it declines. List +//! lengths straddle the threshold at which membership stops being answered by +//! OR-ing one equality per element, and run past it to the sizes the +//! performance claims rest on. Each covered needle type appears on both sides of +//! that threshold, because the two forms have to agree; the null-in-list rows +//! are what the probe refuses, so they exercise the fallback that has to keep +//! answering them. +//! +//! The values are chosen for the cases where a set keyed on anything but the +//! kernel's own equality would answer differently: `NaN`, which matches itself, +//! the two signed zeros, which do not match each other, strings long enough to +//! live in a buffer rather than inline in their view, and lists holding +//! duplicates. +//! +//! To regenerate: `git stash`, run against the parent commit, paste the +//! printed table into `EXPECTED`, then unstash. + +#![expect(clippy::expect_used, clippy::tests_outside_test_module)] + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::ByteBuffer; +use vortex_buffer::buffer; + +/// One row per case: the answers `list_contains` gave, as `T` (true), `.` +/// (false) and `?` (invalid). +fn render(case: &str, list: Scalar, needles: ArrayRef) -> String { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let result = needles + .apply(&list_contains(lit(list), root())) + .expect("apply list_contains") + .execute::(&mut ctx) + .expect("execute list_contains"); + + let bits = result.bit_buffer_view(); + let validity = result + .validity() + .expect("validity") + .execute_mask(bits.len(), &mut ctx) + .expect("validity mask"); + let answers: String = (0..bits.len()) + .map(|i| match (validity.value(i), bits.value(i)) { + (false, _) => '?', + (true, true) => 'T', + (true, false) => '.', + }) + .collect(); + format!("{case:34} {answers}") +} + +/// Integer needles: values in the list, values absent, and nulls landing on +/// both. +fn int_needles() -> ArrayRef { + PrimitiveArray::new( + buffer![0i64, 1, 2, 3, 4, 5, 7, 9, 15, 16, 63, 64, -1, 1_000, 0, 8], + Validity::from_iter((0..16).map(|i| i % 3 != 0)), + ) + .into_array() +} + +fn int_list(len: usize, nullability: Nullability) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + (0..len as i64).map(|v| Scalar::from(v * 2)).collect(), + nullability, + ) +} + +fn main_table() -> String { + let mut rows = Vec::new(); + + // Integer lists across the threshold and up to the sizes the performance + // claims are measured at. + for len in [1usize, 2, 3, 4, 5, 8, 64, 2048] { + for nullability in [Nullability::NonNullable, Nullability::Nullable] { + let tag = if nullability == Nullability::Nullable { + "nullable" + } else { + "nonnull" + }; + rows.push(render( + &format!("i64 m={len} list={tag}"), + int_list(len, nullability), + int_needles(), + )); + } + } + + // A null element: not a set key, so the whole list falls back. + let nullable_i64 = DType::Primitive(PType::I64, Nullability::Nullable); + let mut with_null: Vec = (0..8i64) + .map(|v| { + Scalar::from(v * 2) + .cast(&nullable_i64) + .expect("i64 into nullable i64") + }) + .collect(); + with_null.push(Scalar::null(nullable_i64.clone())); + rows.push(render( + "i64 m=9 one-null-element", + Scalar::list(Arc::new(nullable_i64), with_null, Nullability::Nullable), + int_needles(), + )); + + // Floats key on their bits, which is what `Operator::Eq` compares: `NaN` + // matches itself and `-0.0` does not match `0.0`. + let float_needles = || { + PrimitiveArray::new( + buffer![0.0f64, 1.0, 2.0, f64::NAN, 4.0, -0.0, 8.0, 1e9], + Validity::AllValid, + ) + .into_array() + }; + for len in [1usize, 2, 3, 4, 5, 64] { + rows.push(render( + &format!("f64 m={len}"), + Scalar::list( + Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), + (0..len).map(|v| Scalar::from(v as f64 * 2.0)).collect(), + Nullability::NonNullable, + ), + float_needles(), + )); + } + let mut nan_list: Vec = (0..8).map(|v| Scalar::from(v as f64 * 2.0)).collect(); + nan_list.push(Scalar::from(f64::NAN)); + rows.push(render( + "f64 m=9 nan-in-list", + Scalar::list( + Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), + nan_list, + Nullability::NonNullable, + ), + float_needles(), + )); + + // A signed zero on both sides: the two are distinct values, and each of the + // needles below is in exactly one of the two lists. + for (tag, zero) in [("neg", -0.0f64), ("pos", 0.0f64)] { + rows.push(render( + &format!("f64 m=4 {tag}-zero-in-list"), + Scalar::list( + Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), + [zero, 3.0, 5.0, 7.0].map(Scalar::from).to_vec(), + Nullability::NonNullable, + ), + float_needles(), + )); + } + + // Strings key on their bytes. The needles are short enough to sit inline in + // their view; `utf8 long` below uses values too long for that, which resolve + // through a data buffer instead. + let string_needles = || { + VarBinViewArray::from_iter_str(["k000", "k002", "k004", "zzz", "k126", "k010"]).into_array() + }; + for len in [1usize, 2, 3, 4, 5, 64] { + rows.push(render( + &format!("utf8 m={len}"), + Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + (0..len) + .map(|v| Scalar::from(format!("k{:03}", v * 2))) + .collect(), + Nullability::NonNullable, + ), + string_needles(), + )); + } + + // Values past the 12 bytes a view holds inline, so the probe has to reach + // the data buffer for every one of them. + let long = |v: usize| format!("a string too long to inline {v:04}"); + rows.push(render( + "utf8 m=8 long", + Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + (0..8).map(|v| Scalar::from(long(v * 2))).collect(), + Nullability::NonNullable, + ), + VarBinViewArray::from_iter_str([long(0), long(1), long(4), long(99), long(6)]).into_array(), + )); + + // Nullable needles, including one that is null, against a list of the same + // strings: a null needle answers false rather than null. + rows.push(render( + "utf8 m=4 nullable-needles", + Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + (0..4) + .map(|v| Scalar::from(format!("k{:03}", v * 2))) + .collect(), + Nullability::NonNullable, + ), + VarBinViewArray::from_iter( + [Some("k000"), None, Some("k004"), None, Some("zzz")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(), + )); + + // Byte strings take the same path as `Utf8`, keyed on the same bytes. + rows.push(render( + "binary m=4", + Scalar::list( + Arc::new(DType::Binary(Nullability::NonNullable)), + (0..4u8) + .map(|v| { + Scalar::binary( + ByteBuffer::copy_from([v * 2, 0xff]), + Nullability::NonNullable, + ) + }) + .collect(), + Nullability::NonNullable, + ), + VarBinViewArray::from_iter( + [ + Some([0u8, 0xff].as_slice()), + Some(&[2, 0xff]), + Some(&[1, 0xff]), + None, + Some(&[6, 0xff]), + ], + DType::Binary(Nullability::Nullable), + ) + .into_array(), + )); + + // Integer widths other than the i64 above, including an unsigned one: the + // set is keyed by the needle's own width, so each has to be built for it. + rows.push(render( + "i32 m=8", + Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + (0..8i32).map(|v| Scalar::from(v * 2)).collect(), + Nullability::NonNullable, + ), + PrimitiveArray::new( + buffer![0i32, 3, 4, -1, 14, 15], + Validity::from_iter([true, true, true, false, true, true]), + ) + .into_array(), + )); + rows.push(render( + "u8 m=8", + Scalar::list( + Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), + (0..8u8).map(|v| Scalar::from(v * 2)).collect(), + Nullability::NonNullable, + ), + PrimitiveArray::new(buffer![0u8, 3, 4, 255, 14, 15], Validity::AllValid).into_array(), + )); + + // A list that repeats a value, and one that is empty: both answer exactly + // what the set they key does. + rows.push(render( + "i64 m=8 duplicates", + Scalar::list( + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + [0i64, 0, 2, 2, 2, 4, 4, 4].map(Scalar::from).to_vec(), + Nullability::NonNullable, + ), + int_needles(), + )); + rows.push(render( + "i64 m=0 empty", + Scalar::list_empty( + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + Nullability::NonNullable, + ), + int_needles(), + )); + + rows.join("\n") +} + +/// Recorded from the implementation before the set probe. See the module doc. +const EXPECTED: &str = "\ +i64 m=1 list=nonnull ..............T.\n\ +i64 m=1 list=nullable ..............T.\n\ +i64 m=2 list=nonnull ..T...........T.\n\ +i64 m=2 list=nullable ..T...........T.\n\ +i64 m=3 list=nonnull ..T.T.........T.\n\ +i64 m=3 list=nullable ..T.T.........T.\n\ +i64 m=4 list=nonnull ..T.T.........T.\n\ +i64 m=4 list=nullable ..T.T.........T.\n\ +i64 m=5 list=nonnull ..T.T.........T.\n\ +i64 m=5 list=nullable ..T.T.........T.\n\ +i64 m=8 list=nonnull ..T.T.........T.\n\ +i64 m=8 list=nullable ..T.T.........T.\n\ +i64 m=64 list=nonnull ..T.T......T..T.\n\ +i64 m=64 list=nullable ..T.T......T..T.\n\ +i64 m=2048 list=nonnull ..T.T......T.TT.\n\ +i64 m=2048 list=nullable ..T.T......T.TT.\n\ +i64 m=9 one-null-element ..T.T.........T.\n\ +f64 m=1 T.......\n\ +f64 m=2 T.T.....\n\ +f64 m=3 T.T.T...\n\ +f64 m=4 T.T.T...\n\ +f64 m=5 T.T.T.T.\n\ +f64 m=64 T.T.T.T.\n\ +f64 m=9 nan-in-list T.TTT.T.\n\ +f64 m=4 neg-zero-in-list .....T..\n\ +f64 m=4 pos-zero-in-list T.......\n\ +utf8 m=1 T.....\n\ +utf8 m=2 TT....\n\ +utf8 m=3 TTT...\n\ +utf8 m=4 TTT...\n\ +utf8 m=5 TTT...\n\ +utf8 m=64 TTT.TT\n\ +utf8 m=8 long T.T.T\n\ +utf8 m=4 nullable-needles T.T..\n\ +binary m=4 TT..T\n\ +i32 m=8 T.T.T.\n\ +u8 m=8 T.T.T.\n\ +i64 m=8 duplicates ..T.T.........T.\n\ +i64 m=0 empty ................"; + +#[test] +fn constant_list_contains_answers_are_unchanged() { + let actual = main_table(); + assert_eq!( + actual, EXPECTED, + "\nlist_contains changed an answer. Actual table:\n{actual}\n" + ); +} diff --git a/vortex-layout/tests/list_contains_pruning.rs b/vortex-layout/tests/list_contains_pruning.rs new file mode 100644 index 00000000000..56c21fcf4f3 --- /dev/null +++ b/vortex-layout/tests/list_contains_pruning.rs @@ -0,0 +1,448 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The `list_contains` falsifier must only ever claim that a zone holds no +//! matching row. +//! +//! Deriving the predicate is cheap to check on its own, and the unit tests +//! beside the falsifier do. What they cannot see is whether that predicate, +//! once evaluated against a zone's statistics, excludes a zone that does hold a +//! matching row — which is how an `IN` query comes back missing rows. So this +//! checks the property directly: every case states its zones, derives the +//! statistics from them, prunes, and requires that each excluded zone contains +//! zero matches. +//! +//! The falsifier describes the list as intervals rather than as one equality +//! per element, so it depends on the sorted order of the list agreeing with the +//! order the statistics comparison uses. These cases are built around the +//! places that could disagree: floats (where the order is total, so `NaN` and +//! the two signed zeros have positions), truncated string bounds (which widen a +//! zone rather than narrow it), duplicate elements, and the list lengths either +//! side of the cap on how many interior gaps are emitted. + +#![expect(clippy::expect_used, clippy::panic, clippy::tests_outside_test_module)] + +use std::sync::Arc; +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::aggregate_fn::fns::min::Min; +use vortex_array::aggregate_fn::fns::nan_count::NanCount; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_io::session::RuntimeSession; +use vortex_layout::layouts::zoned::zone_map::ZoneMap; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(|| { + vortex_array::array_session() + .with::() + .with::() +}); + +/// Rows per zone. Every case states its zones as equal-length groups so that +/// the nominal zone length and the row count agree. +const ZONE_LEN: usize = 4; + +/// Prunes `stats` with the falsifier derived for `list`, and fails if a zone it +/// excluded holds a matching row. +/// +/// Returns how many zones were excluded, so a caller can also require that a +/// case prunes at all — a case that prunes nothing satisfies the soundness +/// property without having tested it. +fn zones_pruned( + case: &str, + column_dtype: &DType, + stats: Vec<(AggregateFnRef, ArrayRef)>, + list: Scalar, + zone_has_match: &[bool], +) -> usize { + let zone_count = zone_has_match.len(); + let fields = stats + .iter() + .map(|(agg, values)| (agg.to_string(), values.clone())) + .collect::>(); + let zone_map = ZoneMap::try_new( + column_dtype.clone(), + StructArray::from_fields(&fields).expect("stats table"), + stats.iter().map(|(agg, _)| agg.clone()).collect(), + ZONE_LEN as u64, + (zone_count * ZONE_LEN) as u64, + ) + .expect("zone map"); + + let predicate = list_contains(lit(list), root()) + .falsify(column_dtype, &SESSION) + .expect("falsify") + .unwrap_or_else(|| panic!("{case}: no falsifier derived")); + + let mask = zone_map.prune(&predicate, &SESSION).expect("prune"); + + let mut pruned = 0; + for zone in 0..zone_count { + if !mask.value(zone) { + continue; + } + pruned += 1; + assert!( + !zone_has_match[zone], + "{case}: zone {zone} was pruned but holds a matching row", + ); + } + pruned +} + +/// Statistics for primitive zones, as the writer computes them: the minimum and +/// maximum under the total order skipping NaN, plus the NaN count that a float +/// column's bounds are only usable alongside. +fn primitive_stats(zones: &[Vec]) -> Vec<(AggregateFnRef, ArrayRef)> { + let mut mins = Vec::with_capacity(zones.len()); + let mut maxs = Vec::with_capacity(zones.len()); + let mut nan_counts = Vec::with_capacity(zones.len()); + for zone in zones { + let non_nan = || zone.iter().copied().filter(|v| !v.is_nan()); + // Every case keeps at least one ordinary value in each zone, so the + // bounds are always defined. + mins.push( + non_nan() + .min_by(|a, b| a.total_compare(*b)) + .expect("zone has a non-NaN value"), + ); + maxs.push( + non_nan() + .max_by(|a, b| a.total_compare(*b)) + .expect("zone has a non-NaN value"), + ); + nan_counts.push(zone.iter().filter(|v| v.is_nan()).count() as u64); + } + + let mut stats = vec![ + ( + Min.bind(NumericalAggregateOpts::skip_nans()), + PrimitiveArray::new(Buffer::::copy_from(&mins), Validity::AllValid).into_array(), + ), + ( + Max.bind(NumericalAggregateOpts::skip_nans()), + PrimitiveArray::new(Buffer::::copy_from(&maxs), Validity::AllValid).into_array(), + ), + ]; + // A float column's bounds are only usable alongside its NaN count, and only + // a float column has one to store. + if T::PTYPE.is_float() { + stats.push(( + NanCount.bind(EmptyOptions), + PrimitiveArray::new(Buffer::::copy_from(&nan_counts), Validity::AllValid) + .into_array(), + )); + } + stats +} + +/// Statistics for string zones, taken verbatim so a case can state bounds that +/// are wider than the zone's own values — which is what the writer stores when +/// it truncates a long bound. +fn utf8_stats(mins: &[&str], maxs: &[&str]) -> Vec<(AggregateFnRef, ArrayRef)> { + let column = |values: &[&str]| { + VarBinViewArray::from_iter( + values.iter().map(|value| Some(*value)), + DType::Utf8(Nullability::Nullable), + ) + .into_array() + }; + vec![ + (Min.bind(NumericalAggregateOpts::skip_nans()), column(mins)), + (Max.bind(NumericalAggregateOpts::skip_nans()), column(maxs)), + ] +} + +fn int_list(values: &[i64]) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)), + values.iter().copied().map(Scalar::from).collect(), + Nullability::NonNullable, + ) +} + +fn float_list(values: &[f64]) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), + values.iter().copied().map(Scalar::from).collect(), + Nullability::NonNullable, + ) +} + +fn utf8_list(values: &[&str]) -> Scalar { + Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + values + .iter() + .map(|v| Scalar::utf8(*v, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ) +} + +/// Whether each zone holds a value of `list`, under integer equality. +fn int_matches(zones: &[Vec], list: &[i64]) -> Vec { + zones + .iter() + .map(|zone| zone.iter().any(|v| list.contains(v))) + .collect() +} + +#[test] +fn a_zone_inside_a_gap_is_pruned_and_one_holding_a_value_is_not() { + // The list is two clusters far apart, which is the shape the interior gaps + // exist for: every zone lies inside the list's overall range, so the outer + // bounds alone prove nothing. + let list: Vec = (0..8).chain(1_000..1_008).collect(); + let zones: Vec> = vec![ + vec![2, 3, 4, 5], // holds list values + vec![100, 200, 300, 400], // wholly inside the gap + vec![900, 950, 999, 1_000], // straddles the gap's upper end, holds 1000 + vec![1_004, 1_005, 1_006, 1_007], + vec![20, 30, 40, 50], // inside the gap, nowhere near a value + ]; + let matches = int_matches(&zones, &list); + assert_eq!(matches, vec![true, false, true, true, false]); + + let pruned = zones_pruned( + "two clusters", + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + int_list(&list), + &matches, + ); + assert_eq!(pruned, 2, "both gap-interior zones should prune"); +} + +#[test] +fn duplicate_elements_do_not_open_a_gap() { + // A duplicated value has nothing between it and itself. Treating the pair + // as a gap would exclude every zone whose bounds sit on that value, which + // is exactly the zone that matches. + let list = vec![5i64, 5, 5, 5, 5, 5, 5, 5]; + let zones: Vec> = vec![vec![5, 5, 5, 5], vec![1, 2, 3, 4], vec![6, 7, 8, 9]]; + let matches = int_matches(&zones, &list); + assert_eq!(matches, vec![true, false, false]); + + let pruned = zones_pruned( + "all duplicates", + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + int_list(&list), + &matches, + ); + assert_eq!(pruned, 2, "the zones either side of the value should prune"); +} + +#[test] +fn every_list_length_around_the_gap_cap_stays_sound() { + // Only so many interior gaps are emitted, and which boundaries they fall on + // changes with the list length. Sweeping the lengths either side of that cap + // covers the change of regime; the values are spread so that most zones sit + // inside some gap and pruning is actually exercised. + let mut total_pruned = 0; + for len in [1usize, 2, 3, 4, 5, 30, 31, 32, 33, 34, 35, 64, 129] { + let list: Vec = (0..len as i64).map(|i| i * 100).collect(); + let zones: Vec> = (0..16) + .map(|z| { + let base = z * 61; + vec![base, base + 7, base + 13, base + 29] + }) + .collect(); + let matches = int_matches(&zones, &list); + total_pruned += zones_pruned( + &format!("m={len}"), + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + int_list(&list), + &matches, + ); + } + assert!( + total_pruned > 0, + "the sweep pruned nothing, so proved nothing" + ); +} + +#[test] +fn floats_order_nan_and_signed_zero_the_way_the_statistics_do() { + // Float comparison here is total: `NaN` equals itself and sorts above every + // number, and `-0.0` is distinct from and below `0.0`. The falsifier sorts + // the list with the same order the statistics comparison uses, so these + // three values have positions rather than being unordered — this is the case + // that goes wrong if the two orders ever disagree. + let list = vec![-1.0f64, -0.0, 0.0, 1.0, 100.0, f64::NAN]; + let zones: Vec> = vec![ + vec![-0.0, 0.0, 0.5, 0.75], // holds both zeros + vec![10.0, 20.0, 30.0, 40.0], // inside the gap between 1 and 100 + vec![1.0, 2.0, 3.0, 4.0], // holds 1.0 + vec![f64::NAN, 200.0, 300.0, 400.0], // holds NaN, which is in the list + vec![-100.0, -50.0, -20.0, -10.0], // below the list's lowest value + ]; + let matches: Vec = zones + .iter() + .map(|zone| { + zone.iter() + .any(|v| list.iter().any(|l| l.to_bits() == v.to_bits())) + }) + .collect(); + assert_eq!(matches, vec![true, false, true, true, false]); + + let pruned = zones_pruned( + "floats with NaN and signed zero", + &DType::Primitive(PType::F64, Nullability::NonNullable), + primitive_stats(&zones), + float_list(&list), + &matches, + ); + assert!(pruned > 0, "no float zone pruned, so nothing was tested"); +} + +#[test] +fn truncated_string_bounds_cannot_create_a_false_exclusion() { + // A zone whose extreme value is too long to store is written with a bound + // that is looser than the truth: a minimum below the real minimum and a + // maximum above the real maximum. That widens a zone, which can only stop it + // being pruned, never cause it to be pruned wrongly — this states the case + // in which it would, if the direction were ever reversed. + let list = ["aa", "abzz", "ac", "zz"]; + // Zone 0 holds "abzz", stored truncated to two bytes as min "ab", max "ac". + // Zone 1 holds only values inside the gap between "ac" and "zz". + // Zone 2 holds "zz" with an exact bound. + let mins = ["ab", "b", "zz"]; + let maxs = ["ac", "c", "zz"]; + let matches = vec![true, false, true]; + + let pruned = zones_pruned( + "truncated utf8 bounds", + &DType::Utf8(Nullability::NonNullable), + utf8_stats(&mins, &maxs), + utf8_list(&list), + &matches, + ); + assert_eq!(pruned, 1, "only the gap-interior zone should prune"); +} + +#[test] +fn exhaustive_over_a_small_domain() { + // Every list of up to five values from a small domain, against every + // partition of that domain into three zones. This is the case the + // hand-written ones above cannot cover: it makes no assumption about which + // shape of list or zone is interesting, so a boundary that only misbehaves + // for one arrangement still has to show up. + const DOMAIN: i64 = 7; + let mut checked = 0usize; + let mut pruned_total = 0usize; + + for mask in 1u32..(1 << DOMAIN) { + let list: Vec = (0..DOMAIN).filter(|i| mask & (1 << i) != 0).collect(); + if list.len() > 5 { + continue; + } + for shift in 0..DOMAIN { + let zones: Vec> = (0..3) + .map(|z| { + (0..ZONE_LEN as i64) + .map(|r| (shift + z * 2 + r) % DOMAIN) + .collect() + }) + .collect(); + let matches = int_matches(&zones, &list); + pruned_total += zones_pruned( + &format!("list={list:?} shift={shift}"), + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + int_list(&list), + &matches, + ); + checked += 1; + } + } + + assert!(checked > 500, "expected a wide sweep, ran {checked} cases"); + assert!( + pruned_total > 0, + "the sweep pruned nothing, so proved nothing" + ); +} + +#[test] +fn a_run_of_duplicates_does_not_consume_every_boundary() { + // Only so many boundaries are emitted, and they divide the sorted list + // evenly. A list that is mostly one repeated value has exactly one real gap; + // choosing boundaries before collapsing the duplicates spends all of them + // inside the run and misses it, which prunes nothing at all. + let mut list = vec![0i64; 32]; + list.push(1_000); + let zones: Vec> = vec![ + vec![500, 501, 502, 503], // wholly inside the only gap + vec![0, 1, 2, 3], // holds 0 + vec![997, 998, 999, 1_000], // holds 1000 + ]; + let matches = int_matches(&zones, &list); + assert_eq!(matches, vec![false, true, true]); + + let pruned = zones_pruned( + "a run of duplicates", + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + int_list(&list), + &matches, + ); + assert_eq!(pruned, 1, "the zone inside the only gap should prune"); +} + +#[test] +fn a_null_element_does_not_cost_the_whole_list_its_pruning() { + // A null element cannot make `list_contains` true, so it is not a value a + // zone has to be checked against. Abandoning the list because one is present + // would give an ordinary parameterized `IN` list no pruning whatsoever. + let nullable = DType::Primitive(PType::I64, Nullability::Nullable); + let mut elements: Vec = [0i64, 1, 2, 1_000] + .iter() + .map(|v| { + Scalar::from(*v) + .cast(&nullable) + .expect("i64 into nullable i64") + }) + .collect(); + elements.push(Scalar::null(nullable.clone())); + let list = Scalar::list(Arc::new(nullable), elements, Nullability::Nullable); + + let zones: Vec> = vec![ + vec![500, 501, 502, 503], // inside the gap between 2 and 1000 + vec![0, 1, 2, 3], // holds list values + vec![997, 998, 999, 1_000], // holds 1000 + ]; + let matches = int_matches(&zones, &[0, 1, 2, 1_000]); + assert_eq!(matches, vec![false, true, true]); + + let pruned = zones_pruned( + "a null element", + &DType::Primitive(PType::I64, Nullability::NonNullable), + primitive_stats(&zones), + list, + &matches, + ); + assert_eq!(pruned, 1, "the gap-interior zone should still prune"); +}