jayzhan211 commented on code in PR #27091:
URL: https://github.com/apache/datafusion/pull/27091#discussion_r4499844335
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -722,15 +722,7 @@ impl SharedBuildAccumulator {
}) => {
real_partition_ids.push(partition_id);
keys_have_null |= partition_keys_have_null;
- let mergeable = match &pushdown {
- PushdownStrategy::Map(map, _) => {
- let map: &Map = map;
- matches!(map, Map::HashMap(_))
- }
- PushdownStrategy::InList(_) => true,
- PushdownStrategy::Empty => false,
- };
- members = members.filter(|_| mergeable).map(|mut members| {
+ members = members.map(|mut members| {
Review Comment:
Could you add a test that array map partitions actually reach the merged
filter? Putting the `mergeable` gate back still passes every existing test,
because the fuzz builds the lookup directly. Range routing with array maps
isn't covered either; these two tests fail with the gate restored.
<details>
<summary>Two tests for the `shared_bounds.rs` test module, plus the
`array_map_of` helper</summary>
```rs
/// An array map holding `keys`, as a dense build side gets.
fn array_map_of(keys: &[i32]) -> PushdownStrategy {
use crate::joins::array_map::ArrayMap;
let bound = |key: Option<&i32>| {
ArrayMap::key_to_u64(&ScalarValue::Int32(key.copied())).unwrap()
};
let min = bound(keys.iter().min());
let max = bound(keys.iter().max());
let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec()));
let map = ArrayMap::try_new(&keys, min, max).unwrap();
let map = Arc::new(Map::ArrayMap(map));
PushdownStrategy::Map(Arc::new(ReservedMap::unreserved(map)), None)
}
/// Array map partitions merge into one lookup with hash table partitions.
#[test]
#[cfg(not(feature = "force_hash_collisions"))]
fn partitioned_array_map_and_hash_map_partitions_merge() {
let acc = make_partitioned_expr_accumulator_for_test(4);
let build: Vec<i32> = (0..40).step_by(2).collect();
acc.build_filter(FinalizeInput::Partitioned(routed_partitions(
&acc,
4,
&build,
|partition, keys| {
if partition % 2 == 0 {
array_map_of(keys)
} else {
hash_map_of(keys)
}
},
)))
.unwrap();
assert!(
current_expr(&acc).downcast_ref::<CaseExpr>().is_none(),
"expected a merged lookup, got: {}",
current_expr(&acc)
);
let probe: Vec<i32> = (0..40).collect();
let expected: Vec<bool> = probe.iter().map(|key| key % 2 == 0).collect();
assert_matches(&acc, &probe, &expected);
}
/// Array map partitions merge under Range routing too.
#[test]
#[cfg(not(feature = "force_hash_collisions"))]
fn partitioned_range_array_map_partitions_merge() -> Result<()> {
let mut acc = make_partitioned_expr_accumulator_for_test(4);
acc.probe_range_partitioning = Some(RangePartitioning::try_new(
[PhysicalSortExpr::new(
Arc::clone(&acc.on_right[0]),
Default::default(),
)]
.into(),
vec![
SplitPoint::new(vec![ScalarValue::Int32(Some(10))]),
SplitPoint::new(vec![ScalarValue::Int32(Some(20))]),
SplitPoint::new(vec![ScalarValue::Int32(Some(30))]),
],
)?);
acc.build_filter(FinalizeInput::Partitioned(vec![
reported(array_map_of(&[7, 8, 9]), bounds(7, 9)),
reported(PushdownStrategy::Empty, no_bounds()),
reported(array_map_of(&[27, 28, 29]), bounds(27, 29)),
reported(array_map_of(&[37, 38, 39]), bounds(37, 39)),
]))?;
assert!(current_expr(&acc).downcast_ref::<CaseExpr>().is_none());
let probe: Vec<i32> = (0..50).collect();
let expected: Vec<bool> = probe
.iter()
.map(|key| matches!(key, 7..=9 | 27..=29 | 37..=39))
.collect();
assert_matches(&acc, &probe, &expected);
Ok(())
}
```
</details>
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -1195,7 +1206,12 @@ mod tests {
.partition(
RecordBatch::try_new(schema, vec![Arc::clone(&build)])?,
|partition, batch| {
- maps[partition] = Some(hash_table(batch.column(0),
&random_state)?);
+ // A dense partition builds an array map, others a hash table.
+ maps[partition] = Some(if rng.random_bool(0.5) {
Review Comment:
This test fails with `force_hash_collisions`, which only runs in the merge
queue: every hash collides, so the reference hash table matches every key while
an array map stays exact. Keeping hash tables in that mode fixes it and passes
both ways.
```suggestion
// Half the partitions build an array map, as dense keys
would. Under
// forced collisions a hash table matches every key, so an
array
// map would disagree with the single-table reference below.
let dense =
cfg!(not(feature = "force_hash_collisions")) &&
rng.random_bool(0.5);
maps[partition] = Some(if dense {
```
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -464,17 +454,24 @@ impl HashTableLookupExpr {
for (partition, map) in maps.iter().enumerate() {
// An empty partition holds no key, so its rows stay unmatched.
let Some(map) = map else { continue };
- let map: &Map = map;
- let Map::HashMap(map) = map else {
- return internal_err!(
- "routed lookups only support hash table partitions"
- );
- };
let run = starts[partition]..starts[partition + 1];
if run.is_empty() {
continue;
}
- let hits = map.contain_hashes(&sorted_hashes[run.clone()]);
+ let map: &Map = map;
+ let hits = match map {
+ Map::HashMap(map) =>
map.contain_hashes(&sorted_hashes[run.clone()]),
+ // An array map is keyed by value, not hash, so it probes
+ // the run's own key values.
+ Map::ArrayMap(map) => {
+ let indices =
UInt32Array::from(rows[run.clone()].to_vec());
Review Comment:
When every partition is an array map, `with_hashes` and `sorted_hashes` do
per-row work that nothing reads. Skipping them took lookup cost from 4.8 to 4.1
ns/row at 16 partitions and batch 8192, measured with a local variant on a busy
laptop. It's fine to handle this in a follow-up, since it only saves time and
needs the closure reshaped. The condition to hash:
```rs
let needs_hashes = maps.iter().flatten().any(|map| {
let map: &Map = map;
matches!(map, Map::HashMap(_))
});
```
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -464,17 +454,24 @@ impl HashTableLookupExpr {
for (partition, map) in maps.iter().enumerate() {
// An empty partition holds no key, so its rows stay unmatched.
let Some(map) = map else { continue };
- let map: &Map = map;
- let Map::HashMap(map) = map else {
- return internal_err!(
- "routed lookups only support hash table partitions"
- );
- };
let run = starts[partition]..starts[partition + 1];
if run.is_empty() {
continue;
}
- let hits = map.contain_hashes(&sorted_hashes[run.clone()]);
+ let map: &Map = map;
+ let hits = match map {
+ Map::HashMap(map) =>
map.contain_hashes(&sorted_hashes[run.clone()]),
+ // An array map is keyed by value, not hash, so it probes
+ // the run's own key values.
+ Map::ArrayMap(map) => {
+ let indices =
UInt32Array::from(rows[run.clone()].to_vec());
+ let run_keys = keys
+ .iter()
+ .map(|key| take(key.as_ref(), &indices, None))
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ map.contain_keys(&run_keys)?
Review Comment:
Optional cleanup: `take_arrays` does this collect, with the import changed
to `compute::take_arrays`.
```suggestion
map.contain_keys(&take_arrays(keys, &indices,
None)?)?
```
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -1154,6 +1153,18 @@ mod tests {
Ok(Arc::new(Map::HashMap(Box::new(map))))
}
+ /// An array map of `keys`, as a dense build side gets.
+ fn array_map(keys: &ArrayRef) -> Result<Arc<Map>> {
+ let keys = keys.as_primitive::<arrow::datatypes::Int32Type>();
+ let key = |value: Option<i32>| {
+ ArrayMap::key_to_u64(&ScalarValue::Int32(value)).expect("Int32
key")
+ };
+ let min = key(arrow::compute::min(keys));
+ let max = key(arrow::compute::max(keys));
+ let keys: ArrayRef = Arc::new(keys.clone());
+ Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(&keys, min, max)?)))
Review Comment:
Optional cleanup: the helper can pass its input through instead of
re-wrapping it.
```suggestion
let values = keys.as_primitive::<arrow::datatypes::Int32Type>();
let key = |value: Option<i32>| {
ArrayMap::key_to_u64(&ScalarValue::Int32(value)).expect("Int32
key")
};
let min = key(arrow::compute::min(values));
let max = key(arrow::compute::max(values));
Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(keys, min, max)?)))
```
--
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]
