perf(fastlanes): Evaluate constant list membership

encodings/fastlanes/Cargo.toml

@@ -48,6 +48,10 @@ _test-harness = ["dep:rand"]
 name = "bitpacking_take"
 harness = false
 
+[[bench]]
+name = "bitpacking_list_contains"
+harness = false
+
 [[bench]]
 name = "canonicalize_bench"
 harness = false

encodings/fastlanes/benches/bitpacking_list_contains.rs

@@ -0,0 +1,173 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Compares compressed list membership with the canonical fallback.
+//!
+//! 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`.
+
+#![expect(clippy::cast_possible_truncation)]
+#![expect(clippy::unwrap_used)]
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use divan::Bencher;
+use divan::counter::ItemsCount;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+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::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_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, stride: u32) -> Vec<u32> {
+    (0..count).map(|index| index as u32 * stride).collect()
+}
+
+fn page_aligned(array: BitPackedArray) -> BitPackedArray {
+    let ptype = array.dtype().as_ptype();
+    let parts = BitPacked::into_parts(array);
+    BitPacked::try_new(
+        parts.packed.ensure_aligned(Alignment::new(4_096)).unwrap(),
+        ptype,
+        parts.validity,
+        parts.patches,
+        parts.bit_width,
+        parts.len,
+        parts.offset,
+    )
+    .unwrap()
+}
+
+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(),
+            10,
+            &mut ctx,
+        )
+        .unwrap(),
+    );
+    let member_scalars = members(member_count, member_stride)
+        .into_iter()
+        .map(|value| Scalar::primitive(value, Nullability::NonNullable))
+        .collect();
+    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)
+}
+
+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 = DENSE_CASES)]
+fn compressed_dense(bencher: Bencher, (len, member_count): (usize, usize)) {
+    bench_contains(bencher, len, member_count, 2, true);
+}
+
+#[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);
+}

encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::NativePType;
+use vortex_array::match_each_integer_ptype;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar_fn::fns::list_contains::IntegerMembership;
+use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel;
+use vortex_buffer::BitBuffer;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use super::compare_fused::stream_compare_fused;
+use crate::BitPacked;
+
+impl ListContainsElementKernel for BitPacked {
+    fn list_contains(
+        list: &ArrayRef,
+        element: ArrayView<'_, Self>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        list_contains_compressed(list, element, ctx)
+    }
+}
+
+fn list_contains_compressed(
+    list: &ArrayRef,
+    element: ArrayView<'_, BitPacked>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Option<ArrayRef>> {
+    let Some(list_scalar) = list.as_constant() else {
+        return Ok(None);
+    };
+    let DType::List(member_dtype, _) = list.dtype() else {
+        return Ok(None);
+    };
+    if !member_dtype.eq_ignore_nullability(element.dtype()) {
+        return Ok(None);
+    }
+
+    let nullability = list.dtype().nullability() | element.dtype().nullability();
+    let Some(elements) = list_scalar.as_list().elements() else {
+        return Ok(Some(
+            ConstantArray::new(Scalar::null(DType::Bool(nullability)), element.len()).into_array(),
+        ));
+    };
+
+    let result = match_each_integer_ptype!(element.dtype().as_ptype(), |T| {
+        let members = elements
+            .iter()
+            .map(|value| {
+                value
+                    .as_primitive_opt()
+                    .ok_or_else(|| vortex_err!("List member is not a primitive scalar"))?
+                    .try_typed_value::<T>()
+            })
+            .collect::<VortexResult<Vec<Option<T>>>>()?
+            .into_iter()
+            .flatten()
+            .collect::<Vec<_>>();
+
+        if members.is_empty() && !elements.is_empty() {
+            let validity = element.validity()?.union_nullability(nullability);
+            return Ok(Some(
+                BoolArray::new(BitBuffer::new_unset(element.len()), validity).into_array(),
+            ));
+        }
+        let membership = IntegerMembership::new(members);
+
+        match membership.members() {
+            [] => ConstantArray::new(Scalar::bool(false, nullability), element.len()).into_array(),
+            [member] => {
+                let member = *member;
+                stream_compare_fused::<T, _>(element, member, nullability, NativePType::is_eq, ctx)?
+            }
+            [first, second] => {
+                let (first, second) = (*first, *second);
+                stream_compare_fused::<T, _>(
+                    element,
+                    first,
+                    nullability,
+                    move |value, _| value.is_eq(first) | value.is_eq(second),
+                    ctx,
+                )?
+            }
+            [first, second, third] => {
+                let (first, second, third) = (*first, *second, *third);
+                stream_compare_fused::<T, _>(
+                    element,
+                    first,
+                    nullability,
+                    move |value, _| value.is_eq(first) | value.is_eq(second) | value.is_eq(third),
+                    ctx,
+                )?
+            }
+            [first, second, third, fourth] => {
+                let (first, second, third, fourth) = (*first, *second, *third, *fourth);
+                stream_compare_fused::<T, _>(
+                    element,
+                    first,
+                    nullability,
+                    move |value, _| {
+                        value.is_eq(first)
+                            | value.is_eq(second)
+                            | value.is_eq(third)
+                            | value.is_eq(fourth)
+                    },
+                    ctx,
+                )?
+            }
+            _ => {
+                if membership.uses_dense_table() {
+                    stream_compare_fused::<T, _>(
+                        element,
+                        membership.members()[0],
+                        nullability,
+                        |value, _| membership.contains(value),
+                        ctx,
+                    )?
+                } else {
+                    let primitive = element
+                        .into_owned()
+                        .into_array()
+                        .execute::<PrimitiveArray>(ctx)?;
+                    membership.evaluate_primitive(primitive.as_view(), nullability)?
+                }
+            }
+        }
+    });
+    Ok(Some(result))
+}
+
+#[cfg(test)]
+mod tests;