Skip to content

Commit 3085965

Browse files
geoffreyclaudegeoffreyclaude
committed
Refactor bitmap filters behind shared config
1 parent bde8e5b commit 3085965

2 files changed

Lines changed: 88 additions & 90 deletions

File tree

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

Lines changed: 82 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -29,113 +29,108 @@ use std::hash::{Hash, Hasher};
2929
use super::result::build_in_list_result;
3030
use super::static_filter::{StaticFilter, handle_dictionary};
3131

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],
32+
pub(super) trait BitmapStorage: Send + Sync {
33+
fn new_zeroed() -> Self;
34+
fn set_bit(&mut self, index: usize);
35+
fn get_bit(&self, index: usize) -> bool;
3936
}
4037

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-
})
38+
impl BitmapStorage for [u64; 4] {
39+
#[inline]
40+
fn new_zeroed() -> Self {
41+
[0u64; 4]
42+
}
43+
#[inline]
44+
fn set_bit(&mut self, index: usize) {
45+
self[index / 64] |= 1u64 << (index % 64);
7146
}
47+
#[inline(always)]
48+
fn get_bit(&self, index: usize) -> bool {
49+
(self[index / 64] >> (index % 64)) & 1 != 0
50+
}
51+
}
7252

53+
impl BitmapStorage for Box<[u64; 1024]> {
54+
#[inline]
55+
fn new_zeroed() -> Self {
56+
Box::new([0u64; 1024])
57+
}
58+
#[inline]
59+
fn set_bit(&mut self, index: usize) {
60+
self[index / 64] |= 1u64 << (index % 64);
61+
}
7362
#[inline(always)]
74-
fn check(&self, needle: u8) -> bool {
75-
let index = needle as usize;
76-
(self.bits[index / 64] >> (index % 64)) & 1 != 0
63+
fn get_bit(&self, index: usize) -> bool {
64+
(self[index / 64] >> (index % 64)) & 1 != 0
7765
}
7866
}
7967

80-
impl StaticFilter for UInt8BitmapFilter {
81-
fn null_count(&self) -> usize {
82-
self.null_count
68+
pub(super) trait BitmapFilterConfig: Send + Sync + 'static {
69+
const DATA_TYPE_NAME: &'static str;
70+
71+
type Native: ArrowNativeType + Copy + Send + Sync;
72+
type ArrowType: ArrowPrimitiveType<Native = Self::Native>;
73+
type Storage: BitmapStorage;
74+
75+
fn to_index(v: Self::Native) -> usize;
76+
}
77+
78+
pub(super) enum UInt8BitmapConfig {}
79+
impl BitmapFilterConfig for UInt8BitmapConfig {
80+
const DATA_TYPE_NAME: &'static str = "UInt8";
81+
82+
type Native = u8;
83+
type ArrowType = UInt8Type;
84+
type Storage = [u64; 4];
85+
86+
#[inline(always)]
87+
fn to_index(v: u8) -> usize {
88+
v as usize
8389
}
90+
}
8491

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-
))
92+
pub(super) enum UInt16BitmapConfig {}
93+
impl BitmapFilterConfig for UInt16BitmapConfig {
94+
const DATA_TYPE_NAME: &'static str = "UInt16";
95+
96+
type Native = u16;
97+
type ArrowType = UInt16Type;
98+
type Storage = Box<[u64; 1024]>;
99+
100+
#[inline(always)]
101+
fn to_index(v: u16) -> usize {
102+
v as usize
104103
}
105104
}
106105

