perf: optimize HashTableLookupExpr::evaluate

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

@@ -21,18 +21,18 @@ use std::{any::Any, fmt::Display, hash::Hash, sync::Arc};
 
 use ahash::RandomState;
 use arrow::{
-    array::{BooleanArray, UInt64Array},
-    buffer::MutableBuffer,
+    array::{ArrayRef, UInt64Array},
     datatypes::{DataType, Schema},
-    util::bit_util,
+    record_batch::RecordBatch,
 };
+use datafusion_common::hash_utils::{create_hashes, with_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.
 ///
@@ -114,6 +114,15 @@ impl HashExpr {
     pub fn description(&self) -> &str {
         &self.description
     }
+
+    /// Evaluate the columns to be hashed.
+    fn evaluate_on_columns(&self, batch: &RecordBatch) -> Result<Vec<ArrayRef>> {
+        let num_rows = batch.num_rows();
+        self.on_columns
+            .iter()
+            .map(|c| c.evaluate(batch)?.into_array(num_rows))
+            .collect()
+    }
 }
 
 impl std::fmt::Debug for HashExpr {
@@ -181,18 +190,11 @@ impl PhysicalExpr for HashExpr {
         Ok(false)
     }
 
-    fn evaluate(
-        &self,
-        batch: &arrow::record_batch::RecordBatch,
-    ) -> Result<ColumnarValue> {
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         let num_rows = batch.num_rows();
 
         // Evaluate columns
-        let keys_values = self
-            .on_columns
-            .iter()
-            .map(|c| c.evaluate(batch)?.into_array(num_rows))
-            .collect::<Result<Vec<_>>>()?;
+        let keys_values = self.evaluate_on_columns(batch)?;
 
         // Compute hashes
         let mut hashes_buffer = vec![0; num_rows];
@@ -327,12 +329,24 @@ impl PhysicalExpr for HashTableLookupExpr {
         Ok(false)
     }
 
-    fn evaluate(
-        &self,
-        batch: &arrow::record_batch::RecordBatch,
-    ) -> Result<ColumnarValue> {
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         let num_rows = batch.num_rows();
 
+        // Optimization: if hash_expr is HashExpr, compute hashes directly into callback
+        // to avoid redundant allocations and copies.
+        if let Some(hash_expr) = self.hash_expr.as_any().downcast_ref::<HashExpr>() {
+            let keys_values = hash_expr.evaluate_on_columns(batch)?;
+
+            return with_hashes(
+                &keys_values,
+                hash_expr.random_state.random_state(),
+                |hashes| {
+                    let array = self.hash_map.contain_hashes(hashes);
+                    Ok(ColumnarValue::Array(Arc::new(array)))
+                },
+            );
+        }
+
         // Evaluate hash expression to get hash values
         let hash_array = self.hash_expr.evaluate(batch)?.into_array(num_rows)?;
         let hash_array = hash_array.as_any().downcast_ref::<UInt64Array>().ok_or(
@@ -342,21 +356,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));
-        for (idx, hash_value) in hash_array.values().iter().enumerate() {
-            // Use get_matched_indices to check - if it returns any indices, the hash exists
-            let (matched_indices, _) = self
-                .hash_map
-                .get_matched_indices(Box::new(std::iter::once((idx, hash_value))), None);
-
-            if !matched_indices.is_empty() {
-                bit_util::set_bit(buf.as_slice_mut(), idx);
-            }
-        }
+        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,6 +22,8 @@
 use std::fmt::{self, Debug};
 use std::ops::Sub;
 
+use arrow::array::BooleanArray;
+use arrow::buffer::BooleanBuffer;
 use arrow::datatypes::ArrowNativeType;
 use hashbrown::HashTable;
 use hashbrown::hash_table::Entry::{Occupied, Vacant};
@@ -124,6 +126,9 @@ pub trait JoinHashMapType: Send + Sync {
         match_indices: &mut Vec<u64>,
     ) -> Option<JoinHashMapOffset>;
 
+    /// 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;
 
@@ -196,6 +201,10 @@ impl JoinHashMapType for JoinHashMapU32 {
         )
     }
 
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
+    }
+
     fn is_empty(&self) -> bool {
         self.map.is_empty()
     }
@@ -270,6 +279,10 @@ impl JoinHashMapType for JoinHashMapU64 {
         )
     }
 
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
+    }
+
     fn is_empty(&self) -> bool {
         self.map.is_empty()
     }
@@ -496,3 +509,35 @@ where
     }
     None
 }
+
+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_contain_hashes() {
+        let mut hash_map = JoinHashMapU32::with_capacity(10);
+        hash_map.update_from_iter(Box::new([10u64, 20u64, 30u64].iter().enumerate()), 0);
+
+        let probe_hashes = vec![10, 11, 20, 21, 30, 31];
+        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!(array.value(i), "Hash {hash} should exist in the map");
+            } else {
+                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,
-    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,6 +95,10 @@ impl JoinHashMapType for PruningJoinHashMap {
         )
     }
 
+    fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray {
+        contain_hashes(&self.map, hash_values)
+    }
+
     fn is_empty(&self) -> bool {
         self.map.is_empty()
     }