JoinHashMapType::set_bits_if_exists(&self, hash_values: &[u64], buffer: &mut [u8]) -> contain_hashes(&self, hash_values: &[u64]) -> BooleanArray

Author: UBarney

datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs

@@ -21,18 +21,17 @@ use std::{any::Any, fmt::Display, hash::Hash, sync::Arc};
 
 use ahash::RandomState;
 use arrow::{
-    array::{BooleanArray, UInt64Array},
-    buffer::MutableBuffer,
+    array::UInt64Array,
     datatypes::{DataType, Schema},
-    util::bit_util,
 };
+use datafusion_common::hash_utils::create_hashes;
 use datafusion_common::{Result, internal_datafusion_err, internal_err};
 use datafusion_expr::ColumnarValue;
 use datafusion_physical_expr_common::physical_expr::{
     DynHash, PhysicalExpr, PhysicalExprRef,
 };
 
-use crate::{hash_utils::create_hashes, joins::utils::JoinHashMapType};
+use crate::joins::utils::JoinHashMapType;
 
 /// RandomState wrapper that preserves the seeds used to create it.
 ///
@@ -342,13 +341,9 @@ impl PhysicalExpr for HashTableLookupExpr {
         )?;
 
         // Check each hash against the hash table
-        let mut buf = MutableBuffer::from_len_zeroed(bit_util::ceil(num_rows, 8));
-        self.hash_map
-            .set_bits_if_exists(hash_array.values(), buf.as_slice_mut());
+        let array = self.hash_map.contain_hashes(hash_array.values());
 
-        Ok(ColumnarValue::Array(Arc::new(
-            BooleanArray::new_from_packed(buf, 0, num_rows),
-        )))
+        Ok(ColumnarValue::Array(Arc::new(array)))
     }
 
     fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

datafusion/physical-plan/src/joins/join_hash_map.rs

@@ -22,8 +22,9 @@
 use std::fmt::{self, Debug};
 use std::ops::Sub;
 
+use arrow::array::BooleanArray;
+use arrow::buffer::BooleanBuffer;
 use arrow::datatypes::ArrowNativeType;
-use arrow::util::bit_util;
 use hashbrown::HashTable;
 use hashbrown::hash_table::Entry::{Occupied, Vacant};
 
@@ -125,8 +126,8 @@ pub trait JoinHashMapType: Send + Sync {
         match_indices: &mut Vec<u64>,
     ) -> Option<JoinHashMapOffset>;
 
-    /// Sets bits in the provided buffer if the corresponding hash exists in the map.
-    fn set_bits_if_exists(&self, hash_values: &[u64], buffer: &mut [u8]);
+    /// Returns a BooleanArray indicating which of the provided hashes exist in the map.
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray;
 
     /// Returns `true` if the join hash map contains no entries.
     fn is_empty(&self) -> bool;
@@ -200,8 +201,8 @@ impl JoinHashMapType for JoinHashMapU32 {
         )
     }
 
-    fn set_bits_if_exists(&self, hash_values: &[u64], buffer: &mut [u8]) {
-        set_bits_if_exists::<u32>(&self.map, hash_values, buffer);
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
     }
 
     fn is_empty(&self) -> bool {
@@ -278,8 +279,8 @@ impl JoinHashMapType for JoinHashMapU64 {
         )
     }
 
-    fn set_bits_if_exists(&self, hash_values: &[u64], buffer: &mut [u8]) {
-        set_bits_if_exists::<u64>(&self.map, hash_values, buffer);
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
     }
 
     fn is_empty(&self) -> bool {
@@ -509,44 +510,33 @@ where
     None
 }
 
-pub fn set_bits_if_exists<T>(
-    map: &HashTable<(u64, T)>,
-    hash_values: &[u64],
-    buffer: &mut [u8],
-) {
-    for (i, &hash) in hash_values.iter().enumerate() {
-        if map.find(hash, |(h, _)| hash == *h).is_some() {
-            bit_util::set_bit(buffer, i);
-        }
-    }
+pub fn contain_hashes<T>(map: &HashTable<(u64, T)>, hash_values: &[u64]) -> BooleanArray {
+    let buffer = BooleanBuffer::collect_bool(hash_values.len(), |i| {
+        let hash = hash_values[i];
+        map.find(hash, |(h, _)| hash == *h).is_some()
+    });
+    BooleanArray::new(buffer, None)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_set_bits_if_exists() {
+    fn test_contain_hashes() {
         let mut hash_map = JoinHashMapU32::with_capacity(10);
-        // Build side: insert 10, 20, 30
         hash_map.update_from_iter(Box::new([10u64, 20u64, 30u64].iter().enumerate()), 0);
 
-        // Probe side: test both existing and non-existing hashes
         let probe_hashes = vec![10, 11, 20, 21, 30, 31];
-        let mut buffer = vec![0u8; 1];
-        hash_map.set_bits_if_exists(&probe_hashes, &mut buffer);
+        let array = hash_map.contain_hashes(&probe_hashes);
+
+        assert_eq!(array.len(), probe_hashes.len());
 
         for (i, &hash) in probe_hashes.iter().enumerate() {
             if matches!(hash, 10 | 20 | 30) {
-                assert!(
-                    bit_util::get_bit(&buffer, i),
-                    "Hash {hash} should exist in the map"
-                );
+                assert!(array.value(i), "Hash {hash} should exist in the map");
             } else {
-                assert!(
-                    !bit_util::get_bit(&buffer, i),
-                    "Hash {hash} should NOT exist in the map"
-                );
+                assert!(!array.value(i), "Hash {hash} should NOT exist in the map");
             }
         }
     }

datafusion/physical-plan/src/joins/stream_join_utils.rs

@@ -23,15 +23,16 @@ use std::mem::size_of;
 use std::sync::Arc;
 
 use crate::joins::join_hash_map::{
-    JoinHashMapOffset, get_matched_indices, get_matched_indices_with_limit_offset,
-    set_bits_if_exists, update_from_iter,
+    JoinHashMapOffset, contain_hashes, get_matched_indices,
+    get_matched_indices_with_limit_offset, update_from_iter,
 };
 use crate::joins::utils::{JoinFilter, JoinHashMapType};
 use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder};
 use crate::{ExecutionPlan, metrics};
 
 use arrow::array::{
-    ArrowPrimitiveType, BooleanBufferBuilder, NativeAdapter, PrimitiveArray, RecordBatch,
+    ArrowPrimitiveType, BooleanArray, BooleanBufferBuilder, NativeAdapter,
+    PrimitiveArray, RecordBatch,
 };
 use arrow::compute::concat_batches;
 use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef};
@@ -94,8 +95,8 @@ impl JoinHashMapType for PruningJoinHashMap {
         )
     }
 
-    fn set_bits_if_exists(&self, hash_values: &[u64], buffer: &mut [u8]) {
-        set_bits_if_exists::<u64>(&self.map, hash_values, buffer);
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
     }
 
     fn is_empty(&self) -> bool {