crates/nornir-parquet-delta/src/store.rs
use std::io::Write;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use parquet::file::metadata::ParquetMetaData;
use crate::reader::{Intervals, MirrorReader};
use crate::{plan, Error, Need, Piece, Plan, Presence, Range, FOOTER_TRAILER_LEN, PARQUET_MAGIC};
/// **What the server has to say about its file** before a mirror of it can exist: how
/// long it is, its footer bytes, and a token that changes when the file does.
///
/// The token is deliberately an opaque `String` and not something this crate derives.
/// Every system that has this problem already has an answer to it — an ETag, an object
/// generation, a cut version, a content hash — and inventing a second one here would give
/// one fact two homes that can disagree. This crate only ever asks whether it CHANGED.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Source {
/// The source file's total length in bytes. Every offset in the footer is relative
/// to a file this long, so it is part of the file's identity, not a hint.
pub len: u64,
/// The footer's bytes — the range [`crate::footer_range`] named: the thrift metadata,
/// its `u32` length, and the trailing `PAR1`.
pub footer: Vec<u8>,
/// The source's version token.
pub version: String,
}
const STATE_MAGIC: &[u8; 8] = b"PQDELTA1";
const STATE_EXT: &str = "mirror";
/// **A partially-filled local copy of a remote parquet file.**
///
/// The file on disk is exactly as long as the source and holds the source's bytes at the
/// source's offsets — the parts that have arrived. The rest is a hole, and the filesystem
/// does not store holes, so a mirror that holds 3 % of a 1.5 GB table occupies 3 % of
/// 1.5 GB (see [`Mirror::on_disk_bytes`], which reads the real block count and not the
/// apparent length).
///
/// Beside it sits `<path>.mirror`: the version token and the presence bitmap.
///
/// # The ordering rule, which is the only way to corrupt this
///
/// [`Mirror::accept`] writes bytes and marks a bit **in memory**. [`Mirror::save`]
/// fsyncs the data and only then writes the bitmap. Losing the bitmap after the data
/// landed costs a re-fetch; keeping a bitmap whose data did not land is a claim to hold
/// bytes that are zeros, which [`crate::MirrorReader`] would then serve as a valid page.
/// Under-claiming is free, over-claiming is corruption — so the fsync comes first, always.
pub struct Mirror {
path: PathBuf,
state_path: PathBuf,
file: Arc<std::fs::File>,
meta: Arc<ParquetMetaData>,
presence: Presence,
version: String,
source_len: u64,
footer: Range,
/// `chunk_ranges[rg * columns + col]` — derived once from the footer, because
/// rebuilding the interval set on every `accept` otherwise walks the whole metadata.
chunk_ranges: Vec<Range>,
intervals: Arc<Intervals>,
}
impl Mirror {
/// **Create (or truncate and re-create) a mirror of `src` at `path`.**
///
/// Writes the sparse file, the leading `PAR1` and the footer, and an empty bitmap.
/// Use [`Mirror::open_or_create`] when a mirror may already be there.
pub fn create(path: impl AsRef<Path>, src: &Source) -> Result<Self, Error> {
let path = path.as_ref().to_path_buf();
if src.footer.len() < FOOTER_TRAILER_LEN as usize
|| src.footer[src.footer.len() - 4..] != PARQUET_MAGIC
{
return Err(Error::NotParquet { why: "footer does not end in PAR1".into() });
}
let footer_len = src.footer.len() as u64;
if footer_len + PARQUET_MAGIC.len() as u64 > src.len {
return Err(Error::NotParquet {
why: format!("a {}-byte footer does not fit a {}-byte file", footer_len, src.len),
});
}
if let Some(dir) = path.parent() {
std::fs::create_dir_all(dir)?;
}
let file = std::fs::OpenOptions::new().read(true).write(true).create(true).truncate(true).open(&path)?;
// The hole. `set_len` past EOF is the POSIX way to ask for a sparse file; whether
// the filesystem grants it is the filesystem's business, which is why
// `on_disk_bytes` reports what actually happened rather than assuming.
file.set_len(src.len)?;
// Byte 0 of a sparse file is 0x00, and a parquet reader's first act is to check
// for `PAR1`. Without these four bytes the ordinary reader rejects the mirror
// before it looks at anything else.
file.write_all_at(&PARQUET_MAGIC, 0)?;
let footer = Range { offset: src.len - footer_len, len: footer_len };
file.write_all_at(&src.footer, footer.offset)?;
file.sync_data()?;
let meta = Arc::new(parquet::file::metadata::ParquetMetaDataReader::new().parse_and_finish(&file)?);
let columns = meta.file_metadata().schema_descr().num_columns();
let row_groups = meta.num_row_groups();
let presence = Presence::empty(row_groups, columns);
let chunk_ranges = chunk_ranges(&meta);
let mut m = Self {
state_path: state_path(&path),
path,
file: Arc::new(file),
meta,
presence,
version: src.version.clone(),
source_len: src.len,
footer,
chunk_ranges,
intervals: Arc::new(Intervals::default()),
};
m.rebuild_intervals();
m.save()?;
Ok(m)
}
/// **Open an existing mirror, refusing a stale one.**
///
/// [`Error::StaleSource`] when the token differs — a re-written source moves every
/// byte offset in the footer, so a mirror of the old one is not "mostly right", it is
/// arbitrary bytes at plausible offsets.
pub fn open(path: impl AsRef<Path>, expect_version: &str) -> Result<Self, Error> {
let path = path.as_ref().to_path_buf();
let state_path = state_path(&path);
let state = std::fs::read(&state_path)?;
let st = State::decode(&state)?;
if st.version != expect_version {
return Err(Error::StaleSource { held: st.version, expected: expect_version.to_string() });
}
let file = std::fs::OpenOptions::new().read(true).write(true).open(&path)?;
let on_disk_len = file.metadata()?.len();
if on_disk_len != st.source_len {
return Err(Error::BadState {
why: format!("mirror file is {on_disk_len} bytes, its state says {}", st.source_len),
});
}
let meta = Arc::new(parquet::file::metadata::ParquetMetaDataReader::new().parse_and_finish(&file)?);
let columns = meta.file_metadata().schema_descr().num_columns();
let row_groups = meta.num_row_groups();
// The sidecar's grid must be the FOOTER's grid. A word count does not pin a grid
// (a 20x20 bitmap and a 20x21 one are both seven words), so without this a
// sidecar written against a different schema would come back as a presence that
// maps bits onto the wrong chunks — every one of them a claim to hold bytes that
// are somewhere else entirely.
if (st.row_groups, st.columns) != (row_groups, columns) {
return Err(Error::BadState {
why: format!(
"state file is for a {}x{} grid, this footer is {row_groups}x{columns}",
st.row_groups, st.columns
),
});
}
let presence = Presence::from_words(row_groups, columns, st.words)?;
let chunk_ranges = chunk_ranges(&meta);
let mut m = Self {
state_path,
path,
file: Arc::new(file),
meta,
presence,
version: st.version,
source_len: st.source_len,
footer: Range { offset: st.source_len - st.footer_len, len: st.footer_len },
chunk_ranges,
intervals: Arc::new(Intervals::default()),
};
m.rebuild_intervals();
Ok(m)
}
/// **What version is the mirror at `path` a mirror OF?** — read from its sidecar,
/// without a source and without opening the mirror itself.
///
/// [`Mirror::open`] takes the version it EXPECTS, which is the right shape whenever
/// the source can be asked: a mirror that disagrees with the live source is stale and
/// must be refused. But a consumer that is OFFLINE has no source to ask, and refusing
/// to open the file it already holds because it cannot verify freshness turns "the
/// server is down" into "you have nothing" — which is precisely the situation a
/// mirror exists for.
///
/// So the token is readable on its own, and the caller decides. An offline reader
/// opens on this version and knows what it has assumed; the first answer from a live
/// source re-runs [`Mirror::open`]'s comparison and discards the mirror if the cut
/// moved. The assumption is bounded by the next successful connection, and it is
/// never silent — a caller that does not want it simply does not call this.
pub fn held_version(path: impl AsRef<Path>) -> Result<String, Error> {
let state = std::fs::read(state_path(path.as_ref()))?;
Ok(State::decode(&state)?.version)
}
/// Open the mirror at `path` if it is of this exact source, else start a fresh one.
/// The whole freshness policy, in one call.
pub fn open_or_create(path: impl AsRef<Path>, src: &Source) -> Result<Self, Error> {
match Self::open(path.as_ref(), &src.version) {
Ok(m) => Ok(m),
Err(Error::StaleSource { .. } | Error::Io(_) | Error::BadState { .. }) => {
Self::create(path, src)
}
Err(e) => Err(e),
}
}
/// The source's footer, parsed. This is the same `ParquetMetaData` the server has —
/// a mirror knows the shape of the whole table from its first byte, which is what
/// makes [`Mirror::plan`] answerable without asking anyone.
#[must_use]
pub fn metadata(&self) -> &Arc<ParquetMetaData> {
&self.meta
}
/// Which chunks are held.
#[must_use]
pub fn presence(&self) -> &Presence {
&self.presence
}
/// The source version token this mirror is of.
#[must_use]
pub fn version(&self) -> &str {
&self.version
}
/// Where the mirror file is.
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
/// **What this need still costs.** [`Plan::satisfied`] is the "do I hold what this
/// query needs?" answer with the bytes attached.
pub fn plan(&self, need: &Need) -> Result<Plan, Error> {
plan(&self.meta, need, Some(&self.presence))
}
/// **Do I hold every chunk this need names?** The cheap form of [`Mirror::plan`].
pub fn holds(&self, need: &Need) -> Result<bool, Error> {
let n_rg = self.meta.num_row_groups();
let n_col = self.presence.columns();
for &rg in &need.row_groups {
if rg >= n_rg {
return Err(Error::OutOfRange { what: "row group", index: rg, count: n_rg });
}
for &col in &need.columns {
if col >= n_col {
return Err(Error::OutOfRange { what: "column", index: col, count: n_col });
}
if !self.presence.holds(rg, col) {
return Ok(false);
}
}
}
Ok(true)
}
/// **Take delivery of one chunk.** The bytes must be exactly the length the footer
/// says that chunk is, or they are refused — a short or padded body written at a
/// source offset puts the file permanently out of phase, and nothing downstream would
/// notice.
///
/// Marks the bit in memory only; see [`Mirror::save`] for why.
pub fn accept(&mut self, piece: Piece, body: &[u8]) -> Result<(), Error> {
let want = self.chunk_range(piece.row_group, piece.column)?;
if want != piece.range {
return Err(Error::BadState {
why: format!(
"piece (rg {}, col {}) says {:?}, the footer says {want:?}",
piece.row_group, piece.column, piece.range
),
});
}
if body.len() as u64 != want.len {
return Err(Error::WrongLength {
row_group: piece.row_group,
column: piece.column,
want: want.len,
got: body.len() as u64,
});
}
self.file.write_all_at(body, want.offset)?;
self.presence.set(piece.row_group, piece.column);
self.rebuild_intervals();
Ok(())
}
/// **Make the mirror durable: fsync the data, THEN write the bitmap.**
///
/// That order is the correctness rule of the whole crate. See the type docs.
pub fn save(&self) -> Result<(), Error> {
self.file.sync_data()?;
let tmp = self.state_path.with_extension(format!("{STATE_EXT}.tmp"));
{
let mut f = std::fs::File::create(&tmp)?;
f.write_all(&State::encode(self))?;
f.sync_all()?;
}
std::fs::rename(&tmp, &self.state_path)?;
Ok(())
}
/// A [`parquet::file::reader::ChunkReader`] over what is held. Give it to
/// `ParquetRecordBatchReaderBuilder::try_new`.
#[must_use]
pub fn reader(&self) -> MirrorReader {
MirrorReader::new(
Arc::clone(&self.file),
self.source_len,
Arc::clone(&self.intervals),
self.presence.count(),
self.presence.total(),
)
}
/// The source's length — what the mirror file APPEARS to be.
#[must_use]
pub fn apparent_bytes(&self) -> u64 {
self.source_len
}
/// **What the mirror actually occupies**, from the filesystem's own block count.
///
/// Reported rather than computed on purpose: `set_len` is a request for a hole, not a
/// guarantee of one, and a filesystem that does not do sparse files (or one mounted
/// so it does not) turns this design's central claim into a 1.5 GB file of zeros.
/// The only honest way to know is to ask.
pub fn on_disk_bytes(&self) -> Result<u64, Error> {
use std::os::unix::fs::MetadataExt;
Ok(self.file.metadata()?.blocks() * 512)
}
/// Bytes of the source the mirror can serve — the footer and magic included.
#[must_use]
pub fn held_bytes(&self) -> u64 {
self.intervals.bytes()
}
fn chunk_range(&self, rg: usize, col: usize) -> Result<Range, Error> {
let cols = self.presence.columns();
if rg >= self.presence.row_groups() {
return Err(Error::OutOfRange { what: "row group", index: rg, count: self.presence.row_groups() });
}
if col >= cols {
return Err(Error::OutOfRange { what: "column", index: col, count: cols });
}
Ok(self.chunk_ranges[rg * cols + col])
}
fn rebuild_intervals(&mut self) {
let cols = self.presence.columns();
let mut ranges = Vec::with_capacity(self.presence.count() as usize + 2);
// Always present: the leading magic and the footer.
ranges.push(Range { offset: 0, len: PARQUET_MAGIC.len() as u64 });
ranges.push(self.footer);
for rg in 0..self.presence.row_groups() {
for col in 0..cols {
if self.presence.holds(rg, col) {
ranges.push(self.chunk_ranges[rg * cols + col]);
}
}
}
self.intervals = Arc::new(Intervals::from_ranges(ranges));
}
}
fn state_path(path: &Path) -> PathBuf {
let mut s = path.as_os_str().to_os_string();
s.push(".");
s.push(STATE_EXT);
PathBuf::from(s)
}
fn chunk_ranges(meta: &ParquetMetaData) -> Vec<Range> {
let cols = meta.file_metadata().schema_descr().num_columns();
let mut out = Vec::with_capacity(meta.num_row_groups() * cols);
for rg in 0..meta.num_row_groups() {
let g = meta.row_group(rg);
for col in 0..cols {
let c = g.column(col);
let offset = c.dictionary_page_offset().unwrap_or_else(|| c.data_page_offset());
out.push(Range {
offset: u64::try_from(offset).unwrap_or(0),
len: u64::try_from(c.compressed_size()).unwrap_or(0),
});
}
}
out
}
/// The sidecar, on disk.
struct State {
source_len: u64,
footer_len: u64,
row_groups: usize,
columns: usize,
version: String,
words: Vec<u64>,
}
impl State {
fn encode(m: &Mirror) -> Vec<u8> {
let v = m.version.as_bytes();
let words = m.presence.words();
let mut out = Vec::with_capacity(40 + v.len() + words.len() * 8);
out.extend_from_slice(STATE_MAGIC);
out.extend_from_slice(&m.source_len.to_le_bytes());
out.extend_from_slice(&m.footer.len.to_le_bytes());
out.extend_from_slice(&(m.presence.row_groups() as u32).to_le_bytes());
out.extend_from_slice(&(m.presence.columns() as u32).to_le_bytes());
out.extend_from_slice(&(v.len() as u32).to_le_bytes());
out.extend_from_slice(v);
for w in words {
out.extend_from_slice(&w.to_le_bytes());
}
out
}
fn decode(b: &[u8]) -> Result<Self, Error> {
let bad = |why: &str| Error::BadState { why: why.to_string() };
if b.len() < 36 || &b[..8] != STATE_MAGIC {
return Err(bad("not a PQDELTA1 state file"));
}
let u64at = |o: usize| u64::from_le_bytes(b[o..o + 8].try_into().unwrap());
let u32at = |o: usize| u32::from_le_bytes(b[o..o + 4].try_into().unwrap()) as usize;
let source_len = u64at(8);
let footer_len = u64at(16);
let row_groups = u32at(24);
let columns = u32at(28);
let vlen = u32at(32);
let vstart = 36;
let wstart = vstart + vlen;
if b.len() < wstart {
return Err(bad("state file is shorter than its version token"));
}
let version = String::from_utf8(b[vstart..wstart].to_vec())
.map_err(|_| bad("version token is not utf-8"))?;
let want_words = row_groups.saturating_mul(columns).div_ceil(64);
if b.len() != wstart + want_words * 8 {
return Err(bad("state file bitmap is the wrong length for its grid"));
}
let words = (0..want_words).map(|i| u64at(wstart + i * 8)).collect();
Ok(Self { source_len, footer_len, row_groups, columns, version, words })
}
}