geoffreyclaude commented on code in PR #23011:
URL: https://github.com/apache/datafusion/pull/23011#discussion_r3462308503
##########
datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs:
##########
@@ -15,16 +15,94 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{
- Array, ArrayRef, AsArray, BooleanArray, downcast_array,
downcast_dictionary_array,
-};
+//! Optimized primitive type filters for InList expressions.
+//!
+//! This module provides membership tests for Arrow primitive types.
+
+use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
use arrow::buffer::{BooleanBuffer, NullBuffer};
-use arrow::compute::take;
use arrow::datatypes::*;
+use arrow::util::bit_iterator::BitIndexIterator;
use datafusion_common::{HashSet, Result, exec_datafusion_err};
use std::hash::{Hash, Hasher};
-use super::static_filter::StaticFilter;
+use super::result::build_in_list_result;
+use super::static_filter::{StaticFilter, handle_dictionary};
+
+/// Bitmap filter for O(1) set membership via single bit test.
+///
+/// `UInt8` has only 256 possible values, so the filter stores membership in a
+/// 256-bit bitmap instead of using a hash table.
+pub(super) struct UInt8BitmapFilter {
+ null_count: usize,
+ bits: [u64; 4],
+}
+
+impl UInt8BitmapFilter {
+ pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+ let prim_array =
in_array.as_primitive_opt::<UInt8Type>().ok_or_else(|| {
+ exec_datafusion_err!("UInt8BitmapFilter: expected UInt8 array")
+ })?;
+ let mut bits = [0u64; 4];
+ let mut set_bit = |v: u8| {
+ 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: u8) -> bool {
+ let index = needle as usize;
+ (self.bits[index / 64] >> (index % 64)) & 1 != 0
Review Comment:
Confirmed: rustc/LLVM does optimize this bounds check away in optimized
builds.
Godbolt repro: https://godbolt.org/z/6r1WjvYv4
I used LLVM IR there because retained Rust bounds checks are easy to spot:
they show up as a branch to a `panic` block calling `panic_bounds_check`.
For the `u8` versions (`bitmap_check` / `bitmap_set`), there is no such
branch. The IR widens the `u8` with:
```llvm
%index = zext i8 %needle to i64
```
then computes `index / 64` as a shift and directly loads/stores from the
`[u64; 4]`.
For contrast, the repro also includes the same expression with `index:
usize`; that version *does* emit:
```llvm
icmp ult i64 %index, 256
br i1 ..., label %bb1, label %panic
```
followed by `panic_bounds_check`.
So I think the safe indexing here already compiles to the unchecked access
we want on the hot path, without needing `get_unchecked`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
