parquet 60.1.0

// 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.

//! Which row groups a scan reads next, with their selections and budget.

use crate::arrow::arrow_reader::{RowGroupPlan, RowGroupSelection, RowSelection};
use crate::errors::ParquetError;
use crate::file::metadata::ParquetMetaData;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;

use super::budget::RowBudget;

/// Plan for the next queued row group after row-selection slicing.
#[derive(Debug)]
enum QueuedRowGroupDecision {
    /// Hand this row group to the builder.
    Read(NextRowGroup),
    /// Skip this row group, and keep scanning with the updated budget.
    Skip { remaining_budget: RowBudget },
}

/// Work item handed from [`RowGroupFrontier`] to [`RowGroupReaderBuilder`](crate::arrow::push_decoder::reader_builder::RowGroupReaderBuilder).
#[derive(Debug, Clone)]
pub(crate) struct NextRowGroup {
    pub(crate) row_group_idx: usize,
    pub(crate) row_count: usize,
    /// This row group's selection, or `None` when all rows are selected.
    pub(crate) selection: Option<RowSelection>,
    /// Budget snapshot to apply while decoding this row group.
    pub(crate) budget: RowBudget,
}

/// Row groups and selections that have not yet been handed to the row-group
/// reader builder.
#[derive(Debug, Clone)]
struct QueuedRowGroups {
    queue: Queue,
    /// The number of times each row group is in `queue`, so that
    /// [`Self::contains`] does not scan the queue.
    counts: HashMap<usize, usize>,
}

impl QueuedRowGroups {
    /// Validate and queue a row-group plan for `parquet_metadata`.
    fn try_new(
        parquet_metadata: &ParquetMetaData,
        row_group_plan: RowGroupPlan,
    ) -> Result<Self, ParquetError> {
        let queue = Queue::try_new(parquet_metadata, row_group_plan)?;
        let mut counts = HashMap::new();
        for row_group_idx in queue.row_group_indices() {
            *counts.entry(row_group_idx).or_default() += 1;
        }
        Ok(Self { queue, counts })
    }

    /// Convert the remaining queue back into a builder configuration.
    fn into_plan(self) -> RowGroupPlan {
        self.queue.into_plan()
    }

    fn front(&self) -> Option<usize> {
        self.queue.front()
    }

    /// The row groups in the queue, in no specific order and without
    /// duplicates.
    fn row_groups(&self) -> impl Iterator<Item = usize> + '_ {
        self.counts.keys().copied()
    }

    /// Returns `true` if `row_group_idx` is in the queue.
    fn contains(&self, row_group_idx: usize) -> bool {
        self.counts.contains_key(&row_group_idx)
    }

    fn len(&self) -> usize {
        self.queue.len()
    }

    fn clear(&mut self) {
        self.queue.clear();
        self.counts.clear();
    }

    /// See [`Queue::global_selection_is_exhausted`].
    fn global_selection_is_exhausted(&self) -> bool {
        self.queue.global_selection_is_exhausted()
    }

    /// Remove the front row group and return its local selection.
    fn pop_front_selection(&mut self, row_count: usize) -> Option<RowSelection> {
        if let Some(row_group_idx) = self.queue.front() {
            match self.counts.get_mut(&row_group_idx) {
                Some(1) => {
                    self.counts.remove(&row_group_idx);
                }
                Some(count) => *count -= 1,
                None => debug_assert!(false, "row group {row_group_idx} is queued but not counted"),
            }
        }
        self.queue.pop_front_selection(row_count)
    }
}

/// The queue of [`QueuedRowGroups`].
#[derive(Debug, Clone)]
enum Queue {
    /// One selection cursor spans all queued row groups.
    Global {
        row_groups: VecDeque<usize>,
        selection: Option<RowSelection>,
    },
    /// Selections are already relative to their respective row groups.
    PerRowGroup(VecDeque<RowGroupSelection>),
}

