From 23822817c0a798bd57a1f073cf7c08bfa71bca05 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Jun 2024 15:50:30 +0100 Subject: [PATCH] Use new search-sorted for finding chunk index (#342) --- vortex-array/src/array/chunked/mod.rs | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/vortex-array/src/array/chunked/mod.rs b/vortex-array/src/array/chunked/mod.rs index 8d13e642aa..5c9a7175d7 100644 --- a/vortex-array/src/array/chunked/mod.rs +++ b/vortex-array/src/array/chunked/mod.rs @@ -8,7 +8,7 @@ use vortex_scalar::Scalar; use crate::array::primitive::PrimitiveArray; use crate::compute::scalar_at::scalar_at; use crate::compute::scalar_subtract::{subtract_scalar, SubtractScalarFn}; -use crate::compute::search_sorted::{search_sorted, SearchSortedSide}; +use crate::compute::search_sorted::{search_sorted, SearchResult, SearchSortedSide}; use crate::iter::{ArrayIterator, ArrayIteratorAdapter}; use crate::stream::{ArrayStream, ArrayStreamAdapter}; use crate::validity::Validity::NonNullable; @@ -73,19 +73,14 @@ impl ChunkedArray { pub fn find_chunk_idx(&self, index: usize) -> (usize, usize) { assert!(index <= self.len(), "Index out of bounds of the array"); - // TODO(ngates): migrate to the new search_sorted API to subtract 1 if not exact match. - let mut index_chunk = search_sorted(&self.chunk_ends(), index, SearchSortedSide::Left) - .unwrap() - .to_index(); - let mut chunk_start = + let index_chunk = + match search_sorted(&self.chunk_ends(), index, SearchSortedSide::Left).unwrap() { + SearchResult::Found(i) => i, + SearchResult::NotFound(i) => i - 1, + }; + let chunk_start = usize::try_from(&scalar_at(&self.chunk_ends(), index_chunk).unwrap()).unwrap(); - if chunk_start != index { - index_chunk -= 1; - chunk_start = - usize::try_from(&scalar_at(&self.chunk_ends(), index_chunk).unwrap()).unwrap(); - } - let index_in_chunk = index - chunk_start; (index_chunk, index_in_chunk) }