From f73241661fbe12d73e03050f90fb2c9dbebe8454 Mon Sep 17 00:00:00 2001 From: Ben Chambers Date: Wed, 9 Sep 2026 23:07:19 -0700 Subject: [PATCH 1/6] perf(array): answer a constant IN list by probing a set, and falsify it by interval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A constant `IN (...)` list was answered by building one full-length `Eq` array per list element and OR-reducing them, so a query paid one pass over every row for every element of the list. Deriving the statistics predicate was worse than linear: it emitted an `AND` of one term per element, and the optimize pass compares every pair of conjuncts, so the cost of *deciding whether a file could be skipped* grew with the square of the list length, regardless of how many rows the file held. Membership is now answered by keying a set on the list once and probing it in a single pass over the needles, for primitive, `Utf8` and `Binary` needles. The set's equality is the kernel's own — `NativePType::is_eq`, which compares floats by their bits, so `NaN` matches itself and `-0.0` does not match `0.0`, the answers the OR-of-equalities form gives. Below four elements a comparison still wins outright, so short lists keep that form, which is also what answers everything the probe does not cover. The statistics predicate becomes a single top-level `OR` describing the list as intervals: the scope lies wholly outside the list's range, or wholly inside one of the gaps between two adjacent sorted values, since nothing orders between an adjacent pair. The gaps are what make a clustered list prunable at all — a list split into a low group and a high group leaves most of a column inside its overall range, where the outer bounds prove nothing. Staying a single `OR` is also what keeps the pairwise `between` search off the derivation path. FSST gets a `ListContainsElementKernel`, the first implementation of that trait: compression under a fixed symbol table is deterministic and lossless, so two values are equal exactly when their codes are. Compressing the list once therefore answers `IN` over compressed strings without decompressing any of them — the same property `compare_fsst_constant` already relies on for `Eq`. `ListScalar::element_values` borrows the element values instead of materializing a `Vec`, which the probe pays once per batch. Evidence -------- `vortex-layout/tests/list_contains_pruning.rs` checks the property that makes the interval form safe: for each case it derives the predicate, prunes a zone map built from the zones' own statistics, and requires every excluded zone to hold zero matching rows. It covers floats (where `NaN` and the two signed zeros have positions under the total order), string bounds stored wider than the zone's values as truncation stores them, duplicate elements, the list lengths either side of the gap cap, and an exhaustive sweep over a small domain. Making the gap bounds non-strict fails 5 of its 6 cases; sorting the list the other way fails the same 5; loosening the outer range test fails 2. `vortex-array/tests/in_list_differential.rs` asserts 39 rows of answers recorded by running it against the commit before this one, so a change that makes `list_contains` faster leaves the table untouched and one that makes it answer differently fails. --- encodings/fsst/src/compute/list_contains.rs | 390 +++++++++++++++ encodings/fsst/src/compute/mod.rs | 1 + encodings/fsst/src/kernel.rs | 7 + .../arrays/extension/compute/list_contains.rs | 190 ++++++++ .../src/arrays/extension/compute/mod.rs | 1 + .../src/arrays/extension/vtable/kernel.rs | 7 + vortex-array/src/arrays/varbinview/mod.rs | 2 + .../src/arrays/varbinview/views_side.rs | 57 +++ vortex-array/src/scalar/typed_view/list.rs | 11 + .../src/scalar_fn/fns/binary/compare/bytes.rs | 45 +- .../src/scalar_fn/fns/list_contains/mod.rs | 238 +++++++++- vortex-array/src/stats/rewrite/builtins.rs | 189 ++++++-- vortex-array/tests/in_list_differential.rs | 355 ++++++++++++++ vortex-layout/tests/list_contains_pruning.rs | 448 ++++++++++++++++++ 14 files changed, 1843 insertions(+), 98 deletions(-) create mode 100644 encodings/fsst/src/compute/list_contains.rs create mode 100644 vortex-array/src/arrays/extension/compute/list_contains.rs create mode 100644 vortex-array/src/arrays/varbinview/views_side.rs create mode 100644 vortex-array/tests/in_list_differential.rs create mode 100644 vortex-layout/tests/list_contains_pruning.rs diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs new file mode 100644 index 00000000000..fbbd2e240f1 --- /dev/null +++ b/encodings/fsst/src/compute/list_contains.rs @@ -0,0 +1,390 @@ +// 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); + } + + 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 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("this is a very long string"), + Some("this is a very long strinX"), + Some("hello"), + None, + ]; + + fn answers( + list: Vec<&str>, + needle_nullability: Nullability, + ) -> VortexResult>> { + let mut ctx = SESSION.create_execution_ctx(); + let plain = VarBinArray::from_iter(VALUES, 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"); + Ok(results.remove(1)) + } + + #[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!["this is a very long string", "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<()> { + let mut ctx = SESSION.create_execution_ctx(); + let plain = VarBinArray::from_iter( + ["a", "b", "c", "d"].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 list = Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + ["a", "c", "e", "f"] + .map(|v| Scalar::utf8(v, Nullability::NonNullable)) + .to_vec(), + Nullability::NonNullable, + ); + let len = compressed.len(); + let result = ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(list, len).into_array(), compressed], + )? + .execute::(&mut ctx)?; + + assert_eq!( + bool_answers(&result, &mut ctx), + 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/extension/compute/list_contains.rs b/vortex-array/src/arrays/extension/compute/list_contains.rs new file mode 100644 index 00000000000..4fabd45f905 --- /dev/null +++ b/vortex-array/src/arrays/extension/compute/list_contains.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::Extension; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; +use crate::scalar::Scalar; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_contains::ListContainsElementKernel; + +/// Answers `IN` for an extension needle through its storage. +/// +/// An extension value *is* its storage value — `ScalarValue` has no extension +/// variant — so membership over a timestamp, date or any user-defined extension +/// is membership over the integers or bytes underneath it. `Operator::Eq` +/// already compares extensions this way, so unwrapping both sides gives the same +/// answers while letting the storage type's own set probe do the work. +impl ListContainsElementKernel for Extension { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + 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); + }; + + // Two extensions are only comparable when they agree on the whole + // extension dtype — a timestamp in milliseconds must not be matched + // against one in seconds. Executing as a parent kernel skips the check + // `ListContains` does for itself, so decline a mismatch here and let the + // generic path raise the error. + if !list_scalar + .element_dtype() + .eq_ignore_nullability(element.dtype()) + { + return Ok(None); + } + + let storage = element.storage_array().clone(); + let storage_dtype = storage.dtype().clone(); + let storage_list = Scalar::list( + storage_dtype, + elements + .iter() + .map(|value| { + // SAFETY: these are the values of a validated extension list + // scalar, so each is a valid value of the storage dtype. + unsafe { Scalar::new_unchecked(storage.dtype().clone(), value.clone()) } + }) + .collect(), + list_constant.scalar().dtype().nullability(), + ); + let len = storage.len(); + + ListContains + .try_new_array( + len, + EmptyOptions, + [ConstantArray::new(storage_list, len).into_array(), storage], + )? + .execute(ctx) + .map(Some) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vortex_buffer::buffer; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::ExtensionArray; + use crate::arrays::bool::BoolArrayExt; + use crate::arrays::scalar_fn::ScalarFnFactoryExt; + use crate::dtype::DType; + use crate::dtype::Nullability; + 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::ListContains; + + /// 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"); + let bits = result.bit_buffer_view(); + (0..bits.len()).map(|i| bits.value(i)).collect() + } + + #[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_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/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index 943733eb45d..eacd83b0116 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -4,6 +4,7 @@ mod cast; mod compare; mod filter; +mod list_contains; mod mask; pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index a81863a3ed6..2d9f8587066 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -13,6 +13,8 @@ use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -22,5 +24,10 @@ pub(crate) fn initialize(session: &VortexSession) { CompareExecuteAdaptor(Extension), ); kernels.register_execute_parent_kernel(Cast.id(), Extension, CastExecuteAdaptor(Extension)); + kernels.register_execute_parent_kernel( + ListContains.id(), + Extension, + ListContainsElementExecuteAdaptor(Extension), + ); kernels.register_execute_parent_kernel(Dict.id(), Extension, TakeExecuteAdaptor(Extension)); } 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..e6ecca48be1 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,22 @@ 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_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; @@ -25,18 +29,26 @@ use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBin; +use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; 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 +123,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 +191,55 @@ 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"); + + if element_values.len() >= HASH_PROBE_MIN_ELEMENTS + && 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 +264,155 @@ 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); + } + _ => 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); + }; + probe_rows(needles.as_slice::().iter(), &validity, |needle| { + set.contains(&NativeValue(*needle)) + }) + }); + + 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. + probe_rows(offsets.as_slice::().windows(2), &validity, |pair| { + let (start, end): (usize, usize) = (pair[0].as_(), pair[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); + probe_rows(side.views().iter(), &validity, |view| { + set.contains(side.view_bytes(view)) + }) + }; + + 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. +fn probe_rows( + needles: impl ExactSizeIterator, + validity: &Mask, + hit: impl Fn(N) -> bool, +) -> BitBuffer { + match validity { + Mask::AllTrue(_) => needles.map(hit).collect(), + Mask::AllFalse(len) => BitBuffer::new_unset(*len), + Mask::Values(valid) => needles + .enumerate() + .map(|(idx, needle)| valid.value(idx) && hit(needle)) + .collect(), + } +} + /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. fn list_contains_scalar( array: &ArrayRef, @@ -599,23 +797,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(()) 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"); +} From aff66352b61b0e0fb86af5314d78e1cb932c3ff4 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:31:38 -0700 Subject: [PATCH 2/6] test(fsst): rename two fixture values the spell check reads as a misspelling The pair differed only in a final byte, to show that a code comparison stopping at a shared prefix would answer both alike. `strinX` carried that property but reads as a misspelling of `string`. --- encodings/fsst/src/compute/list_contains.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs index fbbd2e240f1..9a47b9562b6 100644 --- a/encodings/fsst/src/compute/list_contains.rs +++ b/encodings/fsst/src/compute/list_contains.rs @@ -162,8 +162,8 @@ mod tests { None, Some("world"), Some(""), - Some("this is a very long string"), - Some("this is a very long strinX"), + Some("a value long enough to need a buffer 1"), + Some("a value long enough to need a buffer 2"), Some("hello"), None, ]; @@ -253,7 +253,13 @@ mod tests { // comparison that stopped at a prefix would answer both alike. assert_eq!( answers( - vec!["this is a very long string", "q", "qq", "qqq", "qqqq"], + vec![ + "a value long enough to need a buffer 1", + "q", + "qq", + "qqq", + "qqqq" + ], Nullability::Nullable )?, vec![ From af1c5b301a87d8e56115d02a2918af446bf1c027 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:50:09 -0700 Subject: [PATCH 3/6] fix(array): decline a null-bearing extension list instead of panicking An extension list and a needle only have to agree on their dtype ignoring nullability, so a nullable list can reach the extension kernel against a non-nullable column. A null element there has no storage value to unwrap, and building one panicked in `Scalar::new_unchecked`. Such a list now goes back to the generic path, which answers it, matching what the FSST kernel already does. The FSST tests were not reaching the kernel they exist for. It declines below a ratio of rows to list elements, and the fixture was eight rows, so a five-element list needed a hundred and sixty before the kernel would engage. The fixture now fills enough rows for the longest list the tests use, and asserts the ratio it needs rather than leaving it implicit. Removing a matching element from the compressed list now fails five of the seven; the two that still pass are the two that decline before compressing. The probe walks valid rows a word at a time instead of testing validity per row. Measured against each other in one binary over an 8192-row batch, the word-at-a-time form wins at every density: 1.91x at fully valid, where it fast-paths whole words, through 4.06x at a quarter valid to 28x at 0.2%. --- encodings/fsst/src/compute/list_contains.rs | 76 ++++++++++++------- .../arrays/extension/compute/list_contains.rs | 68 +++++++++++++++-- .../src/scalar_fn/fns/list_contains/mod.rs | 45 ++++++----- 3 files changed, 134 insertions(+), 55 deletions(-) diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs index 9a47b9562b6..e77b7170a59 100644 --- a/encodings/fsst/src/compute/list_contains.rs +++ b/encodings/fsst/src/compute/list_contains.rs @@ -130,6 +130,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::MIN_ROWS_PER_LIST_ELEMENT; use crate::fsst_compress; use crate::fsst_train_compressor; @@ -168,12 +169,40 @@ mod tests { 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(VALUES, DType::Utf8(needle_nullability)).into_array(); + 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(); @@ -201,7 +230,19 @@ mod tests { // 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"); - Ok(results.remove(1)) + + // 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] @@ -308,33 +349,12 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn non_nullable_needles() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let plain = VarBinArray::from_iter( - ["a", "b", "c", "d"].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 list = Scalar::list( - Arc::new(DType::Utf8(Nullability::NonNullable)), - ["a", "c", "e", "f"] - .map(|v| Scalar::utf8(v, Nullability::NonNullable)) - .to_vec(), - Nullability::NonNullable, - ); - let len = compressed.len(); - let result = ListContains - .try_new_array( - len, - EmptyOptions, - [ConstantArray::new(list, len).into_array(), compressed], - )? - .execute::(&mut ctx)?; - assert_eq!( - bool_answers(&result, &mut ctx), + 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(()) diff --git a/vortex-array/src/arrays/extension/compute/list_contains.rs b/vortex-array/src/arrays/extension/compute/list_contains.rs index 4fabd45f905..99b8f036d7c 100644 --- a/vortex-array/src/arrays/extension/compute/list_contains.rs +++ b/vortex-array/src/arrays/extension/compute/list_contains.rs @@ -52,16 +52,25 @@ impl ListContainsElementKernel for Extension { let storage = element.storage_array().clone(); let storage_dtype = storage.dtype().clone(); + let mut storage_elements = Vec::with_capacity(elements.len()); + for value in elements { + // A null element has no storage value to carry, and the dtype check + // above ignores nullability, so a nullable list can arrive against a + // non-nullable column. The generic path answers such a list, so hand + // the whole thing back to it rather than building a scalar the + // storage dtype cannot hold. + let Some(value) = value else { + return Ok(None); + }; + // SAFETY: these are the values of a validated extension list scalar, + // so each is a valid non-null value of the storage dtype. + storage_elements + .push(unsafe { Scalar::new_unchecked(storage_dtype.clone(), Some(value.clone())) }); + } + let storage_list = Scalar::list( storage_dtype, - elements - .iter() - .map(|value| { - // SAFETY: these are the values of a validated extension list - // scalar, so each is a valid value of the storage dtype. - unsafe { Scalar::new_unchecked(storage.dtype().clone(), value.clone()) } - }) - .collect(), + storage_elements, list_constant.scalar().dtype().nullability(), ); let len = storage.len(); @@ -152,6 +161,49 @@ mod tests { ); } + #[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"); + let bits = result.bit_buffer_view(); + assert_eq!( + (0..bits.len()).map(|i| bits.value(i)).collect::>(), + 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 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 e6ecca48be1..9d1f14c1352 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -11,6 +11,7 @@ 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; @@ -297,9 +298,8 @@ fn hash_probe_contains( let Some(set) = primitive_key_set::(elements) else { return Ok(None); }; - probe_rows(needles.as_slice::().iter(), &validity, |needle| { - set.contains(&NativeValue(*needle)) - }) + let slice = needles.as_slice::(); + probe_rows(len, &validity, |idx| set.contains(&NativeValue(slice[idx]))) }); Ok(Some(BoolArray::new(bits, nullability.into()).into_array())) @@ -375,8 +375,9 @@ fn bytes_probe_contains( match_each_integer_ptype!(offsets.ptype(), |O| { // One offset per row plus a final end, so adjacent pairs are the // rows in order. - probe_rows(offsets.as_slice::().windows(2), &validity, |pair| { - let (start, end): (usize, usize) = (pair[0].as_(), pair[1].as_()); + 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]) }) }) @@ -386,8 +387,9 @@ fn bytes_probe_contains( // 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); - probe_rows(side.views().iter(), &validity, |view| { - set.contains(side.view_bytes(view)) + let views = side.views(); + probe_rows(len, &validity, |idx| { + set.contains(side.view_bytes(&views[idx])) }) }; @@ -397,19 +399,24 @@ fn bytes_probe_contains( /// 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. -fn probe_rows( - needles: impl ExactSizeIterator, - validity: &Mask, - hit: impl Fn(N) -> bool, -) -> BitBuffer { +/// 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(_) => needles.map(hit).collect(), - Mask::AllFalse(len) => BitBuffer::new_unset(*len), - Mask::Values(valid) => needles - .enumerate() - .map(|(idx, needle)| valid.value(idx) && hit(needle)) - .collect(), + 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() + } } } From 2f1a22ada71507c9d23a96845e7e6fa21e1008e5 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:30:48 -0700 Subject: [PATCH 4/6] perf(array): answer an extension IN list in the generic probe, not a kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The extension kernel rebuilt the whole list as storage `Scalar`s for every batch, then wrapped them in a constant array and a fresh `ListContains` to recurse. Measured against the same query over the storage type directly, that marshalling cost ~12ns an element and did not vary with batch size: +88% at a thousand-element list over an 8192-row batch, +155% at sixty-five thousand. An extension value is its storage value — `ScalarValue` has no extension variant — so the generic probe can key on the list's values unchanged once the needle is unwrapped. Moving it there drops the per-batch cost to nothing, and the dtype check the kernel had to re-derive for itself is one the generic path has already done: executing as a parent kernel is what skipped it. It also removes the `Scalar::new_unchecked` that retagged each element with the column's storage dtype. That call had to assume the two dtypes agreed all the way down, while the check above it compares them ignoring nullability, so a list whose storage was a struct with a nullable field could be retagged against a column whose field was not. Nothing constructs a scalar now, and a storage type the probe does not cover declines to the equality form as before. --- .../arrays/extension/compute/list_contains.rs | 242 ------------------ .../src/arrays/extension/compute/mod.rs | 1 - .../src/arrays/extension/vtable/kernel.rs | 7 - .../src/scalar_fn/fns/list_contains/mod.rs | 153 +++++++++++ 4 files changed, 153 insertions(+), 250 deletions(-) delete mode 100644 vortex-array/src/arrays/extension/compute/list_contains.rs diff --git a/vortex-array/src/arrays/extension/compute/list_contains.rs b/vortex-array/src/arrays/extension/compute/list_contains.rs deleted file mode 100644 index 99b8f036d7c..00000000000 --- a/vortex-array/src/arrays/extension/compute/list_contains.rs +++ /dev/null @@ -1,242 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::Extension; -use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::scalar_fn::ScalarFnFactoryExt; -use crate::scalar::Scalar; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsElementKernel; - -/// Answers `IN` for an extension needle through its storage. -/// -/// An extension value *is* its storage value — `ScalarValue` has no extension -/// variant — so membership over a timestamp, date or any user-defined extension -/// is membership over the integers or bytes underneath it. `Operator::Eq` -/// already compares extensions this way, so unwrapping both sides gives the same -/// answers while letting the storage type's own set probe do the work. -impl ListContainsElementKernel for Extension { - fn list_contains( - list: &ArrayRef, - element: ArrayView<'_, Self>, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - 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); - }; - - // Two extensions are only comparable when they agree on the whole - // extension dtype — a timestamp in milliseconds must not be matched - // against one in seconds. Executing as a parent kernel skips the check - // `ListContains` does for itself, so decline a mismatch here and let the - // generic path raise the error. - if !list_scalar - .element_dtype() - .eq_ignore_nullability(element.dtype()) - { - return Ok(None); - } - - let storage = element.storage_array().clone(); - let storage_dtype = storage.dtype().clone(); - let mut storage_elements = Vec::with_capacity(elements.len()); - for value in elements { - // A null element has no storage value to carry, and the dtype check - // above ignores nullability, so a nullable list can arrive against a - // non-nullable column. The generic path answers such a list, so hand - // the whole thing back to it rather than building a scalar the - // storage dtype cannot hold. - let Some(value) = value else { - return Ok(None); - }; - // SAFETY: these are the values of a validated extension list scalar, - // so each is a valid non-null value of the storage dtype. - storage_elements - .push(unsafe { Scalar::new_unchecked(storage_dtype.clone(), Some(value.clone())) }); - } - - let storage_list = Scalar::list( - storage_dtype, - storage_elements, - list_constant.scalar().dtype().nullability(), - ); - let len = storage.len(); - - ListContains - .try_new_array( - len, - EmptyOptions, - [ConstantArray::new(storage_list, len).into_array(), storage], - )? - .execute(ctx) - .map(Some) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use vortex_buffer::buffer; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::ExtensionArray; - use crate::arrays::bool::BoolArrayExt; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; - use crate::dtype::DType; - use crate::dtype::Nullability; - 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::ListContains; - - /// 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"); - let bits = result.bit_buffer_view(); - (0..bits.len()).map(|i| bits.value(i)).collect() - } - - #[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"); - let bits = result.bit_buffer_view(); - assert_eq!( - (0..bits.len()).map(|i| bits.value(i)).collect::>(), - 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/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index eacd83b0116..943733eb45d 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -4,7 +4,6 @@ mod cast; mod compare; mod filter; -mod list_contains; mod mask; pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index 2d9f8587066..a81863a3ed6 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -13,8 +13,6 @@ use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -24,10 +22,5 @@ pub(crate) fn initialize(session: &VortexSession) { CompareExecuteAdaptor(Extension), ); kernels.register_execute_parent_kernel(Cast.id(), Extension, CastExecuteAdaptor(Extension)); - kernels.register_execute_parent_kernel( - ListContains.id(), - Extension, - ListContainsElementExecuteAdaptor(Extension), - ); kernels.register_execute_parent_kernel(Dict.id(), Extension, TakeExecuteAdaptor(Extension)); } 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 9d1f14c1352..47ec5361727 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -28,11 +28,13 @@ 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; @@ -288,6 +290,16 @@ fn hash_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), }; @@ -599,6 +611,15 @@ fn list_is_not_empty( #[cfg(test)] mod tests { use std::sync::Arc; + + use vortex_buffer::buffer; + + use super::ListContains; + use crate::arrays::ExtensionArray; + use crate::arrays::scalar_fn::ScalarFnFactoryExt; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar_fn::EmptyOptions; use std::sync::LazyLock; use itertools::Itertools; @@ -1141,4 +1162,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" + ); + } } From 95d40c8bbf955e8808cf9eb7678e85775c2e76c1 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:55:34 -0700 Subject: [PATCH 5/6] style: group the lifted test imports the way the repo's rustfmt does Stable rustfmt warns that `group_imports` and `imports_granularity` are nightly-only and formats without them, so a local `cargo fmt` leaves import order untouched while CI enforces it. --- .../src/scalar_fn/fns/list_contains/mod.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) 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 47ec5361727..8b124f11f86 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -611,31 +611,26 @@ fn list_is_not_empty( #[cfg(test)] mod tests { use std::sync::Arc; - - use vortex_buffer::buffer; - - use super::ListContains; - use crate::arrays::ExtensionArray; - use crate::arrays::scalar_fn::ScalarFnFactoryExt; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::scalar_fn::EmptyOptions; use std::sync::LazyLock; use itertools::Itertools; 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; @@ -652,7 +647,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; From 6c9ffc507b81949f229d1f7ee1741958bd6a4f88 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:12:22 -0700 Subject: [PATCH 6/6] perf(array): reject a null-bearing IN list before doing work for it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A null element is not a key on any probe path, so such a list goes back to OR-ing one equality per element. Both probes established that only once they were already committed: the primitive path after canonicalizing the needles, the byte path after allocating and populating the set, and the FSST kernel after compressing every element preceding the null. The caller then ran the equality form anyway, so a nullable list paid all of that on top of the work it did before the probe existed — the one shape this change made slower. The check now happens once, on the borrowed element values, before anything is decoded, allocated or compressed. Answers are unchanged; the differential table's null-element row is what says so. --- encodings/fsst/src/compute/list_contains.rs | 6 ++++++ vortex-array/src/scalar_fn/fns/list_contains/mod.rs | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/encodings/fsst/src/compute/list_contains.rs b/encodings/fsst/src/compute/list_contains.rs index e77b7170a59..e40cb19b854 100644 --- a/encodings/fsst/src/compute/list_contains.rs +++ b/encodings/fsst/src/compute/list_contains.rs @@ -74,6 +74,12 @@ impl ListContainsElementKernel for FSST { 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()); 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 8b124f11f86..8beab939e81 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -236,7 +236,14 @@ fn constant_list_scalar_contains( // 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);