1use crate::columnar;
10use crate::encryption::{setup_run_encryption, Cipher, Kek, RunEncryption};
11use crate::epoch::Epoch;
12use crate::error::{MongrelError, Result};
13use crate::index::pgm::PgmIndex;
14use crate::memtable::{Row, Value};
15use crate::page::{Encoding, PageStat};
16use crate::row_id_set::RowIdSet;
17use crate::rowid::RowId;
18use crate::schema::{Schema, TypeId};
19use serde::{Deserialize, Serialize};
20use sha2::{Digest, Sha256};
21use std::collections::HashMap;
22use std::fs::{File, OpenOptions};
23use std::io::{Read, Seek, SeekFrom, Write};
24use std::path::Path;
25use std::sync::Arc;
26
27pub const RUN_MAGIC: [u8; 8] = *b"MONGRRUN";
28pub const RUN_FORMAT_VERSION: u16 = 1;
29pub const RUN_HEADER_VERSION: u16 = 2;
35pub const RUN_HEADER_PAD: usize = 256;
36
37const ENC_STATS_NONCE_COLUMN: u16 = u16::MAX;
42const ENC_STATS_NONCE_SEQ: u32 = u16::MAX as u32;
43
44type PageMinMax = (Option<Vec<u8>>, Option<Vec<u8>>);
46
47type EncryptedColumnStats = Vec<(u16, Vec<PageMinMax>)>;
54
55const GCM_TAG_LEN: usize = 16;
60
61fn on_disk_len(page_len: usize, encrypted: bool) -> usize {
64 if encrypted {
65 page_len + GCM_TAG_LEN
66 } else {
67 page_len
68 }
69}
70
71pub const RUN_FLAG_ENCRYPTED: u8 = 1 << 0;
72pub const RUN_FLAG_TOMBSTONE_ONLY: u8 = 1 << 1;
73pub const RUN_FLAG_CLEAN: u8 = 1 << 2;
82pub const RUN_FLAG_UNIFORM_EPOCH: u8 = 1 << 3;
90pub const SORT_KEY_ROW_ID: u16 = 0xFFFF;
91
92pub const SYS_ROW_ID: u16 = 0xFFFE;
94pub const SYS_EPOCH: u16 = 0xFFFD;
95pub const SYS_DELETED: u16 = 0xFFFC;
96
97const LEARNED_EPSILON: usize = 64;
100
101fn build_learned_trailer(row_ids: &[Value]) -> Vec<u8> {
105 let points: Vec<(u64, usize)> = row_ids
106 .iter()
107 .enumerate()
108 .filter_map(|(i, v)| match v {
109 Value::Int64(r) => Some((*r as u64, i)),
110 _ => None,
111 })
112 .collect();
113 let pgm = PgmIndex::build(&points, LEARNED_EPSILON);
114 bincode::serialize(&pgm).expect("pgm serialize")
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct RunHeader {
119 pub magic: [u8; 8],
120 pub format_version: u16,
121 pub header_layout_version: u16,
122 pub run_id: u128,
123 pub content_hash: [u8; 32],
124 pub schema_id: u64,
125 pub epoch_created: u64,
126 pub level: u8,
127 pub flags: u8,
128 pub sort_key_column_id: u16,
129 pub row_count: u64,
130 pub min_row_id: u64,
131 pub max_row_id: u64,
132 pub column_count: u64,
133 pub column_dir_offset: u64,
134 pub index_trailer_offset: u64,
135 pub encryption_descriptor_offset: u64,
136 pub footer_offset: u64,
137 pub encrypted_stats_offset: u64,
141 pub encrypted_stats_len: u64,
142}
143
144impl RunHeader {
145 pub fn is_encrypted(&self) -> bool {
146 self.flags & RUN_FLAG_ENCRYPTED != 0
147 }
148 pub fn is_clean(&self) -> bool {
149 self.flags & RUN_FLAG_CLEAN != 0
150 }
151 pub fn is_uniform_epoch(&self) -> bool {
152 self.flags & RUN_FLAG_UNIFORM_EPOCH != 0
153 }
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct ColumnPageHeader {
158 pub column_id: u16,
159 pub type_id_tag: u16,
160 pub encoding: u8,
161 pub flags: u8,
162 pub page_count: u32,
163 pub page_region_offset: u64,
164 pub page_region_len: u64,
165 pub page_stats: Vec<PageStat>,
166}
167
168impl ColumnPageHeader {
169 const PAGE_ENCRYPTED: u8 = 1 << 0;
170}
171
172const RUN_MAC_LEN: usize = 32;
175
176fn compute_run_mac(
180 enc: Option<&RunEncryption>,
181 header_bytes: &[u8],
182 dir_bytes: &[u8],
183) -> Option<[u8; RUN_MAC_LEN]> {
184 #[cfg(feature = "encryption")]
185 {
186 if let Some(e) = enc {
187 if let Some(mac_key) = &e.mac_key {
188 return Some(crate::encryption::run_metadata_mac(
189 mac_key,
190 header_bytes,
191 dir_bytes,
192 &e.descriptor_bytes,
193 ));
194 }
195 }
196 }
197 #[cfg(not(feature = "encryption"))]
198 {
199 let _ = (enc, header_bytes, dir_bytes);
200 }
201 None
202}
203
204pub struct ColumnPayload {
206 pub column_id: u16,
207 pub type_id_tag: u16,
208 pub encoding: Encoding,
209 pub pages: Vec<Vec<u8>>,
210 pub page_stats: Vec<PageStat>,
214}
215
216pub struct RunSpec<'a> {
218 pub run_id: u128,
219 pub schema_id: u64,
220 pub epoch_created: u64,
221 pub level: u8,
222 pub flags: u8,
223 pub sort_key_column_id: u16,
224 pub row_count: u64,
225 pub min_row_id: u64,
226 pub max_row_id: u64,
227 pub columns: &'a [ColumnPayload],
228}
229
230pub fn write_run(path: impl AsRef<Path>, spec: &RunSpec) -> Result<RunHeader> {
232 write_run_with(path, spec, None, &[], None)
233}
234
235pub fn write_run_with(
245 path: impl AsRef<Path>,
246 spec: &RunSpec,
247 kek: Option<&Kek>,
248 indexable_columns: &[(u16, u8)],
249 index_trailer: Option<&[u8]>,
250) -> Result<RunHeader> {
251 let enc: Option<RunEncryption> = match kek {
255 Some(k) => Some(setup_run_encryption(k, indexable_columns)?),
256 None => None,
257 };
258 match write_run_mmap(path.as_ref(), spec, enc.as_ref(), index_trailer) {
259 Ok(h) => Ok(h),
260 Err(e) if is_mmap_unavailable(&e) => write_run_vec(path, spec, enc, index_trailer),
261 Err(e) => Err(e),
262 }
263}
264
265fn is_mmap_unavailable(e: &MongrelError) -> bool {
269 matches!(e, MongrelError::InvalidArgument(m) if m.starts_with("__mmap_unavailable__:"))
270}
271
272fn write_run_mmap(
293 path: &Path,
294 spec: &RunSpec,
295 enc: Option<&RunEncryption>,
296 index_trailer: Option<&[u8]>,
297) -> Result<RunHeader> {
298 let plan = plan_run(spec, enc, index_trailer)?;
299 let file = OpenOptions::new()
300 .create(true)
301 .write(true)
302 .truncate(true)
303 .open(path)?;
304 file.set_len(plan.total as u64)?;
305 let mut mmap = match unsafe { memmap2::MmapMut::map_mut(&file) } {
306 Ok(m) => m,
307 Err(e) => {
308 return Err(MongrelError::InvalidArgument(format!(
309 "__mmap_unavailable__: {e}"
310 )))
311 }
312 };
313 let header = place_run(spec, enc, index_trailer, &plan, &mut mmap[..])?;
314 mmap.flush()?;
315 file.sync_all()?;
316 Ok(header)
317}
318
319struct RunPlan {
325 jobs: Vec<(usize, usize, u64, usize)>, dir_bytes: Vec<u8>,
327 encrypted: bool,
328 column_dir_offset: u64,
329 index_trailer_offset: u64,
330 encryption_descriptor_offset: u64,
331 footer_offset: u64,
332 total: usize,
333 encrypted_stats_plain: Option<Vec<u8>>,
337 encrypted_stats_offset: u64,
338 encrypted_stats_len: u64,
339}
340
341fn plan_run(
345 spec: &RunSpec,
346 enc: Option<&RunEncryption>,
347 index_trailer: Option<&[u8]>,
348) -> Result<RunPlan> {
349 let encrypted = enc.is_some();
350 let columns = spec.columns;
351 let mut jobs: Vec<(usize, usize, u64, usize)> = Vec::new();
352 let mut dir: Vec<ColumnPageHeader> = Vec::with_capacity(columns.len());
353 let mut enc_stats: EncryptedColumnStats = Vec::new();
354 let mut cursor: u64 = RUN_HEADER_PAD as u64;
355 for (ci, col) in columns.iter().enumerate() {
356 let region_offset = cursor;
357 let mut region_len = 0u64;
358 let mut stats = Vec::with_capacity(col.pages.len());
359 let mut col_minmax: Vec<PageMinMax> = Vec::new();
360 for (ps, page) in col.pages.iter().enumerate() {
361 if encrypted && ps >= u16::MAX as usize {
366 return Err(MongrelError::Full(format!(
367 "column {:#x} exceeds 65534 pages; encrypted-run page-seq nonce space exhausted",
368 col.column_id
369 )));
370 }
371 let odl = on_disk_len(page.len(), encrypted);
372 jobs.push((ci, ps, cursor, odl));
373 let mut stat = col.page_stats.get(ps).cloned().unwrap_or(PageStat {
374 first_row_id: 0,
375 last_row_id: 0,
376 null_count: 0,
377 row_count: 0,
378 min: None,
379 max: None,
380 offset: 0,
381 compressed_len: 0,
382 uncompressed_len: 0,
383 });
384 stat.offset = cursor;
385 stat.compressed_len = odl as u32;
386 stat.uncompressed_len = page.len() as u32;
387 if encrypted {
394 col_minmax.push((stat.min.take(), stat.max.take()));
395 }
396 stats.push(stat);
397 cursor += odl as u64;
398 region_len += odl as u64;
399 }
400 if col_minmax
401 .iter()
402 .any(|(mn, mx)| mn.is_some() || mx.is_some())
403 {
404 enc_stats.push((col.column_id, col_minmax));
405 }
406 let page_flags = if encrypted {
407 ColumnPageHeader::PAGE_ENCRYPTED
408 } else {
409 0
410 };
411 dir.push(ColumnPageHeader {
412 column_id: col.column_id,
413 type_id_tag: col.type_id_tag,
414 encoding: col.encoding as u8,
415 flags: page_flags,
416 page_count: col.pages.len() as u32,
417 page_region_offset: region_offset,
418 page_region_len: region_len,
419 page_stats: stats,
420 });
421 }
422 let dir_bytes = bincode::serialize(&dir)?;
423 let column_dir_offset = cursor;
424 cursor += dir_bytes.len() as u64;
425 let index_trailer_offset = match index_trailer {
426 Some(t) => {
427 let off = cursor;
428 cursor += t.len() as u64;
429 off
430 }
431 None => 0,
432 };
433 let encryption_descriptor_offset = match enc {
434 Some(e) => {
435 let off = cursor;
436 cursor += 4 + e.descriptor_bytes.len() as u64;
437 off
438 }
439 None => 0,
440 };
441 let (encrypted_stats_plain, encrypted_stats_offset, encrypted_stats_len) =
442 if encrypted && !enc_stats.is_empty() {
443 let plain = bincode::serialize(&enc_stats)?;
444 let ct_len = on_disk_len(plain.len(), true) as u64;
445 let off = cursor;
446 cursor += ct_len;
447 (Some(plain), off, ct_len)
448 } else {
449 (None, 0, 0)
450 };
451 let footer_offset = cursor;
452 let total = footer_offset as usize + 8 + 8 + 32 + if encrypted { RUN_MAC_LEN } else { 0 };
455 Ok(RunPlan {
456 jobs,
457 dir_bytes,
458 encrypted,
459 column_dir_offset,
460 index_trailer_offset,
461 encryption_descriptor_offset,
462 footer_offset,
463 total,
464 encrypted_stats_plain,
465 encrypted_stats_offset,
466 encrypted_stats_len,
467 })
468}
469
470fn place_run(
477 spec: &RunSpec,
478 enc: Option<&RunEncryption>,
479 index_trailer: Option<&[u8]>,
480 plan: &RunPlan,
481 buf: &mut [u8],
482) -> Result<RunHeader> {
483 use rayon::prelude::*;
484 use std::borrow::Cow;
485 debug_assert_eq!(
486 buf.len(),
487 plan.total,
488 "place_run: buffer must be exactly plan.total bytes"
489 );
490
491 let cipher: Option<&dyn Cipher> = enc.map(|e| e.cipher.as_ref());
492 let nonce_prefix = enc.map(|e| e.nonce_prefix);
493 let columns = spec.columns;
494
495 struct SyncPtr(*mut u8);
500 unsafe impl Send for SyncPtr {}
501 unsafe impl Sync for SyncPtr {}
502 impl SyncPtr {
503 fn get(&self) -> *mut u8 {
504 self.0
505 }
506 }
507 let base = SyncPtr(buf.as_mut_ptr());
508 plan.jobs
509 .par_iter()
510 .map(
511 |&(ci, ps, offset, odl): &(usize, usize, u64, usize)| -> Result<()> {
512 let page = &columns[ci].pages[ps];
513 let dst =
514 unsafe { std::slice::from_raw_parts_mut(base.get().add(offset as usize), odl) };
515 let bytes: Cow<[u8]> = match cipher {
516 Some(c) => {
517 let nonce =
518 page_nonce(nonce_prefix.unwrap(), columns[ci].column_id, ps as u32);
519 let ct = c.encrypt_page(&nonce, page)?;
520 assert_eq!(
523 ct.len(), odl,
524 "ciphertext length {} != predicted {}; GCM tag size assumption is stale",
525 ct.len(), odl
526 );
527 Cow::Owned(ct)
528 }
529 None => {
530 debug_assert_eq!(page.len(), odl);
531 Cow::Borrowed(page.as_slice())
532 }
533 };
534 dst.copy_from_slice(&bytes);
535 Ok(())
536 },
537 )
538 .collect::<Result<()>>()?;
539
540 let content_hash = {
542 let mut h = Sha256::new();
543 h.update(&buf[RUN_HEADER_PAD..plan.column_dir_offset as usize]);
544 h.finalize()
545 };
546
547 let doff = plan.column_dir_offset as usize;
549 buf[doff..doff + plan.dir_bytes.len()].copy_from_slice(&plan.dir_bytes);
550 if let Some(t) = index_trailer {
551 let off = plan.index_trailer_offset as usize;
552 buf[off..off + t.len()].copy_from_slice(t);
553 }
554 if let Some(e) = enc {
555 let off = plan.encryption_descriptor_offset as usize;
556 buf[off..off + 4].copy_from_slice(&(e.descriptor_bytes.len() as u32).to_le_bytes());
557 buf[off + 4..off + 4 + e.descriptor_bytes.len()].copy_from_slice(&e.descriptor_bytes);
558 }
559 if let (Some(e), Some(plain)) = (enc, plan.encrypted_stats_plain.as_ref()) {
560 let nonce = page_nonce(e.nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
561 let ct = e.cipher.encrypt_page(&nonce, plain)?;
562 debug_assert_eq!(ct.len() as u64, plan.encrypted_stats_len);
563 let off = plan.encrypted_stats_offset as usize;
564 buf[off..off + ct.len()].copy_from_slice(&ct);
565 }
566
567 let header_flags = if plan.encrypted {
569 spec.flags | RUN_FLAG_ENCRYPTED
570 } else {
571 spec.flags
572 };
573 let header = RunHeader {
574 magic: RUN_MAGIC,
575 format_version: RUN_FORMAT_VERSION,
576 header_layout_version: RUN_HEADER_VERSION,
577 run_id: spec.run_id,
578 content_hash: content_hash.into(),
579 schema_id: spec.schema_id,
580 epoch_created: spec.epoch_created,
581 level: spec.level,
582 flags: header_flags,
583 sort_key_column_id: spec.sort_key_column_id,
584 row_count: spec.row_count,
585 min_row_id: spec.min_row_id,
586 max_row_id: spec.max_row_id,
587 column_count: columns.len() as u64,
588 column_dir_offset: plan.column_dir_offset,
589 index_trailer_offset: plan.index_trailer_offset,
590 encryption_descriptor_offset: plan.encryption_descriptor_offset,
591 footer_offset: plan.footer_offset,
592 encrypted_stats_offset: plan.encrypted_stats_offset,
593 encrypted_stats_len: plan.encrypted_stats_len,
594 };
595 let header_bytes = bincode::serialize(&header)?;
596 if header_bytes.len() > RUN_HEADER_PAD {
597 return Err(MongrelError::InvalidArgument(format!(
598 "run header too large: {} > {RUN_HEADER_PAD}",
599 header_bytes.len()
600 )));
601 }
602 buf[..header_bytes.len()].copy_from_slice(&header_bytes);
603
604 let checksum = Sha256::digest(&buf[..plan.footer_offset as usize]);
605 let foot = plan.footer_offset as usize;
606 buf[foot..foot + 8].copy_from_slice(&RUN_MAGIC);
607 buf[foot + 8..foot + 16].copy_from_slice(&plan.footer_offset.to_le_bytes());
608 buf[foot + 16..foot + 48].copy_from_slice(&checksum);
609 if let Some(tag) = compute_run_mac(enc, &header_bytes, &plan.dir_bytes) {
611 buf[foot + 48..foot + 48 + RUN_MAC_LEN].copy_from_slice(&tag);
612 }
613 Ok(header)
614}
615
616fn write_run_vec(
619 path: impl AsRef<Path>,
620 spec: &RunSpec,
621 enc: Option<RunEncryption>,
622 index_trailer: Option<&[u8]>,
623) -> Result<RunHeader> {
624 let mut buf: Vec<u8> = vec![0; RUN_HEADER_PAD]; let mut content_hasher = Sha256::new();
626 let mut dir: Vec<ColumnPageHeader> = Vec::with_capacity(spec.columns.len());
627 let mut enc_stats: EncryptedColumnStats = Vec::new();
628
629 for col in spec.columns {
630 let region_offset = buf.len() as u64;
631 let mut region_len = 0u64;
632 let mut stats = Vec::with_capacity(col.pages.len());
633 let mut col_minmax: Vec<PageMinMax> = Vec::new();
634 for (page_seq, page) in col.pages.iter().enumerate() {
635 if enc.is_some() && page_seq >= u16::MAX as usize {
636 return Err(MongrelError::Full(format!(
637 "column {:#x} exceeds 65534 pages; encrypted-run page-seq nonce space exhausted",
638 col.column_id
639 )));
640 }
641 let on_disk: Vec<u8> = match &enc {
642 Some(e) => e.cipher.encrypt_page(
643 &page_nonce(e.nonce_prefix, col.column_id, page_seq as u32),
644 page,
645 )?,
646 None => page.clone(),
647 };
648 let offset = buf.len() as u64;
649 buf.write_all(&on_disk)?;
650 content_hasher.update(&on_disk);
651 region_len += on_disk.len() as u64;
652 let mut stat = if let Some(s) = col.page_stats.get(page_seq) {
653 s.clone()
654 } else {
655 PageStat {
656 first_row_id: 0,
657 last_row_id: 0,
658 null_count: 0,
659 row_count: 0,
660 min: None,
661 max: None,
662 offset: 0,
663 compressed_len: 0,
664 uncompressed_len: 0,
665 }
666 };
667 stat.offset = offset;
668 stat.compressed_len = on_disk.len() as u32;
669 stat.uncompressed_len = page.len() as u32;
670 if enc.is_some() {
675 col_minmax.push((stat.min.take(), stat.max.take()));
676 }
677 stats.push(stat);
678 }
679 if col_minmax
680 .iter()
681 .any(|(mn, mx)| mn.is_some() || mx.is_some())
682 {
683 enc_stats.push((col.column_id, col_minmax));
684 }
685 let page_flags = if enc.is_some() {
686 ColumnPageHeader::PAGE_ENCRYPTED
687 } else {
688 0
689 };
690 dir.push(ColumnPageHeader {
691 column_id: col.column_id,
692 type_id_tag: col.type_id_tag,
693 encoding: col.encoding as u8,
694 flags: page_flags,
695 page_count: col.pages.len() as u32,
696 page_region_offset: region_offset,
697 page_region_len: region_len,
698 page_stats: stats,
699 });
700 }
701
702 let column_dir_offset = buf.len() as u64;
703 let dir_bytes = bincode::serialize(&dir)?;
704 buf.write_all(&dir_bytes)?;
705
706 let index_trailer_offset = match index_trailer {
707 Some(trailer) => {
708 let off = buf.len() as u64;
709 buf.write_all(trailer)?;
710 off
711 }
712 None => 0,
713 };
714 let encryption_descriptor_offset = match &enc {
715 Some(e) => {
716 let off = buf.len() as u64;
717 buf.write_all(&(e.descriptor_bytes.len() as u32).to_le_bytes())?;
718 buf.write_all(&e.descriptor_bytes)?;
719 off
720 }
721 None => 0,
722 };
723 let (encrypted_stats_offset, encrypted_stats_len) = match &enc {
724 Some(e) if !enc_stats.is_empty() => {
725 let plain = bincode::serialize(&enc_stats)?;
726 let nonce = page_nonce(e.nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
727 let ct = e.cipher.encrypt_page(&nonce, &plain)?;
728 let off = buf.len() as u64;
729 let len = ct.len() as u64;
730 buf.write_all(&ct)?;
731 (off, len)
732 }
733 _ => (0, 0),
734 };
735 let footer_offset = buf.len() as u64;
736
737 let header_flags = if enc.is_some() {
738 spec.flags | RUN_FLAG_ENCRYPTED
739 } else {
740 spec.flags
741 };
742 let header = RunHeader {
743 magic: RUN_MAGIC,
744 format_version: RUN_FORMAT_VERSION,
745 header_layout_version: RUN_HEADER_VERSION,
746 run_id: spec.run_id,
747 content_hash: content_hasher.finalize().into(),
748 schema_id: spec.schema_id,
749 epoch_created: spec.epoch_created,
750 level: spec.level,
751 flags: header_flags,
752 sort_key_column_id: spec.sort_key_column_id,
753 row_count: spec.row_count,
754 min_row_id: spec.min_row_id,
755 max_row_id: spec.max_row_id,
756 column_count: spec.columns.len() as u64,
757 column_dir_offset,
758 index_trailer_offset,
759 encryption_descriptor_offset,
760 footer_offset,
761 encrypted_stats_offset,
762 encrypted_stats_len,
763 };
764 let header_bytes = bincode::serialize(&header)?;
765 if header_bytes.len() > RUN_HEADER_PAD {
766 return Err(MongrelError::InvalidArgument(format!(
767 "run header too large: {} > {RUN_HEADER_PAD}",
768 header_bytes.len()
769 )));
770 }
771 buf[..header_bytes.len()].copy_from_slice(&header_bytes);
772
773 let checksum = Sha256::digest(&buf[..footer_offset as usize]);
774 buf.write_all(&RUN_MAGIC)?;
775 buf.write_all(&footer_offset.to_le_bytes())?;
776 buf.write_all(&checksum)?;
777 if let Some(tag) = compute_run_mac(enc.as_ref(), &header_bytes, &dir_bytes) {
780 buf.write_all(&tag)?;
781 }
782
783 let mut file = OpenOptions::new()
784 .create(true)
785 .write(true)
786 .truncate(true)
787 .open(path)?;
788 file.write_all(&buf)?;
789 file.sync_all()?;
790 Ok(header)
791}
792
793fn page_nonce(nonce_prefix: [u8; 12], column_id: u16, page_seq: u32) -> [u8; 12] {
794 let mut n = nonce_prefix;
795 n[8..10].copy_from_slice(&column_id.to_le_bytes());
796 n[10..12].copy_from_slice(&(page_seq as u16).to_le_bytes());
797 n
798}
799
800fn overlay_encrypted_stats(
807 path: &Path,
808 header: &RunHeader,
809 cipher: &dyn Cipher,
810 nonce_prefix: [u8; 12],
811 dir: &mut [ColumnPageHeader],
812) -> Result<()> {
813 let mut file = File::open(path)?;
814 file.seek(SeekFrom::Start(header.encrypted_stats_offset))?;
815 let mut ct = vec![0u8; header.encrypted_stats_len as usize];
816 file.read_exact(&mut ct)?;
817 let nonce = page_nonce(nonce_prefix, ENC_STATS_NONCE_COLUMN, ENC_STATS_NONCE_SEQ);
818 let plain = cipher.decrypt_page(&nonce, &ct)?;
819 let stats: EncryptedColumnStats = bincode::deserialize(&plain)
820 .map_err(|e| MongrelError::Encryption(format!("bad encrypted page-stats envelope: {e}")))?;
821 for (cid, minmax) in stats {
822 let Some(col) = dir.iter_mut().find(|c| c.column_id == cid) else {
823 continue;
824 };
825 for (stat, (mn, mx)) in col.page_stats.iter_mut().zip(minmax) {
826 stat.min = mn;
827 stat.max = mx;
828 }
829 }
830 Ok(())
831}
832
833pub(crate) fn page_cache_key(
840 table_id: u64,
841 run_id: u128,
842 column_id: u16,
843 page_seq: usize,
844) -> [u8; 32] {
845 let mut h = Sha256::new();
846 h.update(table_id.to_be_bytes());
847 h.update(run_id.to_be_bytes());
848 h.update(column_id.to_be_bytes());
849 h.update((page_seq as u64).to_be_bytes());
850 let out = h.finalize();
851 let mut k = [0u8; 32];
852 k.copy_from_slice(&out);
853 k
854}
855
856fn decrypt_or_passthrough(
860 cipher: Option<&dyn Cipher>,
861 nonce_prefix: [u8; 12],
862 column_id: u16,
863 page_seq: usize,
864 encrypted: bool,
865 buf: &[u8],
866) -> Result<Vec<u8>> {
867 if encrypted {
868 match cipher {
869 Some(c) => {
870 Ok(c.decrypt_page(&page_nonce(nonce_prefix, column_id, page_seq as u32), buf)?)
871 }
872 None => Err(MongrelError::Decryption(
873 "encrypted page but no cipher".into(),
874 )),
875 }
876 } else {
877 Ok(buf.to_vec())
878 }
879}
880
881fn read_header_fast(path: impl AsRef<Path>) -> Result<RunHeader> {
891 let mut file = File::open(path)?;
892 let mut header_buf = vec![0u8; RUN_HEADER_PAD];
893 file.read_exact(&mut header_buf)?;
894 let header: RunHeader = bincode::deserialize(&header_buf)
895 .map_err(|e| MongrelError::InvalidArgument(format!("bad run header: {e}")))?;
896 if header.magic != RUN_MAGIC {
897 return Err(MongrelError::MagicMismatch {
898 what: "sorted run",
899 expected: RUN_MAGIC,
900 got: header.magic,
901 });
902 }
903 file.seek(SeekFrom::Start(header.footer_offset))?;
904 let mut footer_magic = [0u8; 8];
905 file.read_exact(&mut footer_magic)?;
906 if footer_magic != RUN_MAGIC {
907 return Err(MongrelError::MagicMismatch {
908 what: "sorted run footer",
909 expected: RUN_MAGIC,
910 got: footer_magic,
911 });
912 }
913 Ok(header)
914}
915
916fn read_header_cached(
926 path: impl AsRef<Path>,
927 verified_runs: &parking_lot::Mutex<std::collections::HashSet<u128>>,
928) -> Result<RunHeader> {
929 let header = read_header_fast(path.as_ref())?;
930 if verified_runs.lock().contains(&header.run_id) {
931 return Ok(header);
932 }
933 let verified = read_header(path)?;
934 verified_runs.lock().insert(verified.run_id);
935 Ok(verified)
936}
937
938pub fn read_header(path: impl AsRef<Path>) -> Result<RunHeader> {
940 let mut file = File::open(path)?;
941 let mut header_buf = vec![0u8; RUN_HEADER_PAD];
942 file.read_exact(&mut header_buf)?;
943 let header: RunHeader = bincode::deserialize(&header_buf)
944 .map_err(|e| MongrelError::InvalidArgument(format!("bad run header: {e}")))?;
945 if header.magic != RUN_MAGIC {
946 return Err(MongrelError::MagicMismatch {
947 what: "sorted run",
948 expected: RUN_MAGIC,
949 got: header.magic,
950 });
951 }
952
953 file.seek(SeekFrom::Start(header.footer_offset))?;
954 let mut footer = [0u8; 8 + 8 + 32];
955 file.read_exact(&mut footer)?;
956 if footer[..8] != RUN_MAGIC {
957 return Err(MongrelError::MagicMismatch {
958 what: "sorted run footer",
959 expected: RUN_MAGIC,
960 got: footer[..8].try_into().unwrap(),
961 });
962 }
963 let mut hasher = Sha256::new();
964 hasher.update(&header_buf);
965 let body_len = header.footer_offset.saturating_sub(RUN_HEADER_PAD as u64);
966 file.seek(SeekFrom::Start(RUN_HEADER_PAD as u64))?;
967 let mut body = vec![0u8; body_len as usize];
968 file.read_exact(&mut body)?;
969 hasher.update(&body);
970 let computed: [u8; 32] = hasher.finalize().into();
971 let stored: [u8; 32] = footer[16..].try_into().unwrap();
972 if computed != stored {
973 return Err(MongrelError::ChecksumMismatch {
974 expected: u64::from_be_bytes(stored[..8].try_into().unwrap()),
975 actual: u64::from_be_bytes(computed[..8].try_into().unwrap()),
976 context: "sorted run footer".into(),
977 });
978 }
979 Ok(header)
980}
981
982pub fn read_column_dir(
984 path: impl AsRef<Path>,
985 header: &RunHeader,
986) -> Result<Vec<ColumnPageHeader>> {
987 let mut file = File::open(path)?;
988 file.seek(SeekFrom::Start(header.column_dir_offset))?;
989 let end = if header.index_trailer_offset != 0 {
990 header.index_trailer_offset
991 } else {
992 header.footer_offset
993 };
994 let len = end.saturating_sub(header.column_dir_offset);
995 let mut buf = vec![0u8; len as usize];
996 file.read_exact(&mut buf)?;
997 let dir: Vec<ColumnPageHeader> = bincode::deserialize(&buf)
998 .map_err(|e| MongrelError::InvalidArgument(format!("bad column dir: {e}")))?;
999 Ok(dir)
1000}
1001
1002pub(crate) fn read_encryption_descriptor_bytes(
1005 path: impl AsRef<Path>,
1006 header: &RunHeader,
1007) -> Result<Vec<u8>> {
1008 let mut file = File::open(path)?;
1009 file.seek(SeekFrom::Start(header.encryption_descriptor_offset))?;
1010 let mut len_buf = [0u8; 4];
1011 file.read_exact(&mut len_buf)?;
1012 let len = u32::from_le_bytes(len_buf) as usize;
1013 const MAX_DESCRIPTOR_BYTES: usize = 65_536;
1018 if len > MAX_DESCRIPTOR_BYTES {
1019 return Err(MongrelError::InvalidArgument(format!(
1020 "encryption descriptor length {len} exceeds {MAX_DESCRIPTOR_BYTES}"
1021 )));
1022 }
1023 let mut buf = vec![0u8; len];
1024 file.read_exact(&mut buf)?;
1025 Ok(buf)
1026}
1027
1028#[cfg(feature = "encryption")]
1035fn verify_run_mac(
1036 path: &Path,
1037 header: &RunHeader,
1038 dir: &[ColumnPageHeader],
1039 kek: &Kek,
1040 desc_bytes: &[u8],
1041) -> Result<()> {
1042 let header_bytes = bincode::serialize(header)?;
1043 let dir_bytes = bincode::serialize(dir)?;
1044 let mac_key = kek.derive_run_mac_key();
1045 let expected =
1046 crate::encryption::run_metadata_mac(&mac_key, &header_bytes, &dir_bytes, desc_bytes);
1047 let mut file = File::open(path)?;
1048 file.seek(SeekFrom::Start(header.footer_offset + 8 + 8 + 32))?;
1049 let mut tag = [0u8; RUN_MAC_LEN];
1050 file.read_exact(&mut tag).map_err(|_| {
1051 MongrelError::Decryption(
1052 "encrypted run is missing or truncated its metadata MAC; cannot \
1053 authenticate metadata"
1054 .into(),
1055 )
1056 })?;
1057 let mut diff = 0u8;
1059 for (x, y) in tag.iter().zip(expected.iter()) {
1060 diff |= x ^ y;
1061 }
1062 if diff != 0 {
1063 return Err(MongrelError::Decryption(
1064 "run metadata authentication failed — tampered run or wrong key".into(),
1065 ));
1066 }
1067 Ok(())
1068}
1069
1070#[cfg(not(feature = "encryption"))]
1071fn verify_run_mac(
1072 _path: &Path,
1073 _header: &RunHeader,
1074 _dir: &[ColumnPageHeader],
1075 _kek: &Kek,
1076 _desc_bytes: &[u8],
1077) -> Result<()> {
1078 Ok(())
1079}
1080
1081pub struct RunWriter<'a> {
1090 schema: &'a Schema,
1091 run_id: u128,
1092 epoch_created: Epoch,
1093 level: u8,
1094 kek: Option<&'a Kek>,
1095 indexable_columns: Vec<(u16, u8)>,
1098 compress: columnar::Compress,
1102 clean: bool,
1107 uniform_epoch: bool,
1112 le: bool,
1119}
1120
1121impl<'a> RunWriter<'a> {
1122 pub fn new(schema: &'a Schema, run_id: u128, epoch_created: Epoch, level: u8) -> Self {
1123 Self {
1124 schema,
1125 run_id,
1126 epoch_created,
1127 level,
1128 kek: None,
1129 indexable_columns: Vec::new(),
1130 compress: columnar::Compress::Zstd(3),
1131 clean: false,
1132 uniform_epoch: false,
1133 le: false,
1134 }
1135 }
1136
1137 pub fn uniform_epoch(mut self, uniform: bool) -> Self {
1143 self.uniform_epoch = uniform;
1144 self
1145 }
1146
1147 pub fn with_encryption(mut self, kek: &'a Kek, indexable_columns: Vec<(u16, u8)>) -> Self {
1151 self.kek = Some(kek);
1152 self.indexable_columns = indexable_columns;
1153 self
1154 }
1155
1156 pub fn with_zstd_level(mut self, level: i32) -> Self {
1159 self.compress = columnar::Compress::Zstd(level);
1160 self
1161 }
1162
1163 pub fn with_lz4(mut self) -> Self {
1167 self.compress = columnar::Compress::Lz4;
1168 self
1169 }
1170
1171 pub fn with_plain(mut self) -> Self {
1174 self.compress = columnar::Compress::Plain;
1175 self
1176 }
1177
1178 pub fn clean(mut self, clean: bool) -> Self {
1184 self.clean = clean;
1185 self
1186 }
1187
1188 pub fn with_native_endian(mut self) -> Self {
1193 if cfg!(target_endian = "little") {
1194 self.le = true;
1195 }
1196 self
1197 }
1198
1199 pub fn write_native(
1206 self,
1207 path: impl AsRef<Path>,
1208 user_columns: &[(u16, columnar::NativeColumn)],
1209 n: usize,
1210 first_row_id: u64,
1211 ) -> Result<RunHeader> {
1212 use columnar::NativeColumn;
1213 let row_id_col = NativeColumn::int64_sequence(first_row_id as i64, n);
1214 let epoch_col = NativeColumn::int64_constant(self.epoch_created.0 as i64, n);
1215 let deleted_col = NativeColumn::bool_constant(false, n);
1216
1217 let learned_trailer = build_learned_trailer_native(&row_id_col);
1218
1219 let row_ids = match &row_id_col {
1222 NativeColumn::Int64 { data, .. } => data.as_slice(),
1223 _ => &[],
1224 };
1225 let bounds = page_bounds(row_ids);
1226 let compress = self.compress;
1227 let le = self.le;
1228 let plain = matches!(compress, columnar::Compress::Plain);
1229 let row_id_enc = if plain {
1232 Encoding::Plain
1233 } else {
1234 Encoding::Delta
1235 };
1236 let sys_enc = if plain {
1237 Encoding::Plain
1238 } else {
1239 Encoding::Zstd
1240 };
1241
1242 let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
1243 let (pages, stats) = native_column_pages(
1244 TypeId::Int64,
1245 &row_id_col,
1246 row_id_enc,
1247 compress,
1248 le,
1249 &bounds,
1250 )?;
1251 columns.push(ColumnPayload {
1252 column_id: SYS_ROW_ID,
1253 type_id_tag: type_tag(&TypeId::Int64),
1254 encoding: row_id_enc,
1255 pages,
1256 page_stats: stats,
1257 });
1258 let (pages, stats) =
1259 native_column_pages(TypeId::Int64, &epoch_col, sys_enc, compress, le, &bounds)?;
1260 columns.push(ColumnPayload {
1261 column_id: SYS_EPOCH,
1262 type_id_tag: type_tag(&TypeId::Int64),
1263 encoding: sys_enc,
1264 pages,
1265 page_stats: stats,
1266 });
1267 let (pages, stats) =
1268 native_column_pages(TypeId::Bool, &deleted_col, sys_enc, compress, le, &bounds)?;
1269 columns.push(ColumnPayload {
1270 column_id: SYS_DELETED,
1271 type_id_tag: type_tag(&TypeId::Bool),
1272 encoding: sys_enc,
1273 pages,
1274 page_stats: stats,
1275 });
1276 use rayon::prelude::*;
1282 let user_cols: Vec<ColumnPayload> = self
1283 .schema
1284 .columns
1285 .par_iter()
1286 .map(|cdef| -> Result<ColumnPayload> {
1287 let col = user_columns
1288 .iter()
1289 .find(|(id, _)| *id == cdef.id)
1290 .map(|(_, c)| c)
1291 .ok_or_else(|| {
1292 MongrelError::ColumnNotFound(format!("user column {}", cdef.id))
1293 })?;
1294 let encoding = if plain {
1295 Encoding::Plain
1296 } else {
1297 choose_encoding_native(&cdef.ty, col)
1298 };
1299 let (pages, stats) =
1300 native_column_pages(cdef.ty, col, encoding, compress, le, &bounds)?;
1301 Ok(ColumnPayload {
1302 column_id: cdef.id,
1303 type_id_tag: type_tag(&cdef.ty),
1304 encoding,
1305 pages,
1306 page_stats: stats,
1307 })
1308 })
1309 .collect::<Result<Vec<_>>>()?;
1310 columns.extend(user_cols);
1311
1312 let flags = if self.uniform_epoch {
1318 RUN_FLAG_UNIFORM_EPOCH
1319 } else {
1320 RUN_FLAG_CLEAN
1321 };
1322 let spec = RunSpec {
1323 run_id: self.run_id,
1324 schema_id: self.schema.schema_id,
1325 epoch_created: self.epoch_created.0,
1326 level: self.level,
1327 flags,
1328 sort_key_column_id: SYS_ROW_ID,
1329 row_count: n as u64,
1330 min_row_id: first_row_id,
1331 max_row_id: first_row_id + n as u64 - 1,
1332 columns: &columns,
1333 };
1334 write_run_with(
1335 path,
1336 &spec,
1337 self.kek,
1338 &self.indexable_columns,
1339 Some(&learned_trailer),
1340 )
1341 }
1342
1343 pub fn write(self, path: impl AsRef<Path>, rows: &[Row]) -> Result<RunHeader> {
1344 let n = rows.len();
1345 let mut row_ids = Vec::with_capacity(n);
1347 let mut epochs = Vec::with_capacity(n);
1348 let mut deleted = Vec::with_capacity(n);
1349 for r in rows {
1350 row_ids.push(Value::Int64(r.row_id.0 as i64));
1351 epochs.push(Value::Int64(r.committed_epoch.0 as i64));
1352 deleted.push(Value::Bool(r.deleted));
1353 }
1354 let learned_trailer = build_learned_trailer(&row_ids);
1355 let (min_rid, max_rid) = row_id_bounds(rows);
1356 let row_id_i64: Vec<i64> = rows.iter().map(|r| r.row_id.0 as i64).collect();
1357 let bounds = page_bounds(&row_id_i64);
1358
1359 let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
1360 let (pages, stats, enc) = value_column_pages(TypeId::Int64, &row_ids, &bounds)?;
1361 columns.push(ColumnPayload {
1362 column_id: SYS_ROW_ID,
1363 type_id_tag: type_tag(&TypeId::Int64),
1364 encoding: enc,
1365 pages,
1366 page_stats: stats,
1367 });
1368 let (pages, stats, enc) = value_column_pages(TypeId::Int64, &epochs, &bounds)?;
1369 columns.push(ColumnPayload {
1370 column_id: SYS_EPOCH,
1371 type_id_tag: type_tag(&TypeId::Int64),
1372 encoding: enc,
1373 pages,
1374 page_stats: stats,
1375 });
1376 let (pages, stats, enc) = value_column_pages(TypeId::Bool, &deleted, &bounds)?;
1377 columns.push(ColumnPayload {
1378 column_id: SYS_DELETED,
1379 type_id_tag: type_tag(&TypeId::Bool),
1380 encoding: enc,
1381 pages,
1382 page_stats: stats,
1383 });
1384 for cdef in &self.schema.columns {
1386 let vals: Vec<Value> = rows
1387 .iter()
1388 .map(|r| r.columns.get(&cdef.id).cloned().unwrap_or(Value::Null))
1389 .collect();
1390 let (pages, stats, encoding) = value_column_pages(cdef.ty, &vals, &bounds)?;
1391 columns.push(ColumnPayload {
1392 column_id: cdef.id,
1393 type_id_tag: type_tag(&cdef.ty),
1394 encoding,
1395 pages,
1396 page_stats: stats,
1397 });
1398 }
1399
1400 let spec = RunSpec {
1401 run_id: self.run_id,
1402 schema_id: self.schema.schema_id,
1403 epoch_created: self.epoch_created.0,
1404 level: self.level,
1405 flags: {
1406 let mut f = if self.clean { RUN_FLAG_CLEAN } else { 0 };
1407 if self.uniform_epoch {
1408 f |= RUN_FLAG_UNIFORM_EPOCH;
1409 }
1410 f
1411 },
1412 sort_key_column_id: SYS_ROW_ID,
1413 row_count: n as u64,
1414 min_row_id: min_rid,
1415 max_row_id: max_rid,
1416 columns: &columns,
1417 };
1418 write_run_with(
1419 path,
1420 &spec,
1421 self.kek,
1422 &self.indexable_columns,
1423 Some(&learned_trailer),
1424 )
1425 }
1426}
1427
1428fn type_tag(ty: &TypeId) -> u16 {
1429 match ty {
1431 TypeId::Bool => 1,
1432 TypeId::Int64 => 8,
1433 TypeId::Float64 => 9,
1434 TypeId::Bytes => 12,
1435 TypeId::Embedding { .. } => 13,
1436 _ => 0,
1437 }
1438}
1439
1440pub(crate) fn be_i64(b: Option<&[u8]>) -> Option<i64> {
1443 let b = b?;
1444 (b.len() == 8).then(|| i64::from_be_bytes(b.try_into().unwrap()))
1445}
1446
1447pub(crate) fn be_f64(b: Option<&[u8]>) -> Option<f64> {
1449 let b = b?;
1450 (b.len() == 8).then(|| f64::from_bits(u64::from_be_bytes(b.try_into().unwrap())))
1451}
1452
1453const PAGE_ROWS: usize = 65_536;
1456
1457fn page_bounds(row_ids: &[i64]) -> Vec<(usize, usize, u64, u64)> {
1460 let n = row_ids.len();
1461 if n == 0 {
1462 return vec![(0, 0, 0, 0)];
1463 }
1464 let mut out = Vec::new();
1465 let mut start = 0;
1466 while start < n {
1467 let end = (start + PAGE_ROWS).min(n);
1468 out.push((start, end, row_ids[start] as u64, row_ids[end - 1] as u64));
1469 start = end;
1470 }
1471 out
1472}
1473
1474fn native_column_pages(
1480 ty: TypeId,
1481 col: &columnar::NativeColumn,
1482 encoding: Encoding,
1483 compress: columnar::Compress,
1484 le: bool,
1485 bounds: &[(usize, usize, u64, u64)],
1486) -> Result<(Vec<Vec<u8>>, Vec<PageStat>)> {
1487 use rayon::prelude::*;
1488 let encode_one =
1489 |&(s, e, frid, lrid): &(usize, usize, u64, u64)| -> Result<(Vec<u8>, PageStat)> {
1490 let chunk = col.slice_range(s, e);
1491 let stat = columnar::page_stat_for(ty, &chunk, frid, lrid);
1492 let page = columnar::encode_page_native(ty, &chunk, encoding, compress, le)?;
1493 Ok((page, stat))
1494 };
1495 let pairs: Vec<(Vec<u8>, PageStat)> = if bounds.len() > 1 {
1497 bounds
1498 .par_iter()
1499 .map(encode_one)
1500 .collect::<Result<Vec<_>>>()?
1501 } else {
1502 bounds.iter().map(encode_one).collect::<Result<Vec<_>>>()?
1503 };
1504 let (pages, stats) = pairs.into_iter().unzip();
1505 Ok((pages, stats))
1506}
1507
1508fn value_column_pages(
1510 ty: TypeId,
1511 vals: &[Value],
1512 bounds: &[(usize, usize, u64, u64)],
1513) -> Result<(Vec<Vec<u8>>, Vec<PageStat>, Encoding)> {
1514 let encoding = choose_encoding(&ty, vals);
1515 let mut pages = Vec::with_capacity(bounds.len());
1516 let mut stats = Vec::with_capacity(bounds.len());
1517 for &(s, e, frid, lrid) in bounds {
1518 let chunk = &vals[s..e];
1519 pages.push(columnar::encode_page(ty, chunk, encoding)?);
1520 let native = columnar::values_to_native(ty, chunk);
1521 stats.push(columnar::page_stat_for(ty, &native, frid, lrid));
1522 }
1523 Ok((pages, stats, encoding))
1524}
1525
1526fn choose_encoding(ty: &TypeId, values: &[Value]) -> Encoding {
1529 use std::collections::HashSet;
1530 if matches!(ty, TypeId::Bytes) {
1531 let n = values.len();
1532 if n > 0 {
1533 let distinct = values
1534 .iter()
1535 .filter(|v| !matches!(v, Value::Null))
1536 .map(|v| v.encode_key())
1537 .collect::<HashSet<_>>()
1538 .len();
1539 if (distinct as f64 / n as f64) < 0.5 {
1540 return Encoding::Dictionary;
1541 }
1542 }
1543 }
1544 Encoding::Zstd
1545}
1546
1547fn choose_encoding_native(ty: &TypeId, col: &columnar::NativeColumn) -> Encoding {
1550 use std::collections::HashSet;
1551 match (ty, col) {
1552 (TypeId::Int64 | TypeId::TimestampNanos, columnar::NativeColumn::Int64 { data, .. }) => {
1553 if data.windows(2).all(|w| w[0] <= w[1]) {
1554 Encoding::Delta
1555 } else {
1556 Encoding::Zstd
1557 }
1558 }
1559 (
1560 TypeId::Bytes,
1561 columnar::NativeColumn::Bytes {
1562 offsets, values, ..
1563 },
1564 ) => {
1565 let n = offsets.len().saturating_sub(1);
1566 if n > 0 {
1567 let distinct: HashSet<&[u8]> = (0..n)
1568 .map(|i| &values[offsets[i] as usize..offsets[i + 1] as usize])
1569 .collect();
1570 if (distinct.len() as f64 / n as f64) < 0.5 {
1571 return Encoding::Dictionary;
1572 }
1573 }
1574 Encoding::Zstd
1575 }
1576 _ => Encoding::Zstd,
1577 }
1578}
1579
1580fn build_learned_trailer_native(col: &columnar::NativeColumn) -> Vec<u8> {
1582 let points: Vec<(u64, usize)> = match col {
1583 columnar::NativeColumn::Int64 { data, .. } => data
1584 .iter()
1585 .enumerate()
1586 .map(|(i, v)| (*v as u64, i))
1587 .collect(),
1588 _ => Vec::new(),
1589 };
1590 let pgm = PgmIndex::build(&points, LEARNED_EPSILON);
1591 bincode::serialize(&pgm).expect("pgm serialize")
1592}
1593
1594fn row_id_bounds(rows: &[Row]) -> (u64, u64) {
1595 match (rows.first(), rows.last()) {
1596 (Some(f), Some(l)) => (f.row_id.0, l.row_id.0),
1597 _ => (0, 0),
1598 }
1599}
1600
1601pub struct RunReader {
1607 file: File,
1608 mmap: Option<memmap2::Mmap>,
1609 header: RunHeader,
1610 dir: Vec<ColumnPageHeader>,
1611 schema: Schema,
1612 table_id: u64,
1614 cipher: Option<Box<dyn Cipher>>,
1616 nonce_prefix: [u8; 12],
1618 col_cache: HashMap<u16, Vec<Value>>,
1619 page_cache: Option<Arc<parking_lot::Mutex<crate::cache::PageCache>>>,
1623 decoded_cache: Option<Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>>,
1627 epoch_override: Option<Epoch>,
1631}
1632
1633impl RunReader {
1634 pub fn open(path: impl AsRef<Path>, schema: Schema, kek: Option<Arc<Kek>>) -> Result<Self> {
1635 Self::open_with_cache(path, schema, kek, None, None, 0, None)
1636 }
1637
1638 #[allow(clippy::too_many_arguments)]
1639 pub(crate) fn open_with_cache(
1640 path: impl AsRef<Path>,
1641 schema: Schema,
1642 kek: Option<Arc<Kek>>,
1643 page_cache: Option<Arc<parking_lot::Mutex<crate::cache::PageCache>>>,
1644 decoded_cache: Option<Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>>,
1645 table_id: u64,
1646 verified_runs: Option<&parking_lot::Mutex<std::collections::HashSet<u128>>>,
1647 ) -> Result<Self> {
1648 let path = path.as_ref().to_path_buf();
1649 let header = match verified_runs {
1650 Some(cache) => read_header_cached(&path, cache)?,
1651 None => read_header(&path)?,
1652 };
1653 let mut dir = read_column_dir(&path, &header)?;
1654 let (cipher, nonce_prefix): (Option<Box<dyn Cipher>>, [u8; 12]) = if header.is_encrypted() {
1657 let kek = kek.as_ref().ok_or_else(|| {
1658 MongrelError::Encryption(
1659 "run is encrypted but no key-encryption key was provided".into(),
1660 )
1661 })?;
1662 if header.encryption_descriptor_offset == 0 {
1663 return Err(MongrelError::Encryption(
1664 "encrypted run has no encryption descriptor".into(),
1665 ));
1666 }
1667 let desc_bytes = read_encryption_descriptor_bytes(&path, &header)?;
1668 verify_run_mac(&path, &header, &dir, kek, &desc_bytes)?;
1674 let enc = crate::encryption::build_run_cipher(kek, &desc_bytes)?;
1675 if header.encrypted_stats_offset != 0 {
1679 overlay_encrypted_stats(
1680 &path,
1681 &header,
1682 enc.cipher.as_ref(),
1683 enc.nonce_prefix,
1684 &mut dir,
1685 )?;
1686 }
1687 (Some(enc.cipher), enc.nonce_prefix)
1688 } else {
1689 (None, [0u8; 12])
1690 };
1691 let file = File::open(&path)?;
1694 let mmap = unsafe { memmap2::MmapOptions::new().map(&file) }.ok();
1701 let _ = path;
1702 Ok(Self {
1703 file,
1704 mmap,
1705 header,
1706 dir,
1707 schema,
1708 table_id,
1709 cipher,
1710 nonce_prefix,
1711 col_cache: HashMap::new(),
1712 epoch_override: None,
1713 page_cache,
1714 decoded_cache,
1715 })
1716 }
1717
1718 pub fn header(&self) -> &RunHeader {
1719 &self.header
1720 }
1721
1722 pub(crate) fn set_uniform_epoch(&mut self, epoch: Epoch) {
1726 if self.header.is_uniform_epoch() {
1727 self.epoch_override = Some(epoch);
1728 self.col_cache.remove(&SYS_EPOCH);
1730 }
1731 }
1732
1733 pub fn is_clean(&self) -> bool {
1736 self.header.is_clean()
1737 }
1738
1739 pub fn row_count(&self) -> usize {
1740 self.header.row_count as usize
1741 }
1742
1743 pub(crate) fn clean_contiguous_row_ids(&self) -> bool {
1744 let n = self.row_count();
1745 n > 0
1746 && self.is_clean()
1747 && self.epoch_override.is_none()
1748 && self.header.max_row_id >= self.header.min_row_id
1749 && self
1750 .header
1751 .max_row_id
1752 .checked_sub(self.header.min_row_id)
1753 .and_then(|span| span.checked_add(1))
1754 == Some(n as u64)
1755 }
1756
1757 pub(crate) fn position_for_row_id_fast(&self, row_id: u64) -> Option<usize> {
1758 if !self.clean_contiguous_row_ids()
1759 || row_id < self.header.min_row_id
1760 || row_id > self.header.max_row_id
1761 {
1762 return None;
1763 }
1764 Some((row_id - self.header.min_row_id) as usize)
1765 }
1766
1767 pub(crate) fn positions_for_row_ids_fast(&self, row_ids: &[u64]) -> Option<Vec<usize>> {
1768 if !self.clean_contiguous_row_ids() {
1769 return None;
1770 }
1771 let mut positions = Vec::with_capacity(row_ids.len());
1772 for &row_id in row_ids {
1773 if let Some(pos) = self.position_for_row_id_fast(row_id) {
1774 positions.push(pos);
1775 }
1776 }
1777 positions.sort_unstable();
1778 Some(positions)
1779 }
1780
1781 fn resolve_type(&self, column_id: u16) -> TypeId {
1782 match column_id {
1783 SYS_ROW_ID | SYS_EPOCH => TypeId::Int64,
1784 SYS_DELETED => TypeId::Bool,
1785 _ => self
1786 .schema
1787 .columns
1788 .iter()
1789 .find(|c| c.id == column_id)
1790 .map(|c| c.ty)
1791 .unwrap_or(TypeId::Bytes),
1792 }
1793 }
1794
1795 fn find_header(&self, column_id: u16) -> Result<&ColumnPageHeader> {
1796 self.dir
1797 .iter()
1798 .find(|h| h.column_id == column_id)
1799 .ok_or_else(|| MongrelError::ColumnNotFound(format!("column {column_id}")))
1800 }
1801
1802 fn col_encrypted(&self, column_id: u16) -> bool {
1808 self.dir
1809 .iter()
1810 .find(|h| h.column_id == column_id)
1811 .map(|h| h.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0)
1812 .unwrap_or(false)
1813 }
1814
1815 pub(crate) fn read_page(&mut self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
1816 let (offset, compressed_len, encrypted) = {
1817 let ch = self.find_header(column_id)?;
1818 let stat = ch
1819 .page_stats
1820 .get(page_seq)
1821 .ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
1822 (
1823 stat.offset,
1824 stat.compressed_len,
1825 ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0,
1826 )
1827 };
1828 let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
1831 if let Some(cache) = &self.page_cache {
1832 if let Some(bytes) = cache.lock().get(
1833 &key,
1834 crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
1835 ) {
1836 return decrypt_or_passthrough(
1837 self.cipher.as_deref(),
1838 self.nonce_prefix,
1839 column_id,
1840 page_seq,
1841 encrypted,
1842 &bytes,
1843 );
1844 }
1845 }
1846 let buf = match &self.mmap {
1847 Some(m) => m[offset as usize..(offset + compressed_len as u64) as usize].to_vec(),
1850 None => {
1851 self.file.seek(SeekFrom::Start(offset))?;
1852 let mut buf = vec![0u8; compressed_len as usize];
1853 self.file.read_exact(&mut buf)?;
1854 buf
1855 }
1856 };
1857 if let Some(cache) = &self.page_cache {
1859 cache.lock().insert(crate::page::CachedPage {
1860 committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
1861 content_hash: key,
1862 bytes: bytes::Bytes::copy_from_slice(&buf),
1863 });
1864 }
1865 decrypt_or_passthrough(
1866 self.cipher.as_deref(),
1867 self.nonce_prefix,
1868 column_id,
1869 page_seq,
1870 encrypted,
1871 &buf,
1872 )
1873 }
1874
1875 fn read_page_shared(&self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
1880 let ch = self.find_header(column_id)?;
1881 let stat = ch
1882 .page_stats
1883 .get(page_seq)
1884 .ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
1885 let encrypted = ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0;
1886 let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
1889 if let Some(cache) = &self.page_cache {
1890 if let Some(guard) = cache.try_lock() {
1891 if let Some(bytes) = guard.try_get(
1892 &key,
1893 crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
1894 ) {
1895 return decrypt_or_passthrough(
1896 self.cipher.as_deref(),
1897 self.nonce_prefix,
1898 column_id,
1899 page_seq,
1900 encrypted,
1901 &bytes,
1902 );
1903 }
1904 }
1905 }
1906 let mmap = self.mmap.as_ref().ok_or_else(|| {
1907 MongrelError::InvalidArgument("parallel page decode requires an mmap-backed run".into())
1908 })?;
1909 let start = stat.offset as usize;
1910 let end = (stat.offset + stat.compressed_len as u64) as usize;
1911 let buf = mmap[start..end].to_vec();
1912 if let Some(cache) = &self.page_cache {
1916 if let Some(mut guard) = cache.try_lock() {
1917 guard.insert(crate::page::CachedPage {
1918 committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
1919 content_hash: key,
1920 bytes: bytes::Bytes::copy_from_slice(&buf),
1921 });
1922 }
1923 }
1924 decrypt_or_passthrough(
1925 self.cipher.as_deref(),
1926 self.nonce_prefix,
1927 column_id,
1928 page_seq,
1929 encrypted,
1930 &buf,
1931 )
1932 }
1933
1934 fn column(&mut self, column_id: u16) -> Result<&[Value]> {
1936 if column_id == SYS_EPOCH {
1939 if let Some(ov) = self.epoch_override {
1940 if !self.col_cache.contains_key(&column_id) {
1941 let n = self.row_count();
1942 self.col_cache
1943 .insert(column_id, vec![Value::Int64(ov.0 as i64); n]);
1944 }
1945 return Ok(self.col_cache.get(&column_id).unwrap().as_slice());
1946 }
1947 }
1948 if !self.col_cache.contains_key(&column_id) {
1949 let ty = self.resolve_type(column_id);
1950 let page_rows: Vec<usize> = {
1951 let ch = self.find_header(column_id)?;
1952 ch.page_stats.iter().map(|s| s.row_count as usize).collect()
1953 };
1954 let mut decoded: Vec<Value> = Vec::with_capacity(self.row_count());
1955 for (seq, &pr) in page_rows.iter().enumerate() {
1956 let page = self.read_page(column_id, seq)?;
1957 decoded.extend(columnar::decode_page(ty, &page, pr)?);
1958 }
1959 self.col_cache.insert(column_id, decoded);
1960 }
1961 Ok(self.col_cache.get(&column_id).unwrap().as_slice())
1962 }
1963
1964 pub fn get_version(&mut self, row_id: RowId, snapshot: Epoch) -> Result<Option<(Epoch, Row)>> {
1978 match self.find_version_page(row_id, snapshot)? {
1979 None => Ok(None),
1980 Some((epoch, seq, local_index)) => Ok(Some((
1981 Epoch(epoch),
1982 self.materialize_in_page(seq, local_index)?,
1983 ))),
1984 }
1985 }
1986
1987 fn find_version_page(
1993 &mut self,
1994 row_id: RowId,
1995 snapshot: Epoch,
1996 ) -> Result<Option<(u64, usize, usize)>> {
1997 let n = self.row_count();
1998 if n == 0 {
1999 return Ok(None);
2000 }
2001 let target = row_id.0 as i64;
2002 let ch = self.find_header(SYS_ROW_ID)?;
2003 let mut page_start = 0u64;
2004 let candidate_pages: Vec<(usize, u64)> = ch .page_stats
2006 .iter()
2007 .enumerate()
2008 .filter_map(|(seq, s)| {
2009 let start = page_start;
2010 page_start += s.row_count as u64;
2011 (s.first_row_id <= row_id.0 && row_id.0 <= s.last_row_id).then_some((seq, start))
2012 })
2013 .collect();
2014 if candidate_pages.is_empty() {
2015 return Ok(None);
2016 }
2017 let ty = self.resolve_type(SYS_ROW_ID);
2018 let mut best: Option<(u64, usize, usize)> = None; for (seq, _page_row_start) in candidate_pages {
2020 let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
2021 let row_ids = match self.decode_page_native_cached(ty, SYS_ROW_ID, seq, page_rows)? {
2022 columnar::NativeColumn::Int64 { data, .. } => data,
2023 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
2024 };
2025 let local = match row_ids.binary_search(&target) {
2026 Ok(i) => i,
2027 Err(_) => continue,
2028 };
2029 let epoch_ty = self.resolve_type(SYS_EPOCH);
2030 let epochs =
2031 match self.decode_page_native_cached(epoch_ty, SYS_EPOCH, seq, page_rows)? {
2032 columnar::NativeColumn::Int64 { data, .. } => data,
2033 _ => return Err(MongrelError::InvalidArgument("sys epoch not int64".into())),
2034 };
2035 let mut lo = local;
2038 while lo > 0 && row_ids[lo - 1] == target {
2039 lo -= 1;
2040 }
2041 let mut hi = local;
2042 while hi + 1 < row_ids.len() && row_ids[hi + 1] == target {
2043 hi += 1;
2044 }
2045 for (i, &epoch) in epochs[lo..=hi].iter().enumerate() {
2046 let epoch = epoch as u64;
2047 if epoch <= snapshot.0 && best.map(|(be, ..)| epoch > be).unwrap_or(true) {
2048 best = Some((epoch, seq, lo + i));
2049 }
2050 }
2051 }
2052 Ok(best)
2053 }
2054
2055 pub fn get_version_column(
2063 &mut self,
2064 row_id: RowId,
2065 snapshot: Epoch,
2066 column_id: u16,
2067 ) -> Result<Option<(Epoch, bool, Option<Value>)>> {
2068 let Some((epoch, seq, local_index)) = self.find_version_page(row_id, snapshot)? else {
2069 return Ok(None);
2070 };
2071 let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
2072 let page_start: usize = self.find_header(SYS_ROW_ID)?.page_stats[..seq]
2073 .iter()
2074 .map(|s| s.row_count as usize)
2075 .sum();
2076 let global_index = page_start + local_index;
2077 let native_at = |slf: &mut Self, cid: u16| -> Result<Option<Value>> {
2078 if !slf.dir.iter().any(|h| h.column_id == cid) {
2079 return Ok(None);
2080 }
2081 let ty = slf.resolve_type(cid);
2082 if !matches!(
2083 ty,
2084 TypeId::Bool
2085 | TypeId::Int8
2086 | TypeId::Int16
2087 | TypeId::Int32
2088 | TypeId::Int64
2089 | TypeId::UInt8
2090 | TypeId::UInt16
2091 | TypeId::UInt32
2092 | TypeId::UInt64
2093 | TypeId::Float32
2094 | TypeId::Float64
2095 | TypeId::TimestampNanos
2096 | TypeId::Date32
2097 | TypeId::Bytes
2098 ) {
2099 return Ok(slf.column(cid)?.get(global_index).cloned());
2100 }
2101 Ok(slf
2102 .decode_page_native_cached(ty, cid, seq, page_rows)?
2103 .value_at(local_index))
2104 };
2105 let deleted = matches!(native_at(self, SYS_DELETED)?, Some(Value::Bool(true)));
2106 let value = native_at(self, column_id)?;
2107 Ok(Some((Epoch(epoch), deleted, value)))
2108 }
2109
2110 fn materialize_in_page(&mut self, seq: usize, local_index: usize) -> Result<Row> {
2119 let page_rows = self.find_header(SYS_ROW_ID)?.page_stats[seq].row_count as usize;
2120 let page_start: usize = self.find_header(SYS_ROW_ID)?.page_stats[..seq]
2121 .iter()
2122 .map(|s| s.row_count as usize)
2123 .sum();
2124 let global_index = page_start + local_index;
2125 let native_at = |slf: &mut Self, column_id: u16| -> Result<Option<Value>> {
2126 if !slf.dir.iter().any(|h| h.column_id == column_id) {
2129 return Ok(None);
2130 }
2131 let ty = slf.resolve_type(column_id);
2132 if !matches!(
2133 ty,
2134 TypeId::Bool
2135 | TypeId::Int8
2136 | TypeId::Int16
2137 | TypeId::Int32
2138 | TypeId::Int64
2139 | TypeId::UInt8
2140 | TypeId::UInt16
2141 | TypeId::UInt32
2142 | TypeId::UInt64
2143 | TypeId::Float32
2144 | TypeId::Float64
2145 | TypeId::TimestampNanos
2146 | TypeId::Date32
2147 | TypeId::Bytes
2148 ) {
2149 return Ok(slf.column(column_id)?.get(global_index).cloned());
2152 }
2153 Ok(slf
2154 .decode_page_native_cached(ty, column_id, seq, page_rows)?
2155 .value_at(local_index))
2156 };
2157 let row_id = RowId(match native_at(self, SYS_ROW_ID)? {
2158 Some(Value::Int64(x)) => x as u64,
2159 _ => 0,
2160 });
2161 let committed_epoch = Epoch(match native_at(self, SYS_EPOCH)? {
2162 Some(Value::Int64(x)) => x as u64,
2163 _ => 0,
2164 });
2165 let deleted = matches!(native_at(self, SYS_DELETED)?, Some(Value::Bool(true)));
2166 let col_ids: Vec<u16> = self.schema.columns.iter().map(|c| c.id).collect();
2167 let mut columns = HashMap::new();
2168 for id in col_ids {
2169 columns.insert(id, native_at(self, id)?.unwrap_or(Value::Null));
2170 }
2171 Ok(Row {
2172 row_id,
2173 committed_epoch,
2174 columns,
2175 deleted,
2176 })
2177 }
2178
2179 pub fn all_rows(&mut self) -> Result<Vec<Row>> {
2182 let n = self.row_count();
2183 let mut out = Vec::with_capacity(n);
2184 for i in 0..n {
2185 out.push(self.materialize(i)?);
2186 }
2187 Ok(out)
2188 }
2189
2190 pub fn visible_indices(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
2196 let n = self.row_count();
2197 if n == 0 {
2198 return Ok(Vec::new());
2199 }
2200 let row_ids = self.column(SYS_ROW_ID)?.to_vec();
2201 let epochs = self.column(SYS_EPOCH)?.to_vec();
2202 let deleted = self.column(SYS_DELETED)?.to_vec();
2203 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
2204 for i in 0..n {
2205 let rid = int_at(&row_ids, i);
2206 let e = int_at(&epochs, i);
2207 if e > snapshot.0 {
2208 continue;
2209 }
2210 best.entry(rid)
2211 .and_modify(|(be, bi)| {
2212 if e > *be {
2213 *be = e;
2214 *bi = i;
2215 }
2216 })
2217 .or_insert((e, i));
2218 }
2219 let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
2220 idxs.retain(|&i| !bool_at(&deleted, i));
2221 idxs.sort_unstable();
2222 Ok(idxs)
2223 }
2224
2225 pub fn gather_column(&mut self, column_id: u16, indices: &[usize]) -> Result<Vec<Value>> {
2230 if !self.dir.iter().any(|h| h.column_id == column_id) {
2231 return Ok(vec![Value::Null; indices.len()]);
2232 }
2233 let col = self.column(column_id)?;
2234 Ok(indices
2235 .iter()
2236 .map(|&i| col.get(i).cloned().unwrap_or(Value::Null))
2237 .collect())
2238 }
2239
2240 pub fn column_native(&mut self, column_id: u16) -> Result<columnar::NativeColumn> {
2245 use rayon::prelude::*;
2246 if column_id == SYS_EPOCH {
2247 if let Some(ov) = self.epoch_override {
2248 return Ok(columnar::NativeColumn::int64_constant(
2249 ov.0 as i64,
2250 self.row_count(),
2251 ));
2252 }
2253 }
2254 let ty = self.resolve_type(column_id);
2255 let n = self.row_count();
2256 let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
2257 return Ok(columnar::null_native(ty, n));
2258 };
2259 let page_count = ch.page_count as usize;
2260 let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
2261 if page_count == 0 {
2262 return Ok(columnar::null_native(ty, n));
2263 }
2264 let parts: Vec<columnar::NativeColumn> = if self.mmap.is_some() && page_count > 1 {
2265 let reader: &RunReader = self;
2269 (0..page_count)
2270 .into_par_iter()
2271 .map(|seq| {
2272 let raw = reader.read_page_shared(column_id, seq)?;
2273 columnar::decode_page_native(ty, &raw, page_rows[seq])
2274 })
2275 .collect::<Result<Vec<_>>>()?
2276 } else {
2277 let mut out = Vec::with_capacity(page_count);
2278 for (seq, &pr) in page_rows.iter().enumerate() {
2279 let page = self.read_page(column_id, seq)?;
2280 out.push(columnar::decode_page_native(ty, &page, pr)?);
2281 }
2282 out
2283 };
2284 Ok(columnar::NativeColumn::concat(&parts))
2285 }
2286
2287 pub fn has_mmap(&self) -> bool {
2291 self.mmap.is_some()
2292 }
2293
2294 pub fn column_native_shared(&self, column_id: u16) -> Result<columnar::NativeColumn> {
2301 use rayon::prelude::*;
2302 if column_id == SYS_EPOCH {
2303 if let Some(ov) = self.epoch_override {
2304 return Ok(columnar::NativeColumn::int64_constant(
2305 ov.0 as i64,
2306 self.row_count(),
2307 ));
2308 }
2309 }
2310 let ty = self.resolve_type(column_id);
2311 let n = self.row_count();
2312 let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
2313 return Ok(columnar::null_native(ty, n));
2314 };
2315 let page_count = ch.page_count as usize;
2316 let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
2317 if page_count == 0 {
2318 return Ok(columnar::null_native(ty, n));
2319 }
2320 #[cfg(unix)]
2324 {
2325 if let (Some(m), Some(first)) = (&self.mmap, ch.page_stats.first()) {
2326 let start = first.offset as usize;
2327 let end = (ch.page_region_offset as usize) + (ch.page_region_len as usize);
2328 if end > start {
2329 let _ = m.advise_range(memmap2::Advice::WillNeed, start, end - start);
2330 }
2331 }
2332 }
2333 let run_id = self.header.run_id;
2334 let mut parts_keys: Vec<(columnar::NativeColumn, Option<[u8; 32]>)> = if page_count > 1 {
2338 (0..page_count)
2339 .into_par_iter()
2340 .map(|seq| self.decode_page_cached(ty, column_id, seq, page_rows[seq], run_id))
2341 .collect::<Result<Vec<_>>>()?
2342 } else {
2343 vec![self.decode_page_cached(ty, column_id, 0, page_rows[0], run_id)?]
2344 };
2345 if let Some(cache) = &self.decoded_cache {
2348 let mut g = cache.lock();
2349 for (col, key) in parts_keys.iter_mut() {
2350 if let Some(k) = key.take() {
2351 g.insert(k, Arc::new(col.clone()));
2352 }
2353 }
2354 }
2355 let parts: Vec<columnar::NativeColumn> = parts_keys.into_iter().map(|(c, _)| c).collect();
2356 Ok(columnar::NativeColumn::concat(&parts))
2357 }
2358
2359 fn decode_page_cached(
2364 &self,
2365 ty: TypeId,
2366 column_id: u16,
2367 seq: usize,
2368 nrows: usize,
2369 run_id: u128,
2370 ) -> Result<(columnar::NativeColumn, Option<[u8; 32]>)> {
2371 let key = page_cache_key(self.table_id, run_id, column_id, seq);
2372 if let Some(cache) = &self.decoded_cache {
2373 if let Some(g) = cache.try_lock() {
2374 if let Some(hit) = g.try_get(&key) {
2375 return Ok(((*hit).clone(), None));
2376 }
2377 }
2378 }
2379 let raw = self.read_page_shared(column_id, seq)?;
2380 let col = columnar::decode_page_native(ty, &raw, nrows)?;
2381 Ok((col, Some(key)))
2382 }
2383
2384 fn decode_page_native_cached(
2395 &mut self,
2396 ty: TypeId,
2397 column_id: u16,
2398 seq: usize,
2399 nrows: usize,
2400 ) -> Result<columnar::NativeColumn> {
2401 let key = page_cache_key(self.table_id, self.header.run_id, column_id, seq);
2402 if let Some(cache) = &self.decoded_cache {
2403 if let Some(hit) = cache.lock().try_get(&key) {
2404 return Ok((*hit).clone());
2405 }
2406 }
2407 let raw = self.read_page(column_id, seq)?;
2408 let col = columnar::decode_page_native(ty, &raw, nrows)?;
2409 if let Some(cache) = &self.decoded_cache {
2410 cache.lock().insert(key, std::sync::Arc::new(col.clone()));
2411 }
2412 Ok(col)
2413 }
2414
2415 pub fn range_row_ids_i64(
2420 &mut self,
2421 column_id: u16,
2422 lo: i64,
2423 hi: i64,
2424 ) -> Result<std::collections::HashSet<u64>> {
2425 Ok(self
2426 .range_row_id_set_i64(column_id, lo, hi)?
2427 .into_sorted_vec()
2428 .into_iter()
2429 .collect())
2430 }
2431
2432 pub(crate) fn range_row_id_set_i64(
2433 &mut self,
2434 column_id: u16,
2435 lo: i64,
2436 hi: i64,
2437 ) -> Result<RowIdSet> {
2438 let info: Vec<(Option<i64>, Option<i64>, usize)> =
2439 match self.dir.iter().find(|h| h.column_id == column_id) {
2440 Some(ch) => ch
2441 .page_stats
2442 .iter()
2443 .map(|s| {
2444 (
2445 be_i64(s.min.as_deref()),
2446 be_i64(s.max.as_deref()),
2447 s.row_count as usize,
2448 )
2449 })
2450 .collect(),
2451 None => return Ok(RowIdSet::empty()),
2452 };
2453 let stats_pruneable =
2459 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2460 let clean_contiguous = self.clean_contiguous_row_ids();
2461 let mut out = Vec::new();
2462 let mut page_start = 0usize;
2463 for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
2464 let current_page_start = page_start;
2465 page_start += nrows;
2466 let skip = stats_pruneable
2468 && match (mn, mx) {
2469 (Some(mn), Some(mx)) => mx < lo || mn > hi,
2470 _ => true,
2471 };
2472 if skip {
2473 continue;
2474 }
2475 let val_page = self.read_page(column_id, seq)?;
2476 let vals =
2477 columnar::decode_page_native(self.resolve_type(column_id), &val_page, nrows)?;
2478 if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
2479 if clean_contiguous {
2480 for (i, val) in v.iter().enumerate() {
2481 if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
2482 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2483 }
2484 }
2485 } else {
2486 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2487 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2488 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2489 for (i, val) in v.iter().enumerate() {
2490 if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
2491 out.push(r[i] as u64);
2492 }
2493 }
2494 }
2495 }
2496 }
2497 }
2498 Ok(RowIdSet::from_unsorted(out))
2499 }
2500
2501 pub fn range_row_ids_f64(
2504 &mut self,
2505 column_id: u16,
2506 lo: f64,
2507 lo_inclusive: bool,
2508 hi: f64,
2509 hi_inclusive: bool,
2510 ) -> Result<std::collections::HashSet<u64>> {
2511 Ok(self
2512 .range_row_id_set_f64(column_id, lo, lo_inclusive, hi, hi_inclusive)?
2513 .into_sorted_vec()
2514 .into_iter()
2515 .collect())
2516 }
2517
2518 pub(crate) fn range_row_id_set_f64(
2519 &mut self,
2520 column_id: u16,
2521 lo: f64,
2522 lo_inclusive: bool,
2523 hi: f64,
2524 hi_inclusive: bool,
2525 ) -> Result<RowIdSet> {
2526 let info: Vec<(Option<f64>, Option<f64>, usize)> =
2527 match self.dir.iter().find(|h| h.column_id == column_id) {
2528 Some(ch) => ch
2529 .page_stats
2530 .iter()
2531 .map(|s| {
2532 (
2533 be_f64(s.min.as_deref()),
2534 be_f64(s.max.as_deref()),
2535 s.row_count as usize,
2536 )
2537 })
2538 .collect(),
2539 None => return Ok(RowIdSet::empty()),
2540 };
2541 let stats_pruneable =
2547 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2548 let clean_contiguous = self.clean_contiguous_row_ids();
2549 let mut out = Vec::new();
2550 let mut page_start = 0usize;
2551 for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
2552 let current_page_start = page_start;
2553 page_start += nrows;
2554 let skip = stats_pruneable
2557 && match (mn, mx) {
2558 (Some(mn), Some(mx)) => {
2559 let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
2560 let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
2561 skip_lo || skip_hi
2562 }
2563 _ => true,
2564 };
2565 if skip {
2566 continue;
2567 }
2568 let val_page = self.read_page(column_id, seq)?;
2569 let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
2570 if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
2571 if clean_contiguous {
2572 for (i, val) in v.iter().enumerate() {
2573 if !columnar::validity_bit(&validity, i) || val.is_nan() {
2574 continue;
2575 }
2576 let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
2577 let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
2578 if ok_lo && ok_hi {
2579 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2580 }
2581 }
2582 } else {
2583 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2584 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2585 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2586 for (i, val) in v.iter().enumerate() {
2587 if !columnar::validity_bit(&validity, i) || val.is_nan() {
2588 continue;
2589 }
2590 let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
2591 let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
2592 if ok_lo && ok_hi {
2593 out.push(r[i] as u64);
2594 }
2595 }
2596 }
2597 }
2598 }
2599 }
2600 Ok(RowIdSet::from_unsorted(out))
2601 }
2602
2603 pub(crate) fn null_row_id_set(&mut self, column_id: u16, want_nulls: bool) -> Result<RowIdSet> {
2608 let stats: Vec<(usize, usize)> = match self.dir.iter().find(|h| h.column_id == column_id) {
2609 Some(ch) => ch
2610 .page_stats
2611 .iter()
2612 .map(|s| (s.null_count as usize, s.row_count as usize))
2613 .collect(),
2614 None => return Ok(RowIdSet::empty()),
2615 };
2616 let ty = self.resolve_type(column_id);
2617 let clean_contiguous = self.clean_contiguous_row_ids();
2618 let mut out = Vec::new();
2619 let mut page_start = 0usize;
2620 for (seq, (null_count, nrows)) in stats.into_iter().enumerate() {
2621 let current_page_start = page_start;
2622 page_start += nrows;
2623 if want_nulls && null_count == 0 {
2625 continue;
2626 }
2627 if !want_nulls && null_count == nrows {
2628 continue;
2629 }
2630 let val_page = self.read_page(column_id, seq)?;
2631 let col = columnar::decode_page_native(ty, &val_page, nrows)?;
2632 let validity = col.validity();
2633 if clean_contiguous {
2634 for i in 0..nrows {
2635 let is_null = !columnar::validity_bit(validity, i);
2636 if is_null == want_nulls {
2637 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2638 }
2639 }
2640 } else {
2641 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2642 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2643 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2644 for (i, &rid) in r.iter().enumerate().take(nrows) {
2645 let is_null = !columnar::validity_bit(validity, i);
2646 if is_null == want_nulls {
2647 out.push(rid as u64);
2648 }
2649 }
2650 }
2651 }
2652 }
2653 Ok(RowIdSet::from_unsorted(out))
2654 }
2655
2656 pub fn visible_indices_native(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
2657 let n = self.row_count();
2658 if n == 0 {
2659 return Ok(Vec::new());
2660 }
2661 let (row_ids, epochs, deleted) = self.system_columns_native()?;
2662 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
2663 for i in 0..n {
2664 let rid = row_ids[i] as u64;
2665 let e = epochs[i] as u64;
2666 if e > snapshot.0 {
2667 continue;
2668 }
2669 best.entry(rid)
2670 .and_modify(|(be, bi)| {
2671 if e > *be {
2672 *be = e;
2673 *bi = i;
2674 }
2675 })
2676 .or_insert((e, i));
2677 }
2678 let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
2679 idxs.retain(|&i| deleted[i] == 0);
2680 idxs.sort_unstable();
2681 Ok(idxs)
2682 }
2683
2684 pub fn range_row_ids_visible_i64(
2693 &mut self,
2694 column_id: u16,
2695 lo: i64,
2696 hi: i64,
2697 snapshot: Epoch,
2698 ) -> Result<Vec<u64>> {
2699 let stats: Vec<(Option<i64>, Option<i64>, usize)> = match self.column_page_stats(column_id)
2700 {
2701 Some(s) => s
2702 .iter()
2703 .map(|st| {
2704 (
2705 be_i64(st.min.as_deref()),
2706 be_i64(st.max.as_deref()),
2707 st.row_count as usize,
2708 )
2709 })
2710 .collect(),
2711 None => return Ok(Vec::new()),
2712 };
2713 let stats_pruneable =
2719 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2720 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2721 let mut out: Vec<u64> = Vec::new();
2722 let mut vis = 0usize;
2723 let mut page_start = 0usize;
2724 for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
2725 let page_end = page_start + nrows;
2726 let skip = stats_pruneable
2728 && match (mn, mx) {
2729 (Some(mn), Some(mx)) => mx < lo || mn > hi,
2730 _ => true, };
2732 if !skip {
2733 let val_page = self.read_page(column_id, seq)?;
2734 let vals = columnar::decode_page_native(TypeId::Int64, &val_page, nrows)?;
2735 if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
2736 while vis < positions.len() && positions[vis] < page_end {
2737 let local = positions[vis] - page_start;
2738 if columnar::validity_bit(&validity, local)
2739 && v[local] >= lo
2740 && v[local] <= hi
2741 {
2742 out.push(rids[vis] as u64);
2743 }
2744 vis += 1;
2745 }
2746 }
2747 } else {
2748 while vis < positions.len() && positions[vis] < page_end {
2749 vis += 1;
2750 }
2751 }
2752 page_start = page_end;
2753 }
2754 Ok(out)
2755 }
2756
2757 pub fn range_row_ids_visible_f64(
2760 &mut self,
2761 column_id: u16,
2762 lo: f64,
2763 lo_inclusive: bool,
2764 hi: f64,
2765 hi_inclusive: bool,
2766 snapshot: Epoch,
2767 ) -> Result<Vec<u64>> {
2768 let stats: Vec<(Option<f64>, Option<f64>, usize)> = match self.column_page_stats(column_id)
2769 {
2770 Some(s) => s
2771 .iter()
2772 .map(|st| {
2773 (
2774 be_f64(st.min.as_deref()),
2775 be_f64(st.max.as_deref()),
2776 st.row_count as usize,
2777 )
2778 })
2779 .collect(),
2780 None => return Ok(Vec::new()),
2781 };
2782 let stats_pruneable =
2788 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2789 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2790 let mut out: Vec<u64> = Vec::new();
2791 let mut vis = 0usize;
2792 let mut page_start = 0usize;
2793 for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
2794 let page_end = page_start + nrows;
2795 let skip = stats_pruneable
2796 && match (mn, mx) {
2797 (Some(mn), Some(mx)) => {
2798 let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
2799 let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
2800 skip_lo || skip_hi
2801 }
2802 _ => true,
2803 };
2804 if !skip {
2805 let val_page = self.read_page(column_id, seq)?;
2806 let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
2807 if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
2808 while vis < positions.len() && positions[vis] < page_end {
2809 let local = positions[vis] - page_start;
2810 if columnar::validity_bit(&validity, local) && !v[local].is_nan() {
2811 let val = v[local];
2812 let ok_lo = if lo_inclusive { val >= lo } else { val > lo };
2813 let ok_hi = if hi_inclusive { val <= hi } else { val < hi };
2814 if ok_lo && ok_hi {
2815 out.push(rids[vis] as u64);
2816 }
2817 }
2818 vis += 1;
2819 }
2820 }
2821 } else {
2822 while vis < positions.len() && positions[vis] < page_end {
2823 vis += 1;
2824 }
2825 }
2826 page_start = page_end;
2827 }
2828 Ok(out)
2829 }
2830
2831 pub fn null_row_ids_visible(
2837 &mut self,
2838 column_id: u16,
2839 want_nulls: bool,
2840 snapshot: Epoch,
2841 ) -> Result<Vec<u64>> {
2842 let stats: Vec<(usize, usize)> = match self.column_page_stats(column_id) {
2843 Some(s) => s
2844 .iter()
2845 .map(|st| (st.null_count as usize, st.row_count as usize))
2846 .collect(),
2847 None => return Ok(Vec::new()),
2848 };
2849 let ty = self.resolve_type(column_id);
2850 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2851 let mut out: Vec<u64> = Vec::new();
2852 let mut vis = 0usize;
2853 let mut page_start = 0usize;
2854 for (seq, &(null_count, nrows)) in stats.iter().enumerate() {
2855 let page_end = page_start + nrows;
2856 let skip = (want_nulls && null_count == 0) || (!want_nulls && null_count == nrows);
2857 if !skip {
2858 let val_page = self.read_page(column_id, seq)?;
2859 let col = columnar::decode_page_native(ty, &val_page, nrows)?;
2860 let validity = col.validity();
2861 while vis < positions.len() && positions[vis] < page_end {
2862 let local = positions[vis] - page_start;
2863 let is_null = !columnar::validity_bit(validity, local);
2864 if is_null == want_nulls {
2865 out.push(rids[vis] as u64);
2866 }
2867 vis += 1;
2868 }
2869 } else {
2870 while vis < positions.len() && positions[vis] < page_end {
2871 vis += 1;
2872 }
2873 }
2874 page_start = page_end;
2875 }
2876 Ok(out)
2877 }
2878
2879 pub fn visible_positions_with_rids(
2883 &mut self,
2884 snapshot: Epoch,
2885 ) -> Result<(Vec<usize>, Vec<i64>)> {
2886 let n = self.row_count();
2887 if n == 0 {
2888 return Ok((Vec::new(), Vec::new()));
2889 }
2890 if self.is_clean() && self.epoch_override.is_none() {
2899 let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
2900 columnar::NativeColumn::Int64 { data, .. } => data,
2901 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
2902 };
2903 let positions: Vec<usize> = (0..n).collect();
2904 return Ok((positions, row_ids));
2905 }
2906 let (row_ids, epochs, deleted) = self.system_columns_native()?;
2907 let mut idxs: Vec<usize> = Vec::new();
2918 let mut i = 0;
2919 while i < n {
2920 let rid = row_ids[i] as u64;
2921 let mut best: Option<usize> = None;
2924 let mut j = i;
2925 while j < n && row_ids[j] as u64 == rid {
2926 if epochs[j] as u64 <= snapshot.0 {
2927 best = Some(j);
2928 }
2929 j += 1;
2930 }
2931 if let Some(b) = best {
2932 if deleted[b] == 0 {
2933 idxs.push(b);
2934 }
2935 }
2936 i = j;
2937 }
2938 let rids: Vec<i64> = idxs.iter().map(|&k| row_ids[k]).collect();
2940 Ok((idxs, rids))
2941 }
2942
2943 pub fn page_row_counts(&self, column_id: u16) -> Result<Vec<usize>> {
2947 Ok(self
2948 .find_header(column_id)?
2949 .page_stats
2950 .iter()
2951 .map(|s| s.row_count as usize)
2952 .collect())
2953 }
2954
2955 pub fn has_column(&self, column_id: u16) -> bool {
2958 self.dir.iter().any(|h| h.column_id == column_id)
2959 }
2960
2961 pub fn column_page_stats(&self, column_id: u16) -> Option<&[crate::page::PageStat]> {
2965 self.dir
2966 .iter()
2967 .find(|h| h.column_id == column_id)
2968 .map(|ch| ch.page_stats.as_slice())
2969 }
2970
2971 pub(crate) fn system_columns_native(&mut self) -> Result<(Vec<i64>, Vec<i64>, Vec<u8>)> {
2976 let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
2977 columnar::NativeColumn::Int64 { data, .. } => data,
2978 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
2979 };
2980 let epochs = match self.column_native_shared(SYS_EPOCH)? {
2981 columnar::NativeColumn::Int64 { data, .. } => data,
2982 _ => return Err(MongrelError::InvalidArgument("sys epoch not int64".into())),
2983 };
2984 let deleted = match self.column_native_shared(SYS_DELETED)? {
2985 columnar::NativeColumn::Bool { data, .. } => data,
2986 _ => return Err(MongrelError::InvalidArgument("sys deleted not bool".into())),
2987 };
2988 Ok((row_ids, epochs, deleted))
2989 }
2990
2991 pub fn visible_versions(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
2995 let n = self.row_count();
2996 if n == 0 {
2997 return Ok(Vec::new());
2998 }
2999 let row_ids = self.column(SYS_ROW_ID)?.to_vec();
3000 let epochs = self.column(SYS_EPOCH)?.to_vec();
3001 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
3002 for i in 0..n {
3003 let rid = int_at(&row_ids, i);
3004 let epoch = int_at(&epochs, i);
3005 if epoch > snapshot.0 {
3006 continue;
3007 }
3008 best.entry(rid)
3009 .and_modify(|e| {
3010 if epoch > e.0 {
3011 *e = (epoch, i);
3012 }
3013 })
3014 .or_insert((epoch, i));
3015 }
3016 let mut picks: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
3017 picks.sort();
3018 let mut out = Vec::with_capacity(picks.len());
3019 for i in picks {
3020 out.push(self.materialize(i)?);
3021 }
3022 Ok(out)
3023 }
3024
3025 pub fn visible_rows(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
3027 Ok(self
3028 .visible_versions(snapshot)?
3029 .into_iter()
3030 .filter(|r| !r.deleted)
3031 .collect())
3032 }
3033
3034 pub(crate) fn materialize(&mut self, index: usize) -> Result<Row> {
3035 let row_id = RowId(int_at(self.column(SYS_ROW_ID)?, index));
3036 let epoch = Epoch(int_at(self.column(SYS_EPOCH)?, index));
3037 let deleted = bool_at(self.column(SYS_DELETED)?, index);
3038 let col_ids: Vec<u16> = self.schema.columns.iter().map(|c| c.id).collect();
3039 let mut columns = HashMap::new();
3040 for id in col_ids {
3041 let val = if self.dir.iter().any(|h| h.column_id == id) {
3042 self.column(id)?.get(index).cloned().unwrap_or(Value::Null)
3043 } else {
3044 Value::Null
3046 };
3047 columns.insert(id, val);
3048 }
3049 Ok(Row {
3050 row_id,
3051 committed_epoch: epoch,
3052 columns,
3053 deleted,
3054 })
3055 }
3056
3057 pub(crate) fn materialize_batch(&mut self, indices: &[usize]) -> Result<Vec<Row>> {
3066 if indices.is_empty() {
3067 return Ok(Vec::new());
3068 }
3069 use std::collections::HashMap;
3070 let rid_col = self.column_native_shared(SYS_ROW_ID)?;
3071 let epoch_col = self.column_native_shared(SYS_EPOCH)?;
3072 let del_col = self.column_native_shared(SYS_DELETED)?;
3073 let mut user: HashMap<u16, columnar::NativeColumn> = HashMap::new();
3074 let present: Vec<u16> = self
3075 .schema
3076 .columns
3077 .iter()
3078 .map(|c| c.id)
3079 .filter(|id| self.dir.iter().any(|h| h.column_id == *id))
3080 .collect();
3081 for id in present {
3082 user.insert(id, self.column_native(id)?);
3083 }
3084 let i64_at = |col: &columnar::NativeColumn, i: usize| -> i64 {
3085 match col {
3086 columnar::NativeColumn::Int64 { data, .. } => data.get(i).copied().unwrap_or(0),
3087 _ => 0,
3088 }
3089 };
3090 let bool_at_native = |col: &columnar::NativeColumn, i: usize| -> bool {
3091 match col {
3092 columnar::NativeColumn::Bool { data, .. } => {
3093 data.get(i).copied().map(|b| b != 0).unwrap_or(false)
3094 }
3095 _ => false,
3096 }
3097 };
3098 let mut rows = Vec::with_capacity(indices.len());
3099 for &idx in indices {
3100 let row_id = RowId(i64_at(&rid_col, idx) as u64);
3101 let epoch = Epoch(i64_at(&epoch_col, idx) as u64);
3102 let deleted = bool_at_native(&del_col, idx);
3103 let mut columns = HashMap::with_capacity(self.schema.columns.len());
3104 for cdef in self.schema.columns.iter() {
3105 let val = match user.get(&cdef.id) {
3106 Some(col) => col.value_at(idx).unwrap_or(Value::Null),
3107 None => Value::Null,
3108 };
3109 columns.insert(cdef.id, val);
3110 }
3111 rows.push(Row {
3112 row_id,
3113 committed_epoch: epoch,
3114 columns,
3115 deleted,
3116 });
3117 }
3118 Ok(rows)
3119 }
3120}
3121
3122fn int_at(vals: &[Value], i: usize) -> u64 {
3123 match vals.get(i) {
3124 Some(Value::Int64(x)) => *x as u64,
3125 _ => 0,
3126 }
3127}
3128
3129fn bool_at(vals: &[Value], i: usize) -> bool {
3130 matches!(vals.get(i), Some(Value::Bool(true)))
3131}
3132
3133#[cfg(test)]
3134mod tests {
3135 use super::*;
3136 use crate::columnar::NativeColumn;
3137 use crate::memtable::Value;
3138 use crate::rowid::RowId;
3139 use crate::schema::{ColumnDef, ColumnFlags};
3140 use tempfile::tempdir;
3141
3142 fn schema() -> Schema {
3143 Schema {
3144 schema_id: 1,
3145 columns: vec![
3146 ColumnDef {
3147 id: 1,
3148 name: "id".into(),
3149 ty: TypeId::Int64,
3150 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
3151 },
3152 ColumnDef {
3153 id: 2,
3154 name: "name".into(),
3155 ty: TypeId::Bytes,
3156 flags: ColumnFlags::empty(),
3157 },
3158 ],
3159 indexes: Vec::new(),
3160 colocation: vec![],
3161 constraints: Default::default(),
3162 }
3163 }
3164
3165 fn rows() -> Vec<Row> {
3166 vec![
3167 Row::new(RowId(1), Epoch(10))
3168 .with_column(1, Value::Int64(1))
3169 .with_column(2, Value::Bytes(b"alice".to_vec())),
3170 Row::new(RowId(2), Epoch(10))
3171 .with_column(1, Value::Int64(2))
3172 .with_column(2, Value::Bytes(b"bob".to_vec())),
3173 Row::new(RowId(3), Epoch(10))
3174 .with_column(1, Value::Int64(3))
3175 .with_column(2, Value::Bytes(b"carol".to_vec())),
3176 ]
3177 }
3178
3179 #[test]
3180 fn flush_then_read_mvcc() {
3181 let dir = tempdir().unwrap();
3182 let path = dir.path().join("r-1.sr");
3183 let header = RunWriter::new(&schema(), 1, Epoch(10), 0)
3184 .write(&path, &rows())
3185 .unwrap();
3186 assert_eq!(header.row_count, 3);
3187
3188 let mut r = RunReader::open(&path, schema(), None).unwrap();
3189 let (e, row) = r.get_version(RowId(2), Epoch(20)).unwrap().unwrap();
3191 assert_eq!(e, Epoch(10));
3192 assert_eq!(row.row_id, RowId(2));
3193 assert!(matches!(row.columns.get(&2), Some(Value::Bytes(_))));
3194 assert!(r.get_version(RowId(99), Epoch(20)).unwrap().is_none());
3196 let all = r.visible_rows(Epoch(20)).unwrap();
3198 assert_eq!(all.len(), 3);
3199 }
3200
3201 #[test]
3202 fn learned_index_was_stored() {
3203 let dir = tempdir().unwrap();
3204 let path = dir.path().join("r-2.sr");
3205 RunWriter::new(&schema(), 2, Epoch(1), 0)
3206 .write(&path, &rows())
3207 .unwrap();
3208 let header = read_header(&path).unwrap();
3209 assert!(
3210 header.index_trailer_offset != 0,
3211 "trailer should be present"
3212 );
3213 }
3214
3215 #[test]
3216 fn low_level_container_still_round_trips() {
3217 let dir = tempdir().unwrap();
3218 let path = dir.path().join("r-3.sr");
3219 let cols = vec![ColumnPayload {
3220 column_id: 1,
3221 type_id_tag: 8,
3222 encoding: Encoding::Plain,
3223 pages: vec![vec![1, 2, 3, 4]],
3224 page_stats: Vec::new(),
3225 }];
3226 let header = write_run(
3227 &path,
3228 &RunSpec {
3229 run_id: 1,
3230 schema_id: 1,
3231 epoch_created: 1,
3232 level: 0,
3233 flags: 0,
3234 sort_key_column_id: SORT_KEY_ROW_ID,
3235 row_count: 1,
3236 min_row_id: 0,
3237 max_row_id: 0,
3238 columns: &cols,
3239 },
3240 )
3241 .unwrap();
3242 let back = read_header(&path).unwrap();
3243 assert_eq!(back.content_hash, header.content_hash);
3244 assert_eq!(read_column_dir(&path, &back).unwrap().len(), 1);
3245 }
3246
3247 #[test]
3248 fn detects_corruption() {
3249 let dir = tempdir().unwrap();
3250 let path = dir.path().join("r-4.sr");
3251 write_run(
3252 &path,
3253 &RunSpec {
3254 run_id: 1,
3255 schema_id: 1,
3256 epoch_created: 1,
3257 level: 0,
3258 flags: 0,
3259 sort_key_column_id: SORT_KEY_ROW_ID,
3260 row_count: 1,
3261 min_row_id: 0,
3262 max_row_id: 0,
3263 columns: &[ColumnPayload {
3264 column_id: 1,
3265 type_id_tag: 8,
3266 encoding: Encoding::Plain,
3267 pages: vec![vec![1, 2, 3]],
3268 page_stats: Vec::new(),
3269 }],
3270 },
3271 )
3272 .unwrap();
3273 let mut bytes = std::fs::read(&path).unwrap();
3274 bytes[300] ^= 0xFF;
3275 std::fs::write(&path, bytes).unwrap();
3276 let err = read_header(&path).unwrap_err();
3277 assert!(
3278 matches!(err, MongrelError::ChecksumMismatch { .. }),
3279 "got {err:?}"
3280 );
3281 }
3282
3283 #[test]
3289 fn mmap_and_vec_writers_are_byte_identical() {
3290 let dir = tempdir().unwrap();
3291 let stats = vec![PageStat {
3292 first_row_id: 0,
3293 last_row_id: 1,
3294 null_count: 0,
3295 row_count: 2,
3296 min: Some(vec![0]),
3297 max: Some(vec![9]),
3298 offset: 0,
3299 compressed_len: 0,
3300 uncompressed_len: 0,
3301 }];
3302 let spec = RunSpec {
3303 run_id: 7,
3304 schema_id: 1,
3305 epoch_created: 10,
3306 level: 0,
3307 flags: 0,
3308 sort_key_column_id: SORT_KEY_ROW_ID,
3309 row_count: 2,
3310 min_row_id: 0,
3311 max_row_id: 1,
3312 columns: &[
3313 ColumnPayload {
3314 column_id: 1,
3315 type_id_tag: 8,
3316 encoding: Encoding::Plain,
3317 pages: vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8]],
3318 page_stats: stats.clone(),
3319 },
3320 ColumnPayload {
3321 column_id: 2,
3322 type_id_tag: 8,
3323 encoding: Encoding::Plain,
3324 pages: vec![vec![10, 20], vec![30, 40]],
3325 page_stats: stats.clone(),
3326 },
3327 ],
3328 };
3329 let trailer = b"learned-trailer-bytes";
3330
3331 let plan = plan_run(&spec, None, Some(trailer)).expect("plan");
3333 let mut buf = vec![0u8; plan.total];
3334 let h_place = place_run(&spec, None, Some(trailer), &plan, &mut buf)
3335 .expect("place_run into Vec succeeds");
3336
3337 let path_vec = dir.path().join("vec.sr");
3339 let h_vec = write_run_vec(&path_vec, &spec, None, Some(trailer)).unwrap();
3340 let bv = std::fs::read(&path_vec).unwrap();
3341
3342 assert_eq!(h_place.content_hash, h_vec.content_hash);
3343 assert_eq!(h_place.footer_offset, h_vec.footer_offset);
3344 assert_eq!(
3345 buf, bv,
3346 "place_run and write_run_vec must be byte-identical"
3347 );
3348 assert!(read_header(&path_vec).is_ok());
3349
3350 let path_mmap = dir.path().join("mmap.sr");
3354 match write_run_mmap(&path_mmap, &spec, None, Some(trailer)) {
3355 Ok(h_mmap) => {
3356 let bm = std::fs::read(&path_mmap).unwrap();
3357 assert_eq!(bm, bv, "mmap run must be byte-identical to vec run");
3358 assert_eq!(h_mmap.content_hash, h_vec.content_hash);
3359 }
3360 Err(e) if is_mmap_unavailable(&e) => {
3361 eprintln!("note: file mmap unavailable here; vec path covered (1)/(2)");
3362 }
3363 Err(e) => panic!("unexpected mmap error: {e:?}"),
3364 }
3365 }
3366
3367 #[test]
3372 fn column_native_shared_matches_column_native() {
3373 let dir = tempdir().unwrap();
3374 let path = dir.path().join("r.sr");
3375 let n = 65_536 * 2 + 9;
3378 let id_col = NativeColumn::int64_sequence(1, n);
3379 let mut offsets = vec![0u32];
3380 let mut values = Vec::new();
3381 for i in 0..n {
3382 values.extend_from_slice(format!("v{}", i % 17).as_bytes());
3383 offsets.push(values.len() as u32);
3384 }
3385 let name_col = NativeColumn::Bytes {
3386 offsets,
3387 values,
3388 validity: vec![0xFF; n.div_ceil(8)],
3389 };
3390 let header = RunWriter::new(&schema(), 9, Epoch(5), 0)
3391 .write_native(&path, &[(1, id_col), (2, name_col)], n, 1)
3392 .unwrap();
3393
3394 let mut reader = RunReader::open_with_cache(&path, schema(), None, None, None, 0, None)
3395 .expect("open reader");
3396 assert!(reader.has_mmap(), "test env must support read-only mmap");
3397 assert_eq!(reader.row_count(), header.row_count as usize);
3398
3399 for cid in [1u16, 2] {
3400 let a = reader.column_native(cid).expect("column_native");
3401 let b = reader
3402 .column_native_shared(cid)
3403 .expect("column_native_shared");
3404 assert_eq!(a.len(), b.len(), "len mismatch col {cid}");
3405 match (&a, &b) {
3407 (
3408 NativeColumn::Int64 {
3409 data: da,
3410 validity: va,
3411 },
3412 NativeColumn::Int64 {
3413 data: db,
3414 validity: vb,
3415 },
3416 ) => {
3417 assert_eq!(da, db, "Int64 data col {cid}");
3418 assert_eq!(va, vb, "Int64 validity col {cid}");
3419 }
3420 (
3421 NativeColumn::Bytes {
3422 offsets: oa,
3423 values: ua,
3424 validity: va,
3425 },
3426 NativeColumn::Bytes {
3427 offsets: ob,
3428 values: ub,
3429 validity: vb,
3430 },
3431 ) => {
3432 assert_eq!(oa, ob, "Bytes offsets col {cid}");
3433 assert_eq!(ua, ub, "Bytes values col {cid}");
3434 assert_eq!(va, vb, "Bytes validity col {cid}");
3435 }
3436 _ => panic!("type mismatch col {cid}: {a:?} vs {b:?}"),
3437 }
3438 }
3439 }
3440
3441 #[test]
3442 fn page_cache_key_distinguishes_tables() {
3443 let a = page_cache_key(1, 5, 2, 3);
3444 let b = page_cache_key(2, 5, 2, 3);
3445 assert_ne!(a, b, "same run/col/page, different table must differ");
3446 assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 6, 2, 3));
3448 assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 5, 2, 4));
3449 }
3450}