ARROW-11005: [Rust] Remove indirection from take kernel

This PR is a small cleanup of the `take` kernel.

1. replaces `&Arc<Array>` by `&Array`. This is a small change in the API, but makes it more obvious that the only requirements for `take` is implementing the trait `Array`.
2. internal concrete implementations of `take` are no longer `dyn`, and instead expect the corresponding array types and return the corresponding array types.
3. Clarified in the documentation that not performing bound checks can lead to undefined behavior.
4. Added equality for `FixedSizeListArray`, that was missing.

The rational for this PR is that it makes it more obvious the design of the module: each array type has its own implementation (e.g. `take_primitive<T>(array: &PrimitiveArray<T>, indices: ...) -> PrimitiveArray<T>`), and there is one `dyn` implementation, `take(array: &Array, indices: ...) -> ArrayRef`, that uses each type-specific implementation and wraps it on an `Arc<Array>` for the `dyn` behavior. This is mostly for simplicity and readability, but we could expose these methods publicly.

Closes #8985 from jorgecarleitao/simpler_take

Authored-by: Jorge C. Leitao <[email protected]>
Signed-off-by: Andrew Lamb <[email protected]>

Author: jorgecarleitao

rust/arrow/benches/take_kernels.rs

@@ -22,8 +22,6 @@ use criterion::Criterion;
 use rand::distributions::{Alphanumeric, Distribution, Standard};
 use rand::Rng;
 
-use std::sync::Arc;
-
 extern crate arrow;
 
 use arrow::array::*;
@@ -32,36 +30,32 @@ use arrow::datatypes::*;
 use arrow::util::test_util::seedable_rng;
 
 // cast array from specified primitive array type to desired data type
-fn create_primitive<T>(size: usize) -> ArrayRef
+fn create_primitive<T>(size: usize) -> PrimitiveArray<T>
 where
     T: ArrowPrimitiveType,
     Standard: Distribution<T::Native>,
     PrimitiveArray<T>: std::convert::From<Vec<T::Native>>,
 {
-    let array: PrimitiveArray<T> = seedable_rng()
+    seedable_rng()
         .sample_iter(&Standard)
         .take(size)
         .map(Some)
-        .collect();
-
-    Arc::new(array) as ArrayRef
+        .collect()
 }
 
 // cast array from specified primitive array type to desired data type
-fn create_boolean(size: usize) -> ArrayRef
+fn create_boolean(size: usize) -> BooleanArray
 where
     Standard: Distribution<bool>,
 {
-    let array: BooleanArray = seedable_rng()
+    seedable_rng()
         .sample_iter(&Standard)
         .take(size)
         .map(Some)
-        .collect();
-
-    Arc::new(array) as ArrayRef
+        .collect()
 }
 
-fn create_strings(size: usize, null_density: f32) -> ArrayRef {
+fn create_strings(size: usize, null_density: f32) -> StringArray {
     let rng = &mut seedable_rng();
 
     let mut builder = StringBuilder::new(size);
@@ -74,7 +68,7 @@ fn create_strings(size: usize, null_density: f32) -> ArrayRef {
             builder.append_null().unwrap()
         }
     }
-    Arc::new(builder.finish())
+    builder.finish()
 }
 
 fn create_random_index(size: usize, null_density: f32) -> UInt32Array {
@@ -91,8 +85,8 @@ fn create_random_index(size: usize, null_density: f32) -> UInt32Array {
     builder.finish()
 }
 
-fn bench_take(values: &ArrayRef, indices: &UInt32Array) {
-    criterion::black_box(take(&values, &indices, None).unwrap());
+fn bench_take(values: &dyn Array, indices: &UInt32Array) {
+    criterion::black_box(take(values, &indices, None).unwrap());
 }
 
 fn add_benchmark(c: &mut Criterion) {

rust/arrow/src/array/equal/mod.rs

@@ -21,8 +21,9 @@
 
 use super::{
     Array, ArrayData, BinaryOffsetSizeTrait, BooleanArray, DecimalArray,
-    FixedSizeBinaryArray, GenericBinaryArray, GenericListArray, GenericStringArray,
-    NullArray, OffsetSizeTrait, PrimitiveArray, StringOffsetSizeTrait, StructArray,
+    FixedSizeBinaryArray, FixedSizeListArray, GenericBinaryArray, GenericListArray,
+    GenericStringArray, NullArray, OffsetSizeTrait, PrimitiveArray,
+    StringOffsetSizeTrait, StructArray,
 };
 
 use crate::{
@@ -116,6 +117,12 @@ impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListArray<OffsetSize> {
     }
 }
 
+impl PartialEq for FixedSizeListArray {
+    fn eq(&self, other: &Self) -> bool {
+        equal(self.data().as_ref(), other.data().as_ref())
+    }
+}
+
 impl PartialEq for StructArray {
     fn eq(&self, other: &Self) -> bool {
         equal(self.data().as_ref(), other.data().as_ref())

rust/arrow/src/compute/kernels/cast.rs

@@ -1123,7 +1123,7 @@ where
                 )
             })?;
 
-    take(&cast_dict_values, u32_indicies, None)
+    take(cast_dict_values.as_ref(), u32_indicies, None)
 }
 
 /// Attempts to encode an array into an `ArrayDictionary` with index

rust/arrow/src/compute/kernels/sort.rs

@@ -38,7 +38,7 @@ use TimeUnit::*;
 ///
 pub fn sort(values: &ArrayRef, options: Option<SortOptions>) -> Result<ArrayRef> {
     let indices = sort_to_indices(values, options)?;
-    take(values, &indices, None)
+    take(values.as_ref(), &indices, None)
 }
 
 // partition indices into non-NaN and NaN
@@ -626,7 +626,7 @@ pub fn lexsort(columns: &[SortColumn]) -> Result<Vec<ArrayRef>> {
     let indices = lexsort_to_indices(columns)?;
     columns
         .iter()
-        .map(|c| take(&c.values, &indices, None))
+        .map(|c| take(c.values.as_ref(), &indices, None))
         .collect()
 }