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]

Reply via email to