diff --git a/vortex-array/src/array/chunked/compute/mod.rs b/vortex-array/src/array/chunked/compute/mod.rs index 9057f619de..8ac65bd5ea 100644 --- a/vortex-array/src/array/chunked/compute/mod.rs +++ b/vortex-array/src/array/chunked/compute/mod.rs @@ -4,10 +4,12 @@ use vortex_scalar::Scalar; use crate::array::chunked::ChunkedArray; use crate::compute::as_contiguous::{as_contiguous, AsContiguousFn}; use crate::compute::scalar_at::{scalar_at, ScalarAtFn}; +use crate::compute::slice::SliceFn; use crate::compute::take::TakeFn; use crate::compute::ArrayCompute; use crate::{Array, OwnedArray, ToStatic}; +mod slice; mod take; impl ArrayCompute for ChunkedArray<'_> { @@ -19,6 +21,10 @@ impl ArrayCompute for ChunkedArray<'_> { Some(self) } + fn slice(&self) -> Option<&dyn SliceFn> { + Some(self) + } + fn take(&self) -> Option<&dyn TakeFn> { Some(self) } diff --git a/vortex-array/src/array/chunked/compute/slice.rs b/vortex-array/src/array/chunked/compute/slice.rs new file mode 100644 index 0000000000..fec66e85a0 --- /dev/null +++ b/vortex-array/src/array/chunked/compute/slice.rs @@ -0,0 +1,40 @@ +use vortex_error::VortexResult; + +use crate::array::chunked::ChunkedArray; +use crate::compute::slice::{slice, SliceFn}; +use crate::{ArrayDType, IntoArray, OwnedArray}; + +impl SliceFn for ChunkedArray<'_> { + fn slice(&self, start: usize, stop: usize) -> VortexResult { + let (offset_chunk, offset_in_first_chunk) = self.find_chunk_idx(start); + let (length_chunk, length_in_last_chunk) = self.find_chunk_idx(stop); + + if length_chunk == offset_chunk { + if let Some(chunk) = self.chunk(offset_chunk) { + return ChunkedArray::try_new( + vec![slice(&chunk, offset_in_first_chunk, length_in_last_chunk)?], + self.dtype().clone(), + ) + .map(|a| a.into_array()); + } + } + + let mut chunks = (offset_chunk..length_chunk + 1) + .map(|i| { + self.chunk(i) + .expect("find_chunk_idx returned an incorrect index") + }) + .collect::>(); + if let Some(c) = chunks.first_mut() { + *c = slice(c, offset_in_first_chunk, c.len())?; + } + + if length_in_last_chunk == 0 { + chunks.pop(); + } else if let Some(c) = chunks.last_mut() { + *c = slice(c, 0, length_in_last_chunk)?; + } + + ChunkedArray::try_new(chunks, self.dtype().clone()).map(|a| a.into_array()) + } +} diff --git a/vortex-array/src/array/chunked/mod.rs b/vortex-array/src/array/chunked/mod.rs index a0f185fd40..407865f88e 100644 --- a/vortex-array/src/array/chunked/mod.rs +++ b/vortex-array/src/array/chunked/mod.rs @@ -145,9 +145,9 @@ mod test { use vortex_dtype::{NativePType, PType}; use crate::array::chunked::{ChunkedArray, OwnedChunkedArray}; + use crate::compute::slice::slice; use crate::{Array, IntoArray}; - #[allow(dead_code)] fn chunked_array() -> OwnedChunkedArray { ChunkedArray::try_new( vec![ @@ -160,7 +160,6 @@ mod test { .unwrap() } - #[allow(dead_code)] fn assert_equal_slices(arr: Array, slice: &[T]) { let mut values = Vec::with_capacity(arr.len()); ChunkedArray::try_from(arr) @@ -171,29 +170,31 @@ mod test { assert_eq!(values, slice); } - // FIXME(ngates): bring back when slicing is a compute function. - // #[test] - // pub fn slice_middle() { - // assert_equal_slices(chunked_array().slice(2, 5).unwrap(), &[3u64, 4, 5]) - // } - // - // #[test] - // pub fn slice_begin() { - // assert_equal_slices(chunked_array().slice(1, 3).unwrap(), &[2u64, 3]); - // } - // - // #[test] - // pub fn slice_aligned() { - // assert_equal_slices(chunked_array().slice(3, 6).unwrap(), &[4u64, 5, 6]); - // } - // - // #[test] - // pub fn slice_many_aligned() { - // assert_equal_slices(chunked_array().slice(0, 6).unwrap(), &[1u64, 2, 3, 4, 5, 6]); - // } - // - // #[test] - // pub fn slice_end() { - // assert_equal_slices(chunked_array().slice(7, 8).unwrap(), &[8u64]); - // } + #[test] + pub fn slice_middle() { + assert_equal_slices(slice(chunked_array().array(), 2, 5).unwrap(), &[3u64, 4, 5]) + } + + #[test] + pub fn slice_begin() { + assert_equal_slices(slice(chunked_array().array(), 1, 3).unwrap(), &[2u64, 3]); + } + + #[test] + pub fn slice_aligned() { + assert_equal_slices(slice(chunked_array().array(), 3, 6).unwrap(), &[4u64, 5, 6]); + } + + #[test] + pub fn slice_many_aligned() { + assert_equal_slices( + slice(chunked_array().array(), 0, 6).unwrap(), + &[1u64, 2, 3, 4, 5, 6], + ); + } + + #[test] + pub fn slice_end() { + assert_equal_slices(slice(chunked_array().array(), 7, 8).unwrap(), &[8u64]); + } }