moraine 0.6.5

//! Which of a file's rows a read decodes, and where each emitted row sits
//! in the file.

use std::sync::Arc;

use object_store::path::Path;
use parquet::{
    arrow::{
        ParquetRecordBatchStreamBuilder, arrow_reader::RowSelection, async_reader::AsyncFileReader,
    },
    file::metadata::{PageIndexPolicy, ParquetMetaData},
};

use crate::{
    data_file::usize_as_u64,
    error::{Error, Result},
};

/// Which of a file's rows a scoped read decodes.
#[derive(Debug, Clone, Copy)]
pub(crate) enum ScopedRows<'a> {
    /// Every row, in file order.
    All,
    /// Only the rows at these physical positions; entries come back in
    /// position order.
    At(&'a RowPositions),
    /// One row group, from the physical position `from` (at or after the
    /// group's first row) through the group's end.
    RowGroup {
        /// Index of the row group in file order.
        group: usize,
        /// Physical position of the first row to read.
        from: u64,
    },
}

impl ScopedRows<'_> {
    /// Whether the read can be answered without touching the store.
    pub(super) fn is_empty(self) -> bool {
        matches!(self, Self::At(wanted) if wanted.is_empty())
    }

    /// A selection reads the page index so it can skip whole pages.
    pub(super) fn page_index_policy(self) -> PageIndexPolicy {
        match self {
            Self::All => PageIndexPolicy::Skip,
            Self::At(_) | Self::RowGroup { .. } => PageIndexPolicy::Optional,
        }
    }
}

/// The row groups and rows a scoped read narrows its reader to.
pub(super) struct RowScope {
    groups: Option<Vec<usize>>,
    rows: Option<RowSelection>,
}

impl RowScope {
    /// Narrows `builder` to this scope.
    pub(super) fn narrow<T: AsyncFileReader>(
        self,
        mut builder: ParquetRecordBatchStreamBuilder<T>,
    ) -> ParquetRecordBatchStreamBuilder<T> {
        if let Some(groups) = self.groups {
            builder = builder.with_row_groups(groups);
        }
        if let Some(rows) = self.rows {
            builder = builder.with_row_selection(rows);
        }
        builder
    }
}

/// Physical Parquet row positions, sorted and duplicate-free.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct RowPositions(Arc<[u64]>);

impl RowPositions {
    /// Sorts and deduplicates `positions`.
    pub(crate) fn from_unsorted(mut positions: Vec<u64>) -> Self {
        positions.sort_unstable();
        positions.dedup();
        Self(positions.into())
    }

    /// The ordered positions.
    pub(crate) fn as_slice(&self) -> &[u64] {
        &self.0
    }

    /// Whether no physical row is selected.
    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Whether the selected set contains `position`.
    pub(crate) fn contains(&self, position: u64) -> bool {
        self.0.binary_search(&position).is_ok()
    }
}

impl FromIterator<u64> for RowPositions {
    fn from_iter<T: IntoIterator<Item = u64>>(iter: T) -> Self {
        Self::from_unsorted(iter.into_iter().collect())
    }
}

/// Maps the nth row a read emits to the file ordinal it sits at.
#[derive(Debug, Clone, Copy)]
pub(super) enum Ordinals<'a> {
    /// Every row was read, so the nth emitted row is the nth row.
    Dense,
    /// A read of consecutive rows starting at this physical position.
    Offset(u64),
    /// Only these positions were read, in this order.
    Selected(&'a [u64]),
}

impl Ordinals<'_> {
    /// The file ordinal of the `nth` row emitted.
    pub(super) fn at(self, nth: usize) -> Result<u64> {
        match self {
            Self::Dense => Ok(usize_as_u64(nth)),
            Self::Offset(start) => Ok(start.saturating_add(usize_as_u64(nth))),
            Self::Selected(positions) => positions.get(nth).copied().ok_or_else(|| {
                Error::Corruption(
                    "scoped read: the reader emitted more rows than the selection named".to_owned(),
                )
            }),
        }
    }
}

