This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch 
gh-readonly-queue/main/pr-11334-4d8b332360e2ad0ca65d95868c506258a5eee0ac
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git

commit 4ef327c91413636931e4967181319516470e0983
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Fri Oct 9 19:19:39 2026 +0000

    perf(parquet): store a ScanPlan cursor as one heap entry (#11334)
    
    Stacked on https://github.com/apache/arrow-rs/pull/10555. Review only
    the commits after ec3b0d2c50.
    
    # Which issue does this PR close?
    
    - Part of https://github.com/apache/arrow-rs/issues/6946.
    
    # Rationale for this change
    
    A caller that reads ahead asks the plan for its first range:
    
    ```rust
    let decoder = 
ParquetPushDecoderBuilder::new_with_metadata(metadata).build()?;
    let first = decoder.scan_plan().next(); // what the first fetch waits for
    ```
    
    Before the first range, the planner starts one cursor for each column of
    the stage. On a wide file this cost is O(columns), and each column used
    about 200 bytes of heap memory. This PR makes each column cost less time
    and less memory. The plan output (order and contents) does not change.
    
    Time of `scan_plan().next()`, one row group of 1,000 rows, every column,
    no selection (criterion, `push_decoder/scan_plan/first_range/all`):
    
    | Columns × pages | Before | After | Change |
    |---|---|---|---|
    | 100 × 10 | 2.08 µs | 0.84 µs | −59% |
    | 1,000 × 10 | 20.2 µs | 6.39 µs | −68% |
    | 10,000 × 10 | 179 µs | 63.9 µs | −64% |
    | 100,000 × 10 | 3.16 ms | 1.83 ms | −41% |
    | 1,000 × 100 | 17.7 µs | 6.77 µs | −62% |
    
    Heap memory of one `scan_plan().next()` call (counting global allocator,
    same files):
    
    | Columns | Allocations before | Allocations after | Bytes before |
    Bytes after |
    |---|---|---|---|---|
    | 100 | 13 | 7 | 22,420 | 4,788 |
    | 1,000 | 16 | 7 | 203,064 | 41,688 |
    | 10,000 | 20 | 7 | 2,419,104 | 410,688 |
    | 100,000 | 23 | 7 | 22,089,136 | 4,100,688 |
    
    At 100,000 columns the time is limited by memory reads of the column
    metadata and the offset index (about 1.2–1.4 ms for only these reads),
    so the gain there is smaller.
    
    # What changes are included in this PR?
    
    Terms:
    
    * **Cursor**: the state that gives the next ranges of one column chunk,
    in file order.
    * **Stage**: one decoding step of a row group (a predicate, or the
    output projection). The planner merges the cursors of a stage with a
    heap, by first row.
    
    | Commit | Change |
    |---|---|
    | bench | Adds a 100,000 column case to the `push_decoder/scan_plan`
    benchmark. |
    | perf | A cursor is now only its heap entry: the sort key of its next
    range, the position of the column, and which range of the column chunk
    is next (dictionary page, data page position, or whole chunk). When an
    entry is popped, the planner computes the end and kind of the range
    again from the page locations. A stage keeps the list of selected pages
    of a column only when the row selection removes pages. |
    
    Before, each cursor held a full `PlannedRange` for its next range (about
    128 bytes), and the heap held a copy of its sort key. The planner also
    looked up the offset index two times for each column.
    
    `ColumnFetch` and `StageSchedule` are still the only source of the bytes
    that the decoder fetches. The planner does not copy decoder logic.
    
    # Are these changes tested?
    
    Yes, by the existing tests. The tests that compare the plan order, that
    check that the plan covers every later decoder request, and the
    lazy-plan test pass. The lazy-plan test now counts heap entries instead
    of cursors, because the cursor list is gone (one entry for each
    projected column, as before).
    
    Benchmarks (`cargo bench -p parquet --features arrow,async --bench
    push_decoder -- scan_plan`, baseline saved on the bench commit, Apple
    M-series laptop). `first_range` is the first `next()`, `whole_scan` is
    `count()`:
    
    | Case | first_range before | first_range after | whole_scan before |
    whole_scan after |
    |---|---|---|---|---|
    | all/100cols_10pages | 2.08 µs | 0.84 µs | 30.1 µs | 22.6 µs |
    | selection/100cols_10pages | 23.7 µs | 21.7 µs | 61.8 µs | 51.2 µs |
    | narrow/100cols_10pages | 574 ns | 349 ns | 2.57 µs | 1.42 µs |
    | all/1000cols_10pages | 20.2 µs | 6.39 µs | 383 µs | 320 µs |
    | selection/1000cols_10pages | 227 µs | 216 µs | 642 µs | 589 µs |
    | narrow/1000cols_10pages | 1.06 µs | 647 ns | 2.70 µs | 1.71 µs |
    | all/10000cols_10pages | 179 µs | 63.9 µs | 5.26 ms | 4.42 ms |
    | selection/10000cols_10pages | 2.27 ms | 2.19 ms | 7.77 ms | 7.05 ms |
    | narrow/10000cols_10pages | 3.50 µs | 3.81 µs | 5.44 µs | 4.95 µs |
    | all/100000cols_10pages | 3.16 ms | 1.83 ms | 67.2 ms | 58.1 ms |
    | selection/100000cols_10pages | 23.7 ms | 22.8 ms | 91.3 ms | 84.3 ms |
    | narrow/100000cols_10pages | 32.1 µs | 32.1 µs | 31.1 µs | 32.3 µs |
    | all/1000cols_100pages | 17.7 µs | 6.77 µs | 3.84 ms | 3.45 ms |
    | selection/1000cols_100pages | 325 µs | 326 µs | 749 µs | 711 µs |
    | narrow/1000cols_100pages | 837 ns | 676 ns | 20.1 µs | 12.3 µs |
    
    The `narrow` cases read 10 columns, so they do not depend on this
    change. Three alternating reruns of `narrow` at 10,000 and 100,000
    columns gave equal or faster times after the change (the +9% in the
    table is noise on this machine).
    
    The `selection` cases stay slow: `ColumnFetch::new` calls
    `RowSelection::scan_ranges` for each column (about 6 allocations per
    column). This PR does not change that.
    
    # Are there any user-facing changes?
    
    No. The public API and the plan output do not change.
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    ---------
    
    Co-authored-by: Claude Opus 5.5 <[email protected]>
---
 parquet/benches/push_decoder.rs                   |   8 +-
 parquet/src/arrow/in_memory_row_group.rs          |   5 +
 parquet/src/arrow/push_decoder/scan_plan/mod.rs   | 332 ++++++++++++----------
 parquet/src/arrow/push_decoder/scan_plan/tests.rs |   2 +-
 4 files changed, 196 insertions(+), 151 deletions(-)

diff --git a/parquet/benches/push_decoder.rs b/parquet/benches/push_decoder.rs
index b43abea9af..72ad9f55dd 100644
--- a/parquet/benches/push_decoder.rs
+++ b/parquet/benches/push_decoder.rs
@@ -210,7 +210,13 @@ fn make_paged_test_file(num_columns: usize, 
pages_per_column: usize) -> Bytes {
 fn bench_scan_plan(c: &mut Criterion) {
     let mut group = c.benchmark_group("push_decoder/scan_plan");
 
-    for (num_cols, pages) in [(100, 10), (1_000, 10), (10_000, 10), (1_000, 
100)] {
+    for (num_cols, pages) in [
+        (100, 10),
+        (1_000, 10),
+        (10_000, 10),
+        (100_000, 10),
+        (1_000, 100),
+    ] {
         let file_data = make_paged_test_file(num_cols, pages);
         let options = 
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
         let metadata = ArrowReaderMetadata::load(&file_data, options).unwrap();
diff --git a/parquet/src/arrow/in_memory_row_group.rs 
b/parquet/src/arrow/in_memory_row_group.rs
index 4a849a0661..d75a81efbf 100644
--- a/parquet/src/arrow/in_memory_row_group.rs
+++ b/parquet/src/arrow/in_memory_row_group.rs
@@ -179,6 +179,7 @@ impl InMemoryRowGroup<'_> {
 ///
 /// A decoding stage of a row group does not fetch a column that an earlier
 /// stage of the same row group read.
+#[inline]
 pub(crate) fn columns_to_fetch<'a>(
     projection: &'a ProjectionMask,
     num_columns: usize,
@@ -190,6 +191,7 @@ pub(crate) fn columns_to_fetch<'a>(
 /// The selection that [`InMemoryRowGroup::fetch_ranges`] uses to choose the
 /// pages of column `idx`: `expanded_selection` (the selection expanded to
 /// batch boundaries) if `cache_mask` includes the column, else `selection`.
+#[inline]
 pub(crate) fn column_selection<'a>(
     selection: Option<&'a RowSelection>,
     expanded_selection: Option<&'a RowSelection>,
@@ -230,6 +232,7 @@ impl<'a> ColumnFetch<'a> {
     /// The fetch of the column chunk at byte range `chunk`, with page
     /// locations `locations` from its offset index, if any, and row
     /// selection `selection`, if any.
+    #[inline]
     pub(crate) fn new(
         chunk: Range<u64>,
         locations: Option<&'a [PageLocation]>,
@@ -278,6 +281,7 @@ impl<'a> ColumnFetch<'a> {
 
 /// Byte range of the dictionary page of a column chunk that starts at
 /// `chunk_start`: the bytes before the first data page, if any.
+#[inline]
 pub(crate) fn dictionary_range(chunk_start: u64, locations: &[PageLocation]) 
-> Option<Range<u64>> {
     match locations.first() {
         Some(first) if first.offset as u64 != chunk_start => 
Some(chunk_start..first.offset as u64),
@@ -286,6 +290,7 @@ pub(crate) fn dictionary_range(chunk_start: u64, locations: 
&[PageLocation]) ->
 }
 
 /// Byte range of a data page.
+#[inline]
 pub(crate) fn page_range(location: &PageLocation) -> Range<u64> {
     let start = location.offset as u64;
     start..start + location.compressed_page_size as u64
diff --git a/parquet/src/arrow/push_decoder/scan_plan/mod.rs 
b/parquet/src/arrow/push_decoder/scan_plan/mod.rs
index d763d69f19..9ec48b332f 100644
--- a/parquet/src/arrow/push_decoder/scan_plan/mod.rs
+++ b/parquet/src/arrow/push_decoder/scan_plan/mod.rs
@@ -29,6 +29,7 @@ use crate::arrow::in_memory_row_group::{
     ColumnFetch, column_selection, columns_to_fetch, dictionary_range, 
page_range,
 };
 use crate::errors::ParquetError;
+use crate::file::metadata::ParquetMetaData;
 use crate::file::metadata::page_index::PageIndexProvider;
 use crate::file::page_index::offset_index::PageLocation;
 
@@ -53,9 +54,11 @@ pub struct PlannedRange {
     /// What the range contains.
     pub kind: PageKind,
     /// First planned row that this range serves, counted from the start of
-    /// the plan. Used only to order the plan.
+    /// the plan. The tests use it to check the order of the plan.
+    #[cfg_attr(not(test), expect(dead_code))]
     pub(crate) first_row: u64,
     /// One past the last planned row that this range serves.
+    #[cfg_attr(not(test), expect(dead_code))]
     pub(crate) last_row: u64,
 }
 
@@ -400,22 +403,14 @@ impl Planner {
             // `columns_to_fetch` filters, so it gives no size hint. Reserve
             // for every column to avoid growing the vector one column at a 
time.
             let mut columns = Vec::with_capacity(num_columns);
-            columns.extend(
-                columns_to_fetch(fetch.projection, num_columns, |idx| 
planned_columns[idx]).map(
-                    |column_idx| {
-                        let (chunk_start, chunk_len) = 
row_group.column(column_idx).byte_range();
-                        StageColumn {
-                            column_idx,
-                            chunk: chunk_start..chunk_start + chunk_len,
-                        }
-                    },
-                ),
-            );
+            columns.extend(columns_to_fetch(fetch.projection, num_columns, 
|idx| {
+                planned_columns[idx]
+            }));
             // The plan keeps `columns` until the end of the row group. Release
             // the capacity that a narrow projection does not use.
             columns.shrink_to_fit();
-            for column in &columns {
-                planned_columns[column.column_idx] = true;
+            for &column_idx in &columns {
+                planned_columns[column_idx] = true;
             }
             stage_plans.push(StagePlan { stage, columns });
         }
@@ -429,6 +424,7 @@ impl Planner {
                 selection,
                 expanded_selection,
                 stages: Arc::clone(&self.columns.stages),
+                metadata: Arc::clone(metadata),
                 page_index,
             },
             stages: stage_plans.into_iter(),
@@ -475,20 +471,13 @@ impl FusedIterator for ScanPlan {}
 #[derive(Debug, Clone)]
 struct StagePlan {
     stage: ScanStage,
-    columns: Vec<StageColumn>,
-}
-
-/// A column chunk that a decoding stage reads.
-#[derive(Debug, Clone)]
-struct StageColumn {
-    column_idx: usize,
-    /// Byte range of the whole column chunk.
-    chunk: Range<u64>,
+    /// Leaf column indexes, in file order.
+    columns: Vec<usize>,
 }
 
 /// The planned ranges of one row group, in decode order.
 ///
-/// Each stage is a merge of one [`ColumnCursor`] per column chunk. A cursor
+/// Each stage is a merge of one cursor per column chunk. A cursor
 /// returns the ranges of its column chunk in file order, and their first
 /// rows do not decrease. A heap picks the cursor whose next range has the
 /// smallest [`RangeOrder`], so the ranges of a stage are ordered by first
@@ -517,56 +506,73 @@ struct RowGroupContext {
     expanded_selection: Option<RowSelection>,
     /// What each decoding stage fetches.
     stages: Arc<StageSchedule>,
+    /// The file metadata, for the byte ranges of the column chunks.
+    metadata: Arc<ParquetMetaData>,
     /// The file's page index, if it has offset indexes.
     page_index: Option<Arc<dyn PageIndexProvider>>,
 }
 
 /// The merge state of one stage.
+///
+/// A cursor of a column chunk is one heap entry: the sort key of its next
+/// range, the column and the position of that range. The other parts of the
+/// range are computed again when the entry is popped. Thus a stage stores
+/// no per-column state other than its heap entry, unless the selection
+/// removes pages.
 #[derive(Debug, Clone)]
 struct StageRanges {
-    cursors: Vec<ColumnCursor>,
-    /// The sort key of the next range of each cursor that has one, smallest
-    /// first.
-    heap: BinaryHeap<Reverse<(RangeOrder, usize)>>,
+    /// Leaf column indexes, in file order. A [`HeapEntry`] refers to a
+    /// column by its position in this list.
+    columns: Vec<usize>,
+    /// For each column in `columns`, the indexes of the data pages that the
+    /// decoder reads, if the selection removes pages. `None` if no column of
+    /// the stage has such a list.
+    selected: Option<Vec<Option<Vec<usize>>>>,
+    /// The next range of each cursor that has one, smallest first.
+    heap: BinaryHeap<Reverse<HeapEntry>>,
 }
 
 /// How the heap orders the next ranges of different column chunks: by first
 /// row, then more rows first (a dictionary page), then by file offset. The
 /// ranges of one column chunk keep file order, also a page with no planned
 /// rows that the predicate cache reads.
+///
+/// The key is `(first row, Reverse(last row), range start)`.
 type RangeOrder = (u64, Reverse<u64>, u64);
 
-fn range_order(range: &PlannedRange) -> RangeOrder {
-    (range.first_row, Reverse(range.last_row), range.range.start)
-}
+/// A heap entry: the sort key of the next range of a cursor, the position
+/// of the column in [`StageRanges::columns`], and the [`HeadPosition`] of
+/// that range. The column position breaks ties, so ranges with equal keys
+/// stay in file order of their columns.
+type HeapEntry = (RangeOrder, u32, HeadPosition);
 
-/// The ranges of one column chunk that are not yet returned.
-#[derive(Debug, Clone)]
-struct ColumnCursor {
-    column_idx: usize,
-    /// The next range of this column chunk.
-    head: Option<PlannedRange>,
-    /// Data pages the decoder reads.
-    pages: PageSet,
-    /// Position in `pages` of the next data page after `head`.
-    next_page: usize,
+/// Which range of a column chunk is next: [`CHUNK`], [`DICTIONARY`], or the
+/// position of a data page in the pages that the decoder reads.
+type HeadPosition = u32;
+
+/// The next range is the whole column chunk. The column has no offset index.
+const CHUNK: HeadPosition = u32::MAX;
+
+/// The next range is the dictionary page.
+const DICTIONARY: HeadPosition = u32::MAX - 1;
+
+/// Converts a column or page position to a heap entry field.
+fn position(position: usize) -> u32 {
+    u32::try_from(position).expect("fewer than u32::MAX - 1 columns and pages")
 }
 
 /// The data pages of a column chunk that the decoder reads, in page order.
-#[derive(Debug, Clone)]
-enum PageSet {
-    /// No offset index, so the column chunk is one range.
-    None,
-    /// Every page: there is no selection.
+#[derive(Debug, Clone, Copy)]
+enum PageSet<'a> {
+    /// Every page of the offset index, which has this many pages.
     All(usize),
     /// The pages at these indexes in the offset index.
-    Selected(Vec<usize>),
+    Selected(&'a [usize]),
 }
 
-impl PageSet {
+impl PageSet<'_> {
     fn len(&self) -> usize {
         match self {
-            Self::None => 0,
             Self::All(len) => *len,
             Self::Selected(pages) => pages.len(),
         }
@@ -575,64 +581,86 @@ impl PageSet {
     /// Index in the offset index of the page at `position`.
     fn page(&self, position: usize) -> usize {
         match self {
-            Self::None => unreachable!("no pages"),
             Self::All(_) => position,
             Self::Selected(pages) => pages[position],
         }
     }
 }
 
+impl StageRanges {
+    /// The data pages that the decoder reads of the column at `slot`, whose
+    /// page locations are `locations`.
+    fn pages(&self, slot: usize, locations: &[PageLocation]) -> PageSet<'_> {
+        match self
+            .selected
+            .as_ref()
+            .and_then(|selected| selected[slot].as_ref())
+        {
+            Some(pages) => PageSet::Selected(pages),
+            None => PageSet::All(locations.len()),
+        }
+    }
+}
+
 impl RowGroupRanges {
     fn next(&mut self) -> Option<PlannedRange> {
         loop {
             if let Some(stage) = self.stage.as_mut()
-                && let Some(Reverse((_, cursor_idx))) = stage.heap.pop()
+                && let Some(Reverse(((first_row, Reverse(last_row), start), 
slot, head))) =
+                    stage.heap.pop()
             {
-                let cursor = &mut stage.cursors[cursor_idx];
-                let range = cursor.head.take().expect("heap entries have a 
head");
-                cursor.head = self.row_group.next_data_page(cursor);
-                if let Some(head) = &cursor.head {
-                    stage.heap.push(Reverse((range_order(head), cursor_idx)));
-                }
-                return Some(range);
+                let row_group = &self.row_group;
+                let column_idx = stage.columns[slot as usize];
+                let (end, kind) = if head == CHUNK {
+                    (row_group.chunk(column_idx).end, PageKind::ColumnChunk)
+                } else {
+                    let locations = row_group
+                        .locations(column_idx)
+                        .expect("pages have locations");
+                    let pages = stage.pages(slot as usize, locations);
+                    let (end, kind, next) = if head == DICTIONARY {
+                        (locations[0].offset as u64, PageKind::Dictionary, 0)
+                    } else {
+                        let page = page_range(&locations[pages.page(head as 
usize)]);
+                        (page.end, PageKind::Data, head as usize + 1)
+                    };
+                    if let Some(order) = row_group.data_page(locations, pages, 
next) {
+                        stage.heap.push(Reverse((order, slot, 
position(next))));
+                    }
+                    (end, kind)
+                };
+                return Some(PlannedRange {
+                    range: start..end,
+                    row_group: row_group.row_group_idx,
+                    column: column_idx,
+                    kind,
+                    first_row,
+                    last_row,
+                });
             }
             let plan = self.stages.next()?;
-            let cursors: Vec<_> = plan
-                .columns
-                .iter()
-                .map(|column| self.row_group.column_cursor(column, plan.stage))
-                .collect();
-            let heap = cursors
-                .iter()
-                .enumerate()
-                .filter_map(|(idx, cursor)| {
-                    let head = cursor.head.as_ref()?;
-                    Some(Reverse((range_order(head), idx)))
-                })
-                .collect();
-            self.stage = Some(StageRanges { cursors, heap });
+            let mut selected = None;
+            let mut heap = Vec::with_capacity(plan.columns.len());
+            for (slot, &column_idx) in plan.columns.iter().enumerate() {
+                let (head, pages) = self.row_group.first_range(column_idx, 
plan.stage);
+                if let Some(pages) = pages {
+                    selected.get_or_insert_with(|| vec![None; 
plan.columns.len()])[slot] =
+                        Some(pages);
+                }
+                if let Some((order, head)) = head {
+                    heap.push(Reverse((order, position(slot), head)));
+                }
+            }
+            self.stage = Some(StageRanges {
+                columns: plan.columns,
+                selected,
+                heap: BinaryHeap::from(heap),
+            });
         }
     }
 }
 
 impl RowGroupContext {
-    fn entry(
-        &self,
-        column_idx: usize,
-        range: Range<u64>,
-        rows: Range<u64>,
-        kind: PageKind,
-    ) -> PlannedRange {
-        PlannedRange {
-            range,
-            first_row: rows.start,
-            last_row: rows.end,
-            row_group: self.row_group_idx,
-            column: column_idx,
-            kind,
-        }
-    }
-
     fn locations(&self, column_idx: usize) -> Option<&[PageLocation]> {
         self.page_index
             .as_ref()?
@@ -640,6 +668,21 @@ impl RowGroupContext {
             .map(|offset_index| offset_index.page_locations().as_slice())
     }
 
+    /// Byte range of the whole column chunk.
+    fn chunk(&self, column_idx: usize) -> Range<u64> {
+        let (start, len) = self
+            .metadata
+            .row_group(self.row_group_idx)
+            .column(column_idx)
+            .byte_range();
+        start..start + len
+    }
+
+    /// Planned rows of the whole row group.
+    fn row_group_rows(&self) -> Range<u64> {
+        self.first_row..self.first_row + 
self.rows.selected_before(self.row_count)
+    }
+
     /// Planned rows of the data page at `idx` in `locations`.
     fn page_rows(&self, locations: &[PageLocation], idx: usize) -> Range<u64> {
         let raw_end = locations
@@ -653,18 +696,32 @@ impl RowGroupContext {
         first_row..self.first_row + self.rows.selected_before(raw_end)
     }
 
-    /// Start the cursor of one column chunk.
+    /// The sort key of the data page at `position` in `pages`, if any.
+    fn data_page(
+        &self,
+        locations: &[PageLocation],
+        pages: PageSet<'_>,
+        position: usize,
+    ) -> Option<RangeOrder> {
+        if position >= pages.len() {
+            return None;
+        }
+        let idx = pages.page(position);
+        let rows = self.page_rows(locations, idx);
+        Some((rows.start, Reverse(rows.end), locations[idx].offset as u64))
+    }
+
+    /// The first range of one column chunk: its sort key and position, or
+    /// `None` if the column chunk has no range. Also returns the data pages
+    /// that the decoder reads, if the selection removes pages.
     ///
     /// The ranges are the same bytes that `InMemoryRowGroup::fetch_ranges`
     /// requests, split at page boundaries when page locations are known.
-    fn column_cursor(&self, column: &StageColumn, stage: ScanStage) -> 
ColumnCursor {
-        let StageColumn {
-            column_idx,
-            ref chunk,
-        } = *column;
-        let row_group_rows =
-            self.first_row..self.first_row + 
self.rows.selected_before(self.row_count);
-
+    fn first_range(
+        &self,
+        column_idx: usize,
+        stage: ScanStage,
+    ) -> (Option<(RangeOrder, HeadPosition)>, Option<Vec<usize>>) {
         let fetch_selection = column_selection(
             self.selection.as_ref(),
             self.expanded_selection.as_ref(),
@@ -672,75 +729,52 @@ impl RowGroupContext {
             column_idx,
         );
         let all_locations = self.locations(column_idx);
-        let fetch = ColumnFetch::new(chunk.clone(), all_locations, 
fetch_selection);
+        let fetch = ColumnFetch::new(self.chunk(column_idx), all_locations, 
fetch_selection);
         // The decoder fetches a whole column chunk as one range. If page
         // locations are known, the plan splits it into the same bytes, page
         // by page.
         let locations = all_locations.filter(|l| !l.is_empty());
-        let (dictionary, pages) = match (fetch, locations) {
-            (ColumnFetch::Chunk { range }, Some(locations)) => (
-                dictionary_range(range.start, locations),
-                PageSet::All(locations.len()),
-            ),
+        let (locations, dictionary, selected) = match (fetch, locations) {
+            (ColumnFetch::Chunk { range }, Some(locations)) => {
+                (locations, dictionary_range(range.start, locations), None)
+            }
             (ColumnFetch::Chunk { range }, None) => {
-                return ColumnCursor {
-                    column_idx,
-                    head: Some(self.entry(
-                        column_idx,
-                        range,
-                        row_group_rows,
-                        PageKind::ColumnChunk,
-                    )),
-                    pages: PageSet::None,
-                    next_page: 0,
-                };
+                let rows = self.row_group_rows();
+                let order = (rows.start, Reverse(rows.end), range.start);
+                return (Some((order, CHUNK)), None);
             }
             (
                 ColumnFetch::Pages {
-                    dictionary, pages, ..
+                    dictionary,
+                    locations,
+                    pages,
                 },
                 _,
-            ) => (dictionary, PageSet::Selected(pages)),
+            ) => (locations, dictionary, Some(pages)),
         };
-
-        let mut cursor = ColumnCursor {
-            column_idx,
-            head: None,
-            pages,
-            next_page: 0,
+        let pages = match &selected {
+            Some(pages) => PageSet::Selected(pages),
+            None => PageSet::All(locations.len()),
         };
-        if let Some(dictionary) = dictionary {
-            let locations = all_locations.expect("pages have locations");
+
+        let head = if let Some(dictionary) = dictionary {
             // The dictionary serves exactly the rows of the data pages read.
-            let rows = match cursor.pages.len() {
-                0 => row_group_rows,
+            let rows = match pages.len() {
+                0 => self.row_group_rows(),
                 len => {
-                    let first = self.page_rows(locations, 
cursor.pages.page(0));
-                    let last = self.page_rows(locations, cursor.pages.page(len 
- 1));
+                    let first = self.page_rows(locations, pages.page(0));
+                    let last = self.page_rows(locations, pages.page(len - 1));
                     first.start..last.end
                 }
             };
-            cursor.head = Some(self.entry(column_idx, dictionary, rows, 
PageKind::Dictionary));
+            Some((
+                (rows.start, Reverse(rows.end), dictionary.start),
+                DICTIONARY,
+            ))
         } else {
-            cursor.head = self.next_data_page(&mut cursor);
-        }
-        cursor
-    }
-
-    /// The next data page of `cursor`, if any.
-    fn next_data_page(&self, cursor: &mut ColumnCursor) -> 
Option<PlannedRange> {
-        if cursor.next_page >= cursor.pages.len() {
-            return None;
-        }
-        let locations = self.locations(cursor.column_idx)?;
-        let idx = cursor.pages.page(cursor.next_page);
-        cursor.next_page += 1;
-        Some(self.entry(
-            cursor.column_idx,
-            page_range(&locations[idx]),
-            self.page_rows(locations, idx),
-            PageKind::Data,
-        ))
+            self.data_page(locations, pages, 0).map(|order| (order, 0))
+        };
+        (head, selected)
     }
 }
 
diff --git a/parquet/src/arrow/push_decoder/scan_plan/tests.rs 
b/parquet/src/arrow/push_decoder/scan_plan/tests.rs
index 56541e4f9a..5a81b02e68 100644
--- a/parquet/src/arrow/push_decoder/scan_plan/tests.rs
+++ b/parquet/src/arrow/push_decoder/scan_plan/tests.rs
@@ -900,7 +900,7 @@ fn plan_is_lazy() {
     assert_eq!(current.row_group.row_group_idx, 0);
     assert_eq!(planner.next_row, 200);
     // One cursor per projected column, not one entry per page.
-    assert_eq!(current.stage.as_ref().unwrap().cursors.len(), 3);
+    assert_eq!(current.stage.as_ref().unwrap().heap.len(), 3);
 }
 
 #[test]

Reply via email to