perf(in_list): add BitmapFilter for u8/u16 types

geoffreyclaude · geoffreyclaude · commit a997d7f4a3dd · 2025-12-19T10:56:45.000+01:00
Adds O(1) bitmap-based set membership for 1-byte and 2-byte integer types.
Also introduces result.rs with shared logic for building BooleanArray
results with correct SQL null propagation.

BitmapFilter stores a bitset where bit N is set if value N is in the set:
- U8Config: 256 bits = 32 bytes (fits in cache line)
- U16Config: 65536 bits = 8 KB (fits in L1 cache)

Lookup is a single bit test: `bits[value / 64] &amp; (1 &lt;&lt; (value % 64))`.
This outperforms both hash lookup and branchless comparison at all list
sizes for these small integer types.
diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs
@@ -18,7 +18,10 @@
 //! Implementation of `InList` expressions: [`InListExpr`]
 
 mod array_filter;
+mod primitive;
+mod result;
 mod strategy;
+mod transform;
 
 use std::any::Any;
 use std::fmt::Debug;
diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive.rs b/datafusion/physical-expr/src/expressions/in_list/primitive.rs
@@ -0,0 +1,173 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Optimized primitive type filters for InList expressions
+//!
+//! This module provides high-performance membership testing for Arrow primitive types.
+
+use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
+use arrow::datatypes::ArrowPrimitiveType;
+use datafusion_common::{Result, exec_datafusion_err};
+
+use super::array_filter::StaticFilter;
+use super::result::{build_in_list_result, handle_dictionary};
+
+// =============================================================================
+// BITMAP FILTERS (O(1) lookup for u8/u16 via bit test)
+// =============================================================================
+
+/// Trait for bitmap storage (stack-allocated for u8, heap-allocated for u16).
+pub(crate) trait BitmapStorage: Send + Sync {
+    fn new_zeroed() -> Self;
+    fn set_bit(&mut self, index: usize);
+    fn get_bit(&self, index: usize) -> bool;
+}
+
+impl BitmapStorage for [u64; 4] {
+    #[inline]
+    fn new_zeroed() -> Self {
+        [0u64; 4]
+    }
+    #[inline]
+    fn set_bit(&mut self, index: usize) {
+        self[index / 64] |= 1u64 << (index % 64);
+    }
+    #[inline(always)]
+    fn get_bit(&self, index: usize) -> bool {
+        (self[index / 64] >> (index % 64)) & 1 != 0
+    }
+}
+
+impl BitmapStorage for Box<[u64; 1024]> {
+    #[inline]
+    fn new_zeroed() -> Self {
+        Box::new([0u64; 1024])
+    }
+    #[inline]
+    fn set_bit(&mut self, index: usize) {
+        self[index / 64] |= 1u64 << (index % 64);
+    }
+    #[inline(always)]
+    fn get_bit(&self, index: usize) -> bool {
+        (self[index / 64] >> (index % 64)) & 1 != 0
+    }
+}
+
+/// Configuration trait for bitmap filters.
+pub(crate) trait BitmapFilterConfig: Send + Sync + 'static {
+    type Native: arrow::datatypes::ArrowNativeType + Copy + Send + Sync;
+    type ArrowType: ArrowPrimitiveType<Native = Self::Native>;
+    type Storage: BitmapStorage;
+
+    fn to_index(v: Self::Native) -> usize;
+}
+
+/// Config for u8 bitmap (256 bits = 32 bytes, fits in cache line).
+pub(crate) enum U8Config {}
+impl BitmapFilterConfig for U8Config {
+    type Native = u8;
+    type ArrowType = arrow::datatypes::UInt8Type;
+    type Storage = [u64; 4];
+
+    #[inline(always)]
+    fn to_index(v: u8) -> usize {
+        v as usize
+    }
+}
+
+/// Config for u16 bitmap (65536 bits = 8 KB, fits in L1 cache).
+pub(crate) enum U16Config {}
+impl BitmapFilterConfig for U16Config {
+    type Native = u16;
+    type ArrowType = arrow::datatypes::UInt16Type;
+    type Storage = Box<[u64; 1024]>;
+
+    #[inline(always)]
+    fn to_index(v: u16) -> usize {
+        v as usize
+    }
+}
+
+/// Bitmap filter for O(1) set membership via single bit test.
+///
+/// For small integer types (u8/u16), bitmap lookup outperforms both branchless
+/// and hashed approaches at all list sizes.
+pub(crate) struct BitmapFilter<C: BitmapFilterConfig> {
+    null_count: usize,
+    bits: C::Storage,
+}
+
+impl<C: BitmapFilterConfig> BitmapFilter<C> {
+    pub(crate) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+        let prim_array =
+            in_array.as_primitive_opt::<C::ArrowType>().ok_or_else(|| {
+                exec_datafusion_err!("BitmapFilter: expected primitive array")
+            })?;
+        let mut bits = C::Storage::new_zeroed();
+        for v in prim_array.iter().flatten() {
+            bits.set_bit(C::to_index(v));
+        }
+        Ok(Self {
+            null_count: prim_array.null_count(),
+            bits,
+        })
+    }
+
+    #[inline(always)]
+    fn check(&self, needle: C::Native) -> bool {
+        self.bits.get_bit(C::to_index(needle))
+    }
+
+    /// Check membership using a raw values slice (zero-copy path for type reinterpretation).
+    #[inline]
+    pub(crate) fn contains_slice(
+        &self,
+        values: &[C::Native],
+        nulls: Option<&arrow::buffer::NullBuffer>,
+        negated: bool,
+    ) -> BooleanArray {
+        build_in_list_result(
+            values.len(),
+            nulls,
+            self.null_count > 0,
+            negated,
+            |i| self.check(unsafe { *values.get_unchecked(i) }),
+        )
+    }
+}
+
+impl<C: BitmapFilterConfig> StaticFilter for BitmapFilter<C> {
+    fn null_count(&self) -> usize {
+        self.null_count
+    }
+
+    fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
+        handle_dictionary!(self, v, negated);
+        let v = v.as_primitive_opt::<C::ArrowType>().ok_or_else(|| {
+            exec_datafusion_err!("BitmapFilter: expected primitive array")
+        })?;
+        let input_values = v.values();
+        Ok(build_in_list_result(
+            v.len(),
+            v.nulls(),
+            self.null_count > 0,
+            negated,
+            #[inline(always)]
+            |i| self.check(unsafe { *input_values.get_unchecked(i) }),
+        ))
+    }
+}
diff --git a/datafusion/physical-expr/src/expressions/in_list/result.rs b/datafusion/physical-expr/src/expressions/in_list/result.rs
@@ -0,0 +1,119 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Result building helpers for InList operations
+//!
+//! This module provides unified logic for building BooleanArray results
+//! from IN list membership tests, handling null propagation correctly
+//! according to SQL three-valued logic.
+
+use arrow::array::BooleanArray;
+use arrow::buffer::{BooleanBuffer, NullBuffer};
+
+// =============================================================================
+// RESULT BUILDER FOR IN LIST OPERATIONS
+// =============================================================================
+//
+// Truth table for (needle_nulls, haystack_has_nulls, negated):
+// (Some, true,  false) → values: valid & contains,           nulls: valid & contains
+// (None, true,  false) → values: contains,                   nulls: contains
+// (Some, true,  true)  → values: valid ^ (valid & contains), nulls: valid & contains
+// (None, true,  true)  → values: !contains,                  nulls: contains
+// (Some, false, false) → values: valid & contains,           nulls: valid
+// (Some, false, true)  → values: valid & !contains,          nulls: valid
+// (None, false, false) → values: contains,                   nulls: none
+// (None, false, true)  → values: !contains,                  nulls: none
+
+/// Builds a BooleanArray result for IN list operations.
+///
+/// This function handles the complex null propagation logic for SQL IN lists:
+/// - If the needle value is null, the result is null
+/// - If the needle is not in the set AND the haystack has nulls, the result is null
+/// - Otherwise, the result is true/false based on membership and negation
+#[inline]
+pub(crate) fn build_in_list_result<C>(
+    len: usize,
+    needle_nulls: Option<&NullBuffer>,
+    haystack_has_nulls: bool,
+    negated: bool,
+    contains: C,
+) -> BooleanArray
+where
+    C: FnMut(usize) -> bool,
+{
+    // Pass closure by value to avoid indirection on each call
+    let contains_buf = BooleanBuffer::collect_bool(len, contains);
+    build_result_from_contains(needle_nulls, haystack_has_nulls, negated, contains_buf)
+}
+
+/// Builds a BooleanArray result from a pre-computed contains buffer.
+#[inline]
+pub(crate) fn build_result_from_contains(
+    needle_nulls: Option<&NullBuffer>,
+    haystack_has_nulls: bool,
+    negated: bool,
+    contains_buf: BooleanBuffer,
+) -> BooleanArray {
+    match (needle_nulls, haystack_has_nulls, negated) {
+        (Some(v), true, false) => {
+            let buf = v.inner() & &contains_buf;
+            BooleanArray::new(buf.clone(), Some(NullBuffer::new(buf)))
+        }
+        (None, true, false) => {
+            BooleanArray::new(contains_buf.clone(), Some(NullBuffer::new(contains_buf)))
+        }
+        (Some(v), true, true) => {
+            let nulls = v.inner() & &contains_buf;
+            BooleanArray::new(v.inner() ^ &nulls, Some(NullBuffer::new(nulls)))
+        }
+        (None, true, true) => {
+            BooleanArray::new(!&contains_buf, Some(NullBuffer::new(contains_buf)))
+        }
+        (Some(v), false, false) => {
+            BooleanArray::new(v.inner() & &contains_buf, Some(v.clone()))
+        }
+        (Some(v), false, true) => {
+            BooleanArray::new(v.inner() & &(!&contains_buf), Some(v.clone()))
+        }
+        (None, false, false) => BooleanArray::new(contains_buf, None),
+        (None, false, true) => BooleanArray::new(!&contains_buf, None),
+    }
+}
+
+// =============================================================================
+// DICTIONARY ARRAY HANDLING
+// =============================================================================
+
+/// Macro to handle dictionary arrays in StaticFilter::contains implementations.
+///
+/// This macro extracts the dictionary values, performs the contains check on
+/// the values array, and then uses `take` to map the results back to the
+/// dictionary keys.
+macro_rules! handle_dictionary {
+    ($self:ident, $v:ident, $negated:ident) => {
+        arrow::array::downcast_dictionary_array! {
+            $v => {
+                let values_contains = $self.contains($v.values().as_ref(), $negated)?;
+                let result = arrow::compute::take(&values_contains, $v.keys(), None)?;
+                return Ok(arrow::array::downcast_array(result.as_ref()))
+            }
+            _ => {}
+        }
+    };
+}
+
+pub(crate) use handle_dictionary;
diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs
@@ -27,19 +27,22 @@ use arrow::datatypes::*;
 use datafusion_common::{HashSet, Result, exec_datafusion_err};
 
 use super::array_filter::{ArrayStaticFilter, StaticFilter};
