Skip to content

Commit d0f8c34

Browse files
IN LIST: add UInt8 bitmap filter (apache#23011)
## Which issue does this PR close? - Part of apache#19241. - Stacked on apache#21927. - Next in stack: apache#23012. - Extracted from apache#19390. ## Rationale for this change `IN LIST` evaluates expressions like `x IN (1, 3, 7)`. The list on the right is fixed, so DataFusion can precompute a small lookup structure once and then reuse it for every input row. For `UInt8`, there are only 256 possible values: 0 through 255. That means the lookup can be a tiny checklist with one bit per possible value: - If the list contains `3`, set bit `3`. - If the list contains `7`, set bit `7`. - To check whether an input value is present, read that one bit. So instead of hashing each input value or comparing it against the list, membership becomes one indexed bit test. The bitmap is only 32 bytes, because 256 bits = 32 bytes. This PR adds the first specialized primitive path in the stack as a concrete `UInt8` filter. The `UInt16` version is added in apache#23012, and the shared bitmap abstraction is introduced only after both concrete implementations are visible in apache#23035. ## What changes are included in this PR? - Adds `UInt8BitmapFilter`, a 32-byte bitmap built from the non-null constants in the `IN` list. - Routes `UInt8` constant-list filtering to that bitmap path. - Keeps the same SQL null behavior as the generic path for both `IN` and `NOT IN`. - Moves shared dictionary-needle handling into `static_filter.rs`, so specialized filters can reuse it consistently. - Adds focused tests for `UInt8` null handling and dictionary-encoded needles. ## Are these changes tested? Yes. - `cargo fmt --all` - `cargo test -p datafusion-physical-expr bitmap_filter_u8 --lib` - `cargo test -p datafusion-physical-expr in_list_int_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 --> ## Local benchmark snapshot Benchmark command: ```bash cargo bench -p datafusion-physical-expr --profile release-nonlto --bench in_list_strategy -- --save-baseline <name> ``` Method: compare adjacent saved baselines using raw Criterion sample minima (`min(time / iters)`). Lower is better; changes within +/-5% are treated as noise. These numbers were not rerun after splitting the bitmap abstraction into apache#23035. Compared baselines: [apache#21927](apache#21927) -> [apache#23011](apache#23011) Relevant scope: UInt8 narrow-integer rows. Summary: 5 relevant rows, 5 faster, 0 slower, 0 within +/-5%. | Benchmark | Before | After | Change | |---|---:|---:|---:| | `narrow_integer/u8/list=16/match=0%` | 20.39 us | 3.94 us | -80.7% (5.18x faster) | | `narrow_integer/u8/list=16/match=50%` | 38.38 us | 3.98 us | -89.6% (9.65x faster) | | `narrow_integer/u8/list=4/match=0%` | 18.18 us | 3.93 us | -78.4% (4.62x faster) | | `narrow_integer/u8/list=4/match=50%` | 34.63 us | 3.96 us | -88.6% (8.75x faster) | | `nulls/narrow_integer/u8/list=16/match=50%/nulls=20%` | 37.12 us | 4.16 us | -88.8% (8.93x faster) | <!-- codex-benchmark-end --> --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 96ff37a commit d0f8c34

3 files changed

Lines changed: 166 additions & 24 deletions

File tree

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

Lines changed: 148 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,94 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::array::{
19-
Array, ArrayRef, AsArray, BooleanArray, downcast_array, downcast_dictionary_array,
20-
};
18+
//! Optimized primitive type filters for InList expressions.
19+
//!
20+
//! This module provides membership tests for Arrow primitive types.
21+
22+
use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
2123
use arrow::buffer::{BooleanBuffer, NullBuffer};
22-
use arrow::compute::take;
2324
use arrow::datatypes::*;
25+
use arrow::util::bit_iterator::BitIndexIterator;
2426
use datafusion_common::{HashSet, Result, exec_datafusion_err};
2527
use std::hash::{Hash, Hasher};
2628

27-
use super::static_filter::StaticFilter;
29+
use super::result::build_in_list_result;
30+
use super::static_filter::{StaticFilter, handle_dictionary};
31+
32+
/// Bitmap filter for O(1) set membership via single bit test.
33+
///
34+
/// `UInt8` has only 256 possible values, so the filter stores membership in a
35+
/// 256-bit bitmap instead of using a hash table.
36+
pub(super) struct UInt8BitmapFilter {
37+
null_count: usize,
38+
bits: [u64; 4],
39+
}
40+
41+
impl UInt8BitmapFilter {
42+
pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
43+
let prim_array = in_array.as_primitive_opt::<UInt8Type>().ok_or_else(|| {
44+
exec_datafusion_err!("UInt8BitmapFilter: expected UInt8 array")
45+
})?;
46+
let mut bits = [0u64; 4];
47+
let mut set_bit = |v: u8| {
48+
let index = usize::from(v);
49+
bits[index / 64] |= 1u64 << (index % 64);
50+
};
51+
52+
let values = prim_array.values();
53+
match prim_array.nulls() {
54+
None => {
55+
for &v in values {
56+
set_bit(v);
57+
}
58+
}
59+
Some(nulls) => {
60+
for i in
61+
BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len())
62+
{
63+
set_bit(values[i]);
64+
}
65+
}
66+
}
67+
Ok(Self {
68+
null_count: prim_array.null_count(),
69+
bits,
70+
})
71+
}
72+
73+
#[inline(always)]
74+
fn check(&self, needle: u8) -> bool {
75+
let index = needle as usize;
76+
(self.bits[index / 64] >> (index % 64)) & 1 != 0
77+
}
78+
}
79+
80+
impl StaticFilter for UInt8BitmapFilter {
81+
fn null_count(&self) -> usize {
82+
self.null_count
83+
}
84+
85+
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
86+
handle_dictionary!(self, v, negated);
87+
let v = v.as_primitive_opt::<UInt8Type>().ok_or_else(|| {
88+
exec_datafusion_err!("UInt8BitmapFilter: expected UInt8 array")
89+
})?;
90+
let input_values = v.values();
91+
Ok(build_in_list_result(
92+
v.len(),
93+
v.nulls(),
94+
self.null_count > 0,
95+
negated,
96+
#[inline(always)]
97+
|i| {
98+
// SAFETY: `build_in_list_result` invokes this closure for
99+
// indices in `0..v.len()`, which matches `input_values.len()`.
100+
let needle = unsafe { *input_values.get_unchecked(i) };
101+
self.check(needle)
102+
},
103+
))
104+
}
105+
}
28106

