Reuse chunk lookup in range selection callers

Author: danking

encodings/pco/src/decoder.rs

@@ -23,7 +23,6 @@ use crate::Pco;
 use crate::array::vortex_err_from_pco;
 use crate::page_index::LogicalPageIndex;
 use crate::page_index::PhysicalPage;
-use crate::page_index::ChunkCursor;
 
 pub(crate) struct PageDecoder<T: Number> {
     file: FileDecompressor,
@@ -114,13 +113,15 @@ impl PageLookup {
 /// descriptor is derived from the chosen mode only when a page is decoded for the first time.
 pub(crate) struct DecodedPages<T: Number> {
     decoder: PageDecoder<T>,
+    lookup: PageLookup,
     pages: BTreeMap<usize, PrimitiveArray>,
 }
 
 impl<T: Number + NativePType> DecodedPages<T> {
     pub fn new(array: ArrayView<'_, Pco>) -> VortexResult<Self> {
         Ok(Self {
             decoder: PageDecoder::new(array.metadata.header.as_slice())?,
+            lookup: PageLookup::default(),
             pages: BTreeMap::new(),
         })
     }
@@ -134,19 +135,19 @@ impl<T: Number + NativePType> DecodedPages<T> {
         let entry = match self.pages.entry(page) {
             Entry::Occupied(entry) => entry.into_mut(),
             Entry::Vacant(entry) => {
-                let info = view.physical(page)?;
+                let (chunk_id, n_values) = self.lookup.resolve(view, array, page)?;
                 let buffer = array
                     .pages
                     .get(page)
                     .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?;
                 let chunk = array
                     .chunk_metas
-                    .get(info.chunk)
+                    .get(chunk_id)
                     .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?;
                 entry.insert(self.decoder.read(
-                    info.chunk,
+                    chunk_id,
                     chunk.as_slice(),
-                    info.n_values,
+                    n_values,
                     buffer.as_slice(),
                     ctx,
                 )?)

encodings/pco/src/take.rs

@@ -672,6 +672,7 @@ fn stream_range<T: Number + NativePType>(
     }
     let mut state = AccessState::default();
     let mut decoder = PageDecoder::<T>::new(array.metadata.header.as_slice())?;
+    let mut lookup = PageLookup::default();
     let stop = array.slice_start() + range.end;
     let mut row = global_start;
     // Descriptor-count iteration bounds work even when optional index contents are bad.
@@ -686,24 +687,24 @@ fn stream_range<T: Number + NativePType>(
             .get_mut(start..start + span.mask.len())
             .ok_or_else(|| vortex_err!("Pco contiguous output exceeds visible range"))?;
         if !span.mask.all_false() {
-            let page = view.physical(span.page)?;
+            let (chunk_id, n_values) = lookup.resolve(view, array, span.page)?;
             let buffer = array
                 .pages
                 .get(span.page)
                 .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?;
             let chunk = array
                 .chunk_metas
-                .get(page.chunk)
+                .get(chunk_id)
                 .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?;
-            let n = page.n_values;
+            let n = n_values;
             if span.mask.all_true() && span.values == (0..n) {
-                decoder.read_into(page.chunk, chunk.as_slice(), n, buffer.as_slice(), output)?;
+                decoder.read_into(chunk_id, chunk.as_slice(), n, buffer.as_slice(), output)?;
             } else {
                 if scratch.len() < n {
                     scratch.push_n(T::default(), n - scratch.len());
                 }
                 decoder.read_into(
-                    page.chunk,
+                    chunk_id,
                     chunk.as_slice(),
                     n,
                     buffer.as_slice(),