impl Queue {
    /// Validate and queue a row-group plan for `parquet_metadata`.
    fn try_new(
        parquet_metadata: &ParquetMetaData,
        row_group_plan: RowGroupPlan,
    ) -> Result<Self, ParquetError> {
        match row_group_plan {
            RowGroupPlan::Global {
                row_groups,
                selection,
            } => Ok(Self::Global {
                row_groups: row_groups
                    .unwrap_or_else(|| (0..parquet_metadata.num_row_groups()).collect())
                    .into(),
                selection,
            }),
            RowGroupPlan::PerRowGroup(row_groups) => {
                for row_group in &row_groups {
                    let row_count =
                        parquet_metadata.row_group_num_rows(row_group.row_group_index)?;
                    if let Some(selection) = &row_group.selection {
                        let selection_rows = selection.total_row_count();
                        if selection_rows > row_count {
                            return Err(ParquetError::General(format!(
                                "Row selection for row group {} contains {selection_rows} rows, but the row group has {row_count}",
                                row_group.row_group_index
                            )));
                        }
                    }
                }
                Ok(Self::PerRowGroup(row_groups.into()))
            }
            RowGroupPlan::Conflicting => Err(RowGroupPlan::conflict_error()),
        }
    }

    /// Convert the remaining queue back into a builder configuration.
    fn into_plan(self) -> RowGroupPlan {
        match self {
            Self::Global {
                row_groups,
                selection,
            } => RowGroupPlan::Global {
                row_groups: Some(Vec::from(row_groups)),
                selection,
            },
            Self::PerRowGroup(row_groups) => RowGroupPlan::PerRowGroup(Vec::from(row_groups)),
        }
    }

    fn front(&self) -> Option<usize> {
        match self {
            Self::Global { row_groups, .. } => row_groups.front().copied(),
            Self::PerRowGroup(row_groups) => row_groups
                .front()
                .map(|row_group| row_group.row_group_index),
        }
    }

    /// The queued row group indexes, in order.
    fn row_group_indices(&self) -> Vec<usize> {
        match self {
            Self::Global { row_groups, .. } => row_groups.iter().copied().collect(),
            Self::PerRowGroup(row_groups) => row_groups
                .iter()
                .map(|row_group| row_group.row_group_index)
                .collect(),
        }
    }

    fn len(&self) -> usize {
        match self {
            Self::Global { row_groups, .. } => row_groups.len(),
            Self::PerRowGroup(row_groups) => row_groups.len(),
        }
    }

    fn clear(&mut self) {
        match self {
            Self::Global {
                row_groups,
                selection,
            } => {
                row_groups.clear();
                *selection = None;
            }
            Self::PerRowGroup(row_groups) => row_groups.clear(),
        }
    }

    /// Returns `true` when a shared global selection has no selected rows left.
    /// Per-row-group selections are independent and are drained one at a time.
    fn global_selection_is_exhausted(&self) -> bool {
        matches!(
            self,
            Self::Global {
                selection: Some(selection),
                ..
            } if selection.row_count() == 0
        )
    }