107-
/// Bitmap filter for O(1) `UInt16` set membership via single bit test.
106+
/// Bitmap filter for O(1) set membership via single bit test.
108107
///
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 {
108+
/// Small integer domains can store membership in a fixed-size bitmap instead
109+
/// of using a hash table.
110+
pub(super) struct BitmapFilter<C: BitmapFilterConfig> {
112111
null_count: usize,
113-
bits: Box<[u64; 1024]>,
112+
bits: C::Storage,
114113
}
115114

116-
impl UInt16BitmapFilter {
115+
impl<C: BitmapFilterConfig> BitmapFilter<C> {
117116
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-
117+
let prim_array =
118+
in_array.as_primitive_opt::<C::ArrowType>().ok_or_else(|| {
119+
exec_datafusion_err!("BitmapFilter: expected {} array", C::DATA_TYPE_NAME)
120+
})?;
121+
let mut bits = C::Storage::new_zeroed();
127122
let values = prim_array.values();
128123
match prim_array.nulls() {
129124
None => {
130125
for &v in values {
131-
set_bit(v);
126+
bits.set_bit(C::to_index(v));
132127
}
133128
}
134129
Some(nulls) => {
135130
for i in
136131
BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len())
137132
{
138-
set_bit(values[i]);
133+
bits.set_bit(C::to_index(values[i]));
139134
}
140135
}
141136
}
@@ -146,21 +141,20 @@ impl UInt16BitmapFilter {
146141
}
147142

148143
#[inline(always)]
149-
fn check(&self, needle: u16) -> bool {
150-
let index = needle as usize;
151-
(self.bits[index / 64] >> (index % 64)) & 1 != 0
144+
fn check(&self, needle: C::Native) -> bool {
145+
self.bits.get_bit(C::to_index(needle))
152146
}
153147
}
154148

155-
impl StaticFilter for UInt16BitmapFilter {
149+
impl<C: BitmapFilterConfig> StaticFilter for BitmapFilter<C> {
156150
fn null_count(&self) -> usize {
157151
self.null_count
158152
}
159153

160154
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
161155
handle_dictionary!(self, v, negated);
162-
let v = v.as_primitive_opt::<UInt16Type>().ok_or_else(|| {
163-
exec_datafusion_err!("UInt16BitmapFilter: expected UInt16 array")
156+
let v = v.as_primitive_opt::<C::ArrowType>().ok_or_else(|| {
157+
exec_datafusion_err!("BitmapFilter: expected {} array", C::DATA_TYPE_NAME)
164158
})?;
165159
let input_values = v.values();
166160
Ok(build_in_list_result(
@@ -406,7 +400,7 @@ mod tests {
406400
#[test]
407401
fn bitmap_filter_u8_handles_nulls() -> Result<()> {
408402
let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)]));
409-
let filter = UInt8BitmapFilter::try_new(&haystack)?;
403+
let filter = BitmapFilter::<UInt8BitmapConfig>::try_new(&haystack)?;
410404
let needles = UInt8Array::from(vec![Some(1), Some(2), None, Some(3)]);
411405

412406
assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?;
@@ -421,7 +415,7 @@ mod tests {
421415
#[test]
422416
fn bitmap_filter_u8_handles_dictionary_needles() -> Result<()> {
423417
let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)]));
424-
let filter = UInt8BitmapFilter::try_new(&haystack)?;
418+
let filter = BitmapFilter::<UInt8BitmapConfig>::try_new(&haystack)?;
425419

426420
let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]);
427421
let values = Arc::new(UInt8Array::from(vec![Some(1), Some(2), Some(3)]));
@@ -438,7 +432,7 @@ mod tests {
438432
Some(1024),
439433
Some(u16::MAX),
440434
]));
441-
let filter = UInt16BitmapFilter::try_new(&haystack)?;
435+
let filter = BitmapFilter::<UInt16BitmapConfig>::try_new(&haystack)?;
442436
let needles =
443437
UInt16Array::from(vec![Some(0), Some(1), Some(1024), Some(u16::MAX), None]);
444438

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,12 @@ 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(UInt8BitmapFilter::try_new(&in_array)?)),
46-
DataType::UInt16 => Ok(Arc::new(UInt16BitmapFilter::try_new(&in_array)?)),
45+
DataType::UInt8 => Ok(Arc::new(BitmapFilter::<UInt8BitmapConfig>::try_new(
46+
&in_array,
47+
)?)),
48+
DataType::UInt16 => Ok(Arc::new(BitmapFilter::<UInt16BitmapConfig>::try_new(
49+
&in_array,
50+
)?)),
4751
DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
4852
DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
4953
// Float primitive types (use ordered wrappers for Hash/Eq)

0 commit comments

Comments
 (0)