From 3566b4d40234b7a5ea9a966c23930db0b2803bc1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 30 Oct 2026 12:19:02 +0000 Subject: [PATCH] Look up prepared integer sets through a bitmap and in batches A set of a few dozen keys spread over a range of 100 000, such as the items DuckDB's join filters push into a fact table scan, was looked up by binary search, which mispredicts about half its comparisons. Look up integer sets that span up to 2^20 values in a bitmap of up to 128 KiB. The bitmap stays in the L2 cache and the lookup does not branch. TPC-DS SF10 q32, q16, q90 and q92 on DuckDB get 6-13% faster. A prepared integer set decoded all of its needles into one array before looking them up. For the 64Ki-row chunks of a scan, this writes a fresh buffer that has left the cache by the time it is read. Decode needles of other encodings in slices of 8192 rows and look each slice up while it is still in the cache. In TPC-DS q25, where the IN lists DuckDB pushes from joins took 15% of the query's CPU, this saves 7% of it. Reduction replaces the literal list of list.contains with a PreparedSet, which the falsifier did not recognise, so no IN list pruned any zone. Recognise the PreparedSet so that IN lists prune zones again. Add list_contains_set benchmarks for i64 sets of 32 to 256 values spread over spans of 2^24 and 2^40, too sparse for a bitmap. A roaring bitmap of the offsets was 1.6-2.1x slower than the binary search on all of them, because a lookup searches the containers and then an array container. Signed-off-by: Robert Kruszewski --- vortex-array/benches/list_contains_set.rs | 27 +++++++- vortex-array/src/array/haystack/mod.rs | 66 ++++++++++++++----- .../fns/list_contains/prepared/tests.rs | 28 ++++++++ vortex-expr/src/stats/rewrite/builtins.rs | 41 ++++++++++-- 4 files changed, 137 insertions(+), 25 deletions(-) diff --git a/vortex-array/benches/list_contains_set.rs b/vortex-array/benches/list_contains_set.rs index 56b486dd404..159f1da36e8 100644 --- a/vortex-array/benches/list_contains_set.rs +++ b/vortex-array/benches/list_contains_set.rs @@ -59,14 +59,20 @@ const NESTED_SET_LENS: &[usize] = &[32]; /// A random set of `len` values, and needles of which about half are members. fn random_i64(len: usize) -> (Vec, Vec) { + random_i64_in_span(len, 40) +} + +/// A random set of `len` values below `2^span_bits`, and needles of which about half are members. +fn random_i64_in_span(len: usize, span_bits: u32) -> (Vec, Vec) { let mut rng = StdRng::seed_from_u64(0); - let set: Vec = (0..len).map(|_| rng.random_range(0..1 << 40)).collect(); + let span = 1i64 << span_bits; + let set: Vec = (0..len).map(|_| rng.random_range(0..span)).collect(); let needles = (0..ROWS) .map(|_| { if rng.random_bool(0.5) { set[rng.random_range(0..len)] } else { - rng.random_range(0..1 << 40) + rng.random_range(0..span) } }) .collect(); @@ -105,6 +111,23 @@ fn i64_random(bencher: Bencher, set_len: usize) { bench_in_set(bencher, i64_set(set), needles.into_array()); } +/// Lengths of sets too sparse for a bitmap over their whole span. +const SPARSE_SET_LENS: &[usize] = &[32, 256]; + +#[divan::bench(args = SPARSE_SET_LENS)] +fn i64_sparse_span_2_24(bencher: Bencher, set_len: usize) { + let (set, needles) = random_i64_in_span(set_len, 24); + let needles = PrimitiveArray::from_iter(needles); + bench_in_set(bencher, i64_set(set), needles.into_array()); +} + +#[divan::bench(args = SPARSE_SET_LENS)] +fn i64_sparse_span_2_40(bencher: Bencher, set_len: usize) { + let (set, needles) = random_i64_in_span(set_len, 40); + let needles = PrimitiveArray::from_iter(needles); + bench_in_set(bencher, i64_set(set), needles.into_array()); +} + #[divan::bench(args = SET_LENS)] fn i64_random_chunked(bencher: Bencher, set_len: usize) { let (set, needles) = random_i64(set_len); diff --git a/vortex-array/src/array/haystack/mod.rs b/vortex-array/src/array/haystack/mod.rs index 5e9065b7113..92623f68b0d 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-array/src/array/haystack/mod.rs @@ -25,6 +25,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::RecursiveCanonical; use crate::arrays::DecimalArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::arrays::decimal::converted_buffer; @@ -44,7 +45,12 @@ use crate::validity::Validity; /// bitmap over the span, bounding the bitmap to a few words per element. const BITMAP_BITS_PER_ELEMENT: u128 = 64; /// A span this narrow is looked up through a bitmap whatever the size of the set. -const BITMAP_MIN_BITS: u128 = 1 << 12; +/// +/// A lookup in a bitmap of up to 128 KiB stays in the L2 cache and does not branch, where a +/// binary search through a sparse set mispredicts about half its comparisons. +const BITMAP_MIN_BITS: u128 = 1 << 20; +/// Needles beyond this many are decoded and looked up a batch of this many at a time. +const NEEDLE_BATCH_LEN: usize = 8192; /// Looks up needles among the elements of an array, for many batches of needles. /// @@ -154,17 +160,25 @@ impl IntegerSet { /// One bit per needle, set when the needle is an element. fn contains(&self, needles: &[T], allocator: &BufferAllocatorRef) -> BitBuffer { match self { - Self::Bitmap { min, bitmap } => collect_bits( - needles, - // A needle below the smallest element wraps past the bitmap, so one comparison - // checks both bounds. - |needle| { - needle - .offset_from(*min) - .is_some_and(|offset| offset < bitmap.len() && bitmap.value(offset)) - }, - allocator, - ), + Self::Bitmap { min, bitmap } => { + let bytes = bitmap.inner().as_slice(); + let (bit_offset, len) = (bitmap.offset(), bitmap.len()); + collect_bits( + needles, + // A needle below the smallest element wraps past the bitmap, so one + // comparison checks both bounds. A needle outside the bitmap reads its + // first bit and is masked off, so the lookup does not branch. + |needle| { + let offset = needle.offset_from(*min).unwrap_or(usize::MAX); + let in_span = offset < len; + let bit = if in_span { offset } else { 0 } + bit_offset; + // SAFETY: `bit` is below the bitmap's offset plus its length. + let byte = unsafe { *bytes.get_unchecked(bit / 8) }; + in_span & ((byte >> (bit % 8)) & 1 != 0) + }, + allocator, + ) + } Self::Sorted(sorted) => collect_bits( needles, |needle| sorted.binary_search(&needle).is_ok(), @@ -187,12 +201,28 @@ impl HaystackState for PrimitiveSet { needles: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult<(BitBuffer, Validity)> { - let needles = needles - .clone() - .execute::(ctx)? - .reinterpret_cast(T::PTYPE); - let bits = self.0.contains(needles.as_slice::(), ctx.allocator()); - Ok((bits, needles.validity()?)) + if needles.len() <= NEEDLE_BATCH_LEN || needles.is::() { + let needles = needles + .clone() + .execute::(ctx)? + .reinterpret_cast(T::PTYPE); + let bits = self.0.contains(needles.as_slice::(), ctx.allocator()); + return Ok((bits, needles.validity()?)); + } + + // Decoding a long compressed array in batches keeps each batch of needles in the cache + // between decoding and looking it up, instead of writing out the whole array first. + let len = needles.len(); + let mut bits = BitBufferMut::with_capacity_in(len, ctx.allocator().clone()); + for start in (0..len).step_by(NEEDLE_BATCH_LEN) { + let end = (start + NEEDLE_BATCH_LEN).min(len); + let batch = needles + .slice(start..end)? + .execute::(ctx)? + .reinterpret_cast(T::PTYPE); + bits.append_buffer(&self.0.contains(batch.as_slice::(), ctx.allocator())); + } + Ok((bits.freeze(), needles.validity()?)) } } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs index e67c1717998..bdc78f6e0b3 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs @@ -299,6 +299,34 @@ fn test_primitive_integers( Ok(()) } +/// Needles longer than one lookup batch, in an encoding other than primitive, are looked up a +/// batch at a time. +#[test] +fn test_compressed_needles_in_batches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let set = PreparedSetData::try_new(list_scalar( + &buffer![3i64, 9000, 20_001].into_array(), + &mut ctx, + )?)?; + let values = |range: std::ops::Range| { + PrimitiveArray::from_option_iter(range.map(|v| (v % 7 != 0).then_some(v))).into_array() + }; + let needles = ChunkedArray::try_new( + vec![values(0..12_000), values(12_000..25_000)], + DType::Primitive(PType::I64, Nullability::Nullable), + )? + .into_array(); + let expected = BoolArray::from_iter( + (0..25_000i64).map(|v| (v % 7 != 0).then_some(matches!(v, 3 | 9000 | 20_001))), + ); + assert_arrays_eq!( + set.contains(&needles, &ListContainsOptions::default(), &mut ctx)?, + expected, + &mut ctx + ); + Ok(()) +} + /// The set `{2, null}` of nullable `i32`. fn set_with_null() -> Scalar { let element = DType::Primitive(PType::I32, Nullability::Nullable); diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index 6db4e3ecf66..d902360c684 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -42,6 +42,8 @@ use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; use vortex_array::scalar_fn::fns::is_null::IsNull; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::PreparedSet; +use vortex_array::scalar_fn::fns::list_contains::PreparedSetData; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; @@ -550,12 +552,11 @@ 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 + // Reduction replaces a literal list with a set prepared from it. + let elements = list .as_opt::() - .and_then(|literal| literal.as_list_opt()) + .or_else(|| list.as_opt::().map(PreparedSetData::list)) + .and_then(|list| list.as_list_opt()) .and_then(|list| list.elements()); let Some(elements) = elements else { return Ok(None); @@ -883,6 +884,10 @@ mod tests { use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; + use vortex_array::scalar_fn::fns::list_contains::ListContains; + use vortex_array::scalar_fn::fns::list_contains::ListContainsOptions; + use vortex_array::scalar_fn::fns::list_contains::PreparedSet; + use vortex_array::scalar_fn::fns::list_contains::PreparedSetData; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; @@ -1236,6 +1241,32 @@ mod tests { Ok(()) } + #[test] + fn rewrites_prepared_set_list_contains_falsifier() -> VortexResult<()> { + let list = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![1.into(), 2.into()], + Nullability::NonNullable, + ); + let set = PreparedSet.new_expr(PreparedSetData::try_new(list)?, Vec::::new()); + let expr = ListContains.new_expr(ListContainsOptions::default(), [set, col("a")]); + + assert_rewrite_eq!( + falsify(&expr)?, + Some(and( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(1i32)), + ), + or( + lt(stat(col("a"), Stat::Max), lit(2i32)), + gt(stat(col("a"), Stat::Min), lit(2i32)), + ), + )) + ); + Ok(()) + } + #[test] fn rewrites_like_falsifier() -> VortexResult<()> { let expr = like(col("s"), lit("prefix%"));