adriangb commented on code in PR #26249:
URL: https://github.com/apache/datafusion/pull/26249#discussion_r4406687502
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -312,16 +369,101 @@ impl HashTableLookupExpr {
Self {
on_columns,
random_state,
- map,
+ lookup: Lookup::Single(map),
description,
pruning_bitmap,
}
}
+ /// Looks each probe row up in the hash table of the build partition
+ /// `router` assigns it to. `maps` holds one entry per partition, each a
+ /// [`Map::HashMap`] or `None` for an empty partition, and `pruning_bitmap`
+ /// must cover every partition's keys.
+ pub(crate) fn new_routed(
+ on_columns: Vec<PhysicalExprRef>,
+ random_state: SeededRandomState,
+ router: PartitionRouter,
+ maps: Vec<Option<Arc<Map>>>,
+ description: String,
+ pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
+ ) -> Result<Self> {
+ if maps
+ .iter()
+ .flatten()
+ .any(|map| !matches!(map.as_ref(), Map::HashMap(_)))
+ {
+ return internal_err!("routed lookups only support hash table
partitions");
+ }
+ Ok(Self {
+ on_columns,
+ random_state,
+ lookup: Lookup::Routed {
+ router,
+ maps: maps.into(),
+ },
+ description,
+ pruning_bitmap,
+ })
+ }
+
/// Which parts of this build side's key range hold keys, for container
pruning.
pub fn pruning_bitmap(&self) -> Option<&Arc<KeyRangeBitmap>> {
self.pruning_bitmap.as_ref()
}
+
+ /// Membership of each row of `keys` in the hash table of its routed
partition.
+ fn evaluate_routed(
+ &self,
+ keys: &[ArrayRef],
+ router: &PartitionRouter,
+ maps: &[Option<Arc<Map>>],
+ ) -> Result<BooleanArray> {
+ let num_rows = keys[0].len();
+ let mut partition_ids = Vec::with_capacity(num_rows);
+ router.route(keys, maps.len(), &mut partition_ids)?;
+
+ with_hashes(keys, self.random_state.random_state(), |hashes| {
Review Comment:
Do you think we could add a proptest against some naive implementation (even
if slow, ideally re-using a known working implementation) to check for
correctness here?
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -803,6 +872,148 @@ impl SharedBuildAccumulator {
.update(self.preserve_probe_nulls(filter_expr, keys_have_null)?)
}
+ /// One membership filter over every partition's build keys, replacing the
+ /// `CASE` over per-partition filters.
Review Comment:
Generally I prefer that comment don't narrate history (that's what git
history is for) unless it's a footgun that would repeat itself if not called
out.
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -312,16 +369,101 @@ impl HashTableLookupExpr {
Self {
on_columns,
random_state,
- map,
+ lookup: Lookup::Single(map),
description,
pruning_bitmap,
}
}
+ /// Looks each probe row up in the hash table of the build partition
+ /// `router` assigns it to. `maps` holds one entry per partition, each a
+ /// [`Map::HashMap`] or `None` for an empty partition, and `pruning_bitmap`
+ /// must cover every partition's keys.
+ pub(crate) fn new_routed(
Review Comment:
Should we deprecate `HashTableLookupExpr::new` or just remove it (it was
never meant to be public anyway)?
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -685,6 +712,20 @@ impl SharedBuildAccumulator {
}) => {
real_partition_ids.push(partition_id);
keys_have_null |= partition_keys_have_null;
+ lookups = match (&pushdown, lookups) {
+ (PushdownStrategy::Map(map, _), Some(mut lookups))
+ if matches!(map.as_ref(), Map::HashMap(_)) =>
+ {
+
lookups.push(PartitionLookup::Map(Arc::clone(map)));
+ Some(lookups)
+ }
+ _ => None,
+ };
+ merged_bounds = match (merged_bounds,
real_partition_ids.len()) {
+ (None, 1) => Some(bounds.clone()),
+ (Some(merged), _) => merged.union(&bounds),
Review Comment:
I'd like to know if it's expected to fail. It might still make sense to not
abort the query and e.g. just log it, but I guess that depends on if it's
expected or impossible. It looks like it's only the case when two partitions
have bounds for a different number of columns, which I think should never
happen. So maybe an unwrap / debug assert is good enough.
##########
datafusion/physical-plan/src/joins/key_range_bitmap.rs:
##########
@@ -101,6 +114,67 @@ impl KeyRangeBitmap {
})
}
+ /// One bitmap over `[min, max]` marking every bucket that any of `parts`
+ /// may hold a key in. `[min, max]` must cover every part.
+ ///
+ /// Coarser than the parts' own bitmaps where its buckets are wider, but
+ /// never clears a bucket a part may hold a key in. `None` when the key
+ /// type has no `u64` ordering, or when nothing can be excluded.
+ pub(crate) fn union(
Review Comment:
Looks like a good candidate for a proptest as well
--
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]
