perf: SIMD short-circuit in JoinHashMap probe

[AURON-2160] Optimize join hash map probe by checking hash_matched
first before computing empty mask. This reduces ~50% SIMD instructions
when hash hit rate is high (typical join scenarios).

Before: Always compute both hash_matched and empty SIMD masks.
After: Only compute empty mask when hash_matched has no hits.

Also add a criterion microbenchmark (benches/join_hash_map.rs) covering
realistic BHJ build sizes (5M/10M/20M keys) × three hit rates (0/50/100%).

Results on Apple M2 Pro (probe_size=4096):

  build size      | hit=0%  | hit=50% | hit=100%
  ----------------+---------+---------+---------
  5M  (~128 MB)   | 6.63 µs | 6.52 µs | 6.35 µs
  10M (~256 MB)   | 6.68 µs | 6.50 µs | 6.36 µs
  20M (~512 MB)   | 6.70 µs | 6.59 µs | 6.36 µs

Latency stays flat because prefetch_read_data (4-step ahead) fully
pipelines cache misses. The hit=100% path is consistently ~4-5% faster,
aligning with the optimization goal. Instruction-count savings can be
confirmed on x86 via: perf stat -e instructions

Run benchmark:
  cargo bench --bench join_hash_map -p datafusion-ext-plans

Author: yew1eb

Cargo.lock

@@ -182,6 +182,7 @@ parking_lot = "0.12.5"
 paste = "1.0.15"
 procfs = "0.18.0"
 prost = "0.14.3"
+criterion = { version = "0.5", features = ["html_reports"] }
 prost-types = "0.14.3"
 prost-reflect = "0.16.3"
 rand = "0.9.3"

Cargo.toml

@@ -73,4 +73,9 @@ sonic-rs = { workspace = true }
 procfs = { workspace = true }
 
 [dev-dependencies]
+criterion = { workspace = true }
 rand = { workspace = true }
+
+[[bench]]
+name = "join_hash_map"
+harness = false

native-engine/datafusion-ext-plans/Cargo.toml

