Add PrimitiveArray slice (apache#3880)

arrow-arith/src/aggregate.rs

@@ -1219,13 +1219,12 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = sliced_input.as_primitive::<Float64Type>();
 
-        assert_eq!(sliced_input, &input);
+        assert_eq!(&sliced_input, &input);
 
-        let actual = min(sliced_input);
+        let actual = min(&sliced_input);
         assert_eq!(actual, expected);
-        let actual = max(sliced_input);
+        let actual = max(&sliced_input);
         assert_eq!(actual, expected);
     }
 

arrow-arith/src/arithmetic.rs

@@ -2043,7 +2043,7 @@ mod tests {
     fn test_primitive_array_add_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let actual = add_scalar(a.as_primitive(), 3).unwrap();
+        let actual = add_scalar(&a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(12), Some(11), None]);
         assert_eq!(actual, expected);
     }
@@ -2073,7 +2073,7 @@ mod tests {
     fn test_primitive_array_subtract_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let actual = subtract_scalar(a.as_primitive(), 3).unwrap();
+        let actual = subtract_scalar(&a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(6), Some(5), None]);
         assert_eq!(actual, expected);
     }
@@ -2103,7 +2103,7 @@ mod tests {
     fn test_primitive_array_multiply_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let actual = multiply_scalar(a.as_primitive(), 3).unwrap();
+        let actual = multiply_scalar(&a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(27), Some(24), None]);
         assert_eq!(actual, expected);
     }
@@ -2223,7 +2223,7 @@ mod tests {
     fn test_primitive_array_divide_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let actual = divide_scalar(a.as_primitive(), 3).unwrap();
+        let actual = divide_scalar(&a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(3), Some(2), None]);
         assert_eq!(actual, expected);
     }
