crates/nornir-parquet-delta/src/reader.rs
use std::io::{Read, Seek, SeekFrom};
use std::sync::Arc;
use crate::{Error, Range};
/// The merged, disjoint, ascending set of source ranges a mirror holds.
///
/// Built from the presence bitmap plus the two regions that are always present (the
/// leading `PAR1` and the footer). Merged, because a row group's column chunks are
/// written back to back and a reader asks for them in one span — an unmerged set would
/// reject that span as a hole even though every byte of it is held.
#[derive(Debug, Default, Clone)]
pub(crate) struct Intervals {
/// `(start, end)` pairs, ascending, disjoint, non-touching.
spans: Vec<(u64, u64)>,
}
impl Intervals {
pub(crate) fn from_ranges(mut ranges: Vec<Range>) -> Self {
ranges.retain(|r| r.len > 0);
ranges.sort_unstable();
let mut spans: Vec<(u64, u64)> = Vec::with_capacity(ranges.len());
for r in ranges {
match spans.last_mut() {
Some(last) if r.offset <= last.1 => last.1 = last.1.max(r.end()),
_ => spans.push((r.offset, r.end())),
}
}
Self { spans }
}
/// Is every byte of `[offset, offset + len)` held? A zero-length read is held by
/// anything — there is nothing to be wrong about.
pub(crate) fn covers(&self, offset: u64, len: u64) -> bool {
if len == 0 {
return true;
}
let end = offset.saturating_add(len);
// The last span that starts at or before `offset`.
let i = self.spans.partition_point(|&(s, _)| s <= offset);
i > 0 && self.spans[i - 1].1 >= end
}
pub(crate) fn bytes(&self) -> u64 {
self.spans.iter().map(|&(s, e)| e - s).sum()
}
}
/// **A [`parquet::file::reader::ChunkReader`] over a partially-filled mirror that turns a
/// hole into an error instead of a page of zeros.**
///
/// Hand it to `ParquetRecordBatchReaderBuilder::try_new` like any file. If the read stays
/// inside the row groups and columns the mirror holds, it behaves exactly like the source
/// file — same bytes at the same offsets, because they ARE the same bytes at the same
/// offsets. If it wanders, [`Error::Hole`] comes back out through
/// `ParquetError::External`.
#[derive(Clone)]
pub struct MirrorReader {
file: Arc<std::fs::File>,
len: u64,
held: Arc<Intervals>,
held_chunks: u64,
total_chunks: u64,
}
impl MirrorReader {
pub(crate) fn new(
file: Arc<std::fs::File>,
len: u64,
held: Arc<Intervals>,
held_chunks: u64,
total_chunks: u64,
) -> Self {
Self { file, len, held, held_chunks, total_chunks }
}
fn hole(&self, offset: u64, len: u64) -> Error {
Error::Hole {
want_offset: offset,
want_len: len,
held_chunks: self.held_chunks,
total_chunks: self.total_chunks,
}
}
fn guard(&self, offset: u64, len: u64) -> Result<(), Error> {
if self.held.covers(offset, len) {
Ok(())
} else {
Err(self.hole(offset, len))
}
}
/// Bytes of the source this mirror can serve.
#[must_use]
pub fn held_bytes(&self) -> u64 {
self.held.bytes()
}
}
fn external(e: Error) -> parquet::errors::ParquetError {
parquet::errors::ParquetError::External(Box::new(e))
}
impl parquet::file::reader::Length for MirrorReader {
fn len(&self) -> u64 {
self.len
}
}
impl parquet::file::reader::ChunkReader for MirrorReader {
type T = GuardedRead;
fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
let mut f = self.file.try_clone()?;
f.seek(SeekFrom::Start(start))?;
Ok(GuardedRead { inner: f, at: start, guard: self.clone() })
}
fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<bytes::Bytes> {
self.guard(start, length as u64).map_err(external)?;
let mut f = self.file.try_clone()?;
f.seek(SeekFrom::Start(start))?;
let mut buf = vec![0u8; length];
f.read_exact(&mut buf)?;
Ok(bytes::Bytes::from(buf))
}
}
/// The streaming half of [`MirrorReader`] — every `read` is checked before it is served.
///
/// It has to be per-`read` and not per-`get_read`: `get_read` is handed a start offset
/// and no length, so the only place the span is known is at the moment bytes are handed
/// over. A guard that only checked the start would pass a read that begins in the footer
/// and runs off into a hole.
pub struct GuardedRead {
inner: std::fs::File,
at: u64,
guard: MirrorReader,
}
impl Read for GuardedRead {
fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
let n = self.inner.read(out)?;
if n > 0 {
self.guard
.guard(self.at, n as u64)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
self.at += n as u64;
}
Ok(n)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn set(spans: &[(u64, u64)]) -> Intervals {
Intervals::from_ranges(spans.iter().map(|&(o, e)| Range { offset: o, len: e - o }).collect())
}
#[test]
fn touching_ranges_merge_so_a_span_across_two_chunks_is_held() {
let i = set(&[(100, 200), (200, 300)]);
assert_eq!(i.spans, vec![(100, 300)]);
assert!(i.covers(150, 100), "a read spanning both chunks is held");
assert_eq!(i.bytes(), 200);
}
/// RED FIRST: with `covers` written as `end <= self.spans[i-1].1` against an UNMERGED
/// set, this passes; with the merge removed it fails. Both halves are the point.
#[test]
fn a_read_that_runs_one_byte_past_the_held_span_is_a_hole() {
let i = set(&[(100, 200)]);
assert!(i.covers(100, 100));
assert!(!i.covers(100, 101), "one byte past the end is not held");
assert!(!i.covers(99, 2), "one byte before the start is not held");
assert!(!i.covers(250, 1), "a byte in the gap is not held");
}
#[test]
fn a_gap_between_two_spans_stays_a_gap() {
let i = set(&[(0, 4), (1000, 1100)]);
assert_eq!(i.spans, vec![(0, 4), (1000, 1100)]);
assert!(i.covers(0, 4));
assert!(i.covers(1000, 100));
assert!(!i.covers(4, 996), "the whole gap");
assert!(!i.covers(3, 2), "straddling the first span's end");
}
#[test]
fn an_empty_mirror_holds_nothing_but_still_answers_a_zero_length_read() {
let i = set(&[]);
assert!(!i.covers(0, 1));
assert!(i.covers(0, 0));
assert_eq!(i.bytes(), 0);
}
}