use crate::columnar;
use crate::encryption::{setup_run_encryption, Cipher, Kek, RunEncryption};
use crate::epoch::Epoch;
use crate::error::{MongrelError, Result};
use crate::index::pgm::PgmIndex;
use crate::memtable::{Row, Value};
use crate::page::{Encoding, PageStat};
use crate::row_id_set::RowIdSet;
use crate::rowid::RowId;
use crate::schema::{Schema, TypeId};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
use std::sync::Arc;
pub const RUN_MAGIC: [u8; 8] = *b"MONGRRUN";
pub const RUN_FORMAT_VERSION: u16 = 1;
pub const RUN_HEADER_VERSION: u16 = 2;
pub const RUN_HEADER_PAD: usize = 256;
const ENC_STATS_NONCE_COLUMN: u16 = u16::MAX;
const ENC_STATS_NONCE_SEQ: u32 = u16::MAX as u32;
type PageMinMax = (Option<Vec<u8>>, Option<Vec<u8>>);
type EncryptedColumnStats = Vec<(u16, Vec<PageMinMax>)>;
const GCM_TAG_LEN: usize = 16;
fn on_disk_len(page_len: usize, encrypted: bool) -> usize {
if encrypted {
page_len + GCM_TAG_LEN
} else {
page_len
}
}
pub const RUN_FLAG_ENCRYPTED: u8 = 1 << 0;
pub const RUN_FLAG_TOMBSTONE_ONLY: u8 = 1 << 1;
pub const RUN_FLAG_CLEAN: u8 = 1 << 2;
pub const RUN_FLAG_UNIFORM_EPOCH: u8 = 1 << 3;
pub const SORT_KEY_ROW_ID: u16 = 0xFFFF;
pub const SYS_ROW_ID: u16 = 0xFFFE;
pub const SYS_EPOCH: u16 = 0xFFFD;
pub const SYS_DELETED: u16 = 0xFFFC;
const LEARNED_EPSILON: usize = 64;
fn build_learned_trailer(row_ids: &[Value]) -> Vec<u8> {
let points: Vec<(u64, usize)> = row_ids
.iter()
.enumerate()
.filter_map(|(i, v)| match v {
Value::Int64(r) => Some((*r as u64, i)),
_ => None,
})
.collect();
let pgm = PgmIndex::build(&points, LEARNED_EPSILON);
bincode::serialize(&pgm).expect("pgm serialize")
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunHeader {
pub magic: [u8; 8],
pub format_version: u16,
pub header_layout_version: u16,
pub run_id: u128,
pub content_hash: [u8; 32],
pub schema_id: u64,
pub epoch_created: u64,
pub level: u8,
pub flags: u8,
pub sort_key_column_id: u16,
pub row_count: u64,
pub min_row_id: u64,
pub max_row_id: u64,
pub column_count: u64,
pub column_dir_offset: u64,
pub index_trailer_offset: u64,
pub encryption_descriptor_offset: u64,
pub footer_offset: u64,
pub encrypted_stats_offset: u64,
pub encrypted_stats_len: u64,
}
impl RunHeader {
pub fn is_encrypted(&self) -> bool {
self.flags & RUN_FLAG_ENCRYPTED != 0
}
pub fn is_clean(&self) -> bool {
self.flags & RUN_FLAG_CLEAN != 0
}
pub fn is_uniform_epoch(&self) -> bool {
self.flags & RUN_FLAG_UNIFORM_EPOCH != 0
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnPageHeader {
pub column_id: u16,
pub type_id_tag: u16,
pub encoding: u8,
pub flags: u8,
pub page_count: u32,
pub page_region_offset: u64,
pub page_region_len: u64,
pub page_stats: Vec<PageStat>,
}
impl ColumnPageHeader {
const PAGE_ENCRYPTED: u8 = 1 << 0;
}
const RUN_MAC_LEN: usize = 32;
fn compute_run_mac(
enc: Option<&RunEncryption>,
header_bytes: &[u8],
dir_bytes: &[u8],
) -> Option<[u8; RUN_MAC_LEN]> {
#[cfg(feature = "encryption")]
{
if let Some(e) = enc {
if let Some(mac_key) = &e.mac_key {
return Some(crate::encryption::run_metadata_mac(
mac_key,
header_bytes,
dir_bytes,
&e.descriptor_bytes,
));
}
}
}
#[cfg(not(feature = "encryption"))]
{
let _ = (enc, header_bytes, dir_bytes);
}
None
}
pub struct ColumnPayload {
pub column_id: u16,
pub type_id_tag: u16,
pub encoding: Encoding,
pub pages: Vec<Vec<u8>>,
pub page_stats: Vec<PageStat>,
}
pub struct RunSpec<'a> {
pub run_id: u128,
pub schema_id: u64,
pub epoch_created: u64,
pub level: u8,
pub flags: u8,
pub sort_key_column_id: u16,
pub row_count: u64,
pub min_row_id: u64,
pub max_row_id: u64,
pub columns: &'a [ColumnPayload],
}
pub fn write_run(path: impl AsRef<Path>, spec: &RunSpec) -> Result<RunHeader> {
write_run_with(path, spec, None, &[], None)
}
pub fn write_run_with(
path: impl AsRef<Path>,
spec: &RunSpec,
kek: Option<&Kek>,
indexable_columns: &[(u16, u8)],
index_trailer: Option<&[u8]>,
) -> Result<RunHeader> {
let enc: Option<RunEncryption> = match kek {
Some(k) => Some(setup_run_encryption(k, indexable_columns)?),
None => None,
};
match write_run_mmap(path.as_ref(), spec, enc.as_ref(), index_trailer) {
Ok(h) => Ok(h),
Err(e) if is_mmap_unavailable(&e) => write_run_vec(path, spec, enc, index_trailer),
Err(e) => Err(e),
}
}
fn is_mmap_unavailable(e: &MongrelError) -> bool {
matches!(e, MongrelError::InvalidArgument(m) if m.starts_with("__mmap_unavailable__:"))
}
fn write_run_mmap(
path: &Path,
spec: &RunSpec,
enc: Option<&RunEncryption>,
index_trailer: Option<&[u8]>,
) -> Result<RunHeader> {
let plan = plan_run(spec, enc, index_trailer)?;
let file = OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(path)?;
file.set_len(plan.total as u64)?;
let mut mmap = match unsafe { memmap2::MmapMut::map_mut(&file) } {
Ok(m) => m,
Err(e) => {
return Err(MongrelError::InvalidArgument(format!(
"__mmap_unavailable__: {e}"
)))
}
};
let header = place_run(spec, enc, index_trailer, &plan, &mut mmap[..])?;
mmap.flush()?;
file.sync_all()?;
Ok(header)
}
struct RunPlan {
jobs: Vec<(usize, usize, u64, usize)>, dir_bytes: Vec<u8>,
encrypted: bool,
column_dir_offset: u64,
index_trailer_offset: u64,
encryption_descriptor_offset: u64,
footer_offset: u64,
total: usize,
encrypted_stats_plain: Option<Vec<u8>>,
encrypted_stats_offset: u64,
encrypted_stats_len: u64,
}
fn plan_run(
spec: &RunSpec,
enc: Option<&RunEncryption>,
index_trailer: Option<&[u8]>,
) -> Result<RunPlan> {
let encrypted = enc.is_some();
let columns = spec.columns;
let mut jobs: Vec<(usize, usize, u64, usize)> = Vec::new();
let mut dir: Vec<ColumnPageHeader> = Vec::with_capacity(columns.len());
let mut enc_stats: EncryptedColumnStats = Vec::new();
let mut cursor: u64 = RUN_HEADER_PAD as u64;
for (ci, col) in columns.iter().enumerate() {
let region_offset = cursor;
let mut region_len = 0u64;
let mut stats = Vec::with_capacity(col.pages.len());
let mut col_minmax: Vec<PageMinMax> = Vec::new();
for (ps, page) in col.pages.iter().enumerate() {
if encrypted && ps >= u16::MAX as usize {
return Err(MongrelError::Full(format!(
"column {:#x} exceeds 65534 pages; encrypted-run page-seq nonce space exhausted",
col.column_id
)));
}
let odl = on_disk_len(page.len(), encrypted);
jobs.push((ci, ps, cursor, odl));
let mut stat = col.page_stats.get(ps).cloned().unwrap_or(PageStat {
first_row_id: 0,
last_row_id: 0,
null_count: 0,
row_count: 0,
min: None,
max: None,
offset: 0,
compressed_len: 0,
uncompressed_len: 0,
});
stat.offset = cursor;
stat.compressed_len = odl as u32;
stat.uncompressed_len = page.len() as u32;
if encrypted {
col_minmax.push((stat.min.take(), stat.max.take()));
}
stats.push(stat);
cursor += odl as u64;
region_len += odl as u64;
}
if col_minmax
.iter()
.any(|(mn, mx)| mn.is_some() || mx.is_some())
{
enc_stats.push((col.column_id, col_minmax));
}
let page_flags = if encrypted {
ColumnPageHeader::PAGE_ENCRYPTED
} else {
0
};
dir.push(ColumnPageHeader {
column_id: col.column_id,
type_id_tag: col.type_id_tag,
encoding: col.encoding as u8,
flags: page_flags,
page_count: col.pages.len() as u32,
page_region_offset: region_offset,
page_region_len: region_len,
page_stats: stats,
});
}
let dir_bytes = bincode::serialize(&dir)?;
let column_dir_offset = cursor;
cursor += dir_bytes.len() as u64;
let index_trailer_offset = match index_trailer {
Some(t) => {
let off = cursor;
cursor += t.len() as u64;
off
}
None => 0,
};
let encryption_descriptor_offset = match enc {
Some(e) => {
let off = cursor;
cursor += 4 + e.descriptor_bytes.len() as u64;
off
}
None => 0,
};
let (encrypted_stats_plain, encrypted_stats_offset, encrypted_stats_len) =
if encrypted && !enc_stats.is_empty() {
let plain = bincode::serialize(&enc_stats)?;
let ct_len = on_disk_len(plain.len(), true) as u64;
let off = cursor;
cursor += ct_len;
(Some(plain), off, ct_len)
} else {
(None, 0, 0)
};
let footer_offset = cursor;
let total = footer_offset as usize + 8 + 8 + 32 + if encrypted { RUN_MAC_LEN } else { 0 };
Ok(RunPlan {
jobs,
dir_bytes,
encrypted,
column_dir_offset,
index_trailer_offset,
encryption_descriptor_offset,
footer_offset,
total,
encrypted_stats_plain,
encrypted_stats_offset,
encrypted_stats_len,
})
}
fn place_run(
spec: &RunSpec,
enc: Option<&RunEncryption>,
index_trailer: Option<&[u8]>,
plan: &RunPlan,
buf: &mut [u8],
) -> Result<RunHeader> {
use rayon::prelude::*;
use std::borrow::Cow;
debug_assert_eq!(
buf.len(),
plan.total,
"place_run: buffer must be exactly plan.total bytes"
);
let cipher: Option<&dyn Cipher> = enc.map(|e| e.cipher.as_ref());
let nonce_prefix = enc.map(|e| e.nonce_prefix);
let columns = spec.columns;
struct SyncPtr(*mut u8);
unsafe impl Send for SyncPtr {}
unsafe impl Sync for SyncPtr {}
impl SyncPtr {
fn get(&self) -> *mut u8 {
self.0
}
}
let base = SyncPtr(buf.as_mut_ptr());
plan.jobs
.par_iter()
.map(
|&(ci, ps, offset, odl): &(usize, usize, u64, usize)| -> Result<()> {
let page = &columns[ci].pages[ps];
let dst =
unsafe { std::slice::from_raw_parts_mut(base.get().add(offset as usize), odl) };
let bytes: Cow<[u8]> = match cipher {
Some(c) => {
let nonce =
page_nonce(nonce_prefix.unwrap(), columns[ci].column_id, ps as u32);
let ct = c.encrypt_page(&nonce, page)?;
assert_eq!(
ct.len(), odl,
"ciphertext length {} != predicted {}; GCM tag size assumption is stale",
ct.len(), odl
);
Cow::Owned(ct)
}
None => {
debug_assert_eq!(page.len(), odl);
Cow::Borrowed(page.as_slice())
}
};
dst.copy_from_slice(&bytes);
Ok(())
},
)
.collect::<Result<()>>()?;
let content_hash = {
let mut h = Sha256::new();
h.update(&buf[RUN_HEADER_PAD..plan.column_dir_offset as usize]);
h.finalize()
};
let doff = plan.column_dir_offset as usize;
buf[doff..doff + plan.dir_bytes.len()].copy_from_slice(&plan.dir_bytes);
if let Some(t) = index_trailer {
let off = plan.index_trailer_offset as usize;
buf[off..off + t.len()].copy_from_slice(t);
}
if let Some(e) = enc {
let off = plan.encryption_descriptor_offset as usize;
buf[off..off + 4].copy_from_slice(&(e.descriptor_bytes.len() as u32).to_le_bytes());
buf[off + 4..off + 4 + e.descriptor_bytes.len()].copy_from_slice(&e.descriptor_bytes);
}
if let (Some(e), Some(plain)) = (enc, plan.encrypted_stats_plain.as_ref()) {
let nonce = page_nonce(e.nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
let ct = e.cipher.encrypt_page(&nonce, plain)?;
debug_assert_eq!(ct.len() as u64, plan.encrypted_stats_len);
let off = plan.encrypted_stats_offset as usize;
buf[off..off + ct.len()].copy_from_slice(&ct);
}
let header_flags = if plan.encrypted {
spec.flags | RUN_FLAG_ENCRYPTED
} else {
spec.flags
};
let header = RunHeader {
magic: RUN_MAGIC,
format_version: RUN_FORMAT_VERSION,
header_layout_version: RUN_HEADER_VERSION,
run_id: spec.run_id,
content_hash: content_hash.into(),
schema_id: spec.schema_id,
epoch_created: spec.epoch_created,
level: spec.level,
flags: header_flags,
sort_key_column_id: spec.sort_key_column_id,
row_count: spec.row_count,
min_row_id: spec.min_row_id,
max_row_id: spec.max_row_id,
column_count: columns.len() as u64,
column_dir_offset: plan.column_dir_offset,
index_trailer_offset: plan.index_trailer_offset,
encryption_descriptor_offset: plan.encryption_descriptor_offset,
footer_offset: plan.footer_offset,
encrypted_stats_offset: plan.encrypted_stats_offset,
encrypted_stats_len: plan.encrypted_stats_len,
};
let header_bytes = bincode::serialize(&header)?;
if header_bytes.len() > RUN_HEADER_PAD {
return Err(MongrelError::InvalidArgument(format!(
"run header too large: {} > {RUN_HEADER_PAD}",
header_bytes.len()
)));
}
buf[..header_bytes.len()].copy_from_slice(&header_bytes);
let checksum = Sha256::digest(&buf[..plan.footer_offset as usize]);
let foot = plan.footer_offset as usize;
buf[foot..foot + 8].copy_from_slice(&RUN_MAGIC);
buf[foot + 8..foot + 16].copy_from_slice(&plan.footer_offset.to_le_bytes());
buf[foot + 16..foot + 48].copy_from_slice(&checksum);
if let Some(tag) = compute_run_mac(enc, &header_bytes, &plan.dir_bytes) {
buf[foot + 48..foot + 48 + RUN_MAC_LEN].copy_from_slice(&tag);
}
Ok(header)
}
fn write_run_vec(
path: impl AsRef<Path>,
spec: &RunSpec,
enc: Option<RunEncryption>,
index_trailer: Option<&[u8]>,
) -> Result<RunHeader> {
let mut buf: Vec<u8> = vec![0; RUN_HEADER_PAD]; let mut content_hasher = Sha256::new();
let mut dir: Vec<ColumnPageHeader> = Vec::with_capacity(spec.columns.len());
let mut enc_stats: EncryptedColumnStats = Vec::new();
for col in spec.columns {
let region_offset = buf.len() as u64;
let mut region_len = 0u64;
let mut stats = Vec::with_capacity(col.pages.len());
let mut col_minmax: Vec<PageMinMax> = Vec::new();
for (page_seq, page) in col.pages.iter().enumerate() {
if enc.is_some() && page_seq >= u16::MAX as usize {
return Err(MongrelError::Full(format!(
"column {:#x} exceeds 65534 pages; encrypted-run page-seq nonce space exhausted",
col.column_id
)));
}
let on_disk: Vec<u8> = match &enc {
Some(e) => e.cipher.encrypt_page(
&page_nonce(e.nonce_prefix, col.column_id, page_seq as u32),
page,
)?,
None => page.clone(),
};
let offset = buf.len() as u64;
buf.write_all(&on_disk)?;
content_hasher.update(&on_disk);
region_len += on_disk.len() as u64;
let mut stat = if let Some(s) = col.page_stats.get(page_seq) {
s.clone()
} else {
PageStat {
first_row_id: 0,
last_row_id: 0,
null_count: 0,
row_count: 0,
min: None,
max: None,
offset: 0,
compressed_len: 0,
uncompressed_len: 0,
}
};
stat.offset = offset;
stat.compressed_len = on_disk.len() as u32;
stat.uncompressed_len = page.len() as u32;
if enc.is_some() {
col_minmax.push((stat.min.take(), stat.max.take()));
}
stats.push(stat);
}
if col_minmax
.iter()
.any(|(mn, mx)| mn.is_some() || mx.is_some())
{
enc_stats.push((col.column_id, col_minmax));
}
let page_flags = if enc.is_some() {
ColumnPageHeader::PAGE_ENCRYPTED
} else {
0
};
dir.push(ColumnPageHeader {
column_id: col.column_id,
type_id_tag: col.type_id_tag,
encoding: col.encoding as u8,
flags: page_flags,
page_count: col.pages.len() as u32,
page_region_offset: region_offset,
page_region_len: region_len,
page_stats: stats,
});
}
let column_dir_offset = buf.len() as u64;
let dir_bytes = bincode::serialize(&dir)?;
buf.write_all(&dir_bytes)?;
let index_trailer_offset = match index_trailer {
Some(trailer) => {
let off = buf.len() as u64;
buf.write_all(trailer)?;
off
}
None => 0,
};
let encryption_descriptor_offset = match &enc {
Some(e) => {
let off = buf.len() as u64;
buf.write_all(&(e.descriptor_bytes.len() as u32).to_le_bytes())?;
buf.write_all(&e.descriptor_bytes)?;
off
}
None => 0,
};
let (encrypted_stats_offset, encrypted_stats_len) = match &enc {
Some(e) if !enc_stats.is_empty() => {
let plain = bincode::serialize(&enc_stats)?;
let nonce = page_nonce(e.nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
let ct = e.cipher.encrypt_page(&nonce, &plain)?;
let off = buf.len() as u64;
let len = ct.len() as u64;
buf.write_all(&ct)?;
(off, len)
}
_ => (0, 0),
};
let footer_offset = buf.len() as u64;
let header_flags = if enc.is_some() {
spec.flags | RUN_FLAG_ENCRYPTED
} else {
spec.flags
};
let header = RunHeader {
magic: RUN_MAGIC,
format_version: RUN_FORMAT_VERSION,
header_layout_version: RUN_HEADER_VERSION,
run_id: spec.run_id,
content_hash: content_hasher.finalize().into(),
schema_id: spec.schema_id,
epoch_created: spec.epoch_created,
level: spec.level,
flags: header_flags,
sort_key_column_id: spec.sort_key_column_id,
row_count: spec.row_count,
min_row_id: spec.min_row_id,
max_row_id: spec.max_row_id,
column_count: spec.columns.len() as u64,
column_dir_offset,
index_trailer_offset,
encryption_descriptor_offset,
footer_offset,
encrypted_stats_offset,
encrypted_stats_len,
};
let header_bytes = bincode::serialize(&header)?;
if header_bytes.len() > RUN_HEADER_PAD {
return Err(MongrelError::InvalidArgument(format!(
"run header too large: {} > {RUN_HEADER_PAD}",
header_bytes.len()
)));
}
buf[..header_bytes.len()].copy_from_slice(&header_bytes);
let checksum = Sha256::digest(&buf[..footer_offset as usize]);
buf.write_all(&RUN_MAGIC)?;
buf.write_all(&footer_offset.to_le_bytes())?;
buf.write_all(&checksum)?;
if let Some(tag) = compute_run_mac(enc.as_ref(), &header_bytes, &dir_bytes) {
buf.write_all(&tag)?;
}
let mut file = OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(path)?;
file.write_all(&buf)?;
file.sync_all()?;
Ok(header)
}
fn page_nonce(nonce_prefix: [u8; 12], column_id: u16, page_seq: u32) -> [u8; 12] {
let mut n = nonce_prefix;
n[8..10].copy_from_slice(&column_id.to_le_bytes());
n[10..12].copy_from_slice(&(page_seq as u16).to_le_bytes());
n
}
fn overlay_encrypted_stats(
path: &Path,
header: &RunHeader,
cipher: &dyn Cipher,
nonce_prefix: [u8; 12],
dir: &mut [ColumnPageHeader],
) -> Result<()> {
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(header.encrypted_stats_offset))?;
let mut ct = vec![0u8; header.encrypted_stats_len as usize];
file.read_exact(&mut ct)?;
let nonce = page_nonce(nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
let plain = cipher.decrypt_page(&nonce, &ct)?;
let stats: EncryptedColumnStats = bincode::deserialize(&plain)
.map_err(|e| MongrelError::Encryption(format!("bad encrypted page-stats envelope: {e}")))?;
for (cid, minmax) in stats {
let Some(col) = dir.iter_mut().find(|c| c.column_id == cid) else {
continue;
};
for (stat, (mn, mx)) in col.page_stats.iter_mut().zip(minmax) {
stat.min = mn;
stat.max = mx;
}
}
Ok(())
}
pub(crate) fn page_cache_key(
table_id: u64,
run_id: u128,
column_id: u16,
page_seq: usize,
) -> [u8; 32] {
let mut h = Sha256::new();
h.update(table_id.to_be_bytes());
h.update(run_id.to_be_bytes());
h.update(column_id.to_be_bytes());
h.update((page_seq as u64).to_be_bytes());
let out = h.finalize();
let mut k = [0u8; 32];
k.copy_from_slice(&out);
k
}
fn decrypt_or_passthrough(
cipher: Option<&dyn Cipher>,
nonce_prefix: [u8; 12],
column_id: u16,
page_seq: usize,
encrypted: bool,
buf: &[u8],
) -> Result<Vec<u8>> {
if encrypted {
match cipher {
Some(c) => {
Ok(c.decrypt_page(&page_nonce(nonce_prefix, column_id, page_seq as u32), buf)?)
}
None => Err(MongrelError::Decryption(
"encrypted page but no cipher".into(),
)),
}
} else {
Ok(buf.to_vec())
}
}
fn read_header_fast(path: impl AsRef<Path>) -> Result<RunHeader> {
let mut file = File::open(path)?;
let mut header_buf = vec![0u8; RUN_HEADER_PAD];
file.read_exact(&mut header_buf)?;
let header: RunHeader = bincode::deserialize(&header_buf)
.map_err(|e| MongrelError::InvalidArgument(format!("bad run header: {e}")))?;
if header.magic != RUN_MAGIC {
return Err(MongrelError::MagicMismatch {
what: "sorted run",
expected: RUN_MAGIC,
got: header.magic,
});
}
file.seek(SeekFrom::Start(header.footer_offset))?;
let mut footer_magic = [0u8; 8];
file.read_exact(&mut footer_magic)?;
if footer_magic != RUN_MAGIC {
return Err(MongrelError::MagicMismatch {
what: "sorted run footer",
expected: RUN_MAGIC,
got: footer_magic,
});
}
Ok(header)
}
fn read_header_cached(
path: impl AsRef<Path>,
verified_runs: &parking_lot::Mutex<std::collections::HashSet<u128>>,
) -> Result<RunHeader> {
let header = read_header_fast(path.as_ref())?;
if verified_runs.lock().contains(&header.run_id) {
return Ok(header);
}
let verified = read_header(path)?;
verified_runs.lock().insert(verified.run_id);
Ok(verified)
}
pub fn read_header(path: impl AsRef<Path>) -> Result<RunHeader> {
let mut file = File::open(path)?;
let mut header_buf = vec![0u8; RUN_HEADER_PAD];
file.read_exact(&mut header_buf)?;
let header: RunHeader = bincode::deserialize(&header_buf)
.map_err(|e| MongrelError::InvalidArgument(format!("bad run header: {e}")))?;
if header.magic != RUN_MAGIC {
return Err(MongrelError::MagicMismatch {
what: "sorted run",
expected: RUN_MAGIC,
got: header.magic,
});
}
file.seek(SeekFrom::Start(header.footer_offset))?;
let mut footer = [0u8; 8 + 8 + 32];
file.read_exact(&mut footer)?;
if footer[..8] != RUN_MAGIC {
return Err(MongrelError::MagicMismatch {
what: "sorted run footer",
expected: RUN_MAGIC,
got: footer[..8].try_into().unwrap(),
});
}
let mut hasher = Sha256::new();
hasher.update(&header_buf);
let body_len = header.footer_offset.saturating_sub(RUN_HEADER_PAD as u64);
file.seek(SeekFrom::Start(RUN_HEADER_PAD as u64))?;
let mut body = vec![0u8; body_len as usize];
file.read_exact(&mut body)?;
hasher.update(&body);
let computed: [u8; 32] = hasher.finalize().into();
let stored: [u8; 32] = footer[16..].try_into().unwrap();
if computed != stored {
return Err(MongrelError::ChecksumMismatch {
expected: u64::from_be_bytes(stored[..8].try_into().unwrap()),
actual: u64::from_be_bytes(computed[..8].try_into().unwrap()),
context: "sorted run footer".into(),
});
}
Ok(header)
}
pub fn read_column_dir(
path: impl AsRef<Path>,
header: &RunHeader,
) -> Result<Vec<ColumnPageHeader>> {
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(header.column_dir_offset))?;
let end = if header.index_trailer_offset != 0 {
header.index_trailer_offset
} else {
header.footer_offset
};
let len = end.saturating_sub(header.column_dir_offset);
let mut buf = vec![0u8; len as usize];
file.read_exact(&mut buf)?;
let dir: Vec<ColumnPageHeader> = bincode::deserialize(&buf)
.map_err(|e| MongrelError::InvalidArgument(format!("bad column dir: {e}")))?;
Ok(dir)
}
pub(crate) fn read_encryption_descriptor_bytes(
path: impl AsRef<Path>,
header: &RunHeader,
) -> Result<Vec<u8>> {
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(header.encryption_descriptor_offset))?;
let mut len_buf = [0u8; 4];
file.read_exact(&mut len_buf)?;
let len = u32::from_le_bytes(len_buf) as usize;
const MAX_DESCRIPTOR_BYTES: usize = 65_536;
if len > MAX_DESCRIPTOR_BYTES {
return Err(MongrelError::InvalidArgument(format!(
"encryption descriptor length {len} exceeds {MAX_DESCRIPTOR_BYTES}"
)));
}
let mut buf = vec![0u8; len];
file.read_exact(&mut buf)?;
Ok(buf)
}
#[cfg(feature = "encryption")]
fn verify_run_mac(
path: &Path,
header: &RunHeader,
dir: &[ColumnPageHeader],
kek: &Kek,
desc_bytes: &[u8],
) -> Result<()> {
let header_bytes = bincode::serialize(header)?;
let dir_bytes = bincode::serialize(dir)?;
let mac_key = kek.derive_run_mac_key();
let expected =
crate::encryption::run_metadata_mac(&mac_key, &header_bytes, &dir_bytes, desc_bytes);
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(header.footer_offset + 8 + 8 + 32))?;
let mut tag = [0u8; RUN_MAC_LEN];
file.read_exact(&mut tag).map_err(|_| {
MongrelError::Decryption(
"encrypted run is missing or truncated its metadata MAC; cannot \
authenticate metadata"
.into(),
)
})?;
let mut diff = 0u8;
for (x, y) in tag.iter().zip(expected.iter()) {
diff |= x ^ y;
}
if diff != 0 {
return Err(MongrelError::Decryption(
"run metadata authentication failed — tampered run or wrong key".into(),
));
}
Ok(())
}
#[cfg(not(feature = "encryption"))]
fn verify_run_mac(
_path: &Path,
_header: &RunHeader,
_dir: &[ColumnPageHeader],
_kek: &Kek,
_desc_bytes: &[u8],
) -> Result<()> {
Ok(())
}
pub struct RunWriter<'a> {
schema: &'a Schema,
run_id: u128,
epoch_created: Epoch,
level: u8,
kek: Option<&'a Kek>,
indexable_columns: Vec<(u16, u8)>,
compress: columnar::Compress,
clean: bool,
uniform_epoch: bool,
le: bool,
}
impl<'a> RunWriter<'a> {
pub fn new(schema: &'a Schema, run_id: u128, epoch_created: Epoch, level: u8) -> Self {
Self {
schema,
run_id,
epoch_created,
level,
kek: None,
indexable_columns: Vec::new(),
compress: columnar::Compress::Zstd(3),
clean: false,
uniform_epoch: false,
le: false,
}
}
pub fn uniform_epoch(mut self, uniform: bool) -> Self {
self.uniform_epoch = uniform;
self
}
pub fn with_encryption(mut self, kek: &'a Kek, indexable_columns: Vec<(u16, u8)>) -> Self {
self.kek = Some(kek);
self.indexable_columns = indexable_columns;
self
}
pub fn with_zstd_level(mut self, level: i32) -> Self {
self.compress = columnar::Compress::Zstd(level);
self
}
pub fn with_lz4(mut self) -> Self {
self.compress = columnar::Compress::Lz4;
self
}
pub fn with_plain(mut self) -> Self {
self.compress = columnar::Compress::Plain;
self
}
pub fn clean(mut self, clean: bool) -> Self {
self.clean = clean;
self
}
pub fn with_native_endian(mut self) -> Self {
if cfg!(target_endian = "little") {
self.le = true;
}
self
}
pub fn write_native(
self,
path: impl AsRef<Path>,
user_columns: &[(u16, columnar::NativeColumn)],
n: usize,
first_row_id: u64,
) -> Result<RunHeader> {
use columnar::NativeColumn;
let row_id_col = NativeColumn::int64_sequence(first_row_id as i64, n);
let epoch_col = NativeColumn::int64_constant(self.epoch_created.0 as i64, n);
let deleted_col = NativeColumn::bool_constant(false, n);
let learned_trailer = build_learned_trailer_native(&row_id_col);
let row_ids = match &row_id_col {
NativeColumn::Int64 { data, .. } => data.as_slice(),
_ => &[],
};
let bounds = page_bounds(row_ids);
let compress = self.compress;
let le = self.le;
let plain = matches!(compress, columnar::Compress::Plain);
let row_id_enc = if plain {
Encoding::Plain
} else {
Encoding::Delta
};
let sys_enc = if plain {
Encoding::Plain
} else {
Encoding::Zstd
};
let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
let (pages, stats) = native_column_pages(
TypeId::Int64,
&row_id_col,
row_id_enc,
compress,
le,
&bounds,
)?;
columns.push(ColumnPayload {
column_id: SYS_ROW_ID,
type_id_tag: type_tag(&TypeId::Int64),
encoding: row_id_enc,
pages,
page_stats: stats,
});
let (pages, stats) =
native_column_pages(TypeId::Int64, &epoch_col, sys_enc, compress, le, &bounds)?;
columns.push(ColumnPayload {
column_id: SYS_EPOCH,
type_id_tag: type_tag(&TypeId::Int64),
encoding: sys_enc,
pages,
page_stats: stats,
});
let (pages, stats) =
native_column_pages(TypeId::Bool, &deleted_col, sys_enc, compress, le, &bounds)?;
columns.push(ColumnPayload {
column_id: SYS_DELETED,
type_id_tag: type_tag(&TypeId::Bool),
encoding: sys_enc,
pages,
page_stats: stats,
});
use rayon::prelude::*;
let user_cols: Vec<ColumnPayload> = self
.schema
.columns
.par_iter()
.map(|cdef| -> Result<ColumnPayload> {
let col = user_columns
.iter()
.find(|(id, _)| *id == cdef.id)
.map(|(_, c)| c)
.ok_or_else(|| {
MongrelError::ColumnNotFound(format!("user column {}", cdef.id))
})?;
let encoding = if plain {
Encoding::Plain
} else {
choose_encoding_native(&cdef.ty, col)
};
let (pages, stats) =
native_column_pages(cdef.ty, col, encoding, compress, le, &bounds)?;
Ok(ColumnPayload {
column_id: cdef.id,
type_id_tag: type_tag(&cdef.ty),
encoding,
pages,
page_stats: stats,
})
})
.collect::<Result<Vec<_>>>()?;
columns.extend(user_cols);
let flags = if self.uniform_epoch {
RUN_FLAG_UNIFORM_EPOCH
} else {
RUN_FLAG_CLEAN
};
let spec = RunSpec {
run_id: self.run_id,
schema_id: self.schema.schema_id,
epoch_created: self.epoch_created.0,
level: self.level,
flags,
sort_key_column_id: SYS_ROW_ID,
row_count: n as u64,
min_row_id: first_row_id,
max_row_id: first_row_id + n as u64 - 1,
columns: &columns,
};
write_run_with(
path,
&spec,
self.kek,
&self.indexable_columns,
Some(&learned_trailer),
)
}
pub fn write(self, path: impl AsRef<Path>, rows: &[Row]) -> Result<RunHeader> {
let n = rows.len();
let mut row_ids = Vec::with_capacity(n);
let mut epochs = Vec::with_capacity(n);
let mut deleted = Vec::with_capacity(n);
for r in rows {
row_ids.push(Value::Int64(r.row_id.0 as i64));
epochs.push(Value::Int64(r.committed_epoch.0 as i64));
deleted.push(Value::Bool(r.deleted));
}
let learned_trailer = build_learned_trailer(&row_ids);
let (min_rid, max_rid) = row_id_bounds(rows);
let row_id_i64: Vec<i64> = rows.iter().map(|r| r.row_id.0 as i64).collect();
let bounds = page_bounds(&row_id_i64);
let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
let (pages, stats, enc) = value_column_pages(TypeId::Int64, &row_ids, &bounds)?;
columns.push(ColumnPayload {
column_id: SYS_ROW_ID,
type_id_tag: type_tag(&TypeId::Int64),
encoding: enc,
pages,
page_stats: stats,
});
let (pages, stats, enc) = value_column_pages(TypeId::Int64, &epochs, &bounds)?;
columns.push(ColumnPayload {
column_id: SYS_EPOCH,
type_id_tag: type_tag(&TypeId::Int64),
encoding: enc,
pages,
page_stats: stats,
});
let (pages, stats, enc) = value_column_pages(TypeId::Bool, &deleted, &bounds)?;
columns.push(ColumnPayload {
column_id: SYS_DELETED,
type_id_tag: type_tag(&TypeId::Bool),
encoding: enc,
pages,
page_stats: stats,
});
for cdef in &self.schema.columns {
let vals: Vec<Value> = rows
.iter()
.map(|r| r.columns.get(&cdef.id).cloned().unwrap_or(Value::Null))
.collect();
let (pages, stats, encoding) = value_column_pages(cdef.ty, &vals, &bounds)?;
columns.push(ColumnPayload {
column_id: cdef.id,
type_id_tag: type_tag(&cdef.ty),
encoding,
pages,
page_stats: stats,
});
}
let spec = RunSpec {
run_id: self.run_id,
schema_id: self.schema.schema_id,
epoch_created: self.epoch_created.0,
level: self.level,
flags: {
let mut f = if self.clean { RUN_FLAG_CLEAN } else { 0 };
if self.uniform_epoch {
f |= RUN_FLAG_UNIFORM_EPOCH;
}
f
},
sort_key_column_id: SYS_ROW_ID,
row_count: n as u64,
min_row_id: min_rid,
max_row_id: max_rid,
columns: &columns,
};
write_run_with(
path,
&spec,
self.kek,
&self.indexable_columns,
Some(&learned_trailer),
)
}
}
fn type_tag(ty: &TypeId) -> u16 {
match ty {
TypeId::Bool => 1,
TypeId::Int64 => 8,
TypeId::Float64 => 9,
TypeId::Bytes => 12,
TypeId::Embedding { .. } => 13,
_ => 0,
}
}
pub(crate) fn be_i64(b: Option<&[u8]>) -> Option<i64> {
let b = b?;
(b.len() == 8).then(|| i64::from_be_bytes(b.try_into().unwrap()))
}
pub(crate) fn be_f64(b: Option<&[u8]>) -> Option<f64> {
let b = b?;
(b.len() == 8).then(|| f64::from_bits(u64::from_be_bytes(b.try_into().unwrap())))
}
const PAGE_ROWS: usize = 65_536;
fn page_bounds(row_ids: &[i64]) -> Vec<(usize, usize, u64, u64)> {
let n = row_ids.len();
if n == 0 {
return vec![(0, 0, 0, 0)];
}
let mut out = Vec::new();
let mut start = 0;
while start < n {
let end = (start + PAGE_ROWS).min(n);
out.push((start, end, row_ids[start] as u64, row_ids[end - 1] as u64));
start = end;
}
out
}
fn native_column_pages(
ty: TypeId,
col: &columnar::NativeColumn,
encoding: Encoding,
compress: columnar::Compress,
le: bool,
bounds: &[(usize, usize, u64, u64)],
) -> Result<(Vec<Vec<u8>>, Vec<PageStat>)> {
use rayon::prelude::*;
let encode_one =
|&(s, e, frid, lrid): &(usize, usize, u64, u64)| -> Result<(Vec<u8>, PageStat)> {
let chunk = col.slice_range(s, e);
let stat = columnar::page_stat_for(ty, &chunk, frid, lrid);
let page = columnar::encode_page_native(ty, &chunk, encoding, compress, le)?;
Ok((page, stat))
};
let pairs: Vec<(Vec<u8>, PageStat)> = if bounds.len() > 1 {
bounds
.par_iter()
.map(encode_one)
.collect::<Result<Vec<_>>>()?
} else {
bounds.iter().map(encode_one).collect::<Result<Vec<_>>>()?
};
let (pages, stats) = pairs.into_iter().unzip();
Ok((pages, stats))
}
fn value_column_pages(
ty: TypeId,
vals: &[Value],
bounds: &[(usize, usize, u64, u64)],
) -> Result<(Vec<Vec<u8>>, Vec<PageStat>, Encoding)> {
let encoding = choose_encoding(&ty, vals);
let mut pages = Vec::with_capacity(bounds.len());
let mut stats = Vec::with_capacity(bounds.len());
for &(s, e, frid, lrid) in bounds {
let chunk = &vals[s..e];
pages.push(columnar::encode_page(ty, chunk, encoding)?);
let native = columnar::values_to_native(ty, chunk);
stats.push(columnar::page_stat_for(ty, &native, frid, lrid));
}
Ok((pages, stats, encoding))
}
fn choose_encoding(ty: &TypeId, values: &[Value]) -> Encoding {
use std::collections::HashSet;
if matches!(ty, TypeId::Bytes) {
let n = values.len();
if n > 0 {
let distinct = values
.iter()
.filter(|v| !matches!(v, Value::Null))
.map(|v| v.encode_key())
.collect::<HashSet<_>>()
.len();
if (distinct as f64 / n as f64) < 0.5 {
return Encoding::Dictionary;
}
}
}
Encoding::Zstd
}
fn choose_encoding_native(ty: &TypeId, col: &columnar::NativeColumn) -> Encoding {
use std::collections::HashSet;
match (ty, col) {
(TypeId::Int64 | TypeId::TimestampNanos, columnar::NativeColumn::Int64 { data, .. }) => {
if data.windows(2).all(|w| w[0] <= w[1]) {
Encoding::Delta
} else {
Encoding::Zstd
}
}
(
TypeId::Bytes,
columnar::NativeColumn::Bytes {
offsets, values, ..
},
) => {
let n = offsets.len().saturating_sub(1);
if n > 0 {
let distinct: HashSet<&[u8]> = (0..n)
.map(|i| &values[offsets[i] as usize..offsets[i + 1] as usize])
.collect();
if (distinct.len() as f64 / n as f64) < 0.5 {
return Encoding::Dictionary;
}
}
Encoding::Zstd
}
_ => Encoding::Zstd,
}
}
fn build_learned_trailer_native(col: &columnar::NativeColumn) -> Vec<u8> {
let points: Vec<(u64, usize)> = match col {
columnar::NativeColumn::Int64 { data, .. } => data
.iter()
.enumerate()
.map(|(i, v)| (*v as u64, i))
.collect(),
_ => Vec::new(),
};
let pgm = PgmIndex::build(&points, LEARNED_EPSILON);
bincode::serialize(&pgm).expect("pgm serialize")
}
fn row_id_bounds(rows: &[Row]) -> (u64, u64) {
match (rows.first(), rows.last()) {
(Some(f), Some(l)) => (f.row_id.0, l.row_id.0),
_ => (0, 0),
}
}
pub struct RunReader {
file: File,
mmap: Option<memmap2::Mmap>,
header: RunHeader,
dir: Vec<ColumnPageHeader>,
schema: Schema,
table_id: u64,
cipher: Option<Box<dyn Cipher>>,
nonce_prefix: [u8; 12],
col_cache: HashMap<u16, Vec<Value>>,
page_cache: Option<Arc<crate::cache::Sharded<crate::cache::PageCache>>>,
decoded_cache: Option<Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>>,
epoch_override: Option<Epoch>,
}
impl RunReader {
pub fn open(path: impl AsRef<Path>, schema: Schema, kek: Option<Arc<Kek>>) -> Result<Self> {
Self::open_with_cache(path, schema, kek, None, None, 0, None)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn open_with_cache(
path: impl AsRef<Path>,
schema: Schema,
kek: Option<Arc<Kek>>,
page_cache: Option<Arc<crate::cache::Sharded<crate::cache::PageCache>>>,
decoded_cache: Option<Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>>,
table_id: u64,
verified_runs: Option<&parking_lot::Mutex<std::collections::HashSet<u128>>>,
) -> Result<Self> {
let path = path.as_ref().to_path_buf();
let header = match verified_runs {
Some(cache) => read_header_cached(&path, cache)?,
None => read_header(&path)?,
};
let mut dir = read_column_dir(&path, &header)?;
let (cipher, nonce_prefix): (Option<Box<dyn Cipher>>, [u8; 12]) = if header.is_encrypted() {
let kek = kek.as_ref().ok_or_else(|| {
MongrelError::Encryption(
"run is encrypted but no key-encryption key was provided".into(),
)
})?;
if header.encryption_descriptor_offset == 0 {
return Err(MongrelError::Encryption(
"encrypted run has no encryption descriptor".into(),
));
}
let desc_bytes = read_encryption_descriptor_bytes(&path, &header)?;
verify_run_mac(&path, &header, &dir, kek, &desc_bytes)?;
let enc = crate::encryption::build_run_cipher(kek, &desc_bytes)?;
if header.encrypted_stats_offset != 0 {
overlay_encrypted_stats(
&path,
&header,
enc.cipher.as_ref(),
enc.nonce_prefix,
&mut dir,
)?;
}
(Some(enc.cipher), enc.nonce_prefix)
} else {
(None, [0u8; 12])
};
let file = File::open(&path)?;
let mmap = unsafe { memmap2::MmapOptions::new().map(&file) }.ok();
let _ = path;
Ok(Self {
file,
mmap,
header,
dir,
schema,
table_id,
cipher,
nonce_prefix,
col_cache: HashMap::new(),
epoch_override: None,
page_cache,
decoded_cache,
})
}
pub fn header(&self) -> &RunHeader {
&self.header
}
pub(crate) fn set_uniform_epoch(&mut self, epoch: Epoch) {
if self.header.is_uniform_epoch() {
self.epoch_override = Some(epoch);
self.col_cache.remove(&SYS_EPOCH);
}
}
pub fn is_clean(&self) -> bool {
self.header.is_clean()
}
pub fn row_count(&self) -> usize {
self.header.row_count as usize
}
pub(crate) fn clean_contiguous_row_ids(&self) -> bool {
let n = self.row_count();
n > 0
&& self.is_clean()
&& self.epoch_override.is_none()
&& self.header.max_row_id >= self.header.min_row_id
&& self
.header
.max_row_id
.checked_sub(self.header.min_row_id)
.and_then(|span| span.checked_add(1))
== Some(n as u64)
}
pub(crate) fn position_for_row_id_fast(&self, row_id: u64) -> Option<usize> {
if !self.clean_contiguous_row_ids()
|| row_id < self.header.min_row_id
|| row_id > self.header.max_row_id
{
return None;
}
Some((row_id - self.header.min_row_id) as usize)
}
pub(crate) fn positions_for_row_ids_fast(&self, row_ids: &[u64]) -> Option<Vec<usize>> {
if !self.clean_contiguous_row_ids() {
return None;
}
let mut positions = Vec::with_capacity(row_ids.len());
for &row_id in row_ids {
if let Some(pos) = self.position_for_row_id_fast(row_id) {
positions.push(pos);
}
}
positions.sort_unstable();
Some(positions)
}
fn resolve_type(&self, column_id: u16) -> TypeId {
match column_id {
SYS_ROW_ID | SYS_EPOCH => TypeId::Int64,
SYS_DELETED => TypeId::Bool,
_ => self
.schema
.columns
.iter()
.find(|c| c.id == column_id)
.map(|c| c.ty)
.unwrap_or(TypeId::Bytes),
}
}
fn find_header(&self, column_id: u16) -> Result<&ColumnPageHeader> {
self.dir
.iter()
.find(|h| h.column_id == column_id)
.ok_or_else(|| MongrelError::ColumnNotFound(format!("column {column_id}")))
}
fn col_encrypted(&self, column_id: u16) -> bool {
self.dir
.iter()
.find(|h| h.column_id == column_id)
.map(|h| h.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0)
.unwrap_or(false)
}
pub(crate) fn read_page(&mut self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
let (offset, compressed_len, encrypted) = {
let ch = self.find_header(column_id)?;
let stat = ch
.page_stats
.get(page_seq)
.ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
(
stat.offset,
stat.compressed_len,
ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0,
)
};
let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
if let Some(cache) = &self.page_cache {
if let Some(bytes) = cache.lock(&key).get(
&key,
crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
) {
return decrypt_or_passthrough(
self.cipher.as_deref(),
self.nonce_prefix,
column_id,
page_seq,
encrypted,
&bytes,
);
}
}
let buf = match &self.mmap {
Some(m) => m[offset as usize..(offset + compressed_len as u64) as usize].to_vec(),
None => {
self.file.seek(SeekFrom::Start(offset))?;
let mut buf = vec![0u8; compressed_len as usize];
self.file.read_exact(&mut buf)?;
buf
}
};
if let Some(cache) = &self.page_cache {
cache.lock(&key).insert(crate::page::CachedPage {
committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
content_hash: key,
bytes: bytes::Bytes::copy_from_slice(&buf),
});
}
decrypt_or_passthrough(
self.cipher.as_deref(),
self.nonce_prefix,
column_id,
page_seq,
encrypted,
&buf,
)
}
fn read_page_shared(&self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
let ch = self.find_header(column_id)?;
let stat = ch
.page_stats
.get(page_seq)
.ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
let encrypted = ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0;
let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
if let Some(cache) = &self.page_cache {
if let Some(guard) = cache.try_lock(&key) {
if let Some(bytes) = guard.try_get(
&key,
crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
) {
return decrypt_or_passthrough(
self.cipher.as_deref(),
self.nonce_prefix,
column_id,
page_seq,
encrypted,
&bytes,
);
}
}
}
let mmap = self.mmap.as_ref().ok_or_else(|| {
MongrelError::InvalidArgument("parallel page decode requires an mmap-backed run".into())
})?;
let start = stat.offset as usize;
let end = (stat.offset + stat.compressed_len as u64) as usize;
let buf = mmap[start..end].to_vec();
if let Some(cache) = &self.page_cache {
if let Some(mut guard) = cache.try_lock(&key) {
guard.insert(crate::page::CachedPage {
committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
content_hash: key,
bytes: bytes::Bytes::copy_from_slice(&buf),
});
}
}
decrypt_or_passthrough(
self.cipher.as_deref(),
self.nonce_prefix,
column_id,
page_seq,
encrypted,
&buf,
)
}
fn column(&mut self, column_id: u16) -> Result<&[Value]> {
if column_id == SYS_EPOCH {
if let Some(ov) = self.epoch_override {
if !self.col_cache.contains_key(&column_id) {
let n = self.row_count();
self.col_cache
.insert(column_id, vec![Value::Int64(ov.0 as i64); n]);
}
return Ok(self.col_cache.get(&column_id).unwrap().as_slice());
}
}
if !self.col_cache.contains_key(&column_id) {
let ty = self.resolve_type(column_id);
let page_rows: Vec<usize> = {
let ch = self.find_header(column_id)?;
ch.page_stats.iter().map(|s| s.row_count as usize).collect()
};
let mut decoded: Vec<Value> = Vec::with_capacity(self.row_count());
for (seq, &pr) in page_rows.iter().enumerate() {
let page = self.read_page(column_id, seq)?;
decoded.extend(columnar::decode_page(ty, &page, pr)?);
}
self.col_cache.insert(column_id, decoded);
}
Ok(self.col_cache.get(&column_id).unwrap().as_slice())
}
pub fn get_version(&mut self, row_id: RowId, snapshot: Epoch) -> Result<Option<(Epoch, Row)>> {
match self.find_version_page(row_id, snapshot)? {
None => Ok(None),
Some((epoch, seq, local_index)) => Ok(Some((
Epoch(epoch),
self.materialize_in_page(seq, local_index)?,
))),
}
}
fn find_version_page(
&mut self,
row_id: RowId,
snapshot: Epoch,
) -> Result<Option<(u64, usize, usize)>> {
let n = self.row_count();
if n == 0 {
return Ok(None);
}
let target = row_id.0 as i64;
let ch = self.find_header(SYS_ROW_ID)?;
let mut page_start = 0u64;
let candidate_pages: Vec<(usize, u64)> = ch .page_stats
.iter()
.enumerate()
.filter_map(|(seq, s)| {
let start = page_start;
page_start += s.row_count as u64;
(s.first_row_id <= row_id.0 && row_id.0 <= s.last_row_id).then_some((seq, start))
})
.collect();
if candidate_pages.is_empty() {
return Ok(None);
}
let ty = self.resolve_type(SYS_ROW_ID);
let mut best: Option<(u64, usize, usize)> = None; for (seq, _page_row_start) in candidate_pages {
let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
let row_ids = match self.decode_page_native_cached(ty, SYS_ROW_ID, seq, page_rows)? {
columnar::NativeColumn::Int64 { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
};
let local = match row_ids.binary_search(&target) {
Ok(i) => i,
Err(_) => continue,
};
let epoch_ty = self.resolve_type(SYS_EPOCH);
let epochs =
match self.decode_page_native_cached(epoch_ty, SYS_EPOCH, seq, page_rows)? {
columnar::NativeColumn::Int64 { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys epoch not int64".into())),
};
let mut lo = local;
while lo > 0 && row_ids[lo - 1] == target {
lo -= 1;
}
let mut hi = local;
while hi + 1 < row_ids.len() && row_ids[hi + 1] == target {
hi += 1;
}
for (i, &epoch) in epochs[lo..=hi].iter().enumerate() {
let epoch = epoch as u64;
if epoch <= snapshot.0 && best.map(|(be, ..)| epoch > be).unwrap_or(true) {
best = Some((epoch, seq, lo + i));
}
}
}
Ok(best)
}
pub fn get_version_column(
&mut self,
row_id: RowId,
snapshot: Epoch,
column_id: u16,
) -> Result<Option<(Epoch, bool, Option<Value>)>> {
let Some((epoch, seq, local_index)) = self.find_version_page(row_id, snapshot)? else {
return Ok(None);
};
let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
let page_start: usize = self.find_header(SYS_ROW_ID)?.page_stats[..seq]
.iter()
.map(|s| s.row_count as usize)
.sum();
let global_index = page_start + local_index;
let native_at = |slf: &mut Self, cid: u16| -> Result<Option<Value>> {
if !slf.dir.iter().any(|h| h.column_id == cid) {
return Ok(None);
}
let ty = slf.resolve_type(cid);
if !matches!(
ty,
TypeId::Bool
| TypeId::Int8
| TypeId::Int16
| TypeId::Int32
| TypeId::Int64
| TypeId::UInt8
| TypeId::UInt16
| TypeId::UInt32
| TypeId::UInt64
| TypeId::Float32
| TypeId::Float64
| TypeId::TimestampNanos
| TypeId::Date32
| TypeId::Bytes
) {
return Ok(slf.column(cid)?.get(global_index).cloned());
}
Ok(slf
.decode_page_native_cached(ty, cid, seq, page_rows)?
.value_at(local_index))
};
let deleted = matches!(native_at(self, SYS_DELETED)?, Some(Value::Bool(true)));
let value = native_at(self, column_id)?;
Ok(Some((Epoch(epoch), deleted, value)))
}
fn materialize_in_page(&mut self, seq: usize, local_index: usize) -> Result<Row> {
let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
let page_start: usize = self.find_header(SYS_ROW_ID)?.page_stats[..seq]
.iter()
.map(|s| s.row_count as usize)
.sum();
let global_index = page_start + local_index;
let native_at = |slf: &mut Self, column_id: u16| -> Result<Option<Value>> {
if !slf.dir.iter().any(|h| h.column_id == column_id) {
return Ok(None);
}
let ty = slf.resolve_type(column_id);
if !matches!(
ty,
TypeId::Bool
| TypeId::Int8
| TypeId::Int16
| TypeId::Int32
| TypeId::Int64
| TypeId::UInt8
| TypeId::UInt16
| TypeId::UInt32
| TypeId::UInt64
| TypeId::Float32
| TypeId::Float64
| TypeId::TimestampNanos
| TypeId::Date32
| TypeId::Bytes
) {
return Ok(slf.column(column_id)?.get(global_index).cloned());
}
Ok(slf
.decode_page_native_cached(ty, column_id, seq, page_rows)?
.value_at(local_index))
};
let row_id = RowId(match native_at(self, SYS_ROW_ID)? {
Some(Value::Int64(x)) => x as u64,
_ => 0,
});
let committed_epoch = Epoch(match native_at(self, SYS_EPOCH)? {
Some(Value::Int64(x)) => x as u64,
_ => 0,
});
let deleted = matches!(native_at(self, SYS_DELETED)?, Some(Value::Bool(true)));
let col_ids: Vec<u16> = self.schema.columns.iter().map(|c| c.id).collect();
let mut columns = HashMap::new();
for id in col_ids {
columns.insert(id, native_at(self, id)?.unwrap_or(Value::Null));
}
Ok(Row {
row_id,
committed_epoch,
columns,
deleted,
})
}
pub fn all_rows(&mut self) -> Result<Vec<Row>> {
let n = self.row_count();
let mut out = Vec::with_capacity(n);
for i in 0..n {
out.push(self.materialize(i)?);
}
Ok(out)
}
pub fn visible_indices(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
let n = self.row_count();
if n == 0 {
return Ok(Vec::new());
}
let row_ids = self.column(SYS_ROW_ID)?.to_vec();
let epochs = self.column(SYS_EPOCH)?.to_vec();
let deleted = self.column(SYS_DELETED)?.to_vec();
let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
for i in 0..n {
let rid = int_at(&row_ids, i);
let e = int_at(&epochs, i);
if e > snapshot.0 {
continue;
}
best.entry(rid)
.and_modify(|(be, bi)| {
if e > *be {
*be = e;
*bi = i;
}
})
.or_insert((e, i));
}
let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
idxs.retain(|&i| !bool_at(&deleted, i));
idxs.sort_unstable();
Ok(idxs)
}
pub fn gather_column(&mut self, column_id: u16, indices: &[usize]) -> Result<Vec<Value>> {
if !self.dir.iter().any(|h| h.column_id == column_id) {
return Ok(vec![Value::Null; indices.len()]);
}
let col = self.column(column_id)?;
Ok(indices
.iter()
.map(|&i| col.get(i).cloned().unwrap_or(Value::Null))
.collect())
}
pub fn column_native(&mut self, column_id: u16) -> Result<columnar::NativeColumn> {
use rayon::prelude::*;
if column_id == SYS_EPOCH {
if let Some(ov) = self.epoch_override {
return Ok(columnar::NativeColumn::int64_constant(
ov.0 as i64,
self.row_count(),
));
}
}
let ty = self.resolve_type(column_id);
let n = self.row_count();
let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
return Ok(columnar::null_native(ty, n));
};
let page_count = ch.page_count as usize;
let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
if page_count == 0 {
return Ok(columnar::null_native(ty, n));
}
let parts: Vec<columnar::NativeColumn> = if self.mmap.is_some() && page_count > 1 {
let reader: &RunReader = self;
(0..page_count)
.into_par_iter()
.map(|seq| {
let raw = reader.read_page_shared(column_id, seq)?;
columnar::decode_page_native(ty, &raw, page_rows[seq])
})
.collect::<Result<Vec<_>>>()?
} else {
let mut out = Vec::with_capacity(page_count);
for (seq, &pr) in page_rows.iter().enumerate() {
let page = self.read_page(column_id, seq)?;
out.push(columnar::decode_page_native(ty, &page, pr)?);
}
out
};
Ok(columnar::NativeColumn::concat(&parts))
}
pub fn has_mmap(&self) -> bool {
self.mmap.is_some()
}
pub fn column_native_shared(&self, column_id: u16) -> Result<columnar::NativeColumn> {
use rayon::prelude::*;
if column_id == SYS_EPOCH {
if let Some(ov) = self.epoch_override {
return Ok(columnar::NativeColumn::int64_constant(
ov.0 as i64,
self.row_count(),
));
}
}
let ty = self.resolve_type(column_id);
let n = self.row_count();
let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
return Ok(columnar::null_native(ty, n));
};
let page_count = ch.page_count as usize;
let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
if page_count == 0 {
return Ok(columnar::null_native(ty, n));
}
#[cfg(unix)]
{
if let (Some(m), Some(first)) = (&self.mmap, ch.page_stats.first()) {
let start = first.offset as usize;
let end = (ch.page_region_offset as usize) + (ch.page_region_len as usize);
if end > start {
let _ = m.advise_range(memmap2::Advice::WillNeed, start, end - start);
}
}
}
let run_id = self.header.run_id;
let mut parts_keys: Vec<(columnar::NativeColumn, Option<[u8; 32]>)> = if page_count > 1 {
(0..page_count)
.into_par_iter()
.map(|seq| self.decode_page_cached(ty, column_id, seq, page_rows[seq], run_id))
.collect::<Result<Vec<_>>>()?
} else {
vec![self.decode_page_cached(ty, column_id, 0, page_rows[0], run_id)?]
};
if let Some(cache) = &self.decoded_cache {
for (col, key) in parts_keys.iter_mut() {
if let Some(k) = key.take() {
cache.lock(&k).insert(k, Arc::new(col.clone()));
}
}
}
let parts: Vec<columnar::NativeColumn> = parts_keys.into_iter().map(|(c, _)| c).collect();
Ok(columnar::NativeColumn::concat(&parts))
}
fn decode_page_cached(
&self,
ty: TypeId,
column_id: u16,
seq: usize,
nrows: usize,
run_id: u128,
) -> Result<(columnar::NativeColumn, Option<[u8; 32]>)> {
let key = page_cache_key(self.table_id, run_id, column_id, seq);
if let Some(cache) = &self.decoded_cache {
if let Some(g) = cache.try_lock(&key) {
if let Some(hit) = g.try_get(&key) {
return Ok(((*hit).clone(), None));
}
}
}
let raw = self.read_page_shared(column_id, seq)?;
let col = columnar::decode_page_native(ty, &raw, nrows)?;
Ok((col, Some(key)))
}
fn decode_page_native_cached(
&mut self,
ty: TypeId,
column_id: u16,
seq: usize,
nrows: usize,
) -> Result<columnar::NativeColumn> {
let key = page_cache_key(self.table_id, self.header.run_id, column_id, seq);
if let Some(cache) = &self.decoded_cache {
if let Some(hit) = cache.lock(&key).try_get(&key) {
return Ok((*hit).clone());
}
}
let raw = self.read_page(column_id, seq)?;
let col = columnar::decode_page_native(ty, &raw, nrows)?;
if let Some(cache) = &self.decoded_cache {
cache
.lock(&key)
.insert(key, std::sync::Arc::new(col.clone()));
}
Ok(col)
}
pub fn range_row_ids_i64(
&mut self,
column_id: u16,
lo: i64,
hi: i64,
) -> Result<std::collections::HashSet<u64>> {
Ok(self
.range_row_id_set_i64(column_id, lo, hi)?
.into_sorted_vec()
.into_iter()
.collect())
}
pub(crate) fn range_row_id_set_i64(
&mut self,
column_id: u16,
lo: i64,
hi: i64,
) -> Result<RowIdSet> {
let info: Vec<(Option<i64>, Option<i64>, usize)> =
match self.dir.iter().find(|h| h.column_id == column_id) {
Some(ch) => ch
.page_stats
.iter()
.map(|s| {
(
be_i64(s.min.as_deref()),
be_i64(s.max.as_deref()),
s.row_count as usize,
)
})
.collect(),
None => return Ok(RowIdSet::empty()),
};
let stats_pruneable =
!self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
let clean_contiguous = self.clean_contiguous_row_ids();
let mut out = Vec::new();
let mut page_start = 0usize;
for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
let current_page_start = page_start;
page_start += nrows;
let skip = stats_pruneable
&& match (mn, mx) {
(Some(mn), Some(mx)) => mx < lo || mn > hi,
_ => true,
};
if skip {
continue;
}
let val_page = self.read_page(column_id, seq)?;
let vals =
columnar::decode_page_native(self.resolve_type(column_id), &val_page, nrows)?;
if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
if clean_contiguous {
for (i, val) in v.iter().enumerate() {
if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
}
}
} else {
let rid_page = self.read_page(SYS_ROW_ID, seq)?;
let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
for (i, val) in v.iter().enumerate() {
if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
out.push(r[i] as u64);
}
}
}
}
}
}
Ok(RowIdSet::from_unsorted(out))
}
pub fn range_row_ids_f64(
&mut self,
column_id: u16,
lo: f64,
lo_inclusive: bool,
hi: f64,
hi_inclusive: bool,
) -> Result<std::collections::HashSet<u64>> {
Ok(self
.range_row_id_set_f64(column_id, lo, lo_inclusive, hi, hi_inclusive)?
.into_sorted_vec()
.into_iter()
.collect())
}
pub(crate) fn range_row_id_set_f64(
&mut self,
column_id: u16,
lo: f64,
lo_inclusive: bool,
hi: f64,
hi_inclusive: bool,
) -> Result<RowIdSet> {
let info: Vec<(Option<f64>, Option<f64>, usize)> =
match self.dir.iter().find(|h| h.column_id == column_id) {
Some(ch) => ch
.page_stats
.iter()
.map(|s| {
(
be_f64(s.min.as_deref()),
be_f64(s.max.as_deref()),
s.row_count as usize,
)
})
.collect(),
None => return Ok(RowIdSet::empty()),
};
let stats_pruneable =
!self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
let clean_contiguous = self.clean_contiguous_row_ids();
let mut out = Vec::new();
let mut page_start = 0usize;
for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
let current_page_start = page_start;
page_start += nrows;
let skip = stats_pruneable
&& match (mn, mx) {
(Some(mn), Some(mx)) => {
let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
skip_lo || skip_hi
}
_ => true,
};
if skip {
continue;
}
let val_page = self.read_page(column_id, seq)?;
let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
if clean_contiguous {
for (i, val) in v.iter().enumerate() {
if !columnar::validity_bit(&validity, i) || val.is_nan() {
continue;
}
let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
if ok_lo && ok_hi {
out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
}
}
} else {
let rid_page = self.read_page(SYS_ROW_ID, seq)?;
let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
for (i, val) in v.iter().enumerate() {
if !columnar::validity_bit(&validity, i) || val.is_nan() {
continue;
}
let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
if ok_lo && ok_hi {
out.push(r[i] as u64);
}
}
}
}
}
}
Ok(RowIdSet::from_unsorted(out))
}
pub(crate) fn null_row_id_set(&mut self, column_id: u16, want_nulls: bool) -> Result<RowIdSet> {
let stats: Vec<(usize, usize)> = match self.dir.iter().find(|h| h.column_id == column_id) {
Some(ch) => ch
.page_stats
.iter()
.map(|s| (s.null_count as usize, s.row_count as usize))
.collect(),
None => return Ok(RowIdSet::empty()),
};
let ty = self.resolve_type(column_id);
let clean_contiguous = self.clean_contiguous_row_ids();
let mut out = Vec::new();
let mut page_start = 0usize;
for (seq, (null_count, nrows)) in stats.into_iter().enumerate() {
let current_page_start = page_start;
page_start += nrows;
if want_nulls && null_count == 0 {
continue;
}
if !want_nulls && null_count == nrows {
continue;
}
let val_page = self.read_page(column_id, seq)?;
let col = columnar::decode_page_native(ty, &val_page, nrows)?;
let validity = col.validity();
if clean_contiguous {
for i in 0..nrows {
let is_null = !columnar::validity_bit(validity, i);
if is_null == want_nulls {
out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
}
}
} else {
let rid_page = self.read_page(SYS_ROW_ID, seq)?;
let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
for (i, &rid) in r.iter().enumerate().take(nrows) {
let is_null = !columnar::validity_bit(validity, i);
if is_null == want_nulls {
out.push(rid as u64);
}
}
}
}
}
Ok(RowIdSet::from_unsorted(out))
}
pub fn visible_indices_native(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
let n = self.row_count();
if n == 0 {
return Ok(Vec::new());
}
let (row_ids, epochs, deleted) = self.system_columns_native()?;
let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
for i in 0..n {
let rid = row_ids[i] as u64;
let e = epochs[i] as u64;
if e > snapshot.0 {
continue;
}
best.entry(rid)
.and_modify(|(be, bi)| {
if e > *be {
*be = e;
*bi = i;
}
})
.or_insert((e, i));
}
let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
idxs.retain(|&i| deleted[i] == 0);
idxs.sort_unstable();
Ok(idxs)
}
pub fn range_row_ids_visible_i64(
&mut self,
column_id: u16,
lo: i64,
hi: i64,
snapshot: Epoch,
) -> Result<Vec<u64>> {
let stats: Vec<(Option<i64>, Option<i64>, usize)> = match self.column_page_stats(column_id)
{
Some(s) => s
.iter()
.map(|st| {
(
be_i64(st.min.as_deref()),
be_i64(st.max.as_deref()),
st.row_count as usize,
)
})
.collect(),
None => return Ok(Vec::new()),
};
let stats_pruneable =
!self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
let mut out: Vec<u64> = Vec::new();
let mut vis = 0usize;
let mut page_start = 0usize;
for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
let page_end = page_start + nrows;
let skip = stats_pruneable
&& match (mn, mx) {
(Some(mn), Some(mx)) => mx < lo || mn > hi,
_ => true, };
if !skip {
let val_page = self.read_page(column_id, seq)?;
let vals = columnar::decode_page_native(TypeId::Int64, &val_page, nrows)?;
if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
while vis < positions.len() && positions[vis] < page_end {
let local = positions[vis] - page_start;
if columnar::validity_bit(&validity, local)
&& v[local] >= lo
&& v[local] <= hi
{
out.push(rids[vis] as u64);
}
vis += 1;
}
}
} else {
while vis < positions.len() && positions[vis] < page_end {
vis += 1;
}
}
page_start = page_end;
}
Ok(out)
}
pub fn range_row_ids_visible_f64(
&mut self,
column_id: u16,
lo: f64,
lo_inclusive: bool,
hi: f64,
hi_inclusive: bool,
snapshot: Epoch,
) -> Result<Vec<u64>> {
let stats: Vec<(Option<f64>, Option<f64>, usize)> = match self.column_page_stats(column_id)
{
Some(s) => s
.iter()
.map(|st| {
(
be_f64(st.min.as_deref()),
be_f64(st.max.as_deref()),
st.row_count as usize,
)
})
.collect(),
None => return Ok(Vec::new()),
};
let stats_pruneable =
!self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
let mut out: Vec<u64> = Vec::new();
let mut vis = 0usize;
let mut page_start = 0usize;
for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
let page_end = page_start + nrows;
let skip = stats_pruneable
&& match (mn, mx) {
(Some(mn), Some(mx)) => {
let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
skip_lo || skip_hi
}
_ => true,
};
if !skip {
let val_page = self.read_page(column_id, seq)?;
let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
while vis < positions.len() && positions[vis] < page_end {
let local = positions[vis] - page_start;
if columnar::validity_bit(&validity, local) && !v[local].is_nan() {
let val = v[local];
let ok_lo = if lo_inclusive { val >= lo } else { val > lo };
let ok_hi = if hi_inclusive { val <= hi } else { val < hi };
if ok_lo && ok_hi {
out.push(rids[vis] as u64);
}
}
vis += 1;
}
}
} else {
while vis < positions.len() && positions[vis] < page_end {
vis += 1;
}
}
page_start = page_end;
}
Ok(out)
}
pub fn null_row_ids_visible(
&mut self,
column_id: u16,
want_nulls: bool,
snapshot: Epoch,
) -> Result<Vec<u64>> {
let stats: Vec<(usize, usize)> = match self.column_page_stats(column_id) {
Some(s) => s
.iter()
.map(|st| (st.null_count as usize, st.row_count as usize))
.collect(),
None => return Ok(Vec::new()),
};
let ty = self.resolve_type(column_id);
let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
let mut out: Vec<u64> = Vec::new();
let mut vis = 0usize;
let mut page_start = 0usize;
for (seq, &(null_count, nrows)) in stats.iter().enumerate() {
let page_end = page_start + nrows;
let skip = (want_nulls && null_count == 0) || (!want_nulls && null_count == nrows);
if !skip {
let val_page = self.read_page(column_id, seq)?;
let col = columnar::decode_page_native(ty, &val_page, nrows)?;
let validity = col.validity();
while vis < positions.len() && positions[vis] < page_end {
let local = positions[vis] - page_start;
let is_null = !columnar::validity_bit(validity, local);
if is_null == want_nulls {
out.push(rids[vis] as u64);
}
vis += 1;
}
} else {
while vis < positions.len() && positions[vis] < page_end {
vis += 1;
}
}
page_start = page_end;
}
Ok(out)
}
pub fn visible_positions_with_rids(
&mut self,
snapshot: Epoch,
) -> Result<(Vec<usize>, Vec<i64>)> {
let n = self.row_count();
if n == 0 {
return Ok((Vec::new(), Vec::new()));
}
if self.is_clean() && self.epoch_override.is_none() {
let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
columnar::NativeColumn::Int64 { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
};
let positions: Vec<usize> = (0..n).collect();
return Ok((positions, row_ids));
}
let (row_ids, epochs, deleted) = self.system_columns_native()?;
let mut idxs: Vec<usize> = Vec::new();
let mut i = 0;
while i < n {
let rid = row_ids[i] as u64;
let mut best: Option<usize> = None;
let mut j = i;
while j < n && row_ids[j] as u64 == rid {
if epochs[j] as u64 <= snapshot.0 {
best = Some(j);
}
j += 1;
}
if let Some(b) = best {
if deleted[b] == 0 {
idxs.push(b);
}
}
i = j;
}
let rids: Vec<i64> = idxs.iter().map(|&k| row_ids[k]).collect();
Ok((idxs, rids))
}
pub fn page_row_counts(&self, column_id: u16) -> Result<Vec<usize>> {
Ok(self
.find_header(column_id)?
.page_stats
.iter()
.map(|s| s.row_count as usize)
.collect())
}
pub fn has_column(&self, column_id: u16) -> bool {
self.dir.iter().any(|h| h.column_id == column_id)
}
pub fn column_page_stats(&self, column_id: u16) -> Option<&[crate::page::PageStat]> {
self.dir
.iter()
.find(|h| h.column_id == column_id)
.map(|ch| ch.page_stats.as_slice())
}
pub(crate) fn system_columns_native(&mut self) -> Result<(Vec<i64>, Vec<i64>, Vec<u8>)> {
let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
columnar::NativeColumn::Int64 { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
};
let epochs = match self.column_native_shared(SYS_EPOCH)? {
columnar::NativeColumn::Int64 { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys epoch not int64".into())),
};
let deleted = match self.column_native_shared(SYS_DELETED)? {
columnar::NativeColumn::Bool { data, .. } => data,
_ => return Err(MongrelError::InvalidArgument("sys deleted not bool".into())),
};
Ok((row_ids, epochs, deleted))
}
pub fn visible_versions(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
let n = self.row_count();
if n == 0 {
return Ok(Vec::new());
}
let row_ids = self.column(SYS_ROW_ID)?.to_vec();
let epochs = self.column(SYS_EPOCH)?.to_vec();
let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
for i in 0..n {
let rid = int_at(&row_ids, i);
let epoch = int_at(&epochs, i);
if epoch > snapshot.0 {
continue;
}
best.entry(rid)
.and_modify(|e| {
if epoch > e.0 {
*e = (epoch, i);
}
})
.or_insert((epoch, i));
}
let mut picks: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
picks.sort();
let mut out = Vec::with_capacity(picks.len());
for i in picks {
out.push(self.materialize(i)?);
}
Ok(out)
}
pub fn visible_rows(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
Ok(self
.visible_versions(snapshot)?
.into_iter()
.filter(|r| !r.deleted)
.collect())
}
pub(crate) fn materialize(&mut self, index: usize) -> Result<Row> {
let row_id = RowId(int_at(self.column(SYS_ROW_ID)?, index));
let epoch = Epoch(int_at(self.column(SYS_EPOCH)?, index));
let deleted = bool_at(self.column(SYS_DELETED)?, index);
let col_ids: Vec<u16> = self.schema.columns.iter().map(|c| c.id).collect();
let mut columns = HashMap::new();
for id in col_ids {
let val = if self.dir.iter().any(|h| h.column_id == id) {
self.column(id)?.get(index).cloned().unwrap_or(Value::Null)
} else {
Value::Null
};
columns.insert(id, val);
}
Ok(Row {
row_id,
committed_epoch: epoch,
columns,
deleted,
})
}
pub(crate) fn materialize_batch(&mut self, indices: &[usize]) -> Result<Vec<Row>> {
if indices.is_empty() {
return Ok(Vec::new());
}
use std::collections::HashMap;
let rid_col = self.column_native_shared(SYS_ROW_ID)?;
let epoch_col = self.column_native_shared(SYS_EPOCH)?;
let del_col = self.column_native_shared(SYS_DELETED)?;
let mut user: HashMap<u16, columnar::NativeColumn> = HashMap::new();
let present: Vec<u16> = self
.schema
.columns
.iter()
.map(|c| c.id)
.filter(|id| self.dir.iter().any(|h| h.column_id == *id))
.collect();
for id in present {
user.insert(id, self.column_native(id)?);
}
let i64_at = |col: &columnar::NativeColumn, i: usize| -> i64 {
match col {
columnar::NativeColumn::Int64 { data, .. } => data.get(i).copied().unwrap_or(0),
_ => 0,
}
};
let bool_at_native = |col: &columnar::NativeColumn, i: usize| -> bool {
match col {
columnar::NativeColumn::Bool { data, .. } => {
data.get(i).copied().map(|b| b != 0).unwrap_or(false)
}
_ => false,
}
};
let mut rows = Vec::with_capacity(indices.len());
for &idx in indices {
let row_id = RowId(i64_at(&rid_col, idx) as u64);
let epoch = Epoch(i64_at(&epoch_col, idx) as u64);
let deleted = bool_at_native(&del_col, idx);
let mut columns = HashMap::with_capacity(self.schema.columns.len());
for cdef in self.schema.columns.iter() {
let val = match user.get(&cdef.id) {
Some(col) => col.value_at(idx).unwrap_or(Value::Null),
None => Value::Null,
};
columns.insert(cdef.id, val);
}
rows.push(Row {
row_id,
committed_epoch: epoch,
columns,
deleted,
});
}
Ok(rows)
}
}
fn int_at(vals: &[Value], i: usize) -> u64 {
match vals.get(i) {
Some(Value::Int64(x)) => *x as u64,
_ => 0,
}
}
fn bool_at(vals: &[Value], i: usize) -> bool {
matches!(vals.get(i), Some(Value::Bool(true)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::columnar::NativeColumn;
use crate::memtable::Value;
use crate::rowid::RowId;
use crate::schema::{ColumnDef, ColumnFlags};
use tempfile::tempdir;
fn schema() -> Schema {
Schema {
schema_id: 1,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "name".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty(),
},
],
indexes: Vec::new(),
colocation: vec![],
constraints: Default::default(),
clustered: false,
}
}
fn rows() -> Vec<Row> {
vec![
Row::new(RowId(1), Epoch(10))
.with_column(1, Value::Int64(1))
.with_column(2, Value::Bytes(b"alice".to_vec())),
Row::new(RowId(2), Epoch(10))
.with_column(1, Value::Int64(2))
.with_column(2, Value::Bytes(b"bob".to_vec())),
Row::new(RowId(3), Epoch(10))
.with_column(1, Value::Int64(3))
.with_column(2, Value::Bytes(b"carol".to_vec())),
]
}
#[test]
fn flush_then_read_mvcc() {
let dir = tempdir().unwrap();
let path = dir.path().join("r-1.sr");
let header = RunWriter::new(&schema(), 1, Epoch(10), 0)
.write(&path, &rows())
.unwrap();
assert_eq!(header.row_count, 3);
let mut r = RunReader::open(&path, schema(), None).unwrap();
let (e, row) = r.get_version(RowId(2), Epoch(20)).unwrap().unwrap();
assert_eq!(e, Epoch(10));
assert_eq!(row.row_id, RowId(2));
assert!(matches!(row.columns.get(&2), Some(Value::Bytes(_))));
assert!(r.get_version(RowId(99), Epoch(20)).unwrap().is_none());
let all = r.visible_rows(Epoch(20)).unwrap();
assert_eq!(all.len(), 3);
}
#[test]
fn learned_index_was_stored() {
let dir = tempdir().unwrap();
let path = dir.path().join("r-2.sr");
RunWriter::new(&schema(), 2, Epoch(1), 0)
.write(&path, &rows())
.unwrap();
let header = read_header(&path).unwrap();
assert!(
header.index_trailer_offset != 0,
"trailer should be present"
);
}
#[test]
fn low_level_container_still_round_trips() {
let dir = tempdir().unwrap();
let path = dir.path().join("r-3.sr");
let cols = vec![ColumnPayload {
column_id: 1,
type_id_tag: 8,
encoding: Encoding::Plain,
pages: vec![vec![1, 2, 3, 4]],
page_stats: Vec::new(),
}];
let header = write_run(
&path,
&RunSpec {
run_id: 1,
schema_id: 1,
epoch_created: 1,
level: 0,
flags: 0,
sort_key_column_id: SORT_KEY_ROW_ID,
row_count: 1,
min_row_id: 0,
max_row_id: 0,
columns: &cols,
},
)
.unwrap();
let back = read_header(&path).unwrap();
assert_eq!(back.content_hash, header.content_hash);
assert_eq!(read_column_dir(&path, &back).unwrap().len(), 1);
}
#[test]
fn detects_corruption() {
let dir = tempdir().unwrap();
let path = dir.path().join("r-4.sr");
write_run(
&path,
&RunSpec {
run_id: 1,
schema_id: 1,
epoch_created: 1,
level: 0,
flags: 0,
sort_key_column_id: SORT_KEY_ROW_ID,
row_count: 1,
min_row_id: 0,
max_row_id: 0,
columns: &[ColumnPayload {
column_id: 1,
type_id_tag: 8,
encoding: Encoding::Plain,
pages: vec![vec![1, 2, 3]],
page_stats: Vec::new(),
}],
},
)
.unwrap();
let mut bytes = std::fs::read(&path).unwrap();
bytes[300] ^= 0xFF;
std::fs::write(&path, bytes).unwrap();
let err = read_header(&path).unwrap_err();
assert!(
matches!(err, MongrelError::ChecksumMismatch { .. }),
"got {err:?}"
);
}
#[test]
fn mmap_and_vec_writers_are_byte_identical() {
let dir = tempdir().unwrap();
let stats = vec![PageStat {
first_row_id: 0,
last_row_id: 1,
null_count: 0,
row_count: 2,
min: Some(vec![0]),
max: Some(vec![9]),
offset: 0,
compressed_len: 0,
uncompressed_len: 0,
}];
let spec = RunSpec {
run_id: 7,
schema_id: 1,
epoch_created: 10,
level: 0,
flags: 0,
sort_key_column_id: SORT_KEY_ROW_ID,
row_count: 2,
min_row_id: 0,
max_row_id: 1,
columns: &[
ColumnPayload {
column_id: 1,
type_id_tag: 8,
encoding: Encoding::Plain,
pages: vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8]],
page_stats: stats.clone(),
},
ColumnPayload {
column_id: 2,
type_id_tag: 8,
encoding: Encoding::Plain,
pages: vec![vec![10, 20], vec![30, 40]],
page_stats: stats.clone(),
},
],
};
let trailer = b"learned-trailer-bytes";
let plan = plan_run(&spec, None, Some(trailer)).expect("plan");
let mut buf = vec![0u8; plan.total];
let h_place = place_run(&spec, None, Some(trailer), &plan, &mut buf)
.expect("place_run into Vec succeeds");
let path_vec = dir.path().join("vec.sr");
let h_vec = write_run_vec(&path_vec, &spec, None, Some(trailer)).unwrap();
let bv = std::fs::read(&path_vec).unwrap();
assert_eq!(h_place.content_hash, h_vec.content_hash);
assert_eq!(h_place.footer_offset, h_vec.footer_offset);
assert_eq!(
buf, bv,
"place_run and write_run_vec must be byte-identical"
);
assert!(read_header(&path_vec).is_ok());
let path_mmap = dir.path().join("mmap.sr");
match write_run_mmap(&path_mmap, &spec, None, Some(trailer)) {
Ok(h_mmap) => {
let bm = std::fs::read(&path_mmap).unwrap();
assert_eq!(bm, bv, "mmap run must be byte-identical to vec run");
assert_eq!(h_mmap.content_hash, h_vec.content_hash);
}
Err(e) if is_mmap_unavailable(&e) => {
eprintln!("note: file mmap unavailable here; vec path covered (1)/(2)");
}
Err(e) => panic!("unexpected mmap error: {e:?}"),
}
}
#[test]
fn column_native_shared_matches_column_native() {
let dir = tempdir().unwrap();
let path = dir.path().join("r.sr");
let n = 65_536 * 2 + 9;
let id_col = NativeColumn::int64_sequence(1, n);
let mut offsets = vec![0u32];
let mut values = Vec::new();
for i in 0..n {
values.extend_from_slice(format!("v{}", i % 17).as_bytes());
offsets.push(values.len() as u32);
}
let name_col = NativeColumn::Bytes {
offsets,
values,
validity: vec![0xFF; n.div_ceil(8)],
};
let header = RunWriter::new(&schema(), 9, Epoch(5), 0)
.write_native(&path, &[(1, id_col), (2, name_col)], n, 1)
.unwrap();
let mut reader = RunReader::open_with_cache(&path, schema(), None, None, None, 0, None)
.expect("open reader");
assert!(reader.has_mmap(), "test env must support read-only mmap");
assert_eq!(reader.row_count(), header.row_count as usize);
for cid in [1u16, 2] {
let a = reader.column_native(cid).expect("column_native");
let b = reader
.column_native_shared(cid)
.expect("column_native_shared");
assert_eq!(a.len(), b.len(), "len mismatch col {cid}");
match (&a, &b) {
(
NativeColumn::Int64 {
data: da,
validity: va,
},
NativeColumn::Int64 {
data: db,
validity: vb,
},
) => {
assert_eq!(da, db, "Int64 data col {cid}");
assert_eq!(va, vb, "Int64 validity col {cid}");
}
(
NativeColumn::Bytes {
offsets: oa,
values: ua,
validity: va,
},
NativeColumn::Bytes {
offsets: ob,
values: ub,
validity: vb,
},
) => {
assert_eq!(oa, ob, "Bytes offsets col {cid}");
assert_eq!(ua, ub, "Bytes values col {cid}");
assert_eq!(va, vb, "Bytes validity col {cid}");
}
_ => panic!("type mismatch col {cid}: {a:?} vs {b:?}"),
}
}
}
#[test]
fn page_cache_key_distinguishes_tables() {
let a = page_cache_key(1, 5, 2, 3);
let b = page_cache_key(2, 5, 2, 3);
assert_ne!(a, b, "same run/col/page, different table must differ");
assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 6, 2, 3));
assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 5, 2, 4));
}
}