@@ -2246,12 +2246,11 @@ mod tests {
     fn test_int_array_modulus_scalar_sliced() {
         let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
         let a = a.slice(1, 4);
-        let a = a.as_primitive();
-        let actual = modulus_scalar(a, 3).unwrap();
+        let actual = modulus_scalar(&a, 3).unwrap();
         let expected = Int32Array::from(vec![None, Some(0), Some(2), None]);
         assert_eq!(actual, expected);
 
-        let actual = modulus_scalar_dyn::<Int32Type>(a, 3).unwrap();
+        let actual = modulus_scalar_dyn::<Int32Type>(&a, 3).unwrap();
         let actual = actual.as_primitive::<Int32Type>();
         let expected = Int32Array::from(vec![None, Some(0), Some(2), None]);
         assert_eq!(actual, &expected);

arrow-arith/src/arity.rs

@@ -499,22 +499,20 @@ where
 mod tests {
     use super::*;
     use arrow_array::builder::*;
-    use arrow_array::cast::*;
     use arrow_array::types::*;
 
     #[test]
     fn test_unary_f64_slice() {
         let input =
             Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
         let input_slice = input.slice(1, 4);
-        let input_slice: &Float64Array = input_slice.as_primitive();
-        let result = unary(input_slice, |n| n.round());
+        let result = unary(&input_slice, |n| n.round());
         assert_eq!(
             result,
             Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
         );
 
-        let result = unary_dyn::<_, Float64Type>(input_slice, |n| n + 1.0).unwrap();
+        let result = unary_dyn::<_, Float64Type>(&input_slice, |n| n + 1.0).unwrap();
 
         assert_eq!(
             result.as_any().downcast_ref::<Float64Array>().unwrap(),

arrow-arith/src/boolean.rs

@@ -775,7 +775,7 @@ mod tests {
         let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
         let a = a.slice(8, 4);
 
-        let res = is_null(a.as_ref()).unwrap();
+        let res = is_null(&a).unwrap();
 
         let expected = BooleanArray::from(vec![false, false, false, false]);
 
@@ -800,7 +800,7 @@ mod tests {
         let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
         let a = a.slice(8, 4);
 
-        let res = is_not_null(a.as_ref()).unwrap();
+        let res = is_not_null(&a).unwrap();
 
         let expected = BooleanArray::from(vec![true, true, true, true]);
 
@@ -843,7 +843,7 @@ mod tests {
         ]);
         let a = a.slice(8, 4);
 
-        let res = is_null(a.as_ref()).unwrap();
+        let res = is_null(&a).unwrap();
 
         let expected = BooleanArray::from(vec![false, true, false, true]);
 
@@ -886,7 +886,7 @@ mod tests {
         ]);
         let a = a.slice(8, 4);
 
-        let res = is_not_null(a.as_ref()).unwrap();
+        let res = is_not_null(&a).unwrap();
 
         let expected = BooleanArray::from(vec![true, false, true, false]);
 

arrow-array/src/array/primitive_array.rs

@@ -408,6 +408,12 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
         indexes.map(|opt_index| opt_index.map(|index| self.value_unchecked(index)))
     }
 
+    /// Returns a zero-copy slice of this array with the indicated offset and length.
+    pub fn slice(&self, offset: usize, length: usize) -> Self {
+        // TODO: Slice buffers directly (#3880)
+        self.data.slice(offset, length).into()
+    }
+
     /// Reinterprets this array's contents as a different data type without copying
     ///
     /// This can be used to efficiently convert between primitive arrays with the
@@ -706,8 +712,7 @@ impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
     }
 
     fn slice(&self, offset: usize, length: usize) -> ArrayRef {
-        // TODO: Slice buffers directly (#3880)
-        Arc::new(Self::from(self.data.slice(offset, length)))
+        Arc::new(self.slice(offset, length))
     }
 
     fn nulls(&self) -> Option<&NullBuffer> {

arrow-ipc/src/writer.rs

@@ -2024,15 +2024,11 @@ mod tests {
         );
 
         let sliced = array.slice(1, 2);
-        let read_sliced: &UInt32Array = sliced.as_primitive();
-        assert_eq!(
-            vec![Some(2), Some(3)],
-            read_sliced.iter().collect::<Vec<_>>()
-        );
+        assert_eq!(vec![Some(2), Some(3)], sliced.iter().collect::<Vec<_>>());
 
         let batch = RecordBatch::try_new(
             Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])),
-            vec![sliced],
+            vec![Arc::new(sliced)],
         )
         .expect("new batch");
 

arrow-ord/src/comparison.rs

@@ -2858,7 +2858,7 @@ mod tests {
             // slice and test if the dynamic array works
             let a = a.slice(0, a.len());
             let b = b.slice(0, b.len());
-            let c = $DYN_KERNEL(a.as_ref(), b.as_ref()).unwrap();
+            let c = $DYN_KERNEL(&a, &b).unwrap();
             assert_eq!(BooleanArray::from($EXPECTED), c);
 
             // test with a larger version of the same data to ensure we cover the chunked part of the comparison
@@ -2995,8 +2995,7 @@ mod tests {
     fn test_primitive_array_eq_scalar_with_slice() {
         let a = Int32Array::from(vec![Some(1), None, Some(2), Some(3)]);
         let a = a.slice(1, 3);
-        let a: &Int32Array = a.as_primitive();
-        let a_eq = eq_scalar(a, 2).unwrap();
+        let a_eq = eq_scalar(&a, 2).unwrap();
         assert_eq!(
             a_eq,
             BooleanArray::from(vec![None, Some(true), Some(false)])

arrow-select/src/concat.rs

@@ -243,7 +243,7 @@ mod tests {
             None,
         ])
         .slice(1, 3);
-        let arr = concat(&[input_1.as_ref(), input_2.as_ref()]).unwrap();
+        let arr = concat(&[&input_1, &input_2]).unwrap();
 
         let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
             Some(-1),

arrow-select/src/filter.rs

@@ -764,13 +764,12 @@ mod tests {
 
     #[test]
     fn test_filter_array_slice() {
-        let a_slice = Int32Array::from(vec![5, 6, 7, 8, 9]).slice(1, 4);
-        let a = a_slice.as_ref();
+        let a = Int32Array::from(vec![5, 6, 7, 8, 9]).slice(1, 4);
         let b = BooleanArray::from(vec![true, false, false, true]);
         // filtering with sliced filter array is not currently supported
         // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
         // let b = b_slice.as_any().downcast_ref().unwrap();
-        let c = filter(a, &b).unwrap();
+        let c = filter(&a, &b).unwrap();
         let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
         assert_eq!(2, d.len());
         assert_eq!(6, d.value(0));
@@ -868,14 +867,13 @@ mod tests {
 
     #[test]
     fn test_filter_array_slice_with_null() {
-        let a_slice =
+        let a =
             Int32Array::from(vec![Some(5), None, Some(7), Some(8), Some(9)]).slice(1, 4);
-        let a = a_slice.as_ref();
         let b = BooleanArray::from(vec![true, false, false, true]);
         // filtering with sliced filter array is not currently supported
         // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
         // let b = b_slice.as_any().downcast_ref().unwrap();
-        let c = filter(a, &b).unwrap();
+        let c = filter(&a, &b).unwrap();
         let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
         assert_eq!(2, d.len());
         assert!(d.is_null(0));
@@ -996,7 +994,7 @@ mod tests {
 
         let mask1 = BooleanArray::from(vec![Some(true), Some(true), None]);
         let out = filter(&a, &mask1).unwrap();
-        assert_eq!(&out, &a.slice(0, 2));
+        assert_eq!(out.as_ref(), &a.slice(0, 2));
     }
 
     #[test]

arrow-select/src/nullif.rs

@@ -500,7 +500,6 @@ mod tests {
 
             for (a_offset, a_length) in a_slices {
                 let a = a.slice(a_offset, a_length);
-                let a = a.as_primitive::<Int32Type>();
 
                 for i in 1..65 {
                     let b_start_offset = rng.gen_range(0..i);
@@ -512,7 +511,7 @@ mod tests {
                     let b = b.slice(b_start_offset, a_length);
                     let b = b.as_boolean();
 
-                    test_nullif(a, b);
+                    test_nullif(&a, b);
                 }
             }
         }

arrow-select/src/take.rs

@@ -1561,14 +1561,8 @@ mod tests {
             StringArray::from(vec![Some("hello"), None, Some("world"), None, Some("hi")]);
         let indices = Int32Array::from(vec![Some(0), Some(1), None, Some(0), Some(2)]);
         let indices_slice = indices.slice(1, 4);
-        let indices_slice = indices_slice
-            .as_ref()
-            .as_any()
-            .downcast_ref::<Int32Array>()
-            .unwrap();
-
         let expected = StringArray::from(vec![None, None, Some("hello"), Some("world")]);
-        let result = take(&strings, indices_slice, None).unwrap();
+        let result = take(&strings, &indices_slice, None).unwrap();
         assert_eq!(result.as_ref(), &expected);
     }