Skip to content

Commit 97db09b

Browse files
adriangbclaude
andcommitted
perf: route probe rows to ArrayMap build partitions in the merged filter
A `Partitioned` hash join merges its partitions' build maps into one routed `hash_lookup` filter. The routed lookup accepted only hash tables, so a join whose partitions built an `ArrayMap` (dense integer keys) fell back to a `CASE` over one filter per partition. The routed lookup now probes an `ArrayMap` partition with the key values of the rows routed to it, because an `ArrayMap` is keyed by value, not hash. Every non-empty partition can now merge. The `CASE` fallback remains only for canceled partitions. The routed lookup fuzz test now builds an `ArrayMap` for about half of the partitions. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 70f2ed0 commit 97db09b

2 files changed

Lines changed: 42 additions & 34 deletions

File tree

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

Lines changed: 41 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
use std::{fmt::Display, hash::Hash, sync::Arc};
88

99
use arrow::{
10-
array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt64Array},
10+
array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt32Array, UInt64Array},
11+
compute::take,
1112
datatypes::{DataType, Schema},
1213
record_batch::RecordBatch,
1314
};
@@ -352,8 +353,7 @@ impl HashTableLookupExpr {
352353
/// for pruning only, or `None`
353354
///
354355
/// A single non-empty map is probed directly without routing, so it may be
355-
/// any [`Map`] and `router` is ignored. Otherwise `router` is required and
356-
/// every map must be a [`Map::HashMap`].
356+
/// any [`Map`] and `router` is ignored. Otherwise `router` is required.
357357
///
358358
/// # Public Only for Internal Use:
359359
/// `datafusion-proto` tests require this constructor, but it is not part of
@@ -375,21 +375,10 @@ impl HashTableLookupExpr {
375375
"a lookup over anything but one non-empty map needs a partition router"
376376
);
377377
}
378-
(Some(router), _) => {
379-
if maps
380-
.iter()
381-
.flatten()
382-
.any(|map| !matches!(map.as_ref(), Map::HashMap(_)))
383-
{
384-
return internal_err!(
385-
"routed lookups only support hash table partitions"
386-
);
387-
}
388-
Lookup::Routed {
389-
router,
390-
maps: maps.into(),
391-
}
392-
}
378+
(Some(router), _) => Lookup::Routed {
379+
router,
380+
maps: maps.into(),
381+
},
393382
};
394383
Ok(Self {
395384
on_columns,
@@ -441,16 +430,23 @@ impl HashTableLookupExpr {
441430
for (partition, map) in maps.iter().enumerate() {
442431
// An empty partition holds no key, so its rows stay unmatched.
443432
let Some(map) = map else { continue };
444-
let Map::HashMap(map) = map.as_ref() else {
445-
return internal_err!(
446-
"routed lookups only support hash table partitions"
447-
);
448-
};
449433
let run = starts[partition]..starts[partition + 1];
450434
if run.is_empty() {
451435
continue;
452436
}
453-
let hits = map.contain_hashes(&sorted_hashes[run.clone()]);
437+
let hits = match map.as_ref() {
438+
Map::HashMap(map) => map.contain_hashes(&sorted_hashes[run.clone()]),
439+
// An array map is keyed by value, not hash, so it probes
440+
// the run's own key values.
441+
Map::ArrayMap(map) => {
442+
let indices = UInt32Array::from(rows[run.clone()].to_vec());
443+
let run_keys = keys
444+
.iter()
445+
.map(|key| take(key.as_ref(), &indices, None))
446+
.collect::<std::result::Result<Vec<_>, _>>()?;
447+
map.contain_keys(&run_keys)?
448+
}
449+
};
454450
for hit in hits.values().set_indices() {
455451
matched.set_bit(rows[run.start + hit] as usize, true);
456452
}
@@ -639,7 +635,9 @@ fn evaluate_columns(
639635
#[cfg(test)]
640636
mod tests {
641637
use super::*;
638+
use crate::joins::array_map::ArrayMap;
642639
use crate::joins::join_hash_map::JoinHashMapU32;
640+
use arrow::array::AsArray;
643641
use datafusion_common::ScalarValue;
644642
use datafusion_physical_expr::expressions::Column;
645643
use std::collections::hash_map::DefaultHasher;
@@ -1111,7 +1109,7 @@ mod tests {
11111109
)]));
11121110
let batch = RecordBatch::try_new(schema, vec![Arc::clone(probe)])?;
11131111
let matched = expr.evaluate(&batch)?.into_array(batch.num_rows())?;
1114-
Ok(arrow::array::AsArray::as_boolean(&matched).clone())
1112+
Ok(matched.as_boolean().clone())
11151113
}
11161114

11171115
/// A hash table of `keys`, hashed with `random_state` as the build side
@@ -1125,6 +1123,18 @@ mod tests {
11251123
Ok(Arc::new(Map::HashMap(Box::new(map))))
11261124
}
11271125

1126+
/// An array map of `keys`, as a dense build side gets.
1127+
fn array_map(keys: &ArrayRef) -> Result<Arc<Map>> {
1128+
let keys = keys.as_primitive::<arrow::datatypes::Int32Type>();
1129+
let key = |value: Option<i32>| {
1130+
ArrayMap::key_to_u64(&ScalarValue::Int32(value)).expect("Int32 key")
1131+
};
1132+
let min = key(arrow::compute::min(keys));
1133+
let max = key(arrow::compute::max(keys));
1134+
let keys: ArrayRef = Arc::new(keys.clone());
1135+
Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(&keys, min, max)?)))
1136+
}
1137+
11281138
/// One randomly generated scenario, see [`routed_lookup_fuzz`].
11291139
fn run_routed_lookup_scenario(seed: u64) -> Result<()> {
11301140
use crate::repartition::{BatchPartitioner, REPARTITION_RANDOM_STATE};
@@ -1166,7 +1176,12 @@ mod tests {
11661176
.partition(
11671177
RecordBatch::try_new(schema, vec![Arc::clone(&build)])?,
11681178
|partition, batch| {
1169-
maps[partition] = Some(hash_table(batch.column(0), &random_state)?);
1179+
// A dense partition builds an array map, others a hash table.
1180+
maps[partition] = Some(if rng.random_bool(0.5) {
1181+
array_map(batch.column(0))?
1182+
} else {
1183+
hash_table(batch.column(0), &random_state)?
1184+
});
11701185
Ok(())
11711186
},
11721187
)?;

‎datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -721,14 +721,7 @@ impl SharedBuildAccumulator {
721721
}) => {
722722
real_partition_ids.push(partition_id);
723723
keys_have_null |= partition_keys_have_null;
724-
let mergeable = match &pushdown {
725-
PushdownStrategy::Map(map, _) => {
726-
matches!(map.as_ref(), Map::HashMap(_))
727-
}
728-
PushdownStrategy::InList(_) => true,
729-
PushdownStrategy::Empty => false,
730-
};
731-
members = members.filter(|_| mergeable).map(|mut members| {
724+
members = members.map(|mut members| {
732725
members.push((partition_id, pushdown.clone(), bounds.clone()));
733726
members
734727
});

0 commit comments

Comments
 (0)