Skip to content

Commit bde8e5b

Browse files
IN LIST: add UInt16 bitmap filter (apache#23012)
## 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. <!-- codex-benchmark-start --> ## 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. <!-- codex-benchmark-end -->
1 parent ff677c4 commit bde8e5b

2 files changed

Lines changed: 103 additions & 4 deletions

File tree

datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs

Lines changed: 102 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,81 @@ impl StaticFilter for UInt8BitmapFilter {
104104
}
105105
}
106106

107+
/// Bitmap filter for O(1) `UInt16` set membership via single bit test.
108+
///
109+
/// `UInt16` has 65,536 possible values, so the filter stores membership in an
110+
/// 8 KiB heap-allocated bitmap instead of using a hash table.
111+
pub(super) struct UInt16BitmapFilter {
112+
null_count: usize,
113+
bits: Box<[u64; 1024]>,
114+
}
115+
116+
impl UInt16BitmapFilter {
117+
pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
118+
let prim_array = in_array.as_primitive_opt::<UInt16Type>().ok_or_else(|| {
119+
exec_datafusion_err!("UInt16BitmapFilter: expected UInt16 array")
120+
})?;
121+
let mut bits = Box::new([0u64; 1024]);
122+
let mut set_bit = |v: u16| {
123+
let index = usize::from(v);
124+
bits[index / 64] |= 1u64 << (index % 64);
125+
};
126+
127+
let values = prim_array.values();
128+
match prim_array.nulls() {
129+
None => {
130+
for &v in values {
131+
set_bit(v);
132+
}
133+
}
134+
Some(nulls) => {
135+
for i in
136+
BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len())
137+
{
138+
set_bit(values[i]);
139+
}
140+
}
141+
}
142+
Ok(Self {
143+
null_count: prim_array.null_count(),
144+
bits,
145+
})
146+
}
147+
148+
#[inline(always)]
149+
fn check(&self, needle: u16) -> bool {
150+
let index = needle as usize;
151+
(self.bits[index / 64] >> (index % 64)) & 1 != 0
152+
}
153+
}
154+
155+
impl StaticFilter for UInt16BitmapFilter {
156+
fn null_count(&self) -> usize {
157+
self.null_count
158+
}
159+
160+
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
161+
handle_dictionary!(self, v, negated);
162+
let v = v.as_primitive_opt::<UInt16Type>().ok_or_else(|| {
163+
exec_datafusion_err!("UInt16BitmapFilter: expected UInt16 array")
164+
})?;
165+
let input_values = v.values();
166+
Ok(build_in_list_result(
167+
v.len(),
168+
v.nulls(),
169+
self.null_count > 0,
170+
negated,
171+
#[inline(always)]
172+
|i| {
173+
// SAFETY: `build_in_list_result` invokes this closure for
174+
// indices in `0..v.len()`, which matches `input_values.len()`.
175+
let needle = unsafe { *input_values.get_unchecked(i) };
176+
self.check(needle)
177+
},
178+
))
179+
}
180+
}
181+
107182
/// Wrapper for f32 that implements Hash and Eq using bit comparison.
108183
/// This treats NaN values as equal to each other when they have the same bit pattern.
109184
#[derive(Clone, Copy)]
@@ -294,7 +369,6 @@ primitive_static_filter!(Int8StaticFilter, Int8Type);
294369
primitive_static_filter!(Int16StaticFilter, Int16Type);
295370
primitive_static_filter!(Int32StaticFilter, Int32Type);
296371
primitive_static_filter!(Int64StaticFilter, Int64Type);
297-
primitive_static_filter!(UInt16StaticFilter, UInt16Type);
298372
primitive_static_filter!(UInt32StaticFilter, UInt32Type);
299373
primitive_static_filter!(UInt64StaticFilter, UInt64Type);
300374

@@ -315,10 +389,10 @@ mod tests {
315389
use super::*;
316390
use std::sync::Arc;
317391

318-
use arrow::array::{DictionaryArray, Int8Array, UInt8Array};
392+
use arrow::array::{DictionaryArray, Int8Array, UInt8Array, UInt16Array};
319393

320394
fn assert_contains(
321-
filter: &UInt8BitmapFilter,
395+
filter: &dyn StaticFilter,
322396
needles: &dyn Array,
323397
expected: Vec<Option<bool>>,
324398
) -> Result<()> {
@@ -355,4 +429,29 @@ mod tests {
355429

356430
assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])
357431
}
432+
433+
#[test]
434+
fn bitmap_filter_u16_handles_boundaries_and_nulls() -> Result<()> {
435+
let haystack: ArrayRef = Arc::new(UInt16Array::from(vec![
436+
Some(0),
437+
None,
438+
Some(1024),
439+
Some(u16::MAX),
440+
]));
441+
let filter = UInt16BitmapFilter::try_new(&haystack)?;
442+
let needles =
443+
UInt16Array::from(vec![Some(0), Some(1), Some(1024), Some(u16::MAX), None]);
444+
445+
assert_contains(
446+
&filter,
447+
&needles,
448+
vec![Some(true), None, Some(true), Some(true), None],
449+
)?;
450+
assert_eq!(
451+
filter.contains(&needles, true)?,
452+
BooleanArray::from(vec![Some(false), None, Some(false), Some(false), None])
453+
);
454+
455+
Ok(())
456+
}
358457
}

datafusion/physical-expr/src/expressions/in_list/strategy.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ pub(super) fn instantiate_static_filter(
4343
DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)),
4444
DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)),
4545
DataType::UInt8 => Ok(Arc::new(UInt8BitmapFilter::try_new(&in_array)?)),
46-
DataType::UInt16 => Ok(Arc::new(UInt16StaticFilter::try_new(&in_array)?)),
46+
DataType::UInt16 => Ok(Arc::new(UInt16BitmapFilter::try_new(&in_array)?)),
4747
DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
4848
DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
4949
// Float primitive types (use ordered wrappers for Hash/Eq)

0 commit comments

Comments
 (0)