parquet/arrow/push_decoder/scan_plan/
budget.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Offset/limit budget shared across the row groups of a scan.
19
20use crate::arrow::arrow_reader::ReadPlanBuilder;
21
22/// Running offset/limit budget shared across row groups.
23#[derive(Debug, Clone, Copy, Eq, PartialEq)]
24pub(crate) struct RowBudget {
25    offset: Option<usize>,
26    limit: Option<usize>,
27}
28
29impl RowBudget {
30    pub(crate) fn new(offset: Option<usize>, limit: Option<usize>) -> Self {
31        Self { offset, limit }
32    }
33
34    pub(crate) fn is_exhausted(self) -> bool {
35        matches!(self.limit, Some(0))
36    }
37
38    /// The offset still to be skipped before the next readable row group.
39    pub(crate) fn offset(self) -> Option<usize> {
40        self.offset
41    }
42
43    /// The number of output rows still permitted across the remaining row groups.
44    pub(crate) fn limit(self) -> Option<usize> {
45        self.limit
46    }
47
48    /// Returns how many selected rows remain after applying this budget.
49    pub(crate) fn rows_after(self, rows_before_budget: usize) -> usize {
50        let rows_after_offset = rows_before_budget.saturating_sub(self.offset.unwrap_or(0));
51        match self.limit {
52            Some(limit) => rows_after_offset.min(limit),
53            None => rows_after_offset,
54        }
55    }
56
57    /// Returns the number of selected rows needed before applying the offset.
58    pub(crate) fn selected_row_limit(self) -> Option<usize> {
59        self.limit
60            .map(|limit| limit.saturating_add(self.offset.unwrap_or(0)))
61    }
62
63    pub(crate) fn apply_to_plan(
64        self,
65        plan_builder: ReadPlanBuilder,
66        row_count: usize,
67    ) -> BudgetedReadPlan {
68        let rows_before_budget = plan_builder.num_rows_selected().unwrap_or(row_count);
69        let plan_builder = plan_builder
70            .limited(row_count)
71            .with_offset(self.offset)
72            .with_limit(self.limit)
73            .build_limited();
74        let rows_after_budget = self.rows_after(rows_before_budget);
75
76        BudgetedReadPlan {
77            plan_builder,
78            rows_before_budget,
79            rows_after_budget,
80            remaining_budget: self.advance(rows_before_budget, rows_after_budget),
81        }
82    }
83
84    /// Advance the budget past one row group.
85    ///
86    /// `rows_before_budget` is the number of rows selected before applying the
87    /// budget, and `rows_after_budget` is the number retained for output from
88    /// this row group.
89    pub(crate) fn advance(mut self, rows_before_budget: usize, rows_after_budget: usize) -> Self {
90        if let Some(offset) = &mut self.offset {
91            // Reduction is either because of offset or limit, as limit is applied
92            // after offset has been "exhausted" can just use saturating sub here.
93            *offset = offset.saturating_sub(rows_before_budget - rows_after_budget);
94        }
95
96        if rows_after_budget != 0
97            && let Some(limit) = &mut self.limit
98        {
99            *limit -= rows_after_budget;
100        }
101
102        self
103    }
104}
105
106#[derive(Debug)]
107pub(crate) struct BudgetedReadPlan {
108    /// Read plan after applying this row group's share of the offset/limit budget.
109    pub(crate) plan_builder: ReadPlanBuilder,
110    /// Number of rows selected by row selection and predicates before applying
111    /// this row group's offset/limit budget.
112    pub(crate) rows_before_budget: usize,
113    /// Number of selected rows that remain to be read after applying this row
114    /// group's offset/limit budget.
115    pub(crate) rows_after_budget: usize,
116    /// Budget remaining for later row groups.
117    pub(crate) remaining_budget: RowBudget,
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::arrow::arrow_reader::{RowSelection, RowSelector};
124
125    #[test]
126    fn test_row_budget_offset_limit_across_row_groups() {
127        let first =
128            RowBudget::new(Some(225), Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
129        assert_eq!(first.rows_before_budget, 200);
130        assert_eq!(first.rows_after_budget, 0);
131        assert_eq!(first.remaining_budget, RowBudget::new(Some(25), Some(20)));
132        assert_eq!(first.plan_builder.num_rows_selected(), Some(0));
133
134        let second = first
135            .remaining_budget
136            .apply_to_plan(ReadPlanBuilder::new(1024), 200);
137        assert_eq!(second.rows_before_budget, 200);
138        assert_eq!(second.rows_after_budget, 20);
139        assert_eq!(second.remaining_budget, RowBudget::new(Some(0), Some(0)));
140        assert_eq!(second.plan_builder.num_rows_selected(), Some(20));
141    }
142
143    #[test]
144    fn test_row_budget_limit_only() {
145        let budgeted =
146            RowBudget::new(None, Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
147        assert_eq!(budgeted.rows_before_budget, 200);
148        assert_eq!(budgeted.rows_after_budget, 20);
149        assert_eq!(budgeted.remaining_budget, RowBudget::new(None, Some(0)));
150        assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(20));
151    }
152
153    #[test]
154    fn test_row_budget_empty_selection() {
155        let empty_selection = RowSelection::from(vec![RowSelector::skip(200)]);
156        let budgeted = RowBudget::new(Some(10), Some(20)).apply_to_plan(
157            ReadPlanBuilder::new(1024).with_selection(Some(empty_selection)),
158            200,
159        );
160        assert_eq!(budgeted.rows_before_budget, 0);
161        assert_eq!(budgeted.rows_after_budget, 0);
162        assert_eq!(
163            budgeted.remaining_budget,
164            RowBudget::new(Some(10), Some(20))
165        );
166        assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(0));
167    }
168}
parquet/arrow/push_decoder/scan_plan/budget.rs

parquet/arrow/push_decoder/scan_plan/
budget.rs
