From e93f29bb4e9ba7bced2ac83d1121c8e6a3feb148 Mon Sep 17 00:00:00 2001
From: "Heres, Daniel" <danielheres@gmail.com>
Date: Thu, 31 Dec 2020 11:02:43 +0100
Subject: [PATCH 1/2] Extend take implementation to more index types

---
 rust/arrow/src/compute/kernels/take.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/rust/arrow/src/compute/kernels/take.rs b/rust/arrow/src/compute/kernels/take.rs
index c8a2a02241e..85567ad8cee 100644
--- a/rust/arrow/src/compute/kernels/take.rs
+++ b/rust/arrow/src/compute/kernels/take.rs
@@ -76,12 +76,16 @@ macro_rules! downcast_dict_take {
 /// # Ok(())
 /// # }
 /// ```
-pub fn take(
+pub fn take<IndexType>(
     values: &Array,
-    indices: &UInt32Array,
+    indices: &PrimitiveArray<IndexType>,
     options: Option<TakeOptions>,
-) -> Result<ArrayRef> {
-    take_impl::<UInt32Type>(values, indices, options)
+) -> Result<ArrayRef>
+where
+    IndexType: ArrowNumericType,
+    IndexType::Native: ToPrimitive,
+{
+    take_impl(values, indices, options)
 }
 
 fn take_impl<IndexType>(

From 571bd653a4202963057b20ddd4a8197564f17c12 Mon Sep 17 00:00:00 2001
From: "Heres, Daniel" <danielheres@gmail.com>
Date: Thu, 31 Dec 2020 11:31:21 +0100
Subject: [PATCH 2/2] Specify type after collect

---
 rust/datafusion/src/physical_plan/hash_join.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/rust/datafusion/src/physical_plan/hash_join.rs b/rust/datafusion/src/physical_plan/hash_join.rs
index 8a3ee4b09ba..ea3b9b00398 100644
--- a/rust/datafusion/src/physical_plan/hash_join.rs
+++ b/rust/datafusion/src/physical_plan/hash_join.rs
@@ -287,7 +287,8 @@ fn build_batch_from_indices(
     // 2. based on the pick, `take` items from the different recordBatches
     let mut columns: Vec<Arc<dyn Array>> = Vec::with_capacity(schema.fields().len());
 
-    let right_indices = indices.iter().map(|(_, join_index)| join_index).collect();
+    let right_indices: UInt32Array =
+        indices.iter().map(|(_, join_index)| join_index).collect();
 
     for field in schema.fields() {
         // pick the column (left or right) based on the field name.