diff --git a/vortex-array/benches/list_contains_set.rs b/vortex-array/benches/list_contains_set.rs index 56b486dd404..b4bc3e89454 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, 4_096]; + +#[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 f1e04f9c686..87c2851364b 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; @@ -45,7 +46,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. /// @@ -155,17 +161,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(), @@ -188,12 +202,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 fe3c6b46da6..cbf503e3c8d 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 d7c65947195..0600ba8bac1 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -41,6 +41,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; @@ -546,12 +548,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); @@ -879,6 +880,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; @@ -1232,6 +1237,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%"));