Skip to content

Commit 012758d

Browse files
Implement Zero-Copy Reinterpretation and enable Int8/Int16 Bitmaps
Introduces zero-copy buffer reinterpretation to allow signed integers and other 1 or 2-byte primitive types (e.g. Float16) to use the high-performance bitmap filters. Triggers for all types with 1-byte or 2-byte width.
1 parent 01bf68c commit 012758d

4 files changed

Lines changed: 171 additions & 8 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ mod primitive_filter;
4141
mod result;
4242
mod static_filter;
4343
mod strategy;
44+
mod transform;
4445

4546
use static_filter::StaticFilter;
4647
use strategy::instantiate_static_filter;

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,22 @@ where
142142
fn check(&self, needle: T::Native) -> bool {
143143
self.bits.get_bit(needle.as_usize())
144144
}
145+
146+
/// Check membership using a raw values slice (zero-copy path for type reinterpretation).
147+
#[inline]
148+
pub(super) fn contains_slice(
149+
&self,
150+
values: &[T::Native],
151+
nulls: Option<&NullBuffer>,
152+
negated: bool,
153+
) -> BooleanArray {
154+
build_in_list_result(values.len(), nulls, self.null_count > 0, negated, |i| {
155+
// SAFETY: `build_in_list_result` invokes this closure for
156+
// indices in `0..values.len()`.
157+
let needle = unsafe { *values.get_unchecked(i) };
158+
self.check(needle)
159+
})
160+
}
145161
}
146162

147163
impl<T> StaticFilter for BitmapFilter<T>
@@ -359,9 +375,6 @@ macro_rules! primitive_static_filter {
359375
};
360376
}
361377

362-
// Generate specialized filters for all integer primitive types
363-
primitive_static_filter!(Int8StaticFilter, Int8Type);
364-
primitive_static_filter!(Int16StaticFilter, Int16Type);
365378
primitive_static_filter!(Int32StaticFilter, Int32Type);
366379
primitive_static_filter!(Int64StaticFilter, Int64Type);
367380
primitive_static_filter!(UInt32StaticFilter, UInt32Type);

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use datafusion_common::Result;
2525
use super::array_static_filter::ArrayStaticFilter;
2626
use super::primitive_filter::*;
2727
use super::static_filter::StaticFilter;
28+
use super::transform::make_bitmap_filter;
2829