/// Owned ordinal mapping retained by a returned stream.
pub(super) enum OwnedOrdinals {
    Dense,
    Offset(u64),
    Selected(RowPositions),
}

impl OwnedOrdinals {
    pub(super) fn borrowed(&self) -> Ordinals<'_> {
        match self {
            Self::Dense => Ordinals::Dense,
            Self::Offset(start) => Ordinals::Offset(*start),
            Self::Selected(positions) => Ordinals::Selected(positions.as_slice()),
        }
    }
}

pub(super) fn scoped_selection(
    rows: ScopedRows<'_>,
    metadata: &ParquetMetaData,
    path: &Path,
) -> Result<(RowScope, OwnedOrdinals)> {
    let total = total_rows(metadata, path)?;
    let whole_file = |rows| RowScope { groups: None, rows };
    match rows {
        ScopedRows::All => Ok((whole_file(None), OwnedOrdinals::Dense)),
        ScopedRows::At(positions) => Ok((
            whole_file(Some(row_selection(positions.as_slice(), total)?)),
            OwnedOrdinals::Selected(positions.clone()),
        )),
        ScopedRows::RowGroup { group, from } => {
            let counts = row_group_row_counts(metadata, path)?;
            let group_rows = counts.get(group).copied().ok_or_else(|| {
                Error::Corruption(format!("scoped read: {path} has no row group {group}"))
            })?;
            let group_start: u64 = counts[..group].iter().sum();
            let offset = from
                .checked_sub(group_start)
                .and_then(|offset| usize::try_from(offset).ok())
                .filter(|offset| usize_as_u64(*offset) <= group_rows)
                .ok_or_else(|| {
                    Error::Corruption(format!(
                        "scoped read: resume row {from} is outside row group {group} of {path}"
                    ))
                })?;
            let group_rows = usize::try_from(group_rows).map_err(|_| {
                Error::Corruption(format!(
                    "scoped read: row group {group} of {path} is larger than memory"
                ))
            })?;
            Ok((
                RowScope {
                    groups: Some(vec![group]),
                    rows: Some(RowSelection::from_consecutive_ranges(
                        std::iter::once(offset..group_rows),
                        group_rows,
                    )),
                },
                OwnedOrdinals::Offset(from),
            ))
        }
    }
}

/// Row counts of the file's row groups, in file order.
pub(super) fn row_group_row_counts(metadata: &ParquetMetaData, path: &Path) -> Result<Vec<u64>> {
    metadata
        .row_groups()
        .iter()
        .map(|group| {
            u64::try_from(group.num_rows()).map_err(|_| {
                Error::Corruption(format!("scoped read: {path} reports a negative row count"))
            })
        })
        .collect()
}

/// The row selection naming `positions` in a file of `total_rows`.
/// Positions must be ascending and unique.
fn row_selection(positions: &[u64], total_rows: usize) -> Result<RowSelection> {
    let ranges = positions
        .iter()
        .map(|&position| {
            usize::try_from(position)
                .ok()
                .filter(|start| *start < total_rows)
                .map(|start| start..start.saturating_add(1))
                .ok_or_else(|| {
                    Error::Corruption(format!(
                        "scoped read: selected row {position} is past the file's {total_rows} rows"
                    ))
                })
        })
        .collect::<Result<Vec<_>>>()?;
    Ok(RowSelection::from_consecutive_ranges(
        ranges.into_iter(),
        total_rows,
    ))
}

/// A file's row count, as the selection builder counts it.
pub(super) fn total_rows(metadata: &ParquetMetaData, path: &Path) -> Result<usize> {
    usize::try_from(metadata.file_metadata().num_rows())
        .map_err(|_| Error::Corruption(format!("scoped read: {path} reports a negative row count")))
}