@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Microbenchmarks for JoinHashMap::lookup_many.
+//!
+//! Sweeps three build-side map sizes to cover different cache regimes:
+//!   - 4K keys  (~2 MB map)  → fits in L2, mostly cache-hot
+//!   - 64K keys (~32 MB map) → spills into L3
+//!   - 1M keys  (~512 MB map)→ well beyond L3, cache-cold
+//!
+//! For each map size, three probe hit rates are measured:
+//!   - 0%   (all misses)
+//!   - 50%  (half hit, half miss)
+//!   - 100% (all hits — typical inner-join scenario)
+//!
+//! Probe batch size is fixed at 4096 to match a typical Spark batch.
+
+use std::sync::Arc;
+
+use arrow::{
+    array::Int32Array,
+    datatypes::{DataType, Field, Schema},
+    record_batch::RecordBatch,
+};
+use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
+use datafusion_ext_plans::joins::join_hash_map::JoinHashMap;
+
+/// Probe batch size — matches the default Spark vectorized batch size.
+const PROBE_SIZE: usize = 4096;
+
+/// Build-side hash base: top bit always set so hashes are non-zero
+/// (matches join_create_hashes convention).
+const BUILD_HASH_BASE: u32 = 0x8000_0001;
+
+fn make_map(build_size: usize) -> (JoinHashMap, Vec<u32>) {
+    let values: Vec<i32> = (0..build_size as i32).collect();
+    let array = Arc::new(Int32Array::from(values));
+    let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)]));
+    let batch = RecordBatch::try_new(schema, vec![array.clone() as _]).expect("build batch");
+
+    let build_hashes: Vec<u32> = (0..build_size as u32)
+        .map(|i| BUILD_HASH_BASE.wrapping_add(i))
+        .collect();
+
+    let map = JoinHashMap::create_from_data_batch_and_hashes(
+        batch,
+        vec![array as _],
+        build_hashes.clone(),
+    )
+    .expect("build map");
+
+    (map, build_hashes)
+}
+
+/// Generate PROBE_SIZE hashes at the given hit rate.
+/// Hits and misses are interleaved evenly (Bresenham) to avoid cache grouping
+/// bias.
+fn make_probe_hashes(build_hashes: &[u32], build_size: usize, hit_rate: f64) -> Vec<u32> {
+    // Miss hashes start well above the build range to avoid accidental collision.
+    let miss_base: u32 = BUILD_HASH_BASE.wrapping_add(build_size as u32 + 0x0010_0000);
+    let build_len = build_hashes.len();
+    (0..PROBE_SIZE)
+        .map(|i| {
+            let cumulative = ((i + 1) as f64 * hit_rate) as usize;
+            let prev = (i as f64 * hit_rate) as usize;
+            if cumulative > prev {
+                build_hashes[i % build_len]
+            } else {
+                miss_base.wrapping_add(i as u32)
+            }
+        })
+        .collect()
+}
+
+fn bench_lookup_many(c: &mut Criterion) {
+    // (label, build_size)
+    // map memory ≈ (build_size * 2 / 8).next_power_of_two() * 64 bytes
+    //   5M  →  ~128 MB  (realistic BHJ, ~50 MB serialized data)
+    //   10M →  ~256 MB  (realistic BHJ, ~100–200 MB serialized data)
+    //   20M →  ~512 MB  (realistic BHJ, ~1 GB serialized data)
+    let build_sizes: &[(&str, usize)] = &[
+        ("build=5M", 5_000_000),
+        ("build=10M", 10_000_000),
+        ("build=20M", 20_000_000),
+    ];
+    let hit_rates: &[(&str, f64)] = &[("hit=0%", 0.0), ("hit=50%", 0.5), ("hit=100%", 1.0)];
+
+    let mut group = c.benchmark_group("JoinHashMap::lookup_many");
+
+    for &(size_label, build_size) in build_sizes {
+        let (map, build_hashes) = make_map(build_size);
+        for &(rate_label, hit_rate) in hit_rates {
+            let probe = make_probe_hashes(&build_hashes, build_size, hit_rate);
+            let label = format!("{size_label}/{rate_label}");
+            group.bench_with_input(BenchmarkId::from_parameter(&label), &label, |b, _| {
+                b.iter(|| {
+                    let result = map.lookup_many(black_box(probe.clone()));
+                    black_box(result)
+                });
+            });
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_lookup_many);
+criterion_main!(benches);

native-engine/datafusion-ext-plans/benches/join_hash_map.rs

@@ -536,8 +536,8 @@ impl AccColumn for AccBytesColumn {
             if array_data.is_valid(i) {
                 let offset_begin = offset_buffer[i] as usize;
                 let offset_end = offset_buffer[i + 1] as usize;
-                self.items.push(Some(AccBytes::from_slice(
-                    &data_buffer[offset_begin..offset_end],
+                self.items.push(Some(AccBytes::from(
+                    data_buffer[offset_begin..offset_end].as_ref(),
                 )));
             } else {
                 self.items.push(None);

native-engine/datafusion-ext-plans/src/agg/acc.rs

@@ -179,7 +179,7 @@ impl AggHashMapKey for &[u8] {
     }
 
     fn into_owned(self) -> OwnedKey {
-        OwnedKey::from_slice(self)
+        OwnedKey::from(self.as_ref())
     }
 }
 

native-engine/datafusion-ext-plans/src/agg/agg_hash_map.rs

@@ -118,7 +118,7 @@ impl Agg for AggFirst {
                     ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
                         if flag_accs.value(acc_idx).is_none() {
                             if partial_arg.is_valid(partial_arg_idx) {
-                                value_accs.set_value(acc_idx, Some(AccBytes::from(partial_arg.value(partial_arg_idx).as_ref())));
+                                value_accs.set_value(acc_idx, Some(AccBytes::from(partial_arg.value(partial_arg_idx).as_ref() as &[u8])));
                                 flag_accs.set_value(acc_idx, Some(true));
                             } else {
                                 value_accs.set_value(acc_idx, None);

native-engine/datafusion-ext-plans/src/agg/first.rs

@@ -107,7 +107,7 @@ impl Agg for AggFirstIgnoresNull {
                 idx_for_zipped! {
                     ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
                         if accs.value(acc_idx).is_none() && partial_arg.is_valid(partial_arg_idx) {
-                            accs.set_value(acc_idx, Some(AccBytes::from(partial_arg.value(partial_arg_idx).as_ref())));
+                            accs.set_value(acc_idx, Some(AccBytes::from(partial_arg.value(partial_arg_idx).as_ref() as &[u8])));
                         }
                     }
                 }

native-engine/datafusion-ext-plans/src/agg/first_ignores_null.rs

@@ -253,15 +253,23 @@ impl Table {
             let mut e = entries![i] as usize;
             loop {
                 let hash_matched = self.map[e].hashes.simd_eq(Simd::splat(hashes[i]));
-                let empty = self.map[e].hashes.simd_eq(Simd::splat(0));
 
-                if let Some(pos) = (hash_matched | empty).first_set() {
+                // Fast path: check hash match first (common case)
+                if let Some(pos) = hash_matched.first_set() {
                     hashes[i] = unsafe {
                         // safety: transmute MapValue(u32) to u32
                         std::mem::transmute(self.map[e].values[pos])
                     };
                     break;
                 }
+
+                // Slow path: check empty slot only when no match
+                let empty = self.map[e].hashes.simd_eq(Simd::splat(0));
+                if empty.any() {
+                    hashes[i] = MapValue::EMPTY.0;
+                    break;
+                }
+
                 e += 1;
                 e %= 1 << self.map_mod_bits;
             }