Bring back slice for ChunkedArray (#280)

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)
     }

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<OwnedArray> {
+        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::<Vec<_>>();
+        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())
+    }
+}

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<T: NativePType>(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]);
+    }
 }