quatzal-storage 0.1.0

// SPDX-License-Identifier: Apache-2.0
//! In-memory sorted table for records not yet flushed to an SSTable.

use std::collections::BTreeMap;

use crate::pax::PaxRecord;

#[derive(Default)]
pub struct MemTable {
    map: BTreeMap<Vec<u8>, PaxRecord>,
    approx_bytes: usize,
}

impl MemTable {
    pub fn new() -> Self {
        Self::default()
    }

    /// Insert a record, resolving `VectorField::Unchanged` against whatever this MemTable
    /// already knows for the key (if anything) so the MemTable never loses the vector
    /// pointer to an in-flight overwrite. If this is the first time the key appears in this
    /// MemTable, `Unchanged` is kept as-is -- the true value lives further down in an
    /// existing L0/L1 table and `resolve_from_records` will find it there at read time.
    pub fn put(&mut self, mut record: PaxRecord) {
        if matches!(record.vector, crate::pax::VectorField::Unchanged)
            && let Some(existing) = self.map.get(&record.key)
        {
            record.vector = existing.vector.clone();
        }
        let size = record.key.len() + record.scalar_blob.len() + 16;
        let size = size
            + match &record.vector {
                crate::pax::VectorField::Set(b) => b.len(),
                _ => 0,
            };
        if let Some(old) = self.map.insert(record.key.clone(), record) {
            self.approx_bytes = self.approx_bytes.saturating_sub(
                old.key.len()
                    + old.scalar_blob.len()
                    + match &old.vector {
                        crate::pax::VectorField::Set(b) => b.len(),
                        _ => 0,
                    }
                    + 16,
            );
        }
        self.approx_bytes += size;
    }

    pub fn get(&self, key: &[u8]) -> Option<&PaxRecord> {
        self.map.get(key)
    }

    pub fn approx_bytes(&self) -> usize {
        self.approx_bytes
    }

    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Records in ascending key order (a `BTreeMap`'s natural iteration order) — this is
    /// exactly the order an SSTable flush wants them in.
    pub fn iter_sorted(&self) -> impl Iterator<Item = &PaxRecord> {
        self.map.values()
    }

    pub fn clear(&mut self) {
        self.map.clear();
        self.approx_bytes = 0;
    }
}