precalculate hashes & remove iterators

Author: korowa

datafusion/physical-plan/src/joins/hash_join.rs

@@ -1039,15 +1039,14 @@ impl RecordBatchStream for HashJoinStream {
 /// Probe indices: 3, 3, 4, 5
 /// ```
 #[allow(clippy::too_many_arguments)]
-fn lookup_join_hashmap<T: JoinHashMapType>(
-    build_hashmap: &T,
+fn lookup_join_hashmap(
+    build_hashmap: &JoinHashMap,
     build_input_buffer: &RecordBatch,
     probe_batch: &RecordBatch,
     build_on: &[Column],
     probe_on: &[Column],
-    random_state: &RandomState,
     null_equals_null: bool,
-    hashes_buffer: &mut Vec<u64>,
+    hashes_buffer: &[u64],
     limit: usize,
     offset: JoinHashMapOffset,
 ) -> Result<(UInt64Array, UInt32Array, Option<JoinHashMapOffset>)> {
@@ -1063,17 +1062,8 @@ fn lookup_join_hashmap<T: JoinHashMapType>(
         })
         .collect::<Result<Vec<_>>>()?;
 
-    hashes_buffer.clear();
-    hashes_buffer.resize(probe_batch.num_rows(), 0);
-    let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?;
-
     let (mut probe_builder, mut build_builder, next_offset) = build_hashmap
-        .get_matched_indices_with_limit_offset(
-            hash_values.iter().enumerate(),
-            None,
-            limit,
-            offset,
-        );
+        .get_matched_indices_with_limit_offset(hashes_buffer, None, limit, offset);
 
     let build_indices: UInt64Array =
         PrimitiveArray::new(build_builder.finish().into(), None);
@@ -1233,6 +1223,17 @@ impl HashJoinStream {
                 self.state = HashJoinStreamState::ExhaustedProbeSide;
             }
             Some(Ok(batch)) => {
+                // Precalculate hash values for fetched batch
+                let keys_values = self
+                    .on_right
+                    .iter()
+                    .map(|c| c.evaluate(&batch)?.into_array(batch.num_rows()))
+                    .collect::<Result<Vec<_>>>()?;
+
+                self.hashes_buffer.clear();
+                self.hashes_buffer.resize(batch.num_rows(), 0);
+                create_hashes(&keys_values, &self.random_state, &mut self.hashes_buffer)?;
+
                 self.state =
                     HashJoinStreamState::ProcessProbeBatch(ProcessProbeBatchState {
                         batch,
@@ -1266,9 +1267,8 @@ impl HashJoinStream {
             &state.batch,
             &self.on_left,
             &self.on_right,
-            &self.random_state,
             self.null_equals_null,
-            &mut self.hashes_buffer,
+            &self.hashes_buffer,
             self.batch_size,
             state.offset,
         )?;
@@ -2935,18 +2935,24 @@ mod tests {
             ("c", &vec![30, 40]),
         );
 
+        // Join key column for both join sides
+        let key_column = Column::new("a", 0);
+
         let join_hash_map = JoinHashMap::new(hashmap_left, next);
-        let mut hashes_buffer = vec![0];
+
+        let right_keys_values =
+            key_column.evaluate(&right)?.into_array(right.num_rows())?;
+        let mut hashes_buffer = vec![0; right.num_rows()];
+        create_hashes(&[right_keys_values], &random_state, &mut hashes_buffer)?;
 
         let (l, r, _) = lookup_join_hashmap(
             &join_hash_map,
             &left,
             &right,
-            &[Column::new("a", 0)],
-            &[Column::new("a", 0)],
-            &random_state,
+            &[key_column.clone()],
+            &[key_column],
             false,
-            &mut hashes_buffer,
+            &hashes_buffer,
             8192,
             (0, None),
         )?;

datafusion/physical-plan/src/joins/utils.rs

@@ -236,9 +236,9 @@ pub trait JoinHashMapType {
     ///
     /// This method only compares hashes, so additional further check for actual values
     /// equality may be required.
-    fn get_matched_indices_with_limit_offset<'a>(
+    fn get_matched_indices_with_limit_offset(
         &self,
-        iter: impl Iterator<Item = (usize, &'a u64)>,
+        iter: &[u64],
         deleted_offset: Option<usize>,
         limit: usize,
         offset: JoinHashMapOffset,
@@ -251,30 +251,52 @@ pub trait JoinHashMapType {
         let mut match_indices = UInt64BufferBuilder::new(0);
 
         let mut output_tuples = 0_usize;
-        let mut next_offset = None;
 
         let hash_map: &RawTable<(u64, u64)> = self.get_map();
         let next_chain = self.get_list();
 
-        let (initial_idx, initial_next_idx) = offset;
-        'probe: for (row_idx, hash_value) in iter.skip(initial_idx) {
-            let index = if initial_next_idx.is_some() && row_idx == initial_idx {
-                // If `initial_next_idx` is zero, then input index has been processed
-                // during previous iteration, and it can be skipped now
-                if let Some(0) = initial_next_idx {
-                    continue;
+        let to_skip = match offset {
+            (initial_idx, None) => initial_idx,
+            (initial_idx, Some(0)) => initial_idx + 1,
+            (initial_idx, Some(initial_next_idx)) => {
+                let mut i = initial_next_idx - 1;
+                loop {
+                    let match_row_idx = if let Some(offset) = deleted_offset {
+                        // This arguments means that we prune the next index way before here.
+                        if i < offset as u64 {
+                            // End of the list due to pruning
+                            break;
+                        }
+                        i - offset as u64
+                    } else {
+                        i
+                    };
+                    match_indices.append(match_row_idx);
+                    input_indices.append(initial_idx as u32);
+                    output_tuples += 1;
+                    // Follow the chain to get the next index value
+                    let next = next_chain[match_row_idx as usize];
+
+                    if output_tuples >= limit {
+                        let next_offset = Some((initial_idx, Some(next)));
+                        return (input_indices, match_indices, next_offset);
+                    }
+                    if next == 0 {
+                        // end of list
+                        break;
+                    }
+                    i = next - 1;
                 }
-                // Otherwise, use `initial_next_idx` as-is
-                initial_next_idx
-            } else if let Some((_, index)) =
+
+                initial_idx + 1
+            }
+        };
+
+        let mut row_idx = to_skip;
+        for hash_value in &iter[to_skip..] {
+            if let Some((_, index)) =
                 hash_map.get(*hash_value, |(hash, _)| *hash_value == *hash)
             {
-                Some(*index)
-            } else {
-                None
-            };
-
-            if let Some(index) = index {
                 let mut i = index - 1;
                 loop {
                     let match_row_idx = if let Some(offset) = deleted_offset {
@@ -294,8 +316,8 @@ pub trait JoinHashMapType {
                     let next = next_chain[match_row_idx as usize];
 
                     if output_tuples >= limit {
-                        next_offset = Some((row_idx, Some(next)));
-                        break 'probe;
+                        let next_offset = Some((row_idx, Some(next)));
+                        return (input_indices, match_indices, next_offset);
                     }
                     if next == 0 {
                         // end of list
@@ -304,9 +326,10 @@ pub trait JoinHashMapType {
                     i = next - 1;
                 }
             }
+            row_idx += 1;
         }
 
-        (input_indices, match_indices, next_offset)
+        (input_indices, match_indices, None)
     }
 }