Optimise in_list lookups for sparse integer sets

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<i64>, Vec<i64>) {
+    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<i64>, Vec<i64>) {
     let mut rng = StdRng::seed_from_u64(0);
-    let set: Vec<i64> = (0..len).map(|_| rng.random_range(0..1 << 40)).collect();
+    let span = 1i64 << span_bits;
+    let set: Vec<i64> = (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);

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<T: SetInteger> IntegerSet<T> {
     /// 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<T: NativePType + SetInteger> HaystackState for PrimitiveSet<T> {
         needles: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<(BitBuffer, Validity)> {
-        let needles = needles
-            .clone()
-            .execute::<PrimitiveArray>(ctx)?
-            .reinterpret_cast(T::PTYPE);
-        let bits = self.0.contains(needles.as_slice::<T>(), ctx.allocator());
-        Ok((bits, needles.validity()?))
+        if needles.len() <= NEEDLE_BATCH_LEN || needles.is::<Primitive>() {
+            let needles = needles
+                .clone()
+                .execute::<PrimitiveArray>(ctx)?
+                .reinterpret_cast(T::PTYPE);
+            let bits = self.0.contains(needles.as_slice::<T>(), 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::<PrimitiveArray>(ctx)?
+                .reinterpret_cast(T::PTYPE);
+            bits.append_buffer(&self.0.contains(batch.as_slice::<T>(), ctx.allocator()));
+        }
+        Ok((bits.freeze(), needles.validity()?))
     }
 }
 

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<i64>| {
+        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);

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<P: NonNanProof>(
     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::<Literal>()
-        .and_then(|literal| literal.as_list_opt())
+        .or_else(|| list.as_opt::<PreparedSet>().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::<Expression>::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%"));