Add Strongly Typed Array Slice (#3929) (#3930)

* Add PrimitiveArray slice (#3880) * Add ByteArray::slice (#3929) * Add strongly typed Array::slice (#3929)
apache · Mar 28, 2023 · a667af8 · a667af8
1 parent b05522f
commit a667af8
Show file tree

Hide file tree

Showing 25 changed files with 127 additions and 99 deletions.
diff --git a/arrow-arith/src/aggregate.rs b/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);
     }
 
@@ -1265,13 +1264,12 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = sliced_input.as_string::<i32>();
 
-        assert_eq!(sliced_input, &input);
+        assert_eq!(&sliced_input, &input);
 
-        let actual = min_string(sliced_input);
+        let actual = min_string(&sliced_input);
         assert_eq!(actual, expected);
-        let actual = max_string(sliced_input);
+        let actual = max_string(&sliced_input);
         assert_eq!(actual, expected);
     }
 
@@ -1288,13 +1286,12 @@ mod tests {
             .into_iter()
             .collect();
         let sliced_input = sliced_input.slice(4, 2);
-        let sliced_input = sliced_input.as_binary::<i32>();
 
-        assert_eq!(sliced_input, &input);
+        assert_eq!(&sliced_input, &input);
 
-        let actual = min_binary(sliced_input);
+        let actual = min_binary(&sliced_input);
         assert_eq!(actual, expected);
-        let actual = max_binary(sliced_input);
+        let actual = max_binary(&sliced_input);
         assert_eq!(actual, expected);
     }
 

diff --git a/arrow-arith/src/arithmetic.rs b/arrow-arith/src/arithmetic.rs
@@ -2095,7 +2095,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);
     }
@@ -2125,7 +2125,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);
     }
@@ -2155,7 +2155,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);
     }
@@ -2275,7 +2275,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);
     }
@@ -2298,12 +2298,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);

diff --git a/arrow-arith/src/arity.rs b/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(),

diff --git a/arrow-arith/src/boolean.rs b/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]);
 

diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs
@@ -134,6 +134,12 @@ impl<T: ByteArrayType> GenericByteArray<T> {
         ArrayIter::new(self)
     }
 
+    /// 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()
+    }
+
     /// Returns `GenericByteBuilder` of this byte array for mutating its values if the underlying
     /// offset and data buffers are not shared by others.
     pub fn into_builder(self) -> Result<GenericByteBuilder<T>, Self> {
@@ -247,8 +253,7 @@ impl<T: ByteArrayType> Array for GenericByteArray<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> {

diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs
@@ -323,6 +323,12 @@ impl<K: ArrowDictionaryKeyType> DictionaryArray<K> {
         self.keys.is_valid(i).then(|| self.keys.value(i).as_usize())
     }
 
+    /// 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()
+    }
+
     /// Downcast this dictionary to a [`TypedDictionaryArray`]
     ///
     /// ```
@@ -601,8 +607,7 @@ impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<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> {
@@ -693,7 +698,7 @@ impl<'a, K: ArrowDictionaryKeyType, V: Sync> Array for TypedDictionaryArray<'a,
     }
 
     fn slice(&self, offset: usize, length: usize) -> ArrayRef {
-        self.dictionary.slice(offset, length)
+        Arc::new(self.dictionary.slice(offset, length))
     }
 
     fn nulls(&self) -> Option<&NullBuffer> {

diff --git a/arrow-array/src/array/fixed_size_binary_array.rs b/arrow-array/src/array/fixed_size_binary_array.rs
@@ -110,6 +110,12 @@ impl FixedSizeBinaryArray {
         self.data.buffers()[0].clone()
     }
 
+    /// 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()
+    }
+
     /// Create an array from an iterable argument of sparse byte slices.
     /// Sparsity means that items returned by the iterator are optional, i.e input argument can
     /// contain `None` items.
@@ -473,8 +479,7 @@ impl Array for FixedSizeBinaryArray {
     }
 
     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> {

diff --git a/arrow-array/src/array/fixed_size_list_array.rs b/arrow-array/src/array/fixed_size_list_array.rs
@@ -106,6 +106,12 @@ impl FixedSizeListArray {
         i as i32 * self.length
     }
 
+    /// 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()
+    }
+
     /// Creates a [`FixedSizeListArray`] from an iterator of primitive values
     /// # Example
     /// ```
@@ -216,8 +222,7 @@ impl Array for FixedSizeListArray {
     }
 
     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> {

diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
@@ -132,6 +132,12 @@ impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
         }
     }
 
+    /// 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()
+    }
+
     /// Creates a [`GenericListArray`] from an iterator of primitive values
     /// # Example
     /// ```
@@ -253,8 +259,7 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
     }
 
     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> {

diff --git a/arrow-array/src/array/map_array.rs b/arrow-array/src/array/map_array.rs
@@ -95,6 +95,12 @@ impl MapArray {
         let offsets = self.value_offsets();
         offsets[i + 1] - offsets[i]
     }
+
+    /// 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()
+    }
 }
 
 impl From<ArrayData> for MapArray {
@@ -222,8 +228,7 @@ impl Array for MapArray {
     }
 
     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> {

diff --git a/arrow-array/src/array/mod.rs b/arrow-array/src/array/mod.rs
@@ -141,7 +141,7 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     /// // Make slice over the values [2, 3, 4]
     /// let array_slice = array.slice(1, 3);
     ///
-    /// assert_eq!(array_slice.as_ref(), &Int32Array::from(vec![2, 3, 4]));
+    /// assert_eq!(&array_slice, &Int32Array::from(vec![2, 3, 4]));
     /// ```
     fn slice(&self, offset: usize, length: usize) -> ArrayRef;
 

diff --git a/arrow-array/src/array/null_array.rs b/arrow-array/src/array/null_array.rs
@@ -54,6 +54,12 @@ impl NullArray {
         let array_data = unsafe { array_data.build_unchecked() };
         NullArray::from(array_data)
     }
+
+    /// 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()
+    }
 }
 
 impl Array for NullArray {
@@ -74,8 +80,7 @@ impl Array for NullArray {
     }
 
     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> {

diff --git a/arrow-array/src/array/primitive_array.rs b/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> {

diff --git a/arrow-array/src/array/run_array.rs b/arrow-array/src/array/run_array.rs
@@ -253,6 +253,12 @@ impl<R: RunEndIndexType> RunArray<R> {
         }
         Ok(physical_indices)
     }
+
+    /// 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()
+    }
 }
 
 impl<R: RunEndIndexType> From<ArrayData> for RunArray<R> {
@@ -307,8 +313,7 @@ impl<T: RunEndIndexType> Array for RunArray<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> {
@@ -505,7 +510,7 @@ impl<'a, R: RunEndIndexType, V: Sync> Array for TypedRunArray<'a, R, V> {
     }
 
     fn slice(&self, offset: usize, length: usize) -> ArrayRef {
-        self.run_array.slice(offset, length)
+        Arc::new(self.run_array.slice(offset, length))
     }
 
     fn nulls(&self) -> Option<&NullBuffer> {