2930
pub(super) fn instantiate_static_filter(
3031
in_array: ArrayRef,
@@ -37,13 +38,10 @@ pub(super) fn instantiate_static_filter(
3738
_ => in_array,
3839
};
3940
match in_array.data_type() {
40-
// Integer primitive types
41-
DataType::Int8 => Ok(Arc::new(Int8StaticFilter::try_new(&in_array)?)),
42-
DataType::Int16 => Ok(Arc::new(Int16StaticFilter::try_new(&in_array)?)),
41+
DataType::Int8 | DataType::UInt8 => make_bitmap_filter::<UInt8Type>(&in_array),
42+
DataType::Int16 | DataType::UInt16 => make_bitmap_filter::<UInt16Type>(&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(BitmapFilter::<UInt8Type>::try_new(&in_array)?)),
46-
DataType::UInt16 => Ok(Arc::new(BitmapFilter::<UInt16Type>::try_new(&in_array)?)),
4745
DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
4846
DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
4947
// Float primitive types (use ordered wrappers for Hash/Eq)
Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Type transformation utilities for InList filters.
19+
//!
20+
//! Some filters only depend on fixed-width value bit patterns. For those cases,
21+
//! compatible primitive arrays can be reinterpreted to the filter's unsigned
22+
//! storage type without copying values.
23+
24+
use std::mem::size_of;
25+
use std::sync::Arc;
26+
27+
use arrow::array::{Array, ArrayRef, BooleanArray, PrimitiveArray};
28+
use arrow::buffer::ScalarBuffer;
29+
use arrow::datatypes::ArrowPrimitiveType;
30+
use datafusion_common::{Result, exec_datafusion_err};
31+
32+
use super::primitive_filter::{BitmapFilter, BitmapFilterType};
33+
use super::static_filter::{StaticFilter, handle_dictionary};
34+
35+
// =============================================================================
36+
// REINTERPRETING FILTERS (zero-copy type conversion)
37+
// =============================================================================
38+
39+
/// Reinterpreting filter for bitmap lookups (u8/u16).
40+
struct ReinterpretedBitmap<T: BitmapFilterType> {
41+
inner: BitmapFilter<T>,
42+
}
43+
44+
impl<T: BitmapFilterType> StaticFilter for ReinterpretedBitmap<T> {
45+
fn null_count(&self) -> usize {
46+
self.inner.null_count()
47+
}
48+
49+
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
50+
handle_dictionary!(self, v, negated);
51+
52+
if v.data_type().primitive_width() != Some(size_of::<T::Native>()) {
53+
return Err(exec_datafusion_err!(
54+
"BitmapFilter: expected {}-byte primitive array, got {}",
55+
size_of::<T::Native>(),
56+
v.data_type()
57+
));
58+
}
59+
60+
let data = v.to_data();
61+
let values: &[T::Native] = &data.buffer::<T::Native>(0)[..v.len()];
62+
63+
Ok(self.inner.contains_slice(values, data.nulls(), negated))
64+
}
65+
}
66+
67+
/// Reinterprets a same-width primitive array as the target primitive type `T`.
68+
///
69+
/// This is a zero-copy operation: the returned array shares the original values
70+
/// buffer and null buffer. Callers must ensure the source array and target type
71+
/// have the same primitive width.
72+
#[inline]
73+
pub(crate) fn reinterpret_any_primitive_to<T: ArrowPrimitiveType>(
74+
array: &dyn Array,
75+
) -> ArrayRef {
76+
let data = array.to_data();
77+
let values = data.buffers()[0].clone();
78+
let buffer = ScalarBuffer::<T::Native>::new(values, data.offset(), data.len());
79+
Arc::new(PrimitiveArray::<T>::new(buffer, array.nulls().cloned()))
80+
}
81+
82+
/// Creates a bitmap filter for u8/u16 types, reinterpreting if needed.
83+
pub(crate) fn make_bitmap_filter<T>(
84+
in_array: &ArrayRef,
85+
) -> Result<Arc<dyn StaticFilter + Send + Sync>>
86+
where
87+
T: BitmapFilterType,
88+
{
89+
if in_array.data_type() == &T::DATA_TYPE {
90+
return Ok(Arc::new(BitmapFilter::<T>::try_new(in_array)?));
91+
}
92+
if in_array.data_type().primitive_width() != Some(size_of::<T::Native>()) {
93+
return Err(exec_datafusion_err!(
94+
"BitmapFilter: expected {}-byte primitive array for {} bitmap, got {}",
95+
size_of::<T::Native>(),
96+
T::DATA_TYPE,
97+
in_array.data_type()
98+
));
99+
}
100+
101+
let reinterpreted = reinterpret_any_primitive_to::<T>(in_array.as_ref());
102+
let inner = BitmapFilter::<T>::try_new(&reinterpreted)?;
103+
Ok(Arc::new(ReinterpretedBitmap { inner }))
104+
}
105+
106+
#[cfg(test)]
107+
mod tests {
108+
use super::*;
109+
use std::sync::Arc;
110+
111+
use arrow::array::{ArrayRef, BooleanArray, Int8Array, Int16Array};
112+
use arrow::datatypes::{UInt8Type, UInt16Type};
113+
114+
#[test]
115+
fn reinterpreted_bitmap_handles_signed_boundaries_and_slices() -> Result<()> {
116+
let haystack: ArrayRef = Arc::new(
117+
Int8Array::from(vec![Some(99), Some(i8::MIN), None, Some(-1), Some(42)])
118+
.slice(1, 3),
119+
);
120+
let filter = make_bitmap_filter::<UInt8Type>(&haystack)?;
121+
let needles =
122+
Int8Array::from(vec![Some(7), Some(i8::MIN), Some(-1), None]).slice(1, 3);
123+
124+
assert_eq!(
125+
filter.contains(&needles, false)?,
126+
BooleanArray::from(vec![Some(true), Some(true), None])
127+
);
128+
129+
let haystack: ArrayRef = Arc::new(
130+
Int16Array::from(vec![
131+
Some(123),
132+
Some(i16::MIN),
133+
None,
134+
Some(-1),
135+
Some(i16::MAX),
136+
])
137+
.slice(1, 4),
138+
);
139+
let filter = make_bitmap_filter::<UInt16Type>(&haystack)?;
140+
let needles =
141+
Int16Array::from(vec![Some(0), Some(i16::MIN), Some(7), Some(i16::MAX)])
142+
.slice(1, 3);
143+
144+
assert_eq!(
145+
filter.contains(&needles, false)?,
146+
BooleanArray::from(vec![Some(true), None, Some(true)])
147+
);
148+
149+
Ok(())
150+
}
151+
}

0 commit comments

Comments
 (0)