optimization: join probe returns first item directly (hydro-project#874)

* optimization: join probe returns first item directly

* update comments

Author: zzlk

hydroflow/src/compiled/pull/half_join_state/mod.rs

@@ -15,7 +15,14 @@ pub use set::HalfSetJoinState;
 pub type SetJoinState<Key, V1, V2> = (HalfSetJoinState<Key, V1, V2>, HalfSetJoinState<Key, V2, V1>);
 
 pub trait HalfJoinState<Key, ValBuild, ValProbe> {
+    /// Insert a key value pair into the join state, currently this is always inserting into a hash table
+    /// If the key-value pair exists then it is implementation defined what hapepns, usually either two copies are stored or only one copy is stored.
     fn build(&mut self, k: Key, v: &ValBuild) -> bool;
-    fn probe(&mut self, k: &Key, v: &ValProbe);
+
+    /// This function does the actual joining part of the join. It looks up a key in the local join state and creates matches
+    /// The first match is return directly to the caller, and any additional matches are stored internally to be retrieved later with `pop_match`
+    fn probe(&mut self, k: &Key, v: &ValProbe) -> Option<(Key, ValProbe, ValBuild)>;
+
+    /// If there are any stored matches from previous calls to probe then this function will remove them one at a time and return it.
     fn pop_match(&mut self) -> Option<(Key, ValProbe, ValBuild)>;
 }

hydroflow/src/compiled/pull/half_join_state/multiset.rs

@@ -57,17 +57,21 @@ where
         true
     }
 
-    fn probe(&mut self, k: &Key, v: &ValProbe) {
-        if let Some(entry) = self.table.get(k) {
-            // TODO: We currently don't free/shrink the self.current_matches vecdeque to save time.
-            // This mean it will grow to eventually become the largest number of matches in a single probe call.
-            // Maybe we should clear this memory at the beginning of every tick/periodically?
-            self.current_matches.extend(
-                entry
-                    .iter()
-                    .map(|valbuild| (k.clone(), v.clone(), valbuild.clone())),
-            );
-        }
+    fn probe(&mut self, k: &Key, v: &ValProbe) -> Option<(Key, ValProbe, ValBuild)> {
+        // TODO: We currently don't free/shrink the self.current_matches vecdeque to save time.
+        // This mean it will grow to eventually become the largest number of matches in a single probe call.
+        // Maybe we should clear this memory at the beginning of every tick/periodically?
+        let mut iter = self
+            .table
+            .get(k)?
+            .iter()
+            .map(|valbuild| (k.clone(), v.clone(), valbuild.clone()));
+
+        let first = iter.next();
+
+        self.current_matches.extend(iter);
+
+        first
     }
 
     fn pop_match(&mut self) -> Option<(Key, ValProbe, ValBuild)> {

hydroflow/src/compiled/pull/half_join_state/set.rs

@@ -62,17 +62,21 @@ where
         false
     }
 
-    fn probe(&mut self, k: &Key, v: &ValProbe) {
-        if let Some(entry) = self.table.get(k) {
-            // TODO: We currently don't free/shrink the self.current_matches vecdeque to save time.
-            // This mean it will grow to eventually become the largest number of matches in a single probe call.
-            // Maybe we should clear this memory at the beginning of every tick/periodically?
-            self.current_matches.extend(
-                entry
-                    .iter()
-                    .map(|valbuild| (k.clone(), v.clone(), valbuild.clone())),
-            );
-        }
+    fn probe(&mut self, k: &Key, v: &ValProbe) -> Option<(Key, ValProbe, ValBuild)> {
+        // TODO: We currently don't free/shrink the self.current_matches vecdeque to save time.
+        // This mean it will grow to eventually become the largest number of matches in a single probe call.
+        // Maybe we should clear this memory at the beginning of every tick/periodically?
+        let mut iter = self
+            .table
+            .get(k)?
+            .iter()
+            .map(|valbuild| (k.clone(), v.clone(), valbuild.clone()));
+
+        let first = iter.next();
+
+        self.current_matches.extend(iter);
+
+        first
     }
 
     fn pop_match(&mut self) -> Option<(Key, ValProbe, ValBuild)> {

hydroflow/src/compiled/pull/symmetric_hash_join.rs

@@ -40,13 +40,17 @@ where
 
             if let Some((k, v1)) = self.lhs.next() {
                 if self.lhs_state.build(k.clone(), &v1) {
-                    self.rhs_state.probe(&k, &v1);
+                    if let Some((k, v1, v2)) = self.rhs_state.probe(&k, &v1) {
+                        return Some((k, (v1, v2)));
+                    }
                 }
                 continue;
             }
             if let Some((k, v2)) = self.rhs.next() {
                 if self.rhs_state.build(k.clone(), &v2) {
-                    self.lhs_state.probe(&k, &v2);
+                    if let Some((k, v2, v1)) = self.lhs_state.probe(&k, &v2) {
+                        return Some((k, (v1, v2)));
+                    }
                 }
                 continue;
             }