perf(fastlanes): Refine constant list membership

Signed-off-by: Will Manning <will@willmanning.io>

Author: lwwmanning

encodings/fastlanes/Cargo.toml

@@ -51,7 +51,6 @@ harness = false
 [[bench]]
 name = "bitpacking_list_contains"
 harness = false
-required-features = ["_test-harness"]
 
 [[bench]]
 name = "canonicalize_bench"

encodings/fastlanes/benches/bitpacking_list_contains.rs

@@ -1,12 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Measures the linear-search and binary-search crossover for constant-list membership.
+//! Compares compressed list membership with the canonical fallback.
 //!
-//! The strategy benchmarks isolate lookup cost across mostly-missing and mixed probes. The kernel
-//! benchmark includes scalar extraction, sorting, FastLanes decoding, and result construction.
-//! The isolated lookup benchmark excludes sorting, so it favors binary search. Use the forced
-//! full-kernel results to select the production strategy.
+//! The specialized session evaluates membership while it decodes FastLanes lanes. The fallback
+//! session decodes the complete array before the generic membership operation.
+//! Density cases stress the 4 KiB lookup-table boundary. Sparse cases exceed that boundary.
 //!
 //! Run with `cargo bench -p vortex-fastlanes --bench bitpacking_list_contains`.
 
@@ -22,100 +21,57 @@ use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
-use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::PrimitiveArray;
 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::scalar_fn::fns::list_contains::ListContainsElementKernel;
+use vortex_array::session::ArraySessionExt;
 use vortex_array::validity::Validity;
 use vortex_buffer::Alignment;
 use vortex_buffer::BufferMut;
 use vortex_fastlanes::BitPacked;
 use vortex_fastlanes::BitPackedArray;
 use vortex_fastlanes::BitPackedData;
-use vortex_fastlanes::list_contains_test_harness;
-use vortex_fastlanes::list_contains_test_harness::MembershipSearch;
-
-const LEN: usize = 64 * 1_024;
-const MEMBER_COUNTS: &[usize] = &[1, 2, 3, 4, 6, 8, 9, 12, 16, 24, 32, 64];
-const STRATEGY_MEMBER_COUNTS: &[usize] = &[3, 4, 6, 8, 9, 12, 16, 24, 32, 64];
+use vortex_session::VortexSession;
+
+const DENSE_CASES: &[(usize, usize)] = &[
+    (64, 1),
+    (64, 4),
+    (64, 8),
+    (64, 32),
+    (64, 64),
+    (1_024, 1),
+    (1_024, 4),
+    (1_024, 8),
+    (1_024, 32),
+    (1_024, 64),
+    (65_536, 1),
+    (65_536, 4),
+    (65_536, 8),
+    (65_536, 32),
+    (65_536, 64),
+];
+const SPARSE_CASES: &[(usize, usize)] = &[(1_024, 8), (1_024, 64), (65_536, 8), (65_536, 64)];
+const DENSITY_CASES: &[(usize, usize, u32)] = &[
+    (64, 5, 1_000),
+    (64, 8, 512),
+    (64, 64, 64),
+    (65_536, 5, 1_000),
+    (65_536, 8, 512),
+    (65_536, 64, 64),
+];
 
 fn main() {
     divan::main();
 }
 