    /// Remove the front row group and return its local selection.
    fn pop_front_selection(&mut self, row_count: usize) -> Option<RowSelection> {
        match self {
            Self::Global {
                row_groups,
                selection,
            } => {
                let popped = row_groups.pop_front();
                debug_assert!(popped.is_some(), "front row group checked before pop");
                selection
                    .as_mut()
                    .map(|selection| selection.split_off(row_count))
            }
            Self::PerRowGroup(row_groups) => {
                row_groups
                    .pop_front()
                    .expect("front row group checked before pop")
                    .selection
            }
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct RowGroupFrontier {
    /// Metadata used to resolve row counts for queued row groups.
    pub(super) parquet_metadata: Arc<ParquetMetaData>,
    /// Row groups not yet handed to the builder.
    queued: QueuedRowGroups,
    /// Offset/limit budget before the next readable row group is planned.
    pub(super) budget: RowBudget,
    /// If predicates are present, row groups with selected rows must be read so
    /// the predicate can decide whether they are actually needed.
    pub(super) has_predicates: bool,
}

impl RowGroupFrontier {
    pub(crate) fn new(
        parquet_metadata: Arc<ParquetMetaData>,
        row_group_plan: RowGroupPlan,
        budget: RowBudget,
        has_predicates: bool,
    ) -> Result<Self, ParquetError> {
        let queued = QueuedRowGroups::try_new(&parquet_metadata, row_group_plan)?;

        Ok(Self {
            parquet_metadata,
            queued,
            budget,
            has_predicates,
        })
    }

    /// The queued row groups, in no specific order and without duplicates.
    pub(crate) fn queued_row_groups(&self) -> impl Iterator<Item = usize> + '_ {
        self.queued.row_groups()
    }

    /// Returns `true` if `row_group_idx` is still queued to be read.
    pub(crate) fn is_queued(&self, row_group_idx: usize) -> bool {
        self.queued.contains(row_group_idx)
    }

    pub(crate) fn update_budget_after_row_group(&mut self, budget: RowBudget) {
        self.budget = budget;
    }

    /// Peek at the next row-group index [`Self::next_readable_row_group`]
    /// would hand out, without mutating any state. Returns `None` if every
    /// remaining row group would be skipped under the current
    /// selection/budget, or if the queue is empty.
    ///
    /// Runs the real [`Self::next_readable_row_group`] advance logic on a
    /// throwaway clone of the frontier, so peek can never drift from the
    /// read path. The clone copies the queued row-group plan and selections;
    /// see
    /// [`RemainingRowGroups::peek_next_row_group`](crate::arrow::push_decoder::remaining::RemainingRowGroups::peek_next_row_group).
    pub(crate) fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
        Ok(self
            .clone()
            .next_readable_row_group()?
            .map(|next_row_group| next_row_group.row_group_idx))
    }

    /// Number of row groups not yet handed to the builder.
    pub(crate) fn row_groups_remaining(&self) -> usize {
        self.queued.len()
    }

    /// Decompose into the metadata, the row groups not yet handed to the
    /// builder, and the remaining offset/limit budget.
    pub(crate) fn into_parts(self) -> (Arc<ParquetMetaData>, RowGroupPlan, RowBudget) {
        let Self {
            parquet_metadata,
            queued,
            budget,
            has_predicates: _,
        } = self;
        (parquet_metadata, queued.into_plan(), budget)
    }

    fn clear_remaining(&mut self) {
        self.queued.clear();
    }

    /// Plan whether a selected row group should be read or skipped.
    ///
    /// Selection-only skips are handled before this method is called. This
    /// method applies the remaining offset/limit budget and predicate
    /// conservatism.
    fn plan_selected_row_group(
        &self,
        next_row_group: NextRowGroup,
        selected_rows: usize,
    ) -> QueuedRowGroupDecision {
        if self.has_predicates {
            return QueuedRowGroupDecision::Read(next_row_group);
        }

        let rows_after_budget = self.budget.rows_after(selected_rows);
        if rows_after_budget != 0 {
            return QueuedRowGroupDecision::Read(next_row_group);
        }

        QueuedRowGroupDecision::Skip {
            remaining_budget: self.budget.advance(selected_rows, rows_after_budget),
        }
    }

    /// Advance queued row groups until one should be handed to the builder.
    pub(crate) fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>, ParquetError> {
        loop {
            let Some(row_group_idx) = self.queued.front() else {
                return Ok(None);
            };
            // A global selection can be exhausted before its row-group queue.
            // Per-row-group selections have no shared cursor to exhaust; empty
            // local selections are discarded by the `selected_rows == 0` path below.
            if self.budget.is_exhausted() || self.queued.global_selection_is_exhausted() {
                self.clear_remaining();
                return Ok(None);
            }

            let row_count = self.parquet_metadata.row_group_num_rows(row_group_idx)?;
            let selection = self.queued.pop_front_selection(row_count);
            let (selection, selected_rows) = match selection {
                Some(selection) => {
                    let selected_rows = selection.row_count();
                    if selected_rows == 0 {
                        continue;
                    }
                    // An all-rows selection is equivalent to no selection
                    (
                        (selected_rows != row_count).then_some(selection),
                        selected_rows,
                    )
                }
                None => (None, row_count),
            };

            let next_row_group = NextRowGroup {
                row_group_idx,
                row_count,
                selection,
                budget: self.budget,
            };

            match self.plan_selected_row_group(next_row_group, selected_rows) {
                QueuedRowGroupDecision::Read(next_row_group) => {
                    return Ok(Some(next_row_group));
                }
                QueuedRowGroupDecision::Skip { remaining_budget } => {
                    self.budget = remaining_budget;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::arrow::arrow_reader::RowSelector;
    use crate::arrow::push_decoder::test::test_file_parquet_metadata;

    fn global_plan(
        row_groups: Option<Vec<usize>>,
        selection: Option<RowSelection>,
    ) -> RowGroupPlan {
        RowGroupPlan::Global {
            row_groups,
            selection,
        }
    }

    #[test]
    fn queued_row_groups_encapsulates_plan_transitions() {
        let metadata = test_file_parquet_metadata();

        let mut all_row_groups =
            QueuedRowGroups::try_new(&metadata, global_plan(None, None)).unwrap();
        assert_eq!(all_row_groups.len(), 2);
        assert_eq!(all_row_groups.front(), Some(0));
        assert!(!all_row_groups.global_selection_is_exhausted());
        assert!(all_row_groups.pop_front_selection(200).is_none());
        assert_eq!(all_row_groups.front(), Some(1));
        all_row_groups.clear();
        assert_eq!(all_row_groups.len(), 0);
        assert!(matches!(
            all_row_groups.into_plan(),
            RowGroupPlan::Global {
                row_groups: Some(row_groups),
                selection: None,
            } if row_groups.is_empty()
        ));

        let global_selection = RowSelection::from(vec![
            RowSelector::skip(10),
            RowSelector::select(5),
            RowSelector::skip(185),
            RowSelector::select(200),
        ]);
        let mut global = QueuedRowGroups::try_new(
            &metadata,
            global_plan(Some(vec![0, 1]), Some(global_selection)),
        )
        .unwrap();
        let first = global.pop_front_selection(200).unwrap();
        assert_eq!(first.row_count(), 5);
        assert!(!global.global_selection_is_exhausted());
        assert!(matches!(
            global.into_plan(),
            RowGroupPlan::Global {
                row_groups: Some(row_groups),
                selection: Some(selection),
            } if row_groups == vec![1] && selection.row_count() == 200
        ));

        let local_selection =
            RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(3)]);
        let mut local = QueuedRowGroups::try_new(
            &metadata,
            RowGroupPlan::PerRowGroup(vec![
                RowGroupSelection::new(1, Some(local_selection)),
                RowGroupSelection::new(0, None),
            ]),
        )
        .unwrap();
        assert_eq!(local.front(), Some(1));
        assert_eq!(local.pop_front_selection(200).unwrap().row_count(), 3);
        assert!(!local.global_selection_is_exhausted());
        assert!(matches!(
            local.into_plan(),
            RowGroupPlan::PerRowGroup(row_groups)
                if row_groups == vec![RowGroupSelection::new(0, None)]
        ));

        let exhausted = QueuedRowGroups::try_new(
            &metadata,
            global_plan(
                Some(vec![0]),
                Some(RowSelection::from(vec![RowSelector::skip(200)])),
            ),
        )
        .unwrap();
        assert!(exhausted.global_selection_is_exhausted());
    }

    #[test]
    fn frontier_handles_global_and_local_exhaustion() {
        let metadata = test_file_parquet_metadata();
        let budget = RowBudget::new(None, None);

        let mut global = RowGroupFrontier::new(
            Arc::clone(&metadata),
            global_plan(
                Some(vec![0, 1]),
                Some(RowSelection::from(vec![RowSelector::skip(400)])),
            ),
            budget,
            false,
        )
        .unwrap();
        assert!(global.next_readable_row_group().unwrap().is_none());
        assert_eq!(global.queued.len(), 0);

        let mut local = RowGroupFrontier::new(
            Arc::clone(&metadata),
            RowGroupPlan::PerRowGroup(vec![
                RowGroupSelection::new(0, Some(RowSelection::from(vec![RowSelector::skip(200)]))),
                RowGroupSelection::new(1, None),
            ]),
            budget,
            false,
        )
        .unwrap();
        let next = local.next_readable_row_group().unwrap().unwrap();
        assert_eq!(next.row_group_idx, 1);
        assert_eq!(next.row_count, 200);
        assert!(next.selection.is_none());

        let mut exhausted_budget = RowGroupFrontier::new(
            metadata,
            RowGroupPlan::PerRowGroup(vec![RowGroupSelection::new(0, None)]),
            RowBudget::new(None, Some(0)),
            false,
        )
        .unwrap();
        assert!(
            exhausted_budget
                .next_readable_row_group()
                .unwrap()
                .is_none()
        );
        assert_eq!(exhausted_budget.queued.len(), 0);
    }

    #[test]
    fn queued_row_groups_count_repeated_row_groups() {
        let metadata = test_file_parquet_metadata();
        let budget = RowBudget::new(None, None);

        // Row group 1 is skipped the first time. A row group stays queued
        // until its last occurrence is popped.
        let mut local = RowGroupFrontier::new(
            Arc::clone(&metadata),
            RowGroupPlan::PerRowGroup(vec![
                RowGroupSelection::new(0, None),
                RowGroupSelection::new(1, Some(RowSelection::from(vec![RowSelector::skip(200)]))),
                RowGroupSelection::new(0, None),
                RowGroupSelection::new(1, None),
            ]),
            budget,
            false,
        )
        .unwrap();
        let queued = |frontier: &RowGroupFrontier| {
            (frontier.queued.contains(0), frontier.queued.contains(1))
        };
        assert_eq!(queued(&local), (true, true));

        assert_eq!(
            local
                .next_readable_row_group()
                .unwrap()
                .unwrap()
                .row_group_idx,
            0
        );
        assert_eq!(queued(&local), (true, true));
        // Peeking does not change the counts of the frontier.
        assert_eq!(local.peek_next_row_group().unwrap(), Some(0));
        assert_eq!(queued(&local), (true, true));

        // Skips the first occurrence of row group 1.
        assert_eq!(
            local
                .next_readable_row_group()
                .unwrap()
                .unwrap()
                .row_group_idx,
            0
        );
        assert_eq!(queued(&local), (false, true));

        assert_eq!(
            local
                .next_readable_row_group()
                .unwrap()
                .unwrap()
                .row_group_idx,
            1
        );
        assert_eq!(queued(&local), (false, false));
        assert!(local.queued.counts.is_empty());
        assert!(local.next_readable_row_group().unwrap().is_none());

        // An exhausted limit clears the queue and its counts.
        let mut global = RowGroupFrontier::new(
            metadata,
            global_plan(Some(vec![0, 1, 0]), None),
            RowBudget::new(None, Some(200)),
            false,
        )
        .unwrap();
        assert_eq!(
            global
                .next_readable_row_group()
                .unwrap()
                .unwrap()
                .row_group_idx,
            0
        );
        assert_eq!(queued(&global), (true, true));
        global.update_budget_after_row_group(RowBudget::new(None, Some(0)));
        assert!(global.next_readable_row_group().unwrap().is_none());
        assert_eq!(queued(&global), (false, false));
        assert!(global.queued.counts.is_empty());
    }

    #[test]
    fn frontier_reports_invalid_global_row_group_while_peeking() {
        let metadata = test_file_parquet_metadata();
        let frontier = RowGroupFrontier::new(
            metadata,
            global_plan(Some(vec![2]), None),
            RowBudget::new(None, None),
            false,
        )
        .unwrap();

        let error = frontier.peek_next_row_group().unwrap_err();
        assert!(
            error
                .to_string()
                .contains("Row group index 2 out of bounds for file with 2 row groups")
        );
    }

    #[test]
    fn metadata_row_count_overflow_is_reported() {
        let metadata = test_file_parquet_metadata();
        let mut builder = metadata.as_ref().clone().into_builder();
        let mut row_groups = builder.take_row_groups();
        let negative_row_group = row_groups
            .remove(0)
            .into_builder()
            .set_num_rows(-1)
            .build()
            .unwrap();
        let metadata = builder.set_row_groups(vec![negative_row_group]).build();

        let error = metadata.row_group_num_rows(0).unwrap_err();
        assert!(error.to_string().contains("Row count overflow"));
    }
}