Implement take_in for primitive types

cargo fmt

arrow-array/src/builder/mod.rs

@@ -232,6 +232,8 @@
 //! assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![true, true, true, true, true, true, true, false]);
 //! ```
 
+mod take_in_utils;
+
 pub use arrow_buffer::BooleanBufferBuilder;
 pub use arrow_buffer::NullBufferBuilder;
 

arrow-array/src/builder/primitive_builder.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::builder::{ArrayBuilder, BufferBuilder};
-use crate::types::*;
+use crate::{types::*, Array};
 use crate::{ArrayRef, PrimitiveArray};
 use arrow_buffer::NullBufferBuilder;
 use arrow_buffer::{Buffer, MutableBuffer};
@@ -25,6 +25,8 @@ use arrow_schema::{ArrowError, DataType};
 use std::any::Any;
 use std::sync::Arc;
 
+use super::take_in_utils;
+
 /// A signed 8-bit integer array builder.
 pub type Int8Builder = PrimitiveBuilder<Int8Type>;
 /// A signed 16-bit integer array builder.
@@ -272,6 +274,23 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
         self.values_builder.append_trusted_len_iter(iter);
     }
 
+    /// Takes values at indices from array into this array
+    ///
+    /// Each index in indices needs to be in-bounds for array or null - Otherwise an incorrect
+    /// result may be returned
+    #[inline]
+    pub fn take_in<I>(&mut self, array: &PrimitiveArray<T>, indices: &PrimitiveArray<I>)
+    where
+        I: ArrowPrimitiveType,
+    {
+        take_in_utils::take_in_nulls(&mut self.null_buffer_builder, array.nulls(), &indices);
+
+        take_in_utils::take_in_native::<T::Native, I>(
+            &mut self.values_builder,
+            &array.values(),
+            &indices,
+        );
+    }
     /// Builds the [`PrimitiveArray`] and reset this builder.
     pub fn finish(&mut self) -> PrimitiveArray<T> {
         let len = self.len();

arrow-array/src/builder/take_in_utils.rs

@@ -0,0 +1,54 @@
+use arrow_buffer::{ArrowNativeType, BufferBuilder, NullBuffer, NullBufferBuilder, ScalarBuffer};
+
+use crate::{Array, ArrowPrimitiveType, PrimitiveArray};
+
+pub(crate) fn take_in_nulls<I>(
+    null_buffer_builder: &mut NullBufferBuilder,
+    array_nulls: Option<&NullBuffer>,
+    indices: &PrimitiveArray<I>,
+) where
+    I: ArrowPrimitiveType,
+{
+    let array_nulls = array_nulls.filter(|n| n.null_count() > 0);
+    let indices_nulls = indices.nulls().filter(|n| n.null_count() > 0);
+
+    match (array_nulls, indices_nulls) {
+        (None, None) => null_buffer_builder.append_n_non_nulls(indices.len()),
+        (None, Some(indices_nulls)) => null_buffer_builder.append_buffer(&indices_nulls),
+        (Some(array_nulls), None) => {
+            let iter = indices
+                .values()
+                .iter()
+                .map(|idx| array_nulls.is_valid(idx.as_usize()));
+            null_buffer_builder.append_iter(iter);
+        }
+        (Some(array_nulls), Some(_indices_nulls)) => {
+            let iter = indices.iter().map(|idx| {
+                idx.map(|idx| array_nulls.is_valid(idx.as_usize()))
+                    .unwrap_or(false)
+            });
+            null_buffer_builder.append_iter(iter);
+        }
+    }
+}
+
+pub(crate) fn take_in_native<T, I>(
+    values_builder: &mut BufferBuilder<T>,
+    array: &ScalarBuffer<T>,
+    indices: &PrimitiveArray<I>,
+) where
+    T: ArrowNativeType,
+    I: ArrowPrimitiveType,
+{
+    values_builder.reserve(indices.len());
+    if indices.null_count() > 0 {
+        values_builder.extend(
+            indices
+                .values()
+                .iter()
+                .map(|index| array.get(index.as_usize()).cloned().unwrap_or(T::default())),
+        )
+    } else {
+        values_builder.extend(indices.values().iter().map(|index| array[index.as_usize()]));
+    }
+}