alamb commented on code in PR #10555: URL: https://github.com/apache/arrow-rs/pull/10555#discussion_r4210301263
########## parquet/src/arrow/push_decoder/scan_plan/mod.rs: ########## @@ -0,0 +1,1778 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScanPlan`]: the byte ranges that a push decoder may still read. + +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::iter::FusedIterator; +use std::ops::Range; +use std::sync::Arc; + +use crate::arrow::ProjectionMask; +use crate::arrow::arrow_reader::{ReadPlanBuilder, RowSelection}; +use crate::errors::ParquetError; +use crate::file::metadata::page_index::PageIndexProvider; +use crate::file::page_index::offset_index::PageLocation; + +mod budget; +mod frontier; + +pub(crate) use budget::{BudgetedReadPlan, RowBudget}; +pub(crate) use frontier::{NextRowGroup, RowGroupFrontier}; + +/// One item of a [`ScanPlan`]. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct PlannedRange { + /// Byte range in the file. + pub range: Range<u64>, + /// Row group index in the file. + pub row_group: usize, + /// Leaf column index in the file. + pub column: usize, + /// 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. + pub(crate) first_row: u64, + /// One past the last planned row that this range serves. + pub(crate) last_row: u64, + /// The decoding stage that first reads this range. + pub(crate) stage: ScanStage, + /// `true` if a predicate result can make this range unnecessary. + pub(crate) conditional: bool, +} + +impl PlannedRange { + /// Number of bytes in this range. + pub fn len(&self) -> u64 { + self.range.end - self.range.start + } + + /// Returns `true` if this range contains no bytes. + pub fn is_empty(&self) -> bool { + self.range.is_empty() + } +} + +/// What a [`PlannedRange`] contains. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum PageKind { + /// The dictionary page of a column chunk. + Dictionary, + /// One data page. + Data, + /// A complete column chunk. The plan uses this when the column has no + /// offset index, so page locations are unknown. + ColumnChunk, +} + +/// The decoding stage that first reads a [`PlannedRange`]. +/// +/// The decoder reads a column once per row group. A column that a predicate +/// reads is not read again for the output. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) enum ScanStage { + /// Evaluation of the predicate at this index in the + /// [`RowFilter`](crate::arrow::arrow_reader::RowFilter). + Predicate(usize), + /// Decoding of the output columns. + Projection, +} + +/// The byte ranges that a push decoder may still read, in the order that +/// decoding needs them. +/// +/// Created by [`ParquetPushDecoder::scan_plan`]. Use it to fetch data before +/// [`DecodeResult::NeedsData`] asks for it. `NeedsData` stays the exact +/// request. +/// +/// * The plan contains every range that the decoder can request after +/// the plan is made. It can contain more, for example ranges that a +/// [`RowFilter`] makes unnecessary. +/// * The plan comes from the current state of the decoder. To plan the +/// whole scan, call [`ParquetPushDecoder::scan_plan`] once after +/// [`ParquetPushDecoderBuilder::build`] and keep the iterator. After +/// [`ParquetPushDecoder::into_builder`], plan again. +/// * Row groups are in read order. In a row group, the columns of the +/// [`RowFilter`] predicates come first, then the other output columns. The +/// ranges are ordered by the first row that they serve, and the ranges of +/// one column chunk stay in file order. +/// * There is one range for each page if the column has an offset index, +/// and one range for the column chunk if not. +/// * The plan does not change decoding, does no I/O and does not depend +/// on pushed data. It is lazy, and its state does not grow with the number +/// of pages. +/// * If a row group is not valid, the plan ends before it. The decoder +/// reports the error. +/// +/// Keep the fetched ranges in your own cache, and push exactly the ranges +/// that `NeedsData` requests. The decoder does not use a requested range +/// that is split over two pushed buffers, and it releases a pushed buffer +/// only if its range is equal to a requested range. +/// +/// # Example +/// +/// ``` +/// # use std::collections::BTreeMap; +/// # use std::ops::Range; +/// # use bytes::Bytes; +/// # use arrow_array::record_batch; +/// # use parquet::DecodeResult; +/// # use parquet::arrow::ArrowWriter; +/// # use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +/// # use parquet::file::metadata::PageIndexPolicy; +/// # use parquet::file::properties::WriterProperties; +/// # let file = { +/// # let mut buffer = vec![]; +/// # let batch = record_batch!(("a", Int32, [1, 2, 3, 4])).unwrap(); +/// # let props = WriterProperties::builder().set_max_row_group_row_count(Some(2)).build(); +/// # let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); +/// # writer.write(&batch).unwrap(); +/// # writer.close().unwrap(); +/// # Bytes::from(buffer) +/// # }; +/// # let fetch = |range: &Range<u64>| file.slice(range.start as usize..range.end as usize); +/// # // Join cached ranges that cover `range`, or fetch it. +/// # let read = |cache: &BTreeMap<u64, Bytes>, range: &Range<u64>| { +/// # let mut data = Vec::new(); +/// # let mut position = range.start; +/// # while position < range.end { +/// # let Some((start, bytes)) = cache.range(..=position).next_back() else { +/// # return fetch(range); +/// # }; +/// # let end = (start + bytes.len() as u64).min(range.end); +/// # if end <= position { +/// # return fetch(range); +/// # } +/// # data.extend_from_slice(&bytes[(position - start) as usize..(end - start) as usize]); +/// # position = end; +/// # } +/// # Bytes::from(data) +/// # }; +/// # let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional); +/// # let metadata = ArrowReaderMetadata::load(&file, options).unwrap(); +/// let mut decoder = ParquetPushDecoderBuilder::new_with_metadata(metadata) +/// .build() +/// .unwrap(); +/// +/// // Read ahead up to 1 MB into a cache +/// let mut cache = BTreeMap::new(); +/// let mut cached = 0; +/// for planned in decoder.scan_plan() { +/// if cached + planned.len() > 1024 * 1024 { +/// break; +/// } +/// cached += planned.len(); +/// cache.insert(planned.range.start, fetch(&planned.range)); +/// } +/// +/// // Answer each request from the cache +/// loop { +/// match decoder.try_decode().unwrap() { +/// DecodeResult::NeedsData(ranges) => { +/// let data = ranges.iter().map(|range| read(&cache, range)).collect(); +/// decoder.push_ranges(ranges, data).unwrap(); +/// } +/// DecodeResult::Data(batch) => println!("{} rows", batch.num_rows()), +/// DecodeResult::Finished => break, +/// } +/// } +/// ``` +/// +/// [`RowFilter`]: crate::arrow::arrow_reader::RowFilter +/// [`DecodeResult::NeedsData`]: crate::DecodeResult::NeedsData +/// [`ParquetPushDecoder::scan_plan`]: super::ParquetPushDecoder::scan_plan +/// [`ParquetPushDecoder::into_builder`]: super::ParquetPushDecoder::into_builder +/// [`ParquetPushDecoderBuilder::build`]: super::ParquetPushDecoderBuilder::build +#[derive(Debug, Clone)] +pub struct ScanPlan { + /// `None` for an empty plan, for example of a finished decoder. + planner: Option<Box<Planner>>, +} + +impl ScanPlan { + /// A plan with no ranges. + pub(crate) fn empty() -> Self { + Self { planner: None } + } +} + +/// The state of a [`ScanPlan`] that is not empty. +#[derive(Debug, Clone)] +struct Planner { + /// The decoder's row-group queue, selections and offset/limit budget, as + /// they were when the plan was made. + frontier: RowGroupFrontier, + /// The row group that the decoder was fetching when the plan was made. + /// It is planned first. + active_row_group: Option<NextRowGroup>, + /// The decoder's projection, predicates and batch size. + columns: Arc<StageColumns>, + /// Whether an output limit is set. + has_limit: bool, + /// Whether no row group has been planned yet. + at_first_row_group: bool, + /// Planned rows before the next row group. + next_row: u64, + /// The row group being planned, if any. + current: Option<RowGroupRanges>, + /// Whether planning has ended. + done: bool, +} + +/// The columns each decoding stage reads, as [`RowGroupReaderBuilder`] +/// uses them to decide which bytes a row group needs. +/// +/// [`RowGroupReaderBuilder`]: super::reader_builder::RowGroupReaderBuilder +#[derive(Debug)] +struct StageColumns { + /// The output batch size, which aligns cached predicate reads. + batch_size: usize, + /// Columns in the output. + projection: ProjectionMask, + /// Columns each [`RowFilter`] predicate reads, in evaluation order. + /// + /// [`RowFilter`]: crate::arrow::arrow_reader::RowFilter + predicate_projections: Vec<ProjectionMask>, + /// Predicate columns whose decoded values are cached for the output, if + /// any. Their selection is expanded to batch boundaries when fetched. + cache_projection: Option<ProjectionMask>, +} + +/// Builds a [`ScanPlan`] from a decoder's row-group frontier and the columns +/// each decoding stage reads. +#[derive(Debug)] +pub(crate) struct ScanPlanBuilder { + frontier: RowGroupFrontier, + active_row_group: Option<NextRowGroup>, + columns: StageColumns, +} + +impl ScanPlanBuilder { + /// Plan the row groups in `frontier`, decoding `projection` in batches of + /// `batch_size` rows, with no predicates. + pub(crate) fn new( + frontier: RowGroupFrontier, + batch_size: usize, + projection: ProjectionMask, + ) -> Self { + Self { + frontier, + active_row_group: None, + columns: StageColumns { + batch_size, + projection, + predicate_projections: vec![], + cache_projection: None, + }, + } + } + + /// Set the columns each predicate reads, in evaluation order. + pub(crate) fn with_predicate_projections( + mut self, + predicate_projections: Vec<ProjectionMask>, + ) -> Self { + self.columns.predicate_projections = predicate_projections; + self + } + + /// Set the predicate columns whose decoded values are cached for the + /// output. + pub(crate) fn with_cache_projection( + mut self, + cache_projection: Option<ProjectionMask>, + ) -> Self { + self.columns.cache_projection = cache_projection; + self + } + + /// Plan `active_row_group` first: the row group that the decoder is + /// fetching, which the frontier has already handed over. + pub(crate) fn with_active_row_group(mut self, active_row_group: Option<NextRowGroup>) -> Self { + self.active_row_group = active_row_group; + self + } + + pub(crate) fn build(self) -> ScanPlan { + let Self { + frontier, + active_row_group, + columns, + } = self; + let has_limit = frontier.budget.limit().is_some(); + let planner = Planner { + frontier, + active_row_group, + columns: Arc::new(columns), + has_limit, + at_first_row_group: true, + next_row: 0, + current: None, + done: false, + }; + ScanPlan { + planner: Some(Box::new(planner)), + } + } +} + +impl Planner { + /// Start planning the next row group the decoder will read. + /// + /// Returns `Ok(false)` when no row group remains. A row group that the + /// offset/limit budget removes is planned with no ranges. + fn plan_next_row_group(&mut self) -> Result<bool, ParquetError> { Review Comment: Specifically, this logic seems very similar to the code in https://github.com/apache/arrow-rs/blob/e1450a346b51f2020c480d7c76c91afc32c63ffe/parquet/src/arrow/push_decoder/reader_builder/mod.rs#L396-L724 I think we should plan over time to try and unify the two paths. I'll file a follow on ticket -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