29107
/// Wrapper for f32 that implements Hash and Eq using bit comparison.
30108
/// This treats NaN values as equal to each other when they have the same bit pattern.
@@ -94,9 +172,13 @@ macro_rules! primitive_static_filter {
94172

95173
impl $Name {
96174
pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
97-
let in_array = in_array
98-
.as_primitive_opt::<$ArrowType>()
99-
.ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?;
175+
let in_array =
176+
in_array.as_primitive_opt::<$ArrowType>().ok_or_else(|| {
177+
exec_datafusion_err!(
178+
"Failed to downcast an array to a '{}' array",
179+
stringify!($ArrowType)
180+
)
181+
})?;
100182

101183
let mut values = HashSet::with_capacity(in_array.len());
102184
let null_count = in_array.null_count();
@@ -115,19 +197,14 @@ macro_rules! primitive_static_filter {
115197
}
116198

117199
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
118-
// Handle dictionary arrays by recursing on the values
119-
downcast_dictionary_array! {
120-
v => {
121-
let values_contains = self.contains(v.values().as_ref(), negated)?;
122-
let result = take(&values_contains, v.keys(), None)?;
123-
return Ok(downcast_array(result.as_ref()))
124-
}
125-
_ => {}
126-
}
200+
handle_dictionary!(self, v, negated);
127201

128-
let v = v
129-
.as_primitive_opt::<$ArrowType>()
130-
.ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?;
202+
let v = v.as_primitive_opt::<$ArrowType>().ok_or_else(|| {
203+
exec_datafusion_err!(
204+
"Failed to downcast an array to a '{}' array",
205+
stringify!($ArrowType)
206+
)
207+
})?;
131208

132209
let haystack_has_nulls = self.null_count > 0;
133210
let needle_values = v.values();
@@ -188,8 +265,10 @@ macro_rules! primitive_static_filter {
188265
}
189266
(true, true) => {
190267
// Both have nulls - combine needle nulls with haystack-induced nulls
191-
let needle_validity = needle_nulls.map(|n| n.inner().clone())
192-
.unwrap_or_else(|| BooleanBuffer::new_set(needle_values.len()));
268+
let needle_validity =
269+
needle_nulls.map(|n| n.inner().clone()).unwrap_or_else(
270+
|| BooleanBuffer::new_set(needle_values.len()),
271+
);
193272

194273
// Valid when original "in set" is true (see above)
195274
let haystack_validity = if negated {
@@ -215,7 +294,6 @@ primitive_static_filter!(Int8StaticFilter, Int8Type);
215294
primitive_static_filter!(Int16StaticFilter, Int16Type);
216295
primitive_static_filter!(Int32StaticFilter, Int32Type);
217296
primitive_static_filter!(Int64StaticFilter, Int64Type);
218-
primitive_static_filter!(UInt8StaticFilter, UInt8Type);
219297
primitive_static_filter!(UInt16StaticFilter, UInt16Type);
220298
primitive_static_filter!(UInt32StaticFilter, UInt32Type);
221299
primitive_static_filter!(UInt64StaticFilter, UInt64Type);
@@ -231,3 +309,50 @@ macro_rules! float_static_filter {
231309
// Generate specialized filters for float types using ordered wrappers
232310
float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32);
233311
float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64);
312+
313+
#[cfg(test)]
314+
mod tests {
315+
use super::*;
316+
use std::sync::Arc;
317+
318+
use arrow::array::{DictionaryArray, Int8Array, UInt8Array};
319+
320+
fn assert_contains(
321+
filter: &UInt8BitmapFilter,
322+
needles: &dyn Array,
323+
expected: Vec<Option<bool>>,
324+
) -> Result<()> {
325+
assert_eq!(
326+
filter.contains(needles, false)?,
327+
BooleanArray::from(expected)
328+
);
329+
Ok(())
330+
}
331+
332+
#[test]
333+
fn bitmap_filter_u8_handles_nulls() -> Result<()> {
334+
let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)]));
335+
let filter = UInt8BitmapFilter::try_new(&haystack)?;
336+
let needles = UInt8Array::from(vec![Some(1), Some(2), None, Some(3)]);
337+
338+
assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?;
339+
assert_eq!(
340+
filter.contains(&needles, true)?,
341+
BooleanArray::from(vec![Some(false), None, None, Some(false)])
342+
);
343+
344+
Ok(())
345+
}
346+
347+
#[test]
348+
fn bitmap_filter_u8_handles_dictionary_needles() -> Result<()> {
349+
let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)]));
350+
let filter = UInt8BitmapFilter::try_new(&haystack)?;
351+
352+
let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]);
353+
let values = Arc::new(UInt8Array::from(vec![Some(1), Some(2), Some(3)]));
354+
let needles = DictionaryArray::try_new(keys, values)?;
355+
356+
assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])
357+
}
358+
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,3 +35,20 @@ pub(super) trait StaticFilter {
3535
/// implementation unwraps the dictionary and operates on its values.
3636
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray>;
3737
}
38+
39+
/// Evaluate dictionary-encoded needles by applying a filter to dictionary
40+
/// values and remapping the result through the keys.
41+
macro_rules! handle_dictionary {
42+
($self:ident, $v:ident, $negated:ident) => {
43+
arrow::array::downcast_dictionary_array! {
44+
$v => {
45+
let values_contains = $self.contains($v.values().as_ref(), $negated)?;
46+
let result = arrow::compute::take(&values_contains, $v.keys(), None)?;
47+
return Ok(arrow::array::downcast_array(result.as_ref()))
48+
}
49+
_ => {}
50+
}
51+
};
52+
}
53+
54+
pub(super) use handle_dictionary;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ pub(super) fn instantiate_static_filter(
4242
DataType::Int16 => Ok(Arc::new(Int16StaticFilter::try_new(&in_array)?)),
4343
DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)),
4444
DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)),
45-
DataType::UInt8 => Ok(Arc::new(UInt8StaticFilter::try_new(&in_array)?)),
45+
DataType::UInt8 => Ok(Arc::new(UInt8BitmapFilter::try_new(&in_array)?)),
4646
DataType::UInt16 => Ok(Arc::new(UInt16StaticFilter::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)?)),

0 commit comments

Comments
 (0)