From e57f0bd255cd48c83ffa6c6100aa380598e33fda Mon Sep 17 00:00:00 2001
From: Robert Kruszewski <github@robertk.io>
Date: Tue, 20 Aug 2024 19:27:40 +0100
Subject: [PATCH] Fix slicing of ChunkedArray if end index == array length
 (#660)

---
 vortex-array/src/array/chunked/compute/mod.rs |  2 +-
 vortex-array/src/array/chunked/mod.rs         | 13 ++++++++++++-
 vortex-array/src/compute/slice.rs             |  2 +-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/vortex-array/src/array/chunked/compute/mod.rs b/vortex-array/src/array/chunked/compute/mod.rs
index dd8b46111c..359116ad72 100644
--- a/vortex-array/src/array/chunked/compute/mod.rs
+++ b/vortex-array/src/array/chunked/compute/mod.rs
@@ -62,7 +62,7 @@ impl CastFn for ChunkedArray {
 impl CompareFn for ChunkedArray {
     fn compare(&self, array: &Array, operator: Operator) -> VortexResult<Array> {
         let mut idx = 0;
-        let mut compare_chunks = Vec::default();
+        let mut compare_chunks = Vec::with_capacity(self.nchunks());
 
         for chunk in self.chunks() {
             let sliced = slice(array, idx, idx + chunk.len())?;
diff --git a/vortex-array/src/array/chunked/mod.rs b/vortex-array/src/array/chunked/mod.rs
index 4c7cd23422..4f51b2b04c 100644
--- a/vortex-array/src/array/chunked/mod.rs
+++ b/vortex-array/src/array/chunked/mod.rs
@@ -97,7 +97,13 @@ impl ChunkedArray {
                 panic!("Search sorted failed in find_chunk_idx: {}", err);
             });
         let index_chunk = match search_result {
-            SearchResult::Found(i) => i,
+            SearchResult::Found(i) => {
+                if i == self.nchunks() {
+                    i - 1
+                } else {
+                    i
+                }
+            }
             SearchResult::NotFound(i) => i - 1,
         };
         let chunk_start = &scalar_at(&self.chunk_ends(), index_chunk)
@@ -242,6 +248,11 @@ mod test {
         assert_equal_slices(slice(chunked_array().array(), 7, 8).unwrap(), &[8u64]);
     }
 
+    #[test]
+    pub fn slice_exactly_end() {
+        assert_equal_slices(slice(chunked_array().array(), 6, 9).unwrap(), &[7u64, 8, 9]);
+    }
+
     #[test]
     fn test_scalar_subtract() {
         let chunked = chunked_array();
diff --git a/vortex-array/src/compute/slice.rs b/vortex-array/src/compute/slice.rs
index 149333bfad..02c1bec983 100644
--- a/vortex-array/src/compute/slice.rs
+++ b/vortex-array/src/compute/slice.rs
@@ -40,7 +40,7 @@ fn check_slice_bounds(array: &Array, start: usize, stop: usize) -> VortexResult<
         vortex_bail!(OutOfBounds: stop, 0, array.len());
     }
     if start > stop {
-        vortex_bail!("start ({}) must be <= stop ({})", start, stop);
+        vortex_bail!("start ({start}) must be <= stop ({stop})");
     }
     Ok(())
 }