From e12f5b61c7b834160c7f55ca0bad7e2172fb44f5 Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Fri, 24 Mar 2023 11:09:09 +0000 Subject: [PATCH] Add PrimitiveArray slice (#3880) --- arrow-arith/src/aggregate.rs | 7 +++---- arrow-arith/src/arithmetic.rs | 13 ++++++------- arrow-arith/src/arity.rs | 6 ++---- arrow-arith/src/boolean.rs | 8 ++++---- arrow-array/src/array/primitive_array.rs | 9 +++++++-- arrow-ipc/src/writer.rs | 8 ++------ arrow-ord/src/comparison.rs | 5 ++--- arrow-select/src/concat.rs | 2 +- arrow-select/src/filter.rs | 12 +++++------- arrow-select/src/nullif.rs | 3 +-- arrow-select/src/take.rs | 8 +------- 11 files changed, 34 insertions(+), 47 deletions(-) diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs index 9e9d9333fdcb..d5dd7a206721 100644 --- 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::(); - 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); } diff --git a/arrow-arith/src/arithmetic.rs b/arrow-arith/src/arithmetic.rs index 7d60d131bf52..b7174711fd5a 100644 --- a/arrow-arith/src/arithmetic.rs +++ b/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::(a, 3).unwrap(); + let actual = modulus_scalar_dyn::(&a, 3).unwrap(); let actual = actual.as_primitive::(); 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 index 501a240f37d5..d69bbde8d056 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -499,7 +499,6 @@ where mod tests { use super::*; use arrow_array::builder::*; - use arrow_array::cast::*; use arrow_array::types::*; #[test] @@ -507,14 +506,13 @@ mod tests { 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::().unwrap(), diff --git a/arrow-arith/src/boolean.rs b/arrow-arith/src/boolean.rs index eaef1378258b..258d683ad71a 100644 --- 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/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index bc62677c738b..92fedfc1ca5d 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -408,6 +408,12 @@ impl PrimitiveArray { 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 Array for PrimitiveArray { } 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-ipc/src/writer.rs b/arrow-ipc/src/writer.rs index 2d859f608387..07d4b0fe9f93 100644 --- a/arrow-ipc/src/writer.rs +++ b/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::>() - ); + assert_eq!(vec![Some(2), Some(3)], sliced.iter().collect::>()); 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"); diff --git a/arrow-ord/src/comparison.rs b/arrow-ord/src/comparison.rs index aa2f1416d83d..ca859ed2a719 100644 --- a/arrow-ord/src/comparison.rs +++ b/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)]) diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs index 7d42584514f1..b87b52d9a5c0 100644 --- a/arrow-select/src/concat.rs +++ b/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::::from(vec![ Some(-1), diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 567aaa58e8bf..f71a3cbc2ab0 100644 --- a/arrow-select/src/filter.rs +++ b/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::().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::().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] diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index a1b9c0e3e183..1d2b6aec589f 100644 --- a/arrow-select/src/nullif.rs +++ b/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::(); 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); } } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 83b58519fdb8..ebac5a58dd7f 100644 --- a/arrow-select/src/take.rs +++ 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::() - .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); }