From afb2f2f2cf185796e9b2d80f946004ead462da1c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20Heres?= <danielheres@gmail.com>
Date: Thu, 22 Aug 2024 22:38:39 +0200
Subject: [PATCH] Use vec instead of builder

---
 .../physical-plan/src/joins/hash_join.rs      |  8 +++---
 .../src/joins/nested_loop_join.rs             | 25 ++++++++-----------
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs
index 69177bb2a1a8..7fac23ad5557 100644
--- a/datafusion/physical-plan/src/joins/hash_join.rs
+++ b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -3166,13 +3166,13 @@ mod tests {
             (0, None),
         )?;
 
-        let mut left_ids = vec![0, 1];
+        let left_ids: UInt64Array = vec![0, 1].into();
 
-        let mut right_ids = vec![0, 1];
+        let right_ids: UInt32Array = vec![0, 1].into();
 
-        assert_eq!(left_ids.finish(), l);
+        assert_eq!(left_ids, l);
 
-        assert_eq!(right_ids.finish(), r);
+        assert_eq!(right_ids, r);
 
         Ok(())
     }
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index 04a025c93288..18de2de03192 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -40,13 +40,12 @@ use crate::{
     RecordBatchStream, SendableRecordBatchStream,
 };
 
-use arrow::array::{
-    BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder,
-};
+use arrow::array::{BooleanBufferBuilder, UInt32Array, UInt64Array};
 use arrow::compute::concat_batches;
-use arrow::datatypes::{Schema, SchemaRef};
+use arrow::datatypes::{Schema, SchemaRef, UInt64Type};
 use arrow::record_batch::RecordBatch;
 use arrow::util::bit_util;
+use arrow_array::PrimitiveArray;
 use datafusion_common::{exec_datafusion_err, JoinSide, Result, Statistics};
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
 use datafusion_execution::TaskContext;
@@ -573,23 +572,21 @@ fn join_left_and_right_batch(
             )
         })?;
 
-    let mut left_indices_builder = UInt64Builder::new();
-    let mut right_indices_builder = UInt32Builder::new();
+    let mut left_indices_builder: Vec<u64> = vec![];
+    let mut right_indices_builder: Vec<u32> = vec![];
     for (left_side, right_side) in indices {
-        left_indices_builder
-            .append_values(left_side.values(), &vec![true; left_side.len()]);
-        right_indices_builder
-            .append_values(right_side.values(), &vec![true; right_side.len()]);
+        left_indices_builder.extend(left_side.values());
+        right_indices_builder.extend(right_side.values());
     }
 
-    let left_side = left_indices_builder.finish();
-    let right_side = right_indices_builder.finish();
+    let left_side: PrimitiveArray<UInt64Type> = left_indices_builder.into();
+    let right_side = right_indices_builder.into();
     // set the left bitmap
     // and only full join need the left bitmap
     if need_produce_result_in_final(join_type) {
         let mut bitmap = visited_left_side.lock();
-        left_side.iter().flatten().for_each(|x| {
-            bitmap.set_bit(x as usize, true);
+        left_side.values().iter().for_each(|x| {
+            bitmap.set_bit(*x as usize, true);
         });
     }
     // adjust the two side indices base on the join type