Reuse indexed pages for compact Take ranges and sliced execution

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

encodings/pco/src/array.rs

@@ -66,6 +66,7 @@ use crate::PcoLogicalRowPrefix;
 use crate::PcoMetadata;
 use crate::PcoPageInfo;
 use crate::page_index::PageIndex;
+use crate::take::slice as decode_slice;
 
 // Overall approach here:
 // Chunk the array into Pco chunks (currently using the default recommended size
@@ -278,6 +279,12 @@ impl VTable for Pco {
     }
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        if array.metadata.logical_row_prefix.is_some()
+            && (array.slice_start() != 0 || array.slice_stop() != array.unsliced_n_rows()) {
+            return Ok(ExecutionResult::done(
+                decode_slice(array.as_view(), ctx)?.into_array(),
+            ));
+        }
         let unsliced_validity = array.unsliced_validity();
         Ok(ExecutionResult::done(
             array

encodings/pco/src/decoder.rs

@@ -9,14 +9,19 @@ use std::collections::btree_map::Entry;
 use pco::data_types::Number;
 use pco::wrapped::ChunkDecompressor;
 use pco::wrapped::FileDecompressor;
+use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::dtype::NativePType;
 use vortex_array::validity::Validity;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 
+use crate::Pco;
 use crate::array::vortex_err_from_pco;
+use crate::page_index::ChunkCursor;
+use crate::page_index::IndexedView;
 
 pub(crate) struct PageDecoder<T: Number> {
     file: FileDecompressor,
@@ -71,3 +76,51 @@ impl<T: Number + NativePType> PageDecoder<T> {
         Ok(())
     }
 }
+
+/// Operation-owned decoded pages for compact repeated/range selections. The physical
+/// 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>,
+    chunks: ChunkCursor,
+    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())?,
+            chunks: ChunkCursor::default(),
+            pages: BTreeMap::new(),
+        })
+    }
+    pub fn values(
+        &mut self,
+        view: IndexedView<'_>,
+        array: ArrayView<'_, Pco>,
+        page: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<&PrimitiveArray> {
+        let entry = match self.pages.entry(page) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => {
+                let info = view.physical(page, &mut self.chunks)?;
+                let buffer = array
+                    .pages
+                    .get(page)
+                    .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?;
+                let chunk = array
+                    .chunk_metas
+                    .get(info.chunk)
+                    .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?;
+                entry.insert(self.decoder.read(
+                    info.chunk,
+                    chunk.as_slice(),
+                    info.n_values,
+                    buffer.as_slice(),
+                    ctx,
+                )?)
+            }
+        };
+        Ok(entry)
+    }
+}