crates/nornir-parquet-delta/src/plan.rs
use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData};
use crate::{Error, Presence};
/// A contiguous byte range **of the source file**. Offsets are source offsets, and the
/// mirror writes them at the same offsets — that identity is the whole design.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Range {
/// First byte.
pub offset: u64,
/// How many bytes.
pub len: u64,
}
impl Range {
/// One past the last byte.
#[must_use]
pub fn end(&self) -> u64 {
self.offset.saturating_add(self.len)
}
}
/// One (row group × leaf column) chunk, and where it lives in the source.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Piece {
/// Row group index.
pub row_group: usize,
/// Leaf column index — the same index [`ParquetMetaData::file_metadata`]'s schema
/// descriptor uses, and the same one a `ProjectionMask::leaves` takes.
pub column: usize,
/// The bytes.
pub range: Range,
}
/// **What a query needs**: some row groups, some leaf columns.
///
/// This crate never decides either. The row groups come from the caller's own statistics
/// prune, the columns from the caller's own projection — the two decisions that are
/// already made, correctly, in whatever reader is asking.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Need {
/// Row groups, any order; duplicates are harmless.
pub row_groups: Vec<usize>,
/// Leaf column indices, any order; duplicates are harmless.
pub columns: Vec<usize>,
}
impl Need {
/// The obvious constructor.
#[must_use]
pub fn new(row_groups: Vec<usize>, columns: Vec<usize>) -> Self {
Self { row_groups, columns }
}
/// **The leaf columns whose dotted path the predicate admits** — `bbox.xmin`,
/// `geometry_z12`, `node_refs.list.item`, whatever the schema calls them.
///
/// A convenience over the caller's closure and not a projection planner: it exists
/// because every consumer otherwise writes the same `(0..num_columns).filter(path)`
/// loop, and one of them writes it against the wrong index space.
pub fn columns_by_path(meta: &ParquetMetaData, admit: impl Fn(&str) -> bool) -> Vec<usize> {
let descr = meta.file_metadata().schema_descr();
(0..descr.num_columns()).filter(|&i| admit(&descr.column(i).path().string())).collect()
}
/// **The row groups the predicate admits.** The caller's prune, applied.
pub fn row_groups_where(meta: &ParquetMetaData, admit: impl Fn(&RowGroupMetaData) -> bool) -> Vec<usize> {
(0..meta.num_row_groups()).filter(|&i| admit(meta.row_group(i))).collect()
}
/// Every row group — the "mirror the whole table" need.
#[must_use]
pub fn all_row_groups(meta: &ParquetMetaData) -> Vec<usize> {
(0..meta.num_row_groups()).collect()
}
}
/// What [`plan`] worked out.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Plan {
/// The chunks to fetch, **in ascending source offset** so a transport can coalesce
/// adjacent ones into a single range request without sorting them again.
pub pieces: Vec<Piece>,
/// Bytes [`Plan::pieces`] adds up to — what this query costs on the wire.
pub fetch_bytes: u64,
/// Bytes of the need the mirror ALREADY holds — what it costs nothing.
pub held_bytes: u64,
}
impl Plan {
/// Everything the need covers, held or not.
#[must_use]
pub fn need_bytes(&self) -> u64 {
self.fetch_bytes.saturating_add(self.held_bytes)
}
/// Nothing left to fetch — the mirror can answer this query on its own.
#[must_use]
pub fn satisfied(&self) -> bool {
self.pieces.is_empty()
}
/// **Adjacent pieces merged into fetchable ranges.** A transport that pays per
/// request (an HTTP range GET, an S3 object read, a gRPC round trip) wants these;
/// one that pays per byte does not care. Contiguity is common: a row group's column
/// chunks are written back to back, so a projection over many columns of one group
/// usually collapses to a handful of ranges.
#[must_use]
pub fn coalesced(&self) -> Vec<Range> {
let mut out: Vec<Range> = Vec::new();
for p in &self.pieces {
match out.last_mut() {
Some(last) if last.end() == p.range.offset => last.len += p.range.len,
_ => out.push(p.range),
}
}
out
}
}
/// **The plan: footer + need + what is held ⇒ the ranges to fetch.**
///
/// `held = None` means an empty mirror, which is the "how big is this query cold?"
/// question a caller asks before it decides to mirror at all.
///
/// Row groups and columns are visited in ascending order and the result is sorted by
/// source offset, so the plan is a function of the SET a need names and not of the order
/// the caller happened to build it in — two callers that need the same chunks get the
/// same plan, byte for byte.
pub fn plan(meta: &ParquetMetaData, need: &Need, held: Option<&Presence>) -> Result<Plan, Error> {
let n_rg = meta.num_row_groups();
let n_col = meta.file_metadata().schema_descr().num_columns();
let mut rgs: Vec<usize> = need.row_groups.clone();
rgs.sort_unstable();
rgs.dedup();
let mut cols: Vec<usize> = need.columns.clone();
cols.sort_unstable();
cols.dedup();
if let Some(&bad) = rgs.iter().find(|&&i| i >= n_rg) {
return Err(Error::OutOfRange { what: "row group", index: bad, count: n_rg });
}
if let Some(&bad) = cols.iter().find(|&&i| i >= n_col) {
return Err(Error::OutOfRange { what: "column", index: bad, count: n_col });
}
let mut out = Plan::default();
for &rg in &rgs {
let group = meta.row_group(rg);
for &col in &cols {
let chunk = group.column(col);
// The chunk starts at its dictionary page when it has one, else at its first
// data page — the same arithmetic `ColumnChunkMetaData::byte_range` does, and
// the reason it is spelled out is that reading from `data_page_offset` on a
// dictionary-encoded column silently drops the dictionary and decodes
// garbage.
let offset = chunk.dictionary_page_offset().unwrap_or_else(|| chunk.data_page_offset());
let len = chunk.compressed_size();
let range = Range {
offset: u64::try_from(offset).unwrap_or(0),
len: u64::try_from(len).unwrap_or(0),
};
if held.is_some_and(|h| h.holds(rg, col)) {
out.held_bytes += range.len;
} else {
out.fetch_bytes += range.len;
out.pieces.push(Piece { row_group: rg, column: col, range });
}
}
}
out.pieces.sort_unstable_by_key(|p| p.range.offset);
Ok(out)
}