devanbenz commented on code in PR #26249:
URL: https://github.com/apache/datafusion/pull/26249#discussion_r4424806901
##########
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:
I think I could remove it, I only see it used in tests now.
##########
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:
Sounds good, adding.
--
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]