⚡ Bolt: Optimize Hash Join with Single-Batch Fast Path

.jules/bolt.md

@@ -0,0 +1,7 @@
+# Bolt's Journal ⚡
+
+This journal is for critical learnings only. Before starting, I read this file. I only add entries for critical learnings that will help me avoid mistakes or make better decisions.
+
+## YYYY-MM-DD - [Title]
+**Learning:** [Insight]
+**Action:** [How to apply next time]

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

@@ -1653,32 +1653,53 @@ async fn collect_left_input(
                 Box::new(JoinHashMapU32::with_capacity(num_rows))
             };
 
-            let mut hashes_buffer = Vec::new();
-            let mut offset = 0;
-
-            let batches_iter = batches.iter().rev();
-
-            // Updating hashmap starting from the last batch
-            for batch in batches_iter.clone() {
-                hashes_buffer.clear();
-                hashes_buffer.resize(batch.num_rows(), 0);
+            let (batch, left_values) = if batches.len() == 1 {
+                // FAST PATH: if there is only one batch on the build side, we can avoid the
+                // expensive `concat_batches` call.
+                let batch = batches.into_iter().next().unwrap();
+                let mut hashes_buffer = vec![0; batch.num_rows()];
                 update_hash(
                     &on_left,
-                    batch,
+                    &batch,
                     &mut *hashmap,
-                    offset,
+                    0, // offset
                     &random_state,
                     &mut hashes_buffer,
                     0,
                     true,
                 )?;
-                offset += batch.num_rows();
-            }
 
-            // Merge all batches into a single batch, so we can directly index into the arrays
-            let batch = concat_batches(&schema, batches_iter.clone())?;
+                let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
+                (batch, left_values)
+            } else {
+                let mut hashes_buffer = Vec::new();
+                let mut offset = 0;
+
+                let batches_iter = batches.iter().rev();
+
+                // Updating hashmap starting from the last batch
+                for batch in batches_iter.clone() {
+                    hashes_buffer.clear();
+                    hashes_buffer.resize(batch.num_rows(), 0);
+                    update_hash(
+                        &on_left,
+                        batch,
+                        &mut *hashmap,
+                        offset,
+                        &random_state,
+                        &mut hashes_buffer,
+                        0,
+                        true,
+                    )?;
+                    offset += batch.num_rows();
+                }
+
+                // Merge all batches into a single batch, so we can directly index into the arrays
+                let batch = concat_batches(&schema, batches.iter().rev())?;
 
-            let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
+                let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
+                (batch, left_values)
+            };
 
             (Map::HashMap(hashmap), batch, left_values)
         };