pigeonhole-db 1.1.0

data/vendor/pigeonhole/src/read/budget.rs

@@ -0,0 +1,584 @@
+//! Byte-budgeted reads (D308, ICR 0103): [`Scan::read_within`](super::Scan::read_within) and
+//! [`KeysRead::read_within`](super::KeysRead::read_within) return the longest prefix of the
+//! read's cells whose size is at most a budget, as owned rows, with where to resume.
+//!
+//! The read runs the engine's bounded read (D296) with a `ByteBudget` reducer over a visitor
+//! that copies each cell's stored value: the scan's own loop is untouched (D296 R3). Sizes
+//! follow D308 §2: each cell its qualifier, its value (as [`CellRef::value`](crate::CellRef)
+//! returns it) and 16 bytes, and each row's key once.
+
+use std::ops::Bound;
+use std::sync::Arc;
+
+use pigeonhole_engine::{
+    BUDGET_ITEM_OVERHEAD, ByteBudget, CellData, CellKey, CellVisitor, ColumnPosition, FamilyId,
+    KeySet, ReadInput, ReduceSpec, TableId, Timestamp, Visit,
+};
+
+use super::{Selection, TableCore};
+use crate::cell::{Row, RowBuf};
+use crate::{Error, ErrorCode, Result};
+
+/// Where a position resumes: after a whole row, after a column inside it, or at a row (from
+/// its first column: an empty page's start).
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+enum Mark {
+    AfterRow,
+    AfterColumn(ColumnPosition),
+    AtRow,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct Position {
+    /// The table the position is from: a read of another table refuses it
+    /// (`InvalidArgument`, perf287's S1 on #1281), whatever its mark.
+    table: TableId,
+    key: Vec<u8>,
+    mark: Mark,
+}
+
+/// The byte form's version.
+const VERSION: u8 = 1;
+const SCAN: u8 = b'S';
+const KEYS: u8 = b'K';
+
+impl Position {
+    /// `version, kind, mark, table (u32 LE), key length (u32 LE), key`, then for a mark
+    /// after a column, the column's own bytes.
+    fn to_bytes(&self, kind: u8) -> Vec<u8> {
+        let mut b = vec![VERSION, kind];
+        b.push(match self.mark {
+            Mark::AfterRow => 0,
+            Mark::AfterColumn(_) => 1,
+            Mark::AtRow => 2,
+        });
+        b.extend_from_slice(&self.table.0.to_le_bytes());
+        b.extend_from_slice(&(self.key.len() as u32).to_le_bytes());
+        b.extend_from_slice(&self.key);
+        if let Mark::AfterColumn(c) = &self.mark {
+            b.extend_from_slice(&c.to_bytes());
+        }
+        b
+    }
+
+    fn from_bytes(b: &[u8], kind: u8, what: &str) -> Result<Self> {
+        let bad = || Error::new(ErrorCode::InvalidArgument, format!("not a {what}"));
+        const HEAD: usize = 11;
+        if b.len() < HEAD || b[0] != VERSION || b[1] != kind {
+            return Err(bad());
+        }
+        let table = TableId(u32::from_le_bytes(b[3..7].try_into().expect("4 bytes")));
+        let len = u32::from_le_bytes(b[7..11].try_into().expect("4 bytes")) as usize;
+        let end = HEAD.checked_add(len).ok_or_else(bad)?;
+        let key = b.get(HEAD..end).ok_or_else(bad)?.to_vec();
+        let rest = &b[end..];
+        let mark = match (b[2], rest.is_empty()) {
+            (0, true) => Mark::AfterRow,
+            (2, true) => Mark::AtRow,
+            (1, false) => {
+                let c = ColumnPosition::from_bytes(rest).map_err(|_| bad())?;
+                if c.table() != table {
+                    return Err(bad());
+                }
+                Mark::AfterColumn(c)
+            }
+            _ => return Err(bad()),
+        };
+        Ok(Self { table, key, mark })
+    }
+
+    /// `InvalidArgument` unless the position is from `table`.
+    fn check(&self, table: TableId) -> Result<()> {
+        if self.table != table {
+            return Err(Error::new(
+                ErrorCode::InvalidArgument,
+                "a position from another table",
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Where a [`Scan::read_within`](super::Scan::read_within) page stopped, to resume from with
+/// [`Scan::resume_after`](super::Scan::resume_after): after its last row, after a column
+/// inside it when the page cut the row, or (an empty page) at its own start. Its bytes
+/// (`to_bytes`) are versioned and stable, so a client can hand it back later.
+/// It names its table: a read of another table refuses it (`InvalidArgument`).
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ScanPosition(Position);
+
+impl ScanPosition {
+    /// The position as bytes.
+    pub fn to_bytes(&self) -> Vec<u8> {
+        self.0.to_bytes(SCAN)
+    }
+
+    /// The position [`to_bytes`](ScanPosition::to_bytes) wrote; anything else is
+    /// `InvalidArgument`.
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
+        Position::from_bytes(bytes, SCAN, "scan position").map(Self)
+    }
+
+    /// The row the position is in or after.
+    pub fn row(&self) -> &[u8] {
+        &self.0.key
+    }
+}
+
+/// Where a [`KeysRead::read_within`](super::KeysRead::read_within) page stopped, to resume
+/// from with [`KeysRead::resume_after`](super::KeysRead::resume_after). It names the last
+/// key returned (and the column when the page cut its row), not an index into the key set,
+/// so a set rebuilt between pages (keys added or removed) resumes at the right key
+/// (cli692's S1 on #1250).
+/// It names its table: a read of another table refuses it (`InvalidArgument`).
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct KeysPosition(Position);
+
+impl KeysPosition {
+    /// The position as bytes.
+    pub fn to_bytes(&self) -> Vec<u8> {
+        self.0.to_bytes(KEYS)
+    }
+
+    /// The position [`to_bytes`](KeysPosition::to_bytes) wrote; anything else is
+    /// `InvalidArgument`.
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
+        Position::from_bytes(bytes, KEYS, "multi-get position").map(Self)
+    }
+
+    /// The key the position is at or after.
+    pub fn key(&self) -> &[u8] {
+        &self.0.key
+    }
+}
+
+macro_rules! page {
+    ($(#[$doc:meta])* $name:ident, $pos:ident) => {
+        $(#[$doc])*
+        #[derive(Debug, Clone)]
+        pub struct $name {
+            rows: Vec<Row>,
+            bytes: u64,
+            next_bytes: Option<u64>,
+            resume: Option<$pos>,
+        }
+
+        impl $name {
+            /// The page's rows, in key order. A row the budget cut holds the columns that fit.
+            pub fn rows(&self) -> &[Row] {
+                &self.rows
+            }
+
+            /// The rows, owned.
+            pub fn into_rows(self) -> Vec<Row> {
+                self.rows
+            }
+
+            /// The page's size by D308's rule: each cell's qualifier, value and 16 bytes,
+            /// and each row's key once. Never more than the budget.
+            pub fn bytes(&self) -> u64 {
+                self.bytes
+            }
+
+            /// Whether the budget stopped the read (a `limit` stopping it first does not
+            /// count).
+            pub fn cut_by_budget(&self) -> bool {
+                self.next_bytes.is_some()
+            }
+
+            /// With [`cut_by_budget`](Self::cut_by_budget), the size of the first column that
+            /// didn't fit, its row key included: a read resumed at this page's
+            /// [`resume`](Self::resume) with a budget of this many bytes returns it.
+            pub fn next_bytes(&self) -> Option<u64> {
+                self.next_bytes
+            }
+
+            /// Where to resume: `Some` whenever the read stopped before its input ended (by
+            /// the budget or by `limit`), `None` once it is exhausted. An empty page cut by
+            /// the budget resumes at its own start.
+            pub fn resume(&self) -> Option<&$pos> {
+                self.resume.as_ref()
+            }
+        }
+    };
+}
+
+page!(
+    /// A [`Scan::read_within`](super::Scan::read_within) result (D308).
+    ScanPage,
+    ScanPosition
+);
+page!(
+    /// A [`KeysRead::read_within`](super::KeysRead::read_within) result (D308).
+    KeysPage,
+    KeysPosition
+);
+
+/// A [`Nearest::run_within`](super::Nearest::run_within) result (D308 PR 3, #1302): the
+/// nearest neighbours, nearest first, that fit a byte budget.
+///
+/// There's no resume position. A nearest search ranks every candidate it reads, so a later
+/// page couldn't continue the ranking; a larger budget, or a larger `k`, returns more.
+#[derive(Debug, Clone)]
+pub struct NearestPage {
+    neighbours: Vec<super::Neighbour>,
+    bytes: u64,
+    next_bytes: Option<u64>,
+}
+
+impl NearestPage {
+    /// The longest prefix of `ranked` (nearest first) whose sizes fit `budget`: a neighbour
+    /// is in whole or not at all, and the first that doesn't fit ends the page.
+    pub(super) fn cut(ranked: Vec<super::Neighbour>, budget: u64) -> Self {
+        let mut bytes = 0u64;
+        let mut keep = ranked.len();
+        let mut next_bytes = None;
+        for (i, n) in ranked.iter().enumerate() {
+            let size = neighbour_size(n);
+            if bytes + size > budget {
+                keep = i;
+                next_bytes = Some(size);
+                break;
+            }
+            bytes += size;
+        }
+        let mut neighbours = ranked;
+        neighbours.truncate(keep);
+        NearestPage {
+            neighbours,
+            bytes,
+            next_bytes,
+        }
+    }
+
+    /// The neighbours that fit, nearest first, in [`run`](super::Nearest::run)'s order.
+    pub fn neighbours(&self) -> &[super::Neighbour] {
+        &self.neighbours
+    }
+
+    /// The neighbours, owned.
+    pub fn into_neighbours(self) -> Vec<super::Neighbour> {
+        self.neighbours
+    }
+
+    /// The page's size by D308's rule: each neighbour's row key, qualifier and 16 bytes.
+    /// Never more than the budget.
+    pub fn bytes(&self) -> u64 {
+        self.bytes
+    }
+
+    /// Whether the budget left out a neighbour the search found (not `k` or the input
+    /// running out).
+    pub fn cut_by_budget(&self) -> bool {
+        self.next_bytes.is_some()
+    }
+
+    /// With [`cut_by_budget`](Self::cut_by_budget), the size of the first neighbour left
+    /// out: a budget of `bytes() + next_bytes()` includes it.
+    pub fn next_bytes(&self) -> Option<u64> {
+        self.next_bytes
+    }
+}
+
+/// A neighbour's size (D308 §2): its row key, its qualifier and the overhead.
+fn neighbour_size(n: &super::Neighbour) -> u64 {
+    n.row.len() as u64 + n.qualifier.len() as u64 + BUDGET_ITEM_OVERHEAD
+}
+
+/// Copies each cell's stored value (tag byte first) for the page.
+struct CopyValue;
+
+impl CellVisitor for CopyValue {
+    type Item = Vec<u8>;
+    fn visit(
+        &mut self,
+        _: &[u8],
+        _: FamilyId,
+        _: &[u8],
+        _: Timestamp,
+        value: &[u8],
+    ) -> Visit<Vec<u8>> {
+        Visit::Emit(value.to_vec())
+    }
+}
+
+/// A cell's own size (D308 §2): its qualifier, its value as `CellRef::value` gives it (the
+/// stored bytes after the tag), and the overhead. The reducer's size function and the
+/// page's row sizes both use it.
+fn own_size(qualifier: &[u8], stored: &[u8]) -> u64 {
+    qualifier.len() as u64 + stored.len().saturating_sub(1) as u64 + BUDGET_ITEM_OVERHEAD
+}
+
+/// [`own_size`] as the reducer's size function. It takes `&Vec<u8>` because the reducer
+/// hands its size function `&Item`, and the item is the copied value.
+#[allow(clippy::ptr_arg)]
+fn size(cell: &pigeonhole_engine::CellRef<'_>, stored: &Vec<u8>) -> u64 {
+    own_size(cell.qualifier, stored)
+}
+
+/// A page's parts before its resume point is set.
+struct Read {
+    rows: Vec<Row>,
+    /// Each row's size, by D308's rule.
+    row_bytes: Vec<u64>,
+    last: Option<CellKey>,
+    overflow: Option<pigeonhole_engine::Overflow>,
+    /// The row limit stopped the read first.
+    cut_by_rows: bool,
+    table: pigeonhole_engine::TableId,
+}
+
+fn read(
+    core: &TableCore,
+    sel: &mut Selection,
+    input: ReadInput,
+    resume: InRow,
+    budget: u64,
+    row_limit: Option<usize>,
+) -> Result<Read> {
+    let (info, snapshot, families) = sel.start(core)?;
+    let mut spec = ReduceSpec::new(input);
+    let mut r = sel.spec.clone();
+    r.families = families.into_vec();
+    spec.read = r;
+    spec.resume = resume;
+    let table = core.info.id;
+    // Two reducer types, so a read without a row limit runs the budget's code alone (#1284).
+    let page = match row_limit {
+        None => core
+            .with(|e| {
+                e.reduce_into(
+                    &snapshot,
+                    table,
+                    &spec,
+                    || CopyValue,
+                    || ByteBudget::new(budget, size),
+                )
+            })?
+            .0
+            .into_page(),
+        Some(n) => core
+            .with(|e| {
+                e.reduce_into(
+                    &snapshot,
+                    table,
+                    &spec,
+                    || CopyValue,
+                    || ByteBudget::new(budget, size).with_rows(n),
+                )
+            })?
+            .0
+            .into_page(),
+    };
+    let last = page.items.last().map(|(k, _)| k.clone());
+    let (mut rows, mut row_bytes) = (Vec::new(), Vec::new());
+    let mut buf: Option<(RowBuf, u64)> = None;
+    for (key, stored) in page.items {
+        let own = own_size(&key.qualifier, &stored);
+        if buf.as_ref().is_none_or(|(b, _)| b.key != key.row) {
+            if let Some((b, n)) = buf.take() {
+                rows.push(Row::new(b));
+                row_bytes.push(n);
+            }
+            let mut b = RowBuf::new(Arc::clone(&info), 1);
+            b.key = key.row.clone();
+            buf = Some((b, key.row.len() as u64));
+        }
+        let (b, n) = buf.as_mut().expect("set above");
+        *n += own;
+        let start = b.qualifiers.len();
+        b.qualifiers.extend_from_slice(&key.qualifier);
+        let range = start..b.qualifiers.len();
+        b.push(key.family, range, CellData::from_stored(key.ts, stored));
+    }
+    if let Some((b, n)) = buf {
+        rows.push(Row::new(b));
+        row_bytes.push(n);
+    }
+    debug_assert_eq!(
+        row_bytes.iter().sum::<u64>(),
+        page.bytes,
+        "the rows' sizes are the reducer's"
+    );
+    Ok(Read {
+        rows,
+        row_bytes,
+        last,
+        overflow: page.overflow,
+        cut_by_rows: page.cut_by_rows,
+        table,
+    })
+}
+
+impl Read {
+    /// The page's resume point when the budget cut it (`None` otherwise), by the pinned
+    /// rules: after the last kept column when the cut is inside its row, after the last kept
+    /// row when it is at a row boundary, and at the page's own start when nothing fit.
+    fn resume(&self, given: Option<&Position>) -> Option<Position> {
+        let o = self.overflow.as_ref()?;
+        Some(match &self.last {
+            Some(l) if l.row == o.key.row => Position {
+                table: self.table,
+                key: l.row.clone(),
+                mark: Mark::AfterColumn(ColumnPosition::new(
+                    self.table,
+                    l.family,
+                    l.qualifier.clone(),
+                )),
+            },
+            Some(l) => Position {
+                table: self.table,
+                key: l.row.clone(),
+                mark: Mark::AfterRow,
+            },
+            None => given.cloned().unwrap_or(Position {
+                table: self.table,
+                key: o.key.row.clone(),
+                mark: Mark::AtRow,
+            }),
+        })
+    }
+}
+
+/// A resume inside a row: the row and the column to resume after.
+type InRow = Option<(Vec<u8>, ColumnPosition)>;
+
+/// The later of two range starts.
+fn later(a: Bound<Vec<u8>>, b: Bound<Vec<u8>>) -> Bound<Vec<u8>> {
+    use Bound::{Excluded, Included, Unbounded};
+    let rank = |x: &Bound<Vec<u8>>| match x {
+        Unbounded => (None, 0),
+        Included(k) => (Some(k.clone()), 0),
+        Excluded(k) => (Some(k.clone()), 1),
+    };
+    if rank(&a) >= rank(&b) { a } else { b }
+}
+
+/// A scan's range and resume after `position` (cli692's pinned rule 3): it starts at the
+/// later of the range's start and the position; a position before the range starts at the
+/// range, one past its end reads nothing. A mid-row position whose row the range excludes
+/// resumes after that row (perf287's N1 on #1264).
+fn clamp(start: Bound<Vec<u8>>, position: Option<&Position>) -> (Bound<Vec<u8>>, InRow) {
+    let Some(p) = position else {
+        return (start, None);
+    };
+    match &p.mark {
+        Mark::AtRow => (later(start, Bound::Included(p.key.clone())), None),
+        Mark::AfterRow => (later(start, Bound::Excluded(p.key.clone())), None),
+        Mark::AfterColumn(c) => {
+            let at = Bound::Included(p.key.clone());
+            if later(start.clone(), at.clone()) == at {
+                (at, Some((p.key.clone(), c.clone())))
+            } else {
+                // The range starts after the position's column (its row excluded, or a
+                // later row): the position is before the range.
+                (start, None)
+            }
+        }
+    }
+}
+
+/// [`Scan::read_within`](super::Scan::read_within).
+pub(super) fn scan(
+    core: &TableCore,
+    sel: &mut Selection,
+    (start, end): (Bound<Vec<u8>>, Bound<Vec<u8>>),
+    limit: Option<u64>,
+    after: Option<&ScanPosition>,
+    budget: u64,
+) -> Result<ScanPage> {
+    let given = after.map(|p| &p.0);
+    if let Some(p) = given {
+        p.check(core.info.id)?;
+    }
+    if limit == Some(0) {
+        return Ok(ScanPage {
+            rows: Vec::new(),
+            bytes: 0,
+            next_bytes: None,
+            resume: None,
+        });
+    }
+    let (lo, resume) = clamp(start, given);
+    // The row limit runs in the reducer (#1280, perf287's N4 on #1281): the read stops at
+    // whichever of the budget and `limit` comes first, reading no further.
+    let row_limit = limit.map(|n| usize::try_from(n).unwrap_or(usize::MAX));
+    let r = read(
+        core,
+        sel,
+        ReadInput::Ranges(vec![(lo, end)]),
+        resume,
+        budget,
+        row_limit,
+    )?;
+    let resume = if r.cut_by_rows {
+        // `limit` stopped the read first: resume after its last row.
+        r.last.as_ref().map(|l| Position {
+            table: core.info.id,
+            key: l.row.clone(),
+            mark: Mark::AfterRow,
+        })
+    } else {
+        r.resume(given)
+    };
+    Ok(ScanPage {
+        resume: resume.map(ScanPosition),
+        next_bytes: r.overflow.as_ref().map(|o| o.bytes),
+        bytes: r.row_bytes.iter().sum(),
+        rows: r.rows,
+    })
+}
+
+/// [`KeysRead::read_within`](super::KeysRead::read_within).
+pub(super) fn keys(
+    core: &TableCore,
+    sel: &mut Selection,
+    set: &KeySet,
+    after: Option<&KeysPosition>,
+    budget: u64,
+) -> Result<KeysPage> {
+    let given = after.map(|p| &p.0);
+    if let Some(p) = given {
+        p.check(core.info.id)?;
+    }
+    // The keys from the position on (cli692's S1: by key, so a rebuilt set resumes at the
+    // right one), and the resume inside the position's row when it names a column.
+    let (from, resume) = match given {
+        None => (set.clone(), None),
+        Some(p) => {
+            let keep = |k: &[u8]| match &p.mark {
+                Mark::AfterRow => k > p.key.as_slice(),
+                Mark::AfterColumn(_) | Mark::AtRow => k >= p.key.as_slice(),
+            };
+            let kept = KeySet::from_sorted(set.iter().filter(|k| keep(k)))
+                .expect("a sorted set's subset is sorted");
+            let resume = match &p.mark {
+                Mark::AfterColumn(c) => Some((p.key.clone(), c.clone())),
+                _ => None,
+            };
+            (kept, resume)
+        }
+    };
+    if from.is_empty() {
+        return Ok(KeysPage {
+            rows: Vec::new(),
+            bytes: 0,
+            next_bytes: None,
+            resume: None,
+        });
+    }
+    let r = read(
+        core,
+        sel,
+        ReadInput::Keys(Arc::new(from)),
+        resume,
+        budget,
+        None,
+    )?;
+    Ok(KeysPage {
+        resume: r.resume(given).map(KeysPosition),
+        next_bytes: r.overflow.as_ref().map(|o| o.bytes),
+        bytes: r.row_bytes.iter().sum(),
+        rows: r.rows,
+    })
+}