IN LIST: add UInt16 bitmap filter (apache#23012)

geoffreyclaude · web-flow · commit bde8e5b24e77 · 2026-06-26T09:21:05.000Z
## Which issue does this PR close? - Part of apache#19241. - Stacked on apache#23011. - Next in stack: apache#23035. - Extracted from apache#19390. ## Rationale for this change apache#23011 uses a bitmap checklist for `UInt8`, where there are 256 possible values. `UInt16` is the same idea with a larger value range: 0 through 65,535. That is still small enough to represent directly. A `UInt16` bitmap needs one bit for each possible value: - 65,536 possible values - 65,536 bits total - 8 KB of memory Then a lookup is still simple: use the input value as the bit position and check whether that bit is set. For example, if the list contains `42`, bit `42` is set, and every input row with value `42` can be recognized with one bit test. This PR keeps the scope narrow: it adds the unsigned 2-byte bitmap path as a concrete `UInt16` filter. apache#23035 then unifies the `UInt8` and `UInt16` implementations, and apache#23013 uses that shared shape for signed same-width reinterpretation. ## What changes are included in this PR? - Adds `UInt16BitmapFilter`, backed by a heap-allocated 65,536-bit bitmap. - Routes `UInt16` constant-list filtering to that bitmap path. - Keeps the same `IN` / `NOT IN` null behavior as the generic path. - Adds focused coverage for `UInt16` boundary values, nulls, and `NOT IN`. ## Are these changes tested? Yes. - `cargo fmt --all` - `cargo test -p datafusion-physical-expr bitmap_filter_u16 --lib` - `cargo test -p datafusion-physical-expr in_list_int_types --lib` - `cargo test -p datafusion-physical-expr test_in_list_from_array_type_combinations --lib` - `cargo test -p datafusion-physical-expr test_in_list_dictionary_types --lib` - `cargo clippy -p datafusion-physical-expr --all-targets --all-features -- -D warnings` ## Are there any user-facing changes? No. This is an internal performance optimization only.  ## Benchmark note No local `in_list_strategy` numbers are included for this PR because the benchmark harness does not currently include a direct `UInt16` case. The available `i16` rows measure the signed reinterpretation path added in apache#23013 after the bitmap unification in apache#23035, not this PR's unsigned `UInt16` bitmap filter.
diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs
@@ -104,6 +104,81 @@ impl StaticFilter for UInt8BitmapFilter {
     }
 }
 
+/// Bitmap filter for O(1) `UInt16` set membership via single bit test.
+///
+/// `UInt16` has 65,536 possible values, so the filter stores membership in an
+/// 8 KiB heap-allocated bitmap instead of using a hash table.
+pub(super) struct UInt16BitmapFilter {
+    null_count: usize,
+    bits: Box<[u64; 1024]>,
+}
+
+impl UInt16BitmapFilter {
+    pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+        let prim_array = in_array.as_primitive_opt::<UInt16Type>().ok_or_else(|| {
+            exec_datafusion_err!("UInt16BitmapFilter: expected UInt16 array")
+        })?;
+        let mut bits = Box::new([0u64; 1024]);
+        let mut set_bit = |v: u16| {
+            let index = usize::from(v);
+            bits[index / 64] |= 1u64 << (index % 64);
+        };
+
+        let values = prim_array.values();
+        match prim_array.nulls() {
+            None => {
+                for &v in values {
+                    set_bit(v);
+                }
+            }
+            Some(nulls) => {
+                for i in
+                    BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len())
+                {
+                    set_bit(values[i]);
+                }
+            }
+        }
+        Ok(Self {
+            null_count: prim_array.null_count(),
+            bits,
+        })
+    }
+
+    #[inline(always)]
+    fn check(&self, needle: u16) -> bool {
+        let index = needle as usize;
+        (self.bits[index / 64] >> (index % 64)) & 1 != 0
+    }
+}
+
+impl StaticFilter for UInt16BitmapFilter {
+    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::<UInt16Type>().ok_or_else(|| {
+            exec_datafusion_err!("UInt16BitmapFilter: expected UInt16 array")
+        })?;
+        let input_values = v.values();
+        Ok(build_in_list_result(
+            v.len(),
+            v.nulls(),
+            self.null_count > 0,
+            negated,
+            #[inline(always)]
+            |i| {
+                // SAFETY: `build_in_list_result` invokes this closure for
+                // indices in `0..v.len()`, which matches `input_values.len()`.
+                let needle = unsafe { *input_values.get_unchecked(i) };
+                self.check(needle)
+            },
+        ))
+    }
+}
+
 /// Wrapper for f32 that implements Hash and Eq using bit comparison.
 /// This treats NaN values as equal to each other when they have the same bit pattern.
 #[derive(Clone, Copy)]
@@ -294,7 +369,6 @@ primitive_static_filter!(Int8StaticFilter, Int8Type);
 primitive_static_filter!(Int16StaticFilter, Int16Type);
 primitive_static_filter!(Int32StaticFilter, Int32Type);
 primitive_static_filter!(Int64StaticFilter, Int64Type);
-primitive_static_filter!(UInt16StaticFilter, UInt16Type);
 primitive_static_filter!(UInt32StaticFilter, UInt32Type);
 primitive_static_filter!(UInt64StaticFilter, UInt64Type);
 
@@ -315,10 +389,10 @@ mod tests {
     use super::*;
     use std::sync::Arc;
 
-    use arrow::array::{DictionaryArray, Int8Array, UInt8Array};
+    use arrow::array::{DictionaryArray, Int8Array, UInt8Array, UInt16Array};
 
     fn assert_contains(
-        filter: &UInt8BitmapFilter,
+        filter: &dyn StaticFilter,
         needles: &dyn Array,
         expected: Vec<Option<bool>>,
     ) -> Result<()> {
@@ -355,4 +429,29 @@ mod tests {
 
         assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])
     }
+
+    #[test]
+    fn bitmap_filter_u16_handles_boundaries_and_nulls() -> Result<()> {
+        let haystack: ArrayRef = Arc::new(UInt16Array::from(vec![
+            Some(0),
+            None,
+            Some(1024),
+            Some(u16::MAX),
+        ]));
+        let filter = UInt16BitmapFilter::try_new(&haystack)?;
+        let needles =
+            UInt16Array::from(vec![Some(0), Some(1), Some(1024), Some(u16::MAX), None]);
+
+        assert_contains(
+            &filter,
+            &needles,
+            vec![Some(true), None, Some(true), Some(true), None],
+        )?;
+        assert_eq!(
+            filter.contains(&needles, true)?,
+            BooleanArray::from(vec![Some(false), None, Some(false), Some(false), None])
+        );
+
+        Ok(())
+    }
 }
diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs
@@ -43,7 +43,7 @@ pub(super) fn instantiate_static_filter(
         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(UInt8BitmapFilter::try_new(&in_array)?)),
-        DataType::UInt16 => Ok(Arc::new(UInt16StaticFilter::try_new(&in_array)?)),
+        DataType::UInt16 => Ok(Arc::new(UInt16BitmapFilter::try_new(&in_array)?)),
         DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
         DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
         // Float primitive types (use ordered wrappers for Hash/Eq)