-fn members(count: usize) -> Vec<u32> {
-    (0..count).map(|index| index as u32 * 2).collect()
-}
-
-fn mostly_missing_values() -> Vec<u32> {
-    (0..LEN)
-        .map(|index| ((index as u32 * 17) % 4_096) | 1)
-        .collect()
-}
-
-fn mixed_values(members: &[u32]) -> Vec<u32> {
-    (0..LEN)
-        .map(|index| {
-            if index.is_multiple_of(2) {
-                members[(index / 2) % members.len()]
-            } else {
-                ((index as u32 * 17) % 4_096) | 1
-            }
-        })
-        .collect()
-}
-
-fn count_linear(values: &[u32], members: &[u32]) -> usize {
-    values
-        .iter()
-        .filter(|value| members.contains(black_box(value)))
-        .count()
-}
-
-fn count_binary(values: &[u32], members: &[u32]) -> usize {
-    values
-        .iter()
-        .filter(|value| members.binary_search(black_box(value)).is_ok())
-        .count()
-}
-
-#[divan::bench(args = MEMBER_COUNTS)]
-fn linear_mostly_missing(bencher: Bencher, member_count: usize) {
-    let values = mostly_missing_values();
-    let members = members(member_count);
-    bencher
-        .counter(ItemsCount::new(LEN))
-        .bench_local(|| black_box(count_linear(&values, &members)));
-}
-
-#[divan::bench(args = MEMBER_COUNTS)]
-fn binary_mostly_missing(bencher: Bencher, member_count: usize) {
-    let values = mostly_missing_values();
-    let members = members(member_count);
-    bencher
-        .counter(ItemsCount::new(LEN))
-        .bench_local(|| black_box(count_binary(&values, &members)));
-}
-
-#[divan::bench(args = MEMBER_COUNTS)]
-fn linear_mixed(bencher: Bencher, member_count: usize) {
-    let members = members(member_count);
-    let values = mixed_values(&members);
-    bencher
-        .counter(ItemsCount::new(LEN))
-        .bench_local(|| black_box(count_linear(&values, &members)));
-}
-
-#[divan::bench(args = MEMBER_COUNTS)]
-fn binary_mixed(bencher: Bencher, member_count: usize) {
-    let members = members(member_count);
-    let values = mixed_values(&members);
-    bencher
-        .counter(ItemsCount::new(LEN))
-        .bench_local(|| black_box(count_binary(&values, &members)));
+fn members(count: usize, stride: u32) -> Vec<u32> {
+    (0..count).map(|index| index as u32 * stride).collect()
 }
 
 fn page_aligned(array: BitPackedArray) -> BitPackedArray {
@@ -133,9 +89,21 @@ fn page_aligned(array: BitPackedArray) -> BitPackedArray {
     .unwrap()
 }
 
-fn kernel_inputs(member_count: usize) -> (BitPackedArray, ArrayRef) {
-    let mut ctx = array_session().create_execution_ctx();
-    let values: BufferMut<u32> = (0..LEN).map(|index| (index as u32 * 17) % 1_024).collect();
+fn benchmark_input(
+    len: usize,
+    member_count: usize,
+    member_stride: u32,
+    specialized: bool,
+) -> (ArrayRef, VortexSession) {
+    let session = array_session();
+    if specialized {
+        vortex_fastlanes::initialize(&session);
+    } else {
+        session.arrays().register(BitPacked);
+    }
+
+    let mut ctx = session.create_execution_ctx();
+    let values: BufferMut<u32> = (0..len).map(|index| (index as u32 * 17) % 1_024).collect();
     let packed = page_aligned(
         BitPackedData::encode(
             &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(),
@@ -144,71 +112,62 @@ fn kernel_inputs(member_count: usize) -> (BitPackedArray, ArrayRef) {
         )
         .unwrap(),
     );
-    let member_scalars = members(member_count)
+    let member_scalars = members(member_count, member_stride)
         .into_iter()
         .map(|value| Scalar::primitive(value, Nullability::NonNullable))
         .collect();
-    let list = ConstantArray::new(
-        Scalar::list(
-            Arc::new(DType::Primitive(PType::U32, Nullability::NonNullable)),
-            member_scalars,
-            Nullability::NonNullable,
-        ),
-        LEN,
-    )
-    .into_array();
-    (packed, list)
+    let list = Scalar::list(
+        Arc::new(DType::Primitive(PType::U32, Nullability::NonNullable)),
+        member_scalars,
+        Nullability::NonNullable,
+    );
+    let contains = packed
+        .into_array()
+        .apply(&list_contains(lit(list), root()))
+        .unwrap();
+    (contains, session)
 }
 
-#[divan::bench(args = MEMBER_COUNTS)]
-fn bitpacked_kernel(bencher: Bencher, member_count: usize) {
-    let (packed, list) = kernel_inputs(member_count);
-    let mut ctx = array_session().create_execution_ctx();
-    bencher.counter(ItemsCount::new(LEN)).bench_local(|| {
-        black_box(
-            <BitPacked as ListContainsElementKernel>::list_contains(
-                &list,
-                packed.as_view(),
-                &mut ctx,
-            )
-            .unwrap()
-            .unwrap(),
-        )
-    });
+fn bench_contains(
+    bencher: Bencher,
+    len: usize,
+    member_count: usize,
+    member_stride: u32,
+    specialized: bool,
+) {
+    let (contains, session) = benchmark_input(len, member_count, member_stride, specialized);
+    let mut ctx = session.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(len))
+        .bench_local(|| black_box(contains.clone().execute::<BoolArray>(&mut ctx).unwrap()));
 }
 
-#[divan::bench(args = STRATEGY_MEMBER_COUNTS)]
-fn bitpacked_linear(bencher: Bencher, member_count: usize) {
-    let (packed, list) = kernel_inputs(member_count);
-    let mut ctx = array_session().create_execution_ctx();
-    bencher.counter(ItemsCount::new(LEN)).bench_local(|| {
-        black_box(
-            list_contains_test_harness::list_contains(
-                &list,
-                packed.as_view(),
-                MembershipSearch::Linear,
-                &mut ctx,
-            )
-            .unwrap()
-            .unwrap(),
-        )
-    });
+#[divan::bench(args = DENSE_CASES)]
+fn compressed_dense(bencher: Bencher, (len, member_count): (usize, usize)) {
+    bench_contains(bencher, len, member_count, 2, true);
 }
 
-#[divan::bench(args = STRATEGY_MEMBER_COUNTS)]
-fn bitpacked_binary(bencher: Bencher, member_count: usize) {
-    let (packed, list) = kernel_inputs(member_count);
-    let mut ctx = array_session().create_execution_ctx();
-    bencher.counter(ItemsCount::new(LEN)).bench_local(|| {
-        black_box(
-            list_contains_test_harness::list_contains(
-                &list,
-                packed.as_view(),
-                MembershipSearch::Binary,
-                &mut ctx,
-            )
-            .unwrap()
-            .unwrap(),
-        )
-    });
+#[divan::bench(args = DENSE_CASES)]
+fn canonical_dense(bencher: Bencher, (len, member_count): (usize, usize)) {
+    bench_contains(bencher, len, member_count, 2, false);
+}
+
+#[divan::bench(args = SPARSE_CASES)]
+fn compressed_sparse(bencher: Bencher, (len, member_count): (usize, usize)) {
+    bench_contains(bencher, len, member_count, 10_000, true);
+}
+
+#[divan::bench(args = SPARSE_CASES)]
+fn canonical_sparse(bencher: Bencher, (len, member_count): (usize, usize)) {
+    bench_contains(bencher, len, member_count, 10_000, false);
+}
+
+#[divan::bench(args = DENSITY_CASES)]
+fn compressed_density(bencher: Bencher, (len, member_count, member_stride): (usize, usize, u32)) {
+    bench_contains(bencher, len, member_count, member_stride, true);
+}
+
+#[divan::bench(args = DENSITY_CASES)]
+fn canonical_density(bencher: Bencher, (len, member_count, member_stride): (usize, usize, u32)) {
+    bench_contains(bencher, len, member_count, member_stride, false);
 }