+use super::primitive::{U8Config, U16Config};
+use super::transform::make_bitmap_filter;
 
 pub(crate) fn instantiate_static_filter(
     in_array: ArrayRef,
 ) -> Result<Arc<dyn StaticFilter + Send + Sync>> {
     match in_array.data_type() {
-        // Integer primitive types
-        DataType::Int8 => Ok(Arc::new(Int8StaticFilter::try_new(&in_array)?)),
-        DataType::Int16 => Ok(Arc::new(Int16StaticFilter::try_new(&in_array)?)),
+        // 1-byte types: use bitmap (256 bits = 32 bytes)
+        DataType::Int8 | DataType::UInt8 => make_bitmap_filter::<U8Config>(&in_array),
+        // 2-byte types: use bitmap (65536 bits = 8 KB)
+        DataType::Int16 | DataType::UInt16 => make_bitmap_filter::<U16Config>(&in_array),
+        // 4-byte integer types
         DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)),
-        DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)),
-        DataType::UInt8 => Ok(Arc::new(UInt8StaticFilter::try_new(&in_array)?)),
-        DataType::UInt16 => Ok(Arc::new(UInt16StaticFilter::try_new(&in_array)?)),
         DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
+        // 8-byte integer types
+        DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)),
         DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
         // Float primitive types (use ordered wrappers for Hash/Eq)
         DataType::Float32 => Ok(Arc::new(Float32StaticFilter::try_new(&in_array)?)),
@@ -228,13 +231,10 @@ macro_rules! primitive_static_filter {
     };
 }
 
-// Generate specialized filters for all integer primitive types
-primitive_static_filter!(Int8StaticFilter, Int8Type);
-primitive_static_filter!(Int16StaticFilter, Int16Type);
+// Generate specialized filters for 4-byte and 8-byte integer primitive types
+// (1-byte and 2-byte types use BitmapFilter instead)
 primitive_static_filter!(Int32StaticFilter, Int32Type);
 primitive_static_filter!(Int64StaticFilter, Int64Type);
-primitive_static_filter!(UInt8StaticFilter, UInt8Type);
-primitive_static_filter!(UInt16StaticFilter, UInt16Type);
 primitive_static_filter!(UInt32StaticFilter, UInt32Type);
 primitive_static_filter!(UInt64StaticFilter, UInt64Type);
 
diff --git a/datafusion/physical-expr/src/expressions/in_list/transform.rs b/datafusion/physical-expr/src/expressions/in_list/transform.rs
diff --git a/scripts/compare_benchmarks.py b/scripts/compare_benchmarks.py
