Stream PCO slices and compact PiecewiseSequence range takes

encodings/pco/src/access.rs

@@ -141,4 +141,30 @@ impl AccessState {
             .value_index::<RANKED>(row - page.rows.start)
             .map(|value| (page_index, value)))
     }
+    pub fn span(
+        &mut self,
+        view: &LogicalPageIndex,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        stop: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<PageSpan> {
+        let (page_index, page) = self.page(view, array, row, ctx)?;
+        let end = stop.min(page.rows.end);
+        let start_in_page = row - page.rows.start;
+        let end_in_page = end - page.rows.start;
+        Ok(PageSpan {
+            page_index,
+            stop: end,
+            values: page.mask.boundary_rank(start_in_page)..page.mask.boundary_rank(end_in_page),
+            mask: page.mask.mask.slice(start_in_page..end_in_page),
+        })
+    }
+}
+pub(crate) struct PageSpan {
+    pub page_index: usize,
+    pub stop: usize,
+    /// Dense-value positions within this physical page, never global value ordinals.
+    pub values: Range<usize>,
+    pub mask: Mask,
 }

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::range_take::slice as decode_slice;
 
 // Overall approach here:
 // Chunk the array into Pco chunks (currently using the default recommended size
@@ -278,6 +279,13 @@ 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

@@ -50,6 +50,17 @@ impl<T: Number + NativePType> PageDecoder<T> {
         unsafe {
             values.set_len(n_values);
         }
+        self.read_into(chunk, chunk_meta, n_values, buffer, &mut values)?;
+        Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable))
+    }
+    pub fn read_into(
+        &mut self,
+        chunk: usize,
+        chunk_meta: &[u8],
+        n_values: usize,
+        buffer: &[u8],
+        output: &mut [T],
+    ) -> VortexResult<()> {
         let chunk = match self.chunks.entry(chunk) {
             Entry::Occupied(entry) => entry.into_mut(),
             Entry::Vacant(entry) => {
@@ -63,8 +74,8 @@ impl<T: Number + NativePType> PageDecoder<T> {
         let mut decoder = chunk
             .page_decompressor(buffer, n_values)
             .map_err(vortex_err_from_pco)?;
-        decoder.read(&mut values).map_err(vortex_err_from_pco)?;
-        Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable))
+        decoder.read(output).map_err(vortex_err_from_pco)?;
+        Ok(())
     }
 }
 
@@ -98,3 +109,51 @@ impl PageLookup {
         Ok((last.chunk, info.n_values as usize))
     }
 }
+
+/// 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>,
+    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(),
+        })
+    }
+    pub fn values(
+        &mut self,
+        view: &LogicalPageIndex,
+        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 (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(chunk_id)
+                    .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?;
+                entry.insert(self.decoder.read(
+                    chunk_id,
+                    chunk.as_slice(),
+                    n_values,
+                    buffer.as_slice(),
+                    ctx,
+                )?)
+            }
+        };
+        Ok(entry)
+    }
+}

encodings/pco/src/lib.rs

@@ -33,6 +33,7 @@ mod compute;
 mod decoder;
 mod page_index;
 mod probe;
+mod range_take;
 mod rules;
 mod slice;
 mod take;

encodings/pco/src/page_index.rs

@@ -43,6 +43,9 @@ pub(crate) struct LogicalPageIndex {
     chunk_ends: Vec<usize>,
 }
 impl LogicalPageIndex {
+    pub fn page_count(&self) -> usize {
+        self.chunk_ends.last().copied().unwrap_or(0)
+    }
     pub fn find_row(&self, row: usize) -> usize {
         self.ends().partition_point(|end| *end <= row as u64)
     }