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::{LearnedIndex, 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
881pub fn read_header(path: impl AsRef<Path>) -> Result<RunHeader> {
883 let mut file = File::open(path)?;
884 let mut header_buf = vec![0u8; RUN_HEADER_PAD];
885 file.read_exact(&mut header_buf)?;
886 let header: RunHeader = bincode::deserialize(&header_buf)
887 .map_err(|e| MongrelError::InvalidArgument(format!("bad run header: {e}")))?;
888 if header.magic != RUN_MAGIC {
889 return Err(MongrelError::MagicMismatch {
890 what: "sorted run",
891 expected: RUN_MAGIC,
892 got: header.magic,
893 });
894 }
895
896 file.seek(SeekFrom::Start(header.footer_offset))?;
897 let mut footer = [0u8; 8 + 8 + 32];
898 file.read_exact(&mut footer)?;
899 if footer[..8] != RUN_MAGIC {
900 return Err(MongrelError::MagicMismatch {
901 what: "sorted run footer",
902 expected: RUN_MAGIC,
903 got: footer[..8].try_into().unwrap(),
904 });
905 }
906 let mut hasher = Sha256::new();
907 hasher.update(&header_buf);
908 let body_len = header.footer_offset.saturating_sub(RUN_HEADER_PAD as u64);
909 file.seek(SeekFrom::Start(RUN_HEADER_PAD as u64))?;
910 let mut body = vec![0u8; body_len as usize];
911 file.read_exact(&mut body)?;
912 hasher.update(&body);
913 let computed: [u8; 32] = hasher.finalize().into();
914 let stored: [u8; 32] = footer[16..].try_into().unwrap();
915 if computed != stored {
916 return Err(MongrelError::ChecksumMismatch {
917 expected: u64::from_be_bytes(stored[..8].try_into().unwrap()),
918 actual: u64::from_be_bytes(computed[..8].try_into().unwrap()),
919 context: "sorted run footer".into(),
920 });
921 }
922 Ok(header)
923}
924
925pub fn read_column_dir(
927 path: impl AsRef<Path>,
928 header: &RunHeader,
929) -> Result<Vec<ColumnPageHeader>> {
930 let mut file = File::open(path)?;
931 file.seek(SeekFrom::Start(header.column_dir_offset))?;
932 let end = if header.index_trailer_offset != 0 {
933 header.index_trailer_offset
934 } else {
935 header.footer_offset
936 };
937 let len = end.saturating_sub(header.column_dir_offset);
938 let mut buf = vec![0u8; len as usize];
939 file.read_exact(&mut buf)?;
940 let dir: Vec<ColumnPageHeader> = bincode::deserialize(&buf)
941 .map_err(|e| MongrelError::InvalidArgument(format!("bad column dir: {e}")))?;
942 Ok(dir)
943}
944
945pub(crate) fn read_encryption_descriptor_bytes(
948 path: impl AsRef<Path>,
949 header: &RunHeader,
950) -> Result<Vec<u8>> {
951 let mut file = File::open(path)?;
952 file.seek(SeekFrom::Start(header.encryption_descriptor_offset))?;
953 let mut len_buf = [0u8; 4];
954 file.read_exact(&mut len_buf)?;
955 let len = u32::from_le_bytes(len_buf) as usize;
956 const MAX_DESCRIPTOR_BYTES: usize = 65_536;
961 if len > MAX_DESCRIPTOR_BYTES {
962 return Err(MongrelError::InvalidArgument(format!(
963 "encryption descriptor length {len} exceeds {MAX_DESCRIPTOR_BYTES}"
964 )));
965 }
966 let mut buf = vec![0u8; len];
967 file.read_exact(&mut buf)?;
968 Ok(buf)
969}
970
971#[cfg(feature = "encryption")]
978fn verify_run_mac(
979 path: &Path,
980 header: &RunHeader,
981 dir: &[ColumnPageHeader],
982 kek: &Kek,
983 desc_bytes: &[u8],
984) -> Result<()> {
985 let header_bytes = bincode::serialize(header)?;
986 let dir_bytes = bincode::serialize(dir)?;
987 let mac_key = kek.derive_run_mac_key();
988 let expected =
989 crate::encryption::run_metadata_mac(&mac_key, &header_bytes, &dir_bytes, desc_bytes);
990 let mut file = File::open(path)?;
991 file.seek(SeekFrom::Start(header.footer_offset + 8 + 8 + 32))?;
992 let mut tag = [0u8; RUN_MAC_LEN];
993 file.read_exact(&mut tag).map_err(|_| {
994 MongrelError::Decryption(
995 "encrypted run is missing or truncated its metadata MAC; cannot \
996 authenticate metadata"
997 .into(),
998 )
999 })?;
1000 let mut diff = 0u8;
1002 for (x, y) in tag.iter().zip(expected.iter()) {
1003 diff |= x ^ y;
1004 }
1005 if diff != 0 {
1006 return Err(MongrelError::Decryption(
1007 "run metadata authentication failed — tampered run or wrong key".into(),
1008 ));
1009 }
1010 Ok(())
1011}
1012
1013#[cfg(not(feature = "encryption"))]
1014fn verify_run_mac(
1015 _path: &Path,
1016 _header: &RunHeader,
1017 _dir: &[ColumnPageHeader],
1018 _kek: &Kek,
1019 _desc_bytes: &[u8],
1020) -> Result<()> {
1021 Ok(())
1022}
1023
1024pub struct RunWriter<'a> {
1033 schema: &'a Schema,
1034 run_id: u128,
1035 epoch_created: Epoch,
1036 level: u8,
1037 kek: Option<&'a Kek>,
1038 indexable_columns: Vec<(u16, u8)>,
1041 compress: columnar::Compress,
1045 clean: bool,
1050 uniform_epoch: bool,
1055 le: bool,
1062}
1063
1064impl<'a> RunWriter<'a> {
1065 pub fn new(schema: &'a Schema, run_id: u128, epoch_created: Epoch, level: u8) -> Self {
1066 Self {
1067 schema,
1068 run_id,
1069 epoch_created,
1070 level,
1071 kek: None,
1072 indexable_columns: Vec::new(),
1073 compress: columnar::Compress::Zstd(3),
1074 clean: false,
1075 uniform_epoch: false,
1076 le: false,
1077 }
1078 }
1079
1080 pub fn uniform_epoch(mut self, uniform: bool) -> Self {
1086 self.uniform_epoch = uniform;
1087 self
1088 }
1089
1090 pub fn with_encryption(mut self, kek: &'a Kek, indexable_columns: Vec<(u16, u8)>) -> Self {
1094 self.kek = Some(kek);
1095 self.indexable_columns = indexable_columns;
1096 self
1097 }
1098
1099 pub fn with_zstd_level(mut self, level: i32) -> Self {
1102 self.compress = columnar::Compress::Zstd(level);
1103 self
1104 }
1105
1106 pub fn with_lz4(mut self) -> Self {
1110 self.compress = columnar::Compress::Lz4;
1111 self
1112 }
1113
1114 pub fn with_plain(mut self) -> Self {
1117 self.compress = columnar::Compress::Plain;
1118 self
1119 }
1120
1121 pub fn clean(mut self, clean: bool) -> Self {
1127 self.clean = clean;
1128 self
1129 }
1130
1131 pub fn with_native_endian(mut self) -> Self {
1136 if cfg!(target_endian = "little") {
1137 self.le = true;
1138 }
1139 self
1140 }
1141
1142 pub fn write_native(
1149 self,
1150 path: impl AsRef<Path>,
1151 user_columns: &[(u16, columnar::NativeColumn)],
1152 n: usize,
1153 first_row_id: u64,
1154 ) -> Result<RunHeader> {
1155 use columnar::NativeColumn;
1156 let row_id_col = NativeColumn::int64_sequence(first_row_id as i64, n);
1157 let epoch_col = NativeColumn::int64_constant(self.epoch_created.0 as i64, n);
1158 let deleted_col = NativeColumn::bool_constant(false, n);
1159
1160 let learned_trailer = build_learned_trailer_native(&row_id_col);
1161
1162 let row_ids = match &row_id_col {
1165 NativeColumn::Int64 { data, .. } => data.as_slice(),
1166 _ => &[],
1167 };
1168 let bounds = page_bounds(row_ids);
1169 let compress = self.compress;
1170 let le = self.le;
1171 let plain = matches!(compress, columnar::Compress::Plain);
1172 let row_id_enc = if plain {
1175 Encoding::Plain
1176 } else {
1177 Encoding::Delta
1178 };
1179 let sys_enc = if plain {
1180 Encoding::Plain
1181 } else {
1182 Encoding::Zstd
1183 };
1184
1185 let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
1186 let (pages, stats) = native_column_pages(
1187 TypeId::Int64,
1188 &row_id_col,
1189 row_id_enc,
1190 compress,
1191 le,
1192 &bounds,
1193 )?;
1194 columns.push(ColumnPayload {
1195 column_id: SYS_ROW_ID,
1196 type_id_tag: type_tag(&TypeId::Int64),
1197 encoding: row_id_enc,
1198 pages,
1199 page_stats: stats,
1200 });
1201 let (pages, stats) =
1202 native_column_pages(TypeId::Int64, &epoch_col, sys_enc, compress, le, &bounds)?;
1203 columns.push(ColumnPayload {
1204 column_id: SYS_EPOCH,
1205 type_id_tag: type_tag(&TypeId::Int64),
1206 encoding: sys_enc,
1207 pages,
1208 page_stats: stats,
1209 });
1210 let (pages, stats) =
1211 native_column_pages(TypeId::Bool, &deleted_col, sys_enc, compress, le, &bounds)?;
1212 columns.push(ColumnPayload {
1213 column_id: SYS_DELETED,
1214 type_id_tag: type_tag(&TypeId::Bool),
1215 encoding: sys_enc,
1216 pages,
1217 page_stats: stats,
1218 });
1219 use rayon::prelude::*;
1225 let user_cols: Vec<ColumnPayload> = self
1226 .schema
1227 .columns
1228 .par_iter()
1229 .map(|cdef| -> Result<ColumnPayload> {
1230 let col = user_columns
1231 .iter()
1232 .find(|(id, _)| *id == cdef.id)
1233 .map(|(_, c)| c)
1234 .ok_or_else(|| {
1235 MongrelError::ColumnNotFound(format!("user column {}", cdef.id))
1236 })?;
1237 let encoding = if plain {
1238 Encoding::Plain
1239 } else {
1240 choose_encoding_native(&cdef.ty, col)
1241 };
1242 let (pages, stats) =
1243 native_column_pages(cdef.ty, col, encoding, compress, le, &bounds)?;
1244 Ok(ColumnPayload {
1245 column_id: cdef.id,
1246 type_id_tag: type_tag(&cdef.ty),
1247 encoding,
1248 pages,
1249 page_stats: stats,
1250 })
1251 })
1252 .collect::<Result<Vec<_>>>()?;
1253 columns.extend(user_cols);
1254
1255 let flags = if self.uniform_epoch {
1261 RUN_FLAG_UNIFORM_EPOCH
1262 } else {
1263 RUN_FLAG_CLEAN
1264 };
1265 let spec = RunSpec {
1266 run_id: self.run_id,
1267 schema_id: self.schema.schema_id,
1268 epoch_created: self.epoch_created.0,
1269 level: self.level,
1270 flags,
1271 sort_key_column_id: SYS_ROW_ID,
1272 row_count: n as u64,
1273 min_row_id: first_row_id,
1274 max_row_id: first_row_id + n as u64 - 1,
1275 columns: &columns,
1276 };
1277 write_run_with(
1278 path,
1279 &spec,
1280 self.kek,
1281 &self.indexable_columns,
1282 Some(&learned_trailer),
1283 )
1284 }
1285
1286 pub fn write(self, path: impl AsRef<Path>, rows: &[Row]) -> Result<RunHeader> {
1287 let n = rows.len();
1288 let mut row_ids = Vec::with_capacity(n);
1290 let mut epochs = Vec::with_capacity(n);
1291 let mut deleted = Vec::with_capacity(n);
1292 for r in rows {
1293 row_ids.push(Value::Int64(r.row_id.0 as i64));
1294 epochs.push(Value::Int64(r.committed_epoch.0 as i64));
1295 deleted.push(Value::Bool(r.deleted));
1296 }
1297 let learned_trailer = build_learned_trailer(&row_ids);
1298 let (min_rid, max_rid) = row_id_bounds(rows);
1299 let row_id_i64: Vec<i64> = rows.iter().map(|r| r.row_id.0 as i64).collect();
1300 let bounds = page_bounds(&row_id_i64);
1301
1302 let mut columns: Vec<ColumnPayload> = Vec::with_capacity(3 + self.schema.columns.len());
1303 let (pages, stats, enc) = value_column_pages(TypeId::Int64, &row_ids, &bounds)?;
1304 columns.push(ColumnPayload {
1305 column_id: SYS_ROW_ID,
1306 type_id_tag: type_tag(&TypeId::Int64),
1307 encoding: enc,
1308 pages,
1309 page_stats: stats,
1310 });
1311 let (pages, stats, enc) = value_column_pages(TypeId::Int64, &epochs, &bounds)?;
1312 columns.push(ColumnPayload {
1313 column_id: SYS_EPOCH,
1314 type_id_tag: type_tag(&TypeId::Int64),
1315 encoding: enc,
1316 pages,
1317 page_stats: stats,
1318 });
1319 let (pages, stats, enc) = value_column_pages(TypeId::Bool, &deleted, &bounds)?;
1320 columns.push(ColumnPayload {
1321 column_id: SYS_DELETED,
1322 type_id_tag: type_tag(&TypeId::Bool),
1323 encoding: enc,
1324 pages,
1325 page_stats: stats,
1326 });
1327 for cdef in &self.schema.columns {
1329 let vals: Vec<Value> = rows
1330 .iter()
1331 .map(|r| r.columns.get(&cdef.id).cloned().unwrap_or(Value::Null))
1332 .collect();
1333 let (pages, stats, encoding) = value_column_pages(cdef.ty, &vals, &bounds)?;
1334 columns.push(ColumnPayload {
1335 column_id: cdef.id,
1336 type_id_tag: type_tag(&cdef.ty),
1337 encoding,
1338 pages,
1339 page_stats: stats,
1340 });
1341 }
1342
1343 let spec = RunSpec {
1344 run_id: self.run_id,
1345 schema_id: self.schema.schema_id,
1346 epoch_created: self.epoch_created.0,
1347 level: self.level,
1348 flags: {
1349 let mut f = if self.clean { RUN_FLAG_CLEAN } else { 0 };
1350 if self.uniform_epoch {
1351 f |= RUN_FLAG_UNIFORM_EPOCH;
1352 }
1353 f
1354 },
1355 sort_key_column_id: SYS_ROW_ID,
1356 row_count: n as u64,
1357 min_row_id: min_rid,
1358 max_row_id: max_rid,
1359 columns: &columns,
1360 };
1361 write_run_with(
1362 path,
1363 &spec,
1364 self.kek,
1365 &self.indexable_columns,
1366 Some(&learned_trailer),
1367 )
1368 }
1369}
1370
1371fn type_tag(ty: &TypeId) -> u16 {
1372 match ty {
1374 TypeId::Bool => 1,
1375 TypeId::Int64 => 8,
1376 TypeId::Float64 => 9,
1377 TypeId::Bytes => 12,
1378 TypeId::Embedding { .. } => 13,
1379 _ => 0,
1380 }
1381}
1382
1383pub(crate) fn be_i64(b: Option<&[u8]>) -> Option<i64> {
1386 let b = b?;
1387 (b.len() == 8).then(|| i64::from_be_bytes(b.try_into().unwrap()))
1388}
1389
1390pub(crate) fn be_f64(b: Option<&[u8]>) -> Option<f64> {
1392 let b = b?;
1393 (b.len() == 8).then(|| f64::from_bits(u64::from_be_bytes(b.try_into().unwrap())))
1394}
1395
1396const PAGE_ROWS: usize = 65_536;
1399
1400fn page_bounds(row_ids: &[i64]) -> Vec<(usize, usize, u64, u64)> {
1403 let n = row_ids.len();
1404 if n == 0 {
1405 return vec![(0, 0, 0, 0)];
1406 }
1407 let mut out = Vec::new();
1408 let mut start = 0;
1409 while start < n {
1410 let end = (start + PAGE_ROWS).min(n);
1411 out.push((start, end, row_ids[start] as u64, row_ids[end - 1] as u64));
1412 start = end;
1413 }
1414 out
1415}
1416
1417fn native_column_pages(
1423 ty: TypeId,
1424 col: &columnar::NativeColumn,
1425 encoding: Encoding,
1426 compress: columnar::Compress,
1427 le: bool,
1428 bounds: &[(usize, usize, u64, u64)],
1429) -> Result<(Vec<Vec<u8>>, Vec<PageStat>)> {
1430 use rayon::prelude::*;
1431 let encode_one =
1432 |&(s, e, frid, lrid): &(usize, usize, u64, u64)| -> Result<(Vec<u8>, PageStat)> {
1433 let chunk = col.slice_range(s, e);
1434 let stat = columnar::page_stat_for(ty, &chunk, frid, lrid);
1435 let page = columnar::encode_page_native(ty, &chunk, encoding, compress, le)?;
1436 Ok((page, stat))
1437 };
1438 let pairs: Vec<(Vec<u8>, PageStat)> = if bounds.len() > 1 {
1440 bounds
1441 .par_iter()
1442 .map(encode_one)
1443 .collect::<Result<Vec<_>>>()?
1444 } else {
1445 bounds.iter().map(encode_one).collect::<Result<Vec<_>>>()?
1446 };
1447 let (pages, stats) = pairs.into_iter().unzip();
1448 Ok((pages, stats))
1449}
1450
1451fn value_column_pages(
1453 ty: TypeId,
1454 vals: &[Value],
1455 bounds: &[(usize, usize, u64, u64)],
1456) -> Result<(Vec<Vec<u8>>, Vec<PageStat>, Encoding)> {
1457 let encoding = choose_encoding(&ty, vals);
1458 let mut pages = Vec::with_capacity(bounds.len());
1459 let mut stats = Vec::with_capacity(bounds.len());
1460 for &(s, e, frid, lrid) in bounds {
1461 let chunk = &vals[s..e];
1462 pages.push(columnar::encode_page(ty, chunk, encoding)?);
1463 let native = columnar::values_to_native(ty, chunk);
1464 stats.push(columnar::page_stat_for(ty, &native, frid, lrid));
1465 }
1466 Ok((pages, stats, encoding))
1467}
1468
1469fn choose_encoding(ty: &TypeId, values: &[Value]) -> Encoding {
1472 use std::collections::HashSet;
1473 if matches!(ty, TypeId::Bytes) {
1474 let n = values.len();
1475 if n > 0 {
1476 let distinct = values
1477 .iter()
1478 .filter(|v| !matches!(v, Value::Null))
1479 .map(|v| v.encode_key())
1480 .collect::<HashSet<_>>()
1481 .len();
1482 if (distinct as f64 / n as f64) < 0.5 {
1483 return Encoding::Dictionary;
1484 }
1485 }
1486 }
1487 Encoding::Zstd
1488}
1489
1490fn choose_encoding_native(ty: &TypeId, col: &columnar::NativeColumn) -> Encoding {
1493 use std::collections::HashSet;
1494 match (ty, col) {
1495 (TypeId::Int64 | TypeId::TimestampNanos, columnar::NativeColumn::Int64 { data, .. }) => {
1496 if data.windows(2).all(|w| w[0] <= w[1]) {
1497 Encoding::Delta
1498 } else {
1499 Encoding::Zstd
1500 }
1501 }
1502 (
1503 TypeId::Bytes,
1504 columnar::NativeColumn::Bytes {
1505 offsets, values, ..
1506 },
1507 ) => {
1508 let n = offsets.len().saturating_sub(1);
1509 if n > 0 {
1510 let distinct: HashSet<&[u8]> = (0..n)
1511 .map(|i| &values[offsets[i] as usize..offsets[i + 1] as usize])
1512 .collect();
1513 if (distinct.len() as f64 / n as f64) < 0.5 {
1514 return Encoding::Dictionary;
1515 }
1516 }
1517 Encoding::Zstd
1518 }
1519 _ => Encoding::Zstd,
1520 }
1521}
1522
1523fn build_learned_trailer_native(col: &columnar::NativeColumn) -> Vec<u8> {
1525 let points: Vec<(u64, usize)> = match col {
1526 columnar::NativeColumn::Int64 { data, .. } => data
1527 .iter()
1528 .enumerate()
1529 .map(|(i, v)| (*v as u64, i))
1530 .collect(),
1531 _ => Vec::new(),
1532 };
1533 let pgm = PgmIndex::build(&points, LEARNED_EPSILON);
1534 bincode::serialize(&pgm).expect("pgm serialize")
1535}
1536
1537fn row_id_bounds(rows: &[Row]) -> (u64, u64) {
1538 match (rows.first(), rows.last()) {
1539 (Some(f), Some(l)) => (f.row_id.0, l.row_id.0),
1540 _ => (0, 0),
1541 }
1542}
1543
1544pub struct RunReader {
1549 file: File,
1550 mmap: Option<memmap2::Mmap>,
1551 header: RunHeader,
1552 dir: Vec<ColumnPageHeader>,
1553 schema: Schema,
1554 table_id: u64,
1556 cipher: Option<Box<dyn Cipher>>,
1558 nonce_prefix: [u8; 12],
1560 col_cache: HashMap<u16, Vec<Value>>,
1561 learned: Option<PgmIndex>,
1562 page_cache: Option<Arc<parking_lot::Mutex<crate::cache::PageCache>>>,
1566 decoded_cache: Option<Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>>,
1570 epoch_override: Option<Epoch>,
1574}
1575
1576impl RunReader {
1577 pub fn open(path: impl AsRef<Path>, schema: Schema, kek: Option<Arc<Kek>>) -> Result<Self> {
1578 Self::open_with_cache(path, schema, kek, None, None, 0)
1579 }
1580
1581 pub(crate) fn open_with_cache(
1582 path: impl AsRef<Path>,
1583 schema: Schema,
1584 kek: Option<Arc<Kek>>,
1585 page_cache: Option<Arc<parking_lot::Mutex<crate::cache::PageCache>>>,
1586 decoded_cache: Option<Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>>,
1587 table_id: u64,
1588 ) -> Result<Self> {
1589 let path = path.as_ref().to_path_buf();
1590 let header = read_header(&path)?;
1591 let mut dir = read_column_dir(&path, &header)?;
1592 let learned = read_learned(&path, &header)?;
1593 let (cipher, nonce_prefix): (Option<Box<dyn Cipher>>, [u8; 12]) = if header.is_encrypted() {
1596 let kek = kek.as_ref().ok_or_else(|| {
1597 MongrelError::Encryption(
1598 "run is encrypted but no key-encryption key was provided".into(),
1599 )
1600 })?;
1601 if header.encryption_descriptor_offset == 0 {
1602 return Err(MongrelError::Encryption(
1603 "encrypted run has no encryption descriptor".into(),
1604 ));
1605 }
1606 let desc_bytes = read_encryption_descriptor_bytes(&path, &header)?;
1607 verify_run_mac(&path, &header, &dir, kek, &desc_bytes)?;
1613 let enc = crate::encryption::build_run_cipher(kek, &desc_bytes)?;
1614 if header.encrypted_stats_offset != 0 {
1618 overlay_encrypted_stats(
1619 &path,
1620 &header,
1621 enc.cipher.as_ref(),
1622 enc.nonce_prefix,
1623 &mut dir,
1624 )?;
1625 }
1626 (Some(enc.cipher), enc.nonce_prefix)
1627 } else {
1628 (None, [0u8; 12])
1629 };
1630 let file = File::open(&path)?;
1633 let mmap = unsafe { memmap2::MmapOptions::new().map(&file) }.ok();
1640 let _ = path;
1641 Ok(Self {
1642 file,
1643 mmap,
1644 header,
1645 dir,
1646 schema,
1647 table_id,
1648 cipher,
1649 nonce_prefix,
1650 col_cache: HashMap::new(),
1651 learned,
1652 epoch_override: None,
1653 page_cache,
1654 decoded_cache,
1655 })
1656 }
1657
1658 pub fn header(&self) -> &RunHeader {
1659 &self.header
1660 }
1661
1662 pub(crate) fn set_uniform_epoch(&mut self, epoch: Epoch) {
1666 if self.header.is_uniform_epoch() {
1667 self.epoch_override = Some(epoch);
1668 self.col_cache.remove(&SYS_EPOCH);
1670 }
1671 }
1672
1673 pub fn is_clean(&self) -> bool {
1676 self.header.is_clean()
1677 }
1678
1679 pub fn row_count(&self) -> usize {
1680 self.header.row_count as usize
1681 }
1682
1683 pub(crate) fn clean_contiguous_row_ids(&self) -> bool {
1684 let n = self.row_count();
1685 n > 0
1686 && self.is_clean()
1687 && self.epoch_override.is_none()
1688 && self.header.max_row_id >= self.header.min_row_id
1689 && self
1690 .header
1691 .max_row_id
1692 .checked_sub(self.header.min_row_id)
1693 .and_then(|span| span.checked_add(1))
1694 == Some(n as u64)
1695 }
1696
1697 pub(crate) fn position_for_row_id_fast(&self, row_id: u64) -> Option<usize> {
1698 if !self.clean_contiguous_row_ids()
1699 || row_id < self.header.min_row_id
1700 || row_id > self.header.max_row_id
1701 {
1702 return None;
1703 }
1704 Some((row_id - self.header.min_row_id) as usize)
1705 }
1706
1707 pub(crate) fn positions_for_row_ids_fast(&self, row_ids: &[u64]) -> Option<Vec<usize>> {
1708 if !self.clean_contiguous_row_ids() {
1709 return None;
1710 }
1711 let mut positions = Vec::with_capacity(row_ids.len());
1712 for &row_id in row_ids {
1713 if let Some(pos) = self.position_for_row_id_fast(row_id) {
1714 positions.push(pos);
1715 }
1716 }
1717 positions.sort_unstable();
1718 Some(positions)
1719 }
1720
1721 fn resolve_type(&self, column_id: u16) -> TypeId {
1722 match column_id {
1723 SYS_ROW_ID | SYS_EPOCH => TypeId::Int64,
1724 SYS_DELETED => TypeId::Bool,
1725 _ => self
1726 .schema
1727 .columns
1728 .iter()
1729 .find(|c| c.id == column_id)
1730 .map(|c| c.ty)
1731 .unwrap_or(TypeId::Bytes),
1732 }
1733 }
1734
1735 fn find_header(&self, column_id: u16) -> Result<&ColumnPageHeader> {
1736 self.dir
1737 .iter()
1738 .find(|h| h.column_id == column_id)
1739 .ok_or_else(|| MongrelError::ColumnNotFound(format!("column {column_id}")))
1740 }
1741
1742 fn col_encrypted(&self, column_id: u16) -> bool {
1748 self.dir
1749 .iter()
1750 .find(|h| h.column_id == column_id)
1751 .map(|h| h.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0)
1752 .unwrap_or(false)
1753 }
1754
1755 pub(crate) fn read_page(&mut self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
1756 let (offset, compressed_len, encrypted) = {
1757 let ch = self.find_header(column_id)?;
1758 let stat = ch
1759 .page_stats
1760 .get(page_seq)
1761 .ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
1762 (
1763 stat.offset,
1764 stat.compressed_len,
1765 ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0,
1766 )
1767 };
1768 let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
1771 if let Some(cache) = &self.page_cache {
1772 if let Some(bytes) = cache.lock().get(
1773 &key,
1774 crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
1775 ) {
1776 return decrypt_or_passthrough(
1777 self.cipher.as_deref(),
1778 self.nonce_prefix,
1779 column_id,
1780 page_seq,
1781 encrypted,
1782 &bytes,
1783 );
1784 }
1785 }
1786 let buf = match &self.mmap {
1787 Some(m) => m[offset as usize..(offset + compressed_len as u64) as usize].to_vec(),
1790 None => {
1791 self.file.seek(SeekFrom::Start(offset))?;
1792 let mut buf = vec![0u8; compressed_len as usize];
1793 self.file.read_exact(&mut buf)?;
1794 buf
1795 }
1796 };
1797 if let Some(cache) = &self.page_cache {
1799 cache.lock().insert(crate::page::CachedPage {
1800 committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
1801 content_hash: key,
1802 bytes: bytes::Bytes::copy_from_slice(&buf),
1803 });
1804 }
1805 decrypt_or_passthrough(
1806 self.cipher.as_deref(),
1807 self.nonce_prefix,
1808 column_id,
1809 page_seq,
1810 encrypted,
1811 &buf,
1812 )
1813 }
1814
1815 fn read_page_shared(&self, column_id: u16, page_seq: usize) -> Result<Vec<u8>> {
1820 let ch = self.find_header(column_id)?;
1821 let stat = ch
1822 .page_stats
1823 .get(page_seq)
1824 .ok_or_else(|| MongrelError::InvalidArgument("page seq out of range".into()))?;
1825 let encrypted = ch.flags & ColumnPageHeader::PAGE_ENCRYPTED != 0;
1826 let key = page_cache_key(self.table_id, self.header.run_id, column_id, page_seq);
1829 if let Some(cache) = &self.page_cache {
1830 if let Some(guard) = cache.try_lock() {
1831 if let Some(bytes) = guard.try_get(
1832 &key,
1833 crate::epoch::Snapshot::at(crate::epoch::Epoch(u64::MAX)),
1834 ) {
1835 return decrypt_or_passthrough(
1836 self.cipher.as_deref(),
1837 self.nonce_prefix,
1838 column_id,
1839 page_seq,
1840 encrypted,
1841 &bytes,
1842 );
1843 }
1844 }
1845 }
1846 let mmap = self.mmap.as_ref().ok_or_else(|| {
1847 MongrelError::InvalidArgument("parallel page decode requires an mmap-backed run".into())
1848 })?;
1849 let start = stat.offset as usize;
1850 let end = (stat.offset + stat.compressed_len as u64) as usize;
1851 let buf = mmap[start..end].to_vec();
1852 if let Some(cache) = &self.page_cache {
1856 if let Some(mut guard) = cache.try_lock() {
1857 guard.insert(crate::page::CachedPage {
1858 committed_epoch: crate::epoch::Epoch(self.header.epoch_created),
1859 content_hash: key,
1860 bytes: bytes::Bytes::copy_from_slice(&buf),
1861 });
1862 }
1863 }
1864 decrypt_or_passthrough(
1865 self.cipher.as_deref(),
1866 self.nonce_prefix,
1867 column_id,
1868 page_seq,
1869 encrypted,
1870 &buf,
1871 )
1872 }
1873
1874 fn column(&mut self, column_id: u16) -> Result<&[Value]> {
1876 if column_id == SYS_EPOCH {
1879 if let Some(ov) = self.epoch_override {
1880 if !self.col_cache.contains_key(&column_id) {
1881 let n = self.row_count();
1882 self.col_cache
1883 .insert(column_id, vec![Value::Int64(ov.0 as i64); n]);
1884 }
1885 return Ok(self.col_cache.get(&column_id).unwrap().as_slice());
1886 }
1887 }
1888 if !self.col_cache.contains_key(&column_id) {
1889 let ty = self.resolve_type(column_id);
1890 let page_rows: Vec<usize> = {
1891 let ch = self.find_header(column_id)?;
1892 ch.page_stats.iter().map(|s| s.row_count as usize).collect()
1893 };
1894 let mut decoded: Vec<Value> = Vec::with_capacity(self.row_count());
1895 for (seq, &pr) in page_rows.iter().enumerate() {
1896 let page = self.read_page(column_id, seq)?;
1897 decoded.extend(columnar::decode_page(ty, &page, pr)?);
1898 }
1899 self.col_cache.insert(column_id, decoded);
1900 }
1901 Ok(self.col_cache.get(&column_id).unwrap().as_slice())
1902 }
1903
1904 pub fn get_version(&mut self, row_id: RowId, snapshot: Epoch) -> Result<Option<(Epoch, Row)>> {
1907 let n = self.row_count();
1908 if n == 0 {
1909 return Ok(None);
1910 }
1911 let target = row_id.0;
1912 let row_ids = self.column(SYS_ROW_ID)?.to_vec();
1913 let window = self.predict_window(target, n);
1914 let mut start = None;
1917 let mut probe = window;
1918 if probe.is_empty() {
1919 probe = 0..n;
1920 }
1921 for i in probe.clone() {
1922 if int_at(&row_ids, i) == target {
1923 start = Some(i);
1924 break;
1925 }
1926 }
1927 let mut idx = start;
1928 if idx.is_none() {
1929 match row_ids.binary_search_by_key(&target, value_row_id) {
1930 Ok(i) => idx = Some(i),
1931 Err(_) => return Ok(None),
1932 }
1933 }
1934 let mut best: Option<(u64, usize)> = None; let mut i = idx.unwrap();
1936 while i < n && int_at(&row_ids, i) == target {
1937 let epoch = int_at(self.column(SYS_EPOCH)?, i);
1938 if epoch <= snapshot.0 && best.map(|(be, _)| epoch > be).unwrap_or(true) {
1939 best = Some((epoch, i));
1940 }
1941 i += 1;
1942 }
1943 match best {
1944 None => Ok(None),
1945 Some((epoch, index)) => Ok(Some((Epoch(epoch), self.materialize(index)?))),
1946 }
1947 }
1948
1949 pub fn all_rows(&mut self) -> Result<Vec<Row>> {
1952 let n = self.row_count();
1953 let mut out = Vec::with_capacity(n);
1954 for i in 0..n {
1955 out.push(self.materialize(i)?);
1956 }
1957 Ok(out)
1958 }
1959
1960 pub fn visible_indices(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
1966 let n = self.row_count();
1967 if n == 0 {
1968 return Ok(Vec::new());
1969 }
1970 let row_ids = self.column(SYS_ROW_ID)?.to_vec();
1971 let epochs = self.column(SYS_EPOCH)?.to_vec();
1972 let deleted = self.column(SYS_DELETED)?.to_vec();
1973 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
1974 for i in 0..n {
1975 let rid = int_at(&row_ids, i);
1976 let e = int_at(&epochs, i);
1977 if e > snapshot.0 {
1978 continue;
1979 }
1980 best.entry(rid)
1981 .and_modify(|(be, bi)| {
1982 if e > *be {
1983 *be = e;
1984 *bi = i;
1985 }
1986 })
1987 .or_insert((e, i));
1988 }
1989 let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
1990 idxs.retain(|&i| !bool_at(&deleted, i));
1991 idxs.sort_unstable();
1992 Ok(idxs)
1993 }
1994
1995 pub fn gather_column(&mut self, column_id: u16, indices: &[usize]) -> Result<Vec<Value>> {
2000 if !self.dir.iter().any(|h| h.column_id == column_id) {
2001 return Ok(vec![Value::Null; indices.len()]);
2002 }
2003 let col = self.column(column_id)?;
2004 Ok(indices
2005 .iter()
2006 .map(|&i| col.get(i).cloned().unwrap_or(Value::Null))
2007 .collect())
2008 }
2009
2010 pub fn column_native(&mut self, column_id: u16) -> Result<columnar::NativeColumn> {
2015 use rayon::prelude::*;
2016 if column_id == SYS_EPOCH {
2017 if let Some(ov) = self.epoch_override {
2018 return Ok(columnar::NativeColumn::int64_constant(
2019 ov.0 as i64,
2020 self.row_count(),
2021 ));
2022 }
2023 }
2024 let ty = self.resolve_type(column_id);
2025 let n = self.row_count();
2026 let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
2027 return Ok(columnar::null_native(ty, n));
2028 };
2029 let page_count = ch.page_count as usize;
2030 let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
2031 if page_count == 0 {
2032 return Ok(columnar::null_native(ty, n));
2033 }
2034 let parts: Vec<columnar::NativeColumn> = if self.mmap.is_some() && page_count > 1 {
2035 let reader: &RunReader = self;
2039 (0..page_count)
2040 .into_par_iter()
2041 .map(|seq| {
2042 let raw = reader.read_page_shared(column_id, seq)?;
2043 columnar::decode_page_native(ty, &raw, page_rows[seq])
2044 })
2045 .collect::<Result<Vec<_>>>()?
2046 } else {
2047 let mut out = Vec::with_capacity(page_count);
2048 for (seq, &pr) in page_rows.iter().enumerate() {
2049 let page = self.read_page(column_id, seq)?;
2050 out.push(columnar::decode_page_native(ty, &page, pr)?);
2051 }
2052 out
2053 };
2054 Ok(columnar::NativeColumn::concat(&parts))
2055 }
2056
2057 pub fn has_mmap(&self) -> bool {
2061 self.mmap.is_some()
2062 }
2063
2064 pub fn column_native_shared(&self, column_id: u16) -> Result<columnar::NativeColumn> {
2071 use rayon::prelude::*;
2072 if column_id == SYS_EPOCH {
2073 if let Some(ov) = self.epoch_override {
2074 return Ok(columnar::NativeColumn::int64_constant(
2075 ov.0 as i64,
2076 self.row_count(),
2077 ));
2078 }
2079 }
2080 let ty = self.resolve_type(column_id);
2081 let n = self.row_count();
2082 let Some(ch) = self.dir.iter().find(|h| h.column_id == column_id) else {
2083 return Ok(columnar::null_native(ty, n));
2084 };
2085 let page_count = ch.page_count as usize;
2086 let page_rows: Vec<usize> = ch.page_stats.iter().map(|s| s.row_count as usize).collect();
2087 if page_count == 0 {
2088 return Ok(columnar::null_native(ty, n));
2089 }
2090 if let (Some(m), Some(first)) = (&self.mmap, ch.page_stats.first()) {
2094 let start = first.offset as usize;
2095 let end = (ch.page_region_offset as usize) + (ch.page_region_len as usize);
2096 if end > start {
2097 let _ = m.advise_range(memmap2::Advice::WillNeed, start, end - start);
2098 }
2099 }
2100 let run_id = self.header.run_id;
2101 let mut parts_keys: Vec<(columnar::NativeColumn, Option<[u8; 32]>)> = if page_count > 1 {
2105 (0..page_count)
2106 .into_par_iter()
2107 .map(|seq| self.decode_page_cached(ty, column_id, seq, page_rows[seq], run_id))
2108 .collect::<Result<Vec<_>>>()?
2109 } else {
2110 vec![self.decode_page_cached(ty, column_id, 0, page_rows[0], run_id)?]
2111 };
2112 if let Some(cache) = &self.decoded_cache {
2115 let mut g = cache.lock();
2116 for (col, key) in parts_keys.iter_mut() {
2117 if let Some(k) = key.take() {
2118 g.insert(k, Arc::new(col.clone()));
2119 }
2120 }
2121 }
2122 let parts: Vec<columnar::NativeColumn> = parts_keys.into_iter().map(|(c, _)| c).collect();
2123 Ok(columnar::NativeColumn::concat(&parts))
2124 }
2125
2126 fn decode_page_cached(
2131 &self,
2132 ty: TypeId,
2133 column_id: u16,
2134 seq: usize,
2135 nrows: usize,
2136 run_id: u128,
2137 ) -> Result<(columnar::NativeColumn, Option<[u8; 32]>)> {
2138 let key = page_cache_key(self.table_id, run_id, column_id, seq);
2139 if let Some(cache) = &self.decoded_cache {
2140 if let Some(g) = cache.try_lock() {
2141 if let Some(hit) = g.try_get(&key) {
2142 return Ok(((*hit).clone(), None));
2143 }
2144 }
2145 }
2146 let raw = self.read_page_shared(column_id, seq)?;
2147 let col = columnar::decode_page_native(ty, &raw, nrows)?;
2148 Ok((col, Some(key)))
2149 }
2150
2151 pub fn range_row_ids_i64(
2156 &mut self,
2157 column_id: u16,
2158 lo: i64,
2159 hi: i64,
2160 ) -> Result<std::collections::HashSet<u64>> {
2161 Ok(self
2162 .range_row_id_set_i64(column_id, lo, hi)?
2163 .into_sorted_vec()
2164 .into_iter()
2165 .collect())
2166 }
2167
2168 pub(crate) fn range_row_id_set_i64(
2169 &mut self,
2170 column_id: u16,
2171 lo: i64,
2172 hi: i64,
2173 ) -> Result<RowIdSet> {
2174 let info: Vec<(Option<i64>, Option<i64>, usize)> =
2175 match self.dir.iter().find(|h| h.column_id == column_id) {
2176 Some(ch) => ch
2177 .page_stats
2178 .iter()
2179 .map(|s| {
2180 (
2181 be_i64(s.min.as_deref()),
2182 be_i64(s.max.as_deref()),
2183 s.row_count as usize,
2184 )
2185 })
2186 .collect(),
2187 None => return Ok(RowIdSet::empty()),
2188 };
2189 let stats_pruneable =
2195 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2196 let clean_contiguous = self.clean_contiguous_row_ids();
2197 let mut out = Vec::new();
2198 let mut page_start = 0usize;
2199 for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
2200 let current_page_start = page_start;
2201 page_start += nrows;
2202 let skip = stats_pruneable
2204 && match (mn, mx) {
2205 (Some(mn), Some(mx)) => mx < lo || mn > hi,
2206 _ => true,
2207 };
2208 if skip {
2209 continue;
2210 }
2211 let val_page = self.read_page(column_id, seq)?;
2212 let vals =
2213 columnar::decode_page_native(self.resolve_type(column_id), &val_page, nrows)?;
2214 if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
2215 if clean_contiguous {
2216 for (i, val) in v.iter().enumerate() {
2217 if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
2218 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2219 }
2220 }
2221 } else {
2222 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2223 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2224 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2225 for (i, val) in v.iter().enumerate() {
2226 if columnar::validity_bit(&validity, i) && *val >= lo && *val <= hi {
2227 out.push(r[i] as u64);
2228 }
2229 }
2230 }
2231 }
2232 }
2233 }
2234 Ok(RowIdSet::from_unsorted(out))
2235 }
2236
2237 pub fn range_row_ids_f64(
2240 &mut self,
2241 column_id: u16,
2242 lo: f64,
2243 lo_inclusive: bool,
2244 hi: f64,
2245 hi_inclusive: bool,
2246 ) -> Result<std::collections::HashSet<u64>> {
2247 Ok(self
2248 .range_row_id_set_f64(column_id, lo, lo_inclusive, hi, hi_inclusive)?
2249 .into_sorted_vec()
2250 .into_iter()
2251 .collect())
2252 }
2253
2254 pub(crate) fn range_row_id_set_f64(
2255 &mut self,
2256 column_id: u16,
2257 lo: f64,
2258 lo_inclusive: bool,
2259 hi: f64,
2260 hi_inclusive: bool,
2261 ) -> Result<RowIdSet> {
2262 let info: Vec<(Option<f64>, Option<f64>, usize)> =
2263 match self.dir.iter().find(|h| h.column_id == column_id) {
2264 Some(ch) => ch
2265 .page_stats
2266 .iter()
2267 .map(|s| {
2268 (
2269 be_f64(s.min.as_deref()),
2270 be_f64(s.max.as_deref()),
2271 s.row_count as usize,
2272 )
2273 })
2274 .collect(),
2275 None => return Ok(RowIdSet::empty()),
2276 };
2277 let stats_pruneable =
2283 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2284 let clean_contiguous = self.clean_contiguous_row_ids();
2285 let mut out = Vec::new();
2286 let mut page_start = 0usize;
2287 for (seq, (mn, mx, nrows)) in info.into_iter().enumerate() {
2288 let current_page_start = page_start;
2289 page_start += nrows;
2290 let skip = stats_pruneable
2293 && match (mn, mx) {
2294 (Some(mn), Some(mx)) => {
2295 let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
2296 let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
2297 skip_lo || skip_hi
2298 }
2299 _ => true,
2300 };
2301 if skip {
2302 continue;
2303 }
2304 let val_page = self.read_page(column_id, seq)?;
2305 let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
2306 if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
2307 if clean_contiguous {
2308 for (i, val) in v.iter().enumerate() {
2309 if !columnar::validity_bit(&validity, i) || val.is_nan() {
2310 continue;
2311 }
2312 let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
2313 let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
2314 if ok_lo && ok_hi {
2315 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2316 }
2317 }
2318 } else {
2319 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2320 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2321 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2322 for (i, val) in v.iter().enumerate() {
2323 if !columnar::validity_bit(&validity, i) || val.is_nan() {
2324 continue;
2325 }
2326 let ok_lo = if lo_inclusive { *val >= lo } else { *val > lo };
2327 let ok_hi = if hi_inclusive { *val <= hi } else { *val < hi };
2328 if ok_lo && ok_hi {
2329 out.push(r[i] as u64);
2330 }
2331 }
2332 }
2333 }
2334 }
2335 }
2336 Ok(RowIdSet::from_unsorted(out))
2337 }
2338
2339 pub(crate) fn null_row_id_set(&mut self, column_id: u16, want_nulls: bool) -> Result<RowIdSet> {
2344 let stats: Vec<(usize, usize)> = match self.dir.iter().find(|h| h.column_id == column_id) {
2345 Some(ch) => ch
2346 .page_stats
2347 .iter()
2348 .map(|s| (s.null_count as usize, s.row_count as usize))
2349 .collect(),
2350 None => return Ok(RowIdSet::empty()),
2351 };
2352 let ty = self.resolve_type(column_id);
2353 let clean_contiguous = self.clean_contiguous_row_ids();
2354 let mut out = Vec::new();
2355 let mut page_start = 0usize;
2356 for (seq, (null_count, nrows)) in stats.into_iter().enumerate() {
2357 let current_page_start = page_start;
2358 page_start += nrows;
2359 if want_nulls && null_count == 0 {
2361 continue;
2362 }
2363 if !want_nulls && null_count == nrows {
2364 continue;
2365 }
2366 let val_page = self.read_page(column_id, seq)?;
2367 let col = columnar::decode_page_native(ty, &val_page, nrows)?;
2368 let validity = col.validity();
2369 if clean_contiguous {
2370 for i in 0..nrows {
2371 let is_null = !columnar::validity_bit(validity, i);
2372 if is_null == want_nulls {
2373 out.push(self.header.min_row_id + current_page_start as u64 + i as u64);
2374 }
2375 }
2376 } else {
2377 let rid_page = self.read_page(SYS_ROW_ID, seq)?;
2378 let rids = columnar::decode_page_native(TypeId::Int64, &rid_page, nrows)?;
2379 if let columnar::NativeColumn::Int64 { data: r, .. } = rids {
2380 for (i, &rid) in r.iter().enumerate().take(nrows) {
2381 let is_null = !columnar::validity_bit(validity, i);
2382 if is_null == want_nulls {
2383 out.push(rid as u64);
2384 }
2385 }
2386 }
2387 }
2388 }
2389 Ok(RowIdSet::from_unsorted(out))
2390 }
2391
2392 pub fn visible_indices_native(&mut self, snapshot: Epoch) -> Result<Vec<usize>> {
2393 let n = self.row_count();
2394 if n == 0 {
2395 return Ok(Vec::new());
2396 }
2397 let (row_ids, epochs, deleted) = self.system_columns_native()?;
2398 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
2399 for i in 0..n {
2400 let rid = row_ids[i] as u64;
2401 let e = epochs[i] as u64;
2402 if e > snapshot.0 {
2403 continue;
2404 }
2405 best.entry(rid)
2406 .and_modify(|(be, bi)| {
2407 if e > *be {
2408 *be = e;
2409 *bi = i;
2410 }
2411 })
2412 .or_insert((e, i));
2413 }
2414 let mut idxs: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
2415 idxs.retain(|&i| deleted[i] == 0);
2416 idxs.sort_unstable();
2417 Ok(idxs)
2418 }
2419
2420 pub fn range_row_ids_visible_i64(
2429 &mut self,
2430 column_id: u16,
2431 lo: i64,
2432 hi: i64,
2433 snapshot: Epoch,
2434 ) -> Result<Vec<u64>> {
2435 let stats: Vec<(Option<i64>, Option<i64>, usize)> = match self.column_page_stats(column_id)
2436 {
2437 Some(s) => s
2438 .iter()
2439 .map(|st| {
2440 (
2441 be_i64(st.min.as_deref()),
2442 be_i64(st.max.as_deref()),
2443 st.row_count as usize,
2444 )
2445 })
2446 .collect(),
2447 None => return Ok(Vec::new()),
2448 };
2449 let stats_pruneable =
2455 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2456 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2457 let mut out: Vec<u64> = Vec::new();
2458 let mut vis = 0usize;
2459 let mut page_start = 0usize;
2460 for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
2461 let page_end = page_start + nrows;
2462 let skip = stats_pruneable
2464 && match (mn, mx) {
2465 (Some(mn), Some(mx)) => mx < lo || mn > hi,
2466 _ => true, };
2468 if !skip {
2469 let val_page = self.read_page(column_id, seq)?;
2470 let vals = columnar::decode_page_native(TypeId::Int64, &val_page, nrows)?;
2471 if let columnar::NativeColumn::Int64 { data: v, validity } = vals {
2472 while vis < positions.len() && positions[vis] < page_end {
2473 let local = positions[vis] - page_start;
2474 if columnar::validity_bit(&validity, local)
2475 && v[local] >= lo
2476 && v[local] <= hi
2477 {
2478 out.push(rids[vis] as u64);
2479 }
2480 vis += 1;
2481 }
2482 }
2483 } else {
2484 while vis < positions.len() && positions[vis] < page_end {
2485 vis += 1;
2486 }
2487 }
2488 page_start = page_end;
2489 }
2490 Ok(out)
2491 }
2492
2493 pub fn range_row_ids_visible_f64(
2496 &mut self,
2497 column_id: u16,
2498 lo: f64,
2499 lo_inclusive: bool,
2500 hi: f64,
2501 hi_inclusive: bool,
2502 snapshot: Epoch,
2503 ) -> Result<Vec<u64>> {
2504 let stats: Vec<(Option<f64>, Option<f64>, usize)> = match self.column_page_stats(column_id)
2505 {
2506 Some(s) => s
2507 .iter()
2508 .map(|st| {
2509 (
2510 be_f64(st.min.as_deref()),
2511 be_f64(st.max.as_deref()),
2512 st.row_count as usize,
2513 )
2514 })
2515 .collect(),
2516 None => return Ok(Vec::new()),
2517 };
2518 let stats_pruneable =
2524 !self.col_encrypted(column_id) || self.header.encrypted_stats_offset != 0;
2525 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2526 let mut out: Vec<u64> = Vec::new();
2527 let mut vis = 0usize;
2528 let mut page_start = 0usize;
2529 for (seq, &(mn, mx, nrows)) in stats.iter().enumerate() {
2530 let page_end = page_start + nrows;
2531 let skip = stats_pruneable
2532 && match (mn, mx) {
2533 (Some(mn), Some(mx)) => {
2534 let skip_lo = mx < lo || (!lo_inclusive && mx == lo);
2535 let skip_hi = mn > hi || (!hi_inclusive && mn == hi);
2536 skip_lo || skip_hi
2537 }
2538 _ => true,
2539 };
2540 if !skip {
2541 let val_page = self.read_page(column_id, seq)?;
2542 let vals = columnar::decode_page_native(TypeId::Float64, &val_page, nrows)?;
2543 if let columnar::NativeColumn::Float64 { data: v, validity } = vals {
2544 while vis < positions.len() && positions[vis] < page_end {
2545 let local = positions[vis] - page_start;
2546 if columnar::validity_bit(&validity, local) && !v[local].is_nan() {
2547 let val = v[local];
2548 let ok_lo = if lo_inclusive { val >= lo } else { val > lo };
2549 let ok_hi = if hi_inclusive { val <= hi } else { val < hi };
2550 if ok_lo && ok_hi {
2551 out.push(rids[vis] as u64);
2552 }
2553 }
2554 vis += 1;
2555 }
2556 }
2557 } else {
2558 while vis < positions.len() && positions[vis] < page_end {
2559 vis += 1;
2560 }
2561 }
2562 page_start = page_end;
2563 }
2564 Ok(out)
2565 }
2566
2567 pub fn null_row_ids_visible(
2573 &mut self,
2574 column_id: u16,
2575 want_nulls: bool,
2576 snapshot: Epoch,
2577 ) -> Result<Vec<u64>> {
2578 let stats: Vec<(usize, usize)> = match self.column_page_stats(column_id) {
2579 Some(s) => s
2580 .iter()
2581 .map(|st| (st.null_count as usize, st.row_count as usize))
2582 .collect(),
2583 None => return Ok(Vec::new()),
2584 };
2585 let ty = self.resolve_type(column_id);
2586 let (positions, rids) = self.visible_positions_with_rids(snapshot)?;
2587 let mut out: Vec<u64> = Vec::new();
2588 let mut vis = 0usize;
2589 let mut page_start = 0usize;
2590 for (seq, &(null_count, nrows)) in stats.iter().enumerate() {
2591 let page_end = page_start + nrows;
2592 let skip = (want_nulls && null_count == 0) || (!want_nulls && null_count == nrows);
2593 if !skip {
2594 let val_page = self.read_page(column_id, seq)?;
2595 let col = columnar::decode_page_native(ty, &val_page, nrows)?;
2596 let validity = col.validity();
2597 while vis < positions.len() && positions[vis] < page_end {
2598 let local = positions[vis] - page_start;
2599 let is_null = !columnar::validity_bit(validity, local);
2600 if is_null == want_nulls {
2601 out.push(rids[vis] as u64);
2602 }
2603 vis += 1;
2604 }
2605 } else {
2606 while vis < positions.len() && positions[vis] < page_end {
2607 vis += 1;
2608 }
2609 }
2610 page_start = page_end;
2611 }
2612 Ok(out)
2613 }
2614
2615 pub fn visible_positions_with_rids(
2619 &mut self,
2620 snapshot: Epoch,
2621 ) -> Result<(Vec<usize>, Vec<i64>)> {
2622 let n = self.row_count();
2623 if n == 0 {
2624 return Ok((Vec::new(), Vec::new()));
2625 }
2626 if self.is_clean() && self.epoch_override.is_none() {
2635 let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
2636 columnar::NativeColumn::Int64 { data, .. } => data,
2637 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
2638 };
2639 let positions: Vec<usize> = (0..n).collect();
2640 return Ok((positions, row_ids));
2641 }
2642 let (row_ids, epochs, deleted) = self.system_columns_native()?;
2643 let mut idxs: Vec<usize> = Vec::new();
2654 let mut i = 0;
2655 while i < n {
2656 let rid = row_ids[i] as u64;
2657 let mut best: Option<usize> = None;
2660 let mut j = i;
2661 while j < n && row_ids[j] as u64 == rid {
2662 if epochs[j] as u64 <= snapshot.0 {
2663 best = Some(j);
2664 }
2665 j += 1;
2666 }
2667 if let Some(b) = best {
2668 if deleted[b] == 0 {
2669 idxs.push(b);
2670 }
2671 }
2672 i = j;
2673 }
2674 let rids: Vec<i64> = idxs.iter().map(|&k| row_ids[k]).collect();
2676 Ok((idxs, rids))
2677 }
2678
2679 pub fn page_row_counts(&self, column_id: u16) -> Result<Vec<usize>> {
2683 Ok(self
2684 .find_header(column_id)?
2685 .page_stats
2686 .iter()
2687 .map(|s| s.row_count as usize)
2688 .collect())
2689 }
2690
2691 pub fn has_column(&self, column_id: u16) -> bool {
2694 self.dir.iter().any(|h| h.column_id == column_id)
2695 }
2696
2697 pub fn column_page_stats(&self, column_id: u16) -> Option<&[crate::page::PageStat]> {
2701 self.dir
2702 .iter()
2703 .find(|h| h.column_id == column_id)
2704 .map(|ch| ch.page_stats.as_slice())
2705 }
2706
2707 pub(crate) fn system_columns_native(&mut self) -> Result<(Vec<i64>, Vec<i64>, Vec<u8>)> {
2712 let row_ids = match self.column_native_shared(SYS_ROW_ID)? {
2713 columnar::NativeColumn::Int64 { data, .. } => data,
2714 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
2715 };
2716 let epochs = match self.column_native_shared(SYS_EPOCH)? {
2717 columnar::NativeColumn::Int64 { data, .. } => data,
2718 _ => return Err(MongrelError::InvalidArgument("sys epoch not int64".into())),
2719 };
2720 let deleted = match self.column_native_shared(SYS_DELETED)? {
2721 columnar::NativeColumn::Bool { data, .. } => data,
2722 _ => return Err(MongrelError::InvalidArgument("sys deleted not bool".into())),
2723 };
2724 Ok((row_ids, epochs, deleted))
2725 }
2726
2727 pub fn visible_versions(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
2731 let n = self.row_count();
2732 if n == 0 {
2733 return Ok(Vec::new());
2734 }
2735 let row_ids = self.column(SYS_ROW_ID)?.to_vec();
2736 let epochs = self.column(SYS_EPOCH)?.to_vec();
2737 let mut best: HashMap<u64, (u64, usize)> = HashMap::new();
2738 for i in 0..n {
2739 let rid = int_at(&row_ids, i);
2740 let epoch = int_at(&epochs, i);
2741 if epoch > snapshot.0 {
2742 continue;
2743 }
2744 best.entry(rid)
2745 .and_modify(|e| {
2746 if epoch > e.0 {
2747 *e = (epoch, i);
2748 }
2749 })
2750 .or_insert((epoch, i));
2751 }
2752 let mut picks: Vec<usize> = best.into_values().map(|(_, i)| i).collect();
2753 picks.sort();
2754 let mut out = Vec::with_capacity(picks.len());
2755 for i in picks {
2756 out.push(self.materialize(i)?);
2757 }
2758 Ok(out)
2759 }
2760
2761 pub fn visible_rows(&mut self, snapshot: Epoch) -> Result<Vec<Row>> {
2763 Ok(self
2764 .visible_versions(snapshot)?
2765 .into_iter()
2766 .filter(|r| !r.deleted)
2767 .collect())
2768 }
2769
2770 fn predict_window(&self, target: u64, n: usize) -> std::ops::Range<usize> {
2771 let Some(idx) = &self.learned else {
2772 return 0..n;
2773 };
2774 let (lo, hi) = idx.predict(target);
2775 let lo = lo.min(n);
2776 let hi = if hi == usize::MAX { n } else { hi.min(n) };
2777 lo..hi
2778 }
2779
2780 pub(crate) fn materialize(&mut self, index: usize) -> Result<Row> {
2781 let row_id = RowId(int_at(self.column(SYS_ROW_ID)?, index));
2782 let epoch = Epoch(int_at(self.column(SYS_EPOCH)?, index));
2783 let deleted = bool_at(self.column(SYS_DELETED)?, index);
2784 let col_ids: Vec<u16> = self.schema.columns.iter().map(|c| c.id).collect();
2785 let mut columns = HashMap::new();
2786 for id in col_ids {
2787 let val = if self.dir.iter().any(|h| h.column_id == id) {
2788 self.column(id)?.get(index).cloned().unwrap_or(Value::Null)
2789 } else {
2790 Value::Null
2792 };
2793 columns.insert(id, val);
2794 }
2795 Ok(Row {
2796 row_id,
2797 committed_epoch: epoch,
2798 columns,
2799 deleted,
2800 })
2801 }
2802
2803 pub(crate) fn materialize_batch(&mut self, indices: &[usize]) -> Result<Vec<Row>> {
2812 if indices.is_empty() {
2813 return Ok(Vec::new());
2814 }
2815 use std::collections::HashMap;
2816 let rid_col = self.column_native_shared(SYS_ROW_ID)?;
2817 let epoch_col = self.column_native_shared(SYS_EPOCH)?;
2818 let del_col = self.column_native_shared(SYS_DELETED)?;
2819 let mut user: HashMap<u16, columnar::NativeColumn> = HashMap::new();
2820 let present: Vec<u16> = self
2821 .schema
2822 .columns
2823 .iter()
2824 .map(|c| c.id)
2825 .filter(|id| self.dir.iter().any(|h| h.column_id == *id))
2826 .collect();
2827 for id in present {
2828 user.insert(id, self.column_native(id)?);
2829 }
2830 let i64_at = |col: &columnar::NativeColumn, i: usize| -> i64 {
2831 match col {
2832 columnar::NativeColumn::Int64 { data, .. } => data.get(i).copied().unwrap_or(0),
2833 _ => 0,
2834 }
2835 };
2836 let bool_at_native = |col: &columnar::NativeColumn, i: usize| -> bool {
2837 match col {
2838 columnar::NativeColumn::Bool { data, .. } => {
2839 data.get(i).copied().map(|b| b != 0).unwrap_or(false)
2840 }
2841 _ => false,
2842 }
2843 };
2844 let mut rows = Vec::with_capacity(indices.len());
2845 for &idx in indices {
2846 let row_id = RowId(i64_at(&rid_col, idx) as u64);
2847 let epoch = Epoch(i64_at(&epoch_col, idx) as u64);
2848 let deleted = bool_at_native(&del_col, idx);
2849 let mut columns = HashMap::with_capacity(self.schema.columns.len());
2850 for cdef in self.schema.columns.iter() {
2851 let val = match user.get(&cdef.id) {
2852 Some(col) => col.value_at(idx).unwrap_or(Value::Null),
2853 None => Value::Null,
2854 };
2855 columns.insert(cdef.id, val);
2856 }
2857 rows.push(Row {
2858 row_id,
2859 committed_epoch: epoch,
2860 columns,
2861 deleted,
2862 });
2863 }
2864 Ok(rows)
2865 }
2866}
2867
2868fn read_learned(path: &Path, header: &RunHeader) -> Result<Option<PgmIndex>> {
2869 if header.index_trailer_offset == 0 {
2870 return Ok(None);
2871 }
2872 let mut file = File::open(path)?;
2873 file.seek(SeekFrom::Start(header.index_trailer_offset))?;
2874 let len = header
2875 .footer_offset
2876 .saturating_sub(header.index_trailer_offset);
2877 let mut buf = vec![0u8; len as usize];
2878 file.read_exact(&mut buf)?;
2879 let pgm: PgmIndex = bincode::deserialize(&buf)
2880 .map_err(|e| MongrelError::InvalidArgument(format!("bad learned trailer: {e}")))?;
2881 Ok(Some(pgm))
2882}
2883
2884fn int_at(vals: &[Value], i: usize) -> u64 {
2885 match vals.get(i) {
2886 Some(Value::Int64(x)) => *x as u64,
2887 _ => 0,
2888 }
2889}
2890
2891fn bool_at(vals: &[Value], i: usize) -> bool {
2892 matches!(vals.get(i), Some(Value::Bool(true)))
2893}
2894
2895fn value_row_id(v: &Value) -> u64 {
2896 match v {
2897 Value::Int64(x) => *x as u64,
2898 _ => u64::MAX,
2899 }
2900}
2901
2902#[cfg(test)]
2903mod tests {
2904 use super::*;
2905 use crate::columnar::NativeColumn;
2906 use crate::memtable::Value;
2907 use crate::rowid::RowId;
2908 use crate::schema::{ColumnDef, ColumnFlags};
2909 use tempfile::tempdir;
2910
2911 fn schema() -> Schema {
2912 Schema {
2913 schema_id: 1,
2914 columns: vec![
2915 ColumnDef {
2916 id: 1,
2917 name: "id".into(),
2918 ty: TypeId::Int64,
2919 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
2920 },
2921 ColumnDef {
2922 id: 2,
2923 name: "name".into(),
2924 ty: TypeId::Bytes,
2925 flags: ColumnFlags::empty(),
2926 },
2927 ],
2928 indexes: Vec::new(),
2929 colocation: vec![],
2930 constraints: Default::default(),
2931 }
2932 }
2933
2934 fn rows() -> Vec<Row> {
2935 vec![
2936 Row::new(RowId(1), Epoch(10))
2937 .with_column(1, Value::Int64(1))
2938 .with_column(2, Value::Bytes(b"alice".to_vec())),
2939 Row::new(RowId(2), Epoch(10))
2940 .with_column(1, Value::Int64(2))
2941 .with_column(2, Value::Bytes(b"bob".to_vec())),
2942 Row::new(RowId(3), Epoch(10))
2943 .with_column(1, Value::Int64(3))
2944 .with_column(2, Value::Bytes(b"carol".to_vec())),
2945 ]
2946 }
2947
2948 #[test]
2949 fn flush_then_read_mvcc() {
2950 let dir = tempdir().unwrap();
2951 let path = dir.path().join("r-1.sr");
2952 let header = RunWriter::new(&schema(), 1, Epoch(10), 0)
2953 .write(&path, &rows())
2954 .unwrap();
2955 assert_eq!(header.row_count, 3);
2956
2957 let mut r = RunReader::open(&path, schema(), None).unwrap();
2958 let (e, row) = r.get_version(RowId(2), Epoch(20)).unwrap().unwrap();
2960 assert_eq!(e, Epoch(10));
2961 assert_eq!(row.row_id, RowId(2));
2962 assert!(matches!(row.columns.get(&2), Some(Value::Bytes(_))));
2963 assert!(r.get_version(RowId(99), Epoch(20)).unwrap().is_none());
2965 let all = r.visible_rows(Epoch(20)).unwrap();
2967 assert_eq!(all.len(), 3);
2968 }
2969
2970 #[test]
2971 fn learned_index_was_stored() {
2972 let dir = tempdir().unwrap();
2973 let path = dir.path().join("r-2.sr");
2974 RunWriter::new(&schema(), 2, Epoch(1), 0)
2975 .write(&path, &rows())
2976 .unwrap();
2977 let header = read_header(&path).unwrap();
2978 assert!(
2979 header.index_trailer_offset != 0,
2980 "trailer should be present"
2981 );
2982 }
2983
2984 #[test]
2985 fn low_level_container_still_round_trips() {
2986 let dir = tempdir().unwrap();
2987 let path = dir.path().join("r-3.sr");
2988 let cols = vec![ColumnPayload {
2989 column_id: 1,
2990 type_id_tag: 8,
2991 encoding: Encoding::Plain,
2992 pages: vec![vec![1, 2, 3, 4]],
2993 page_stats: Vec::new(),
2994 }];
2995 let header = write_run(
2996 &path,
2997 &RunSpec {
2998 run_id: 1,
2999 schema_id: 1,
3000 epoch_created: 1,
3001 level: 0,
3002 flags: 0,
3003 sort_key_column_id: SORT_KEY_ROW_ID,
3004 row_count: 1,
3005 min_row_id: 0,
3006 max_row_id: 0,
3007 columns: &cols,
3008 },
3009 )
3010 .unwrap();
3011 let back = read_header(&path).unwrap();
3012 assert_eq!(back.content_hash, header.content_hash);
3013 assert_eq!(read_column_dir(&path, &back).unwrap().len(), 1);
3014 }
3015
3016 #[test]
3017 fn detects_corruption() {
3018 let dir = tempdir().unwrap();
3019 let path = dir.path().join("r-4.sr");
3020 write_run(
3021 &path,
3022 &RunSpec {
3023 run_id: 1,
3024 schema_id: 1,
3025 epoch_created: 1,
3026 level: 0,
3027 flags: 0,
3028 sort_key_column_id: SORT_KEY_ROW_ID,
3029 row_count: 1,
3030 min_row_id: 0,
3031 max_row_id: 0,
3032 columns: &[ColumnPayload {
3033 column_id: 1,
3034 type_id_tag: 8,
3035 encoding: Encoding::Plain,
3036 pages: vec![vec![1, 2, 3]],
3037 page_stats: Vec::new(),
3038 }],
3039 },
3040 )
3041 .unwrap();
3042 let mut bytes = std::fs::read(&path).unwrap();
3043 bytes[300] ^= 0xFF;
3044 std::fs::write(&path, bytes).unwrap();
3045 let err = read_header(&path).unwrap_err();
3046 assert!(
3047 matches!(err, MongrelError::ChecksumMismatch { .. }),
3048 "got {err:?}"
3049 );
3050 }
3051
3052 #[test]
3058 fn mmap_and_vec_writers_are_byte_identical() {
3059 let dir = tempdir().unwrap();
3060 let stats = vec![PageStat {
3061 first_row_id: 0,
3062 last_row_id: 1,
3063 null_count: 0,
3064 row_count: 2,
3065 min: Some(vec![0]),
3066 max: Some(vec![9]),
3067 offset: 0,
3068 compressed_len: 0,
3069 uncompressed_len: 0,
3070 }];
3071 let spec = RunSpec {
3072 run_id: 7,
3073 schema_id: 1,
3074 epoch_created: 10,
3075 level: 0,
3076 flags: 0,
3077 sort_key_column_id: SORT_KEY_ROW_ID,
3078 row_count: 2,
3079 min_row_id: 0,
3080 max_row_id: 1,
3081 columns: &[
3082 ColumnPayload {
3083 column_id: 1,
3084 type_id_tag: 8,
3085 encoding: Encoding::Plain,
3086 pages: vec![vec![1, 2, 3, 4], vec![5, 6, 7, 8]],
3087 page_stats: stats.clone(),
3088 },
3089 ColumnPayload {
3090 column_id: 2,
3091 type_id_tag: 8,
3092 encoding: Encoding::Plain,
3093 pages: vec![vec![10, 20], vec![30, 40]],
3094 page_stats: stats.clone(),
3095 },
3096 ],
3097 };
3098 let trailer = b"learned-trailer-bytes";
3099
3100 let plan = plan_run(&spec, None, Some(trailer)).expect("plan");
3102 let mut buf = vec![0u8; plan.total];
3103 let h_place = place_run(&spec, None, Some(trailer), &plan, &mut buf)
3104 .expect("place_run into Vec succeeds");
3105
3106 let path_vec = dir.path().join("vec.sr");
3108 let h_vec = write_run_vec(&path_vec, &spec, None, Some(trailer)).unwrap();
3109 let bv = std::fs::read(&path_vec).unwrap();
3110
3111 assert_eq!(h_place.content_hash, h_vec.content_hash);
3112 assert_eq!(h_place.footer_offset, h_vec.footer_offset);
3113 assert_eq!(
3114 buf, bv,
3115 "place_run and write_run_vec must be byte-identical"
3116 );
3117 assert!(read_header(&path_vec).is_ok());
3118
3119 let path_mmap = dir.path().join("mmap.sr");
3123 match write_run_mmap(&path_mmap, &spec, None, Some(trailer)) {
3124 Ok(h_mmap) => {
3125 let bm = std::fs::read(&path_mmap).unwrap();
3126 assert_eq!(bm, bv, "mmap run must be byte-identical to vec run");
3127 assert_eq!(h_mmap.content_hash, h_vec.content_hash);
3128 }
3129 Err(e) if is_mmap_unavailable(&e) => {
3130 eprintln!("note: file mmap unavailable here; vec path covered (1)/(2)");
3131 }
3132 Err(e) => panic!("unexpected mmap error: {e:?}"),
3133 }
3134 }
3135
3136 #[test]
3141 fn column_native_shared_matches_column_native() {
3142 let dir = tempdir().unwrap();
3143 let path = dir.path().join("r.sr");
3144 let n = 65_536 * 2 + 9;
3147 let id_col = NativeColumn::int64_sequence(1, n);
3148 let mut offsets = vec![0u32];
3149 let mut values = Vec::new();
3150 for i in 0..n {
3151 values.extend_from_slice(format!("v{}", i % 17).as_bytes());
3152 offsets.push(values.len() as u32);
3153 }
3154 let name_col = NativeColumn::Bytes {
3155 offsets,
3156 values,
3157 validity: vec![0xFF; n.div_ceil(8)],
3158 };
3159 let header = RunWriter::new(&schema(), 9, Epoch(5), 0)
3160 .write_native(&path, &[(1, id_col), (2, name_col)], n, 1)
3161 .unwrap();
3162
3163 let mut reader =
3164 RunReader::open_with_cache(&path, schema(), None, None, None, 0).expect("open reader");
3165 assert!(reader.has_mmap(), "test env must support read-only mmap");
3166 assert_eq!(reader.row_count(), header.row_count as usize);
3167
3168 for cid in [1u16, 2] {
3169 let a = reader.column_native(cid).expect("column_native");
3170 let b = reader
3171 .column_native_shared(cid)
3172 .expect("column_native_shared");
3173 assert_eq!(a.len(), b.len(), "len mismatch col {cid}");
3174 match (&a, &b) {
3176 (
3177 NativeColumn::Int64 {
3178 data: da,
3179 validity: va,
3180 },
3181 NativeColumn::Int64 {
3182 data: db,
3183 validity: vb,
3184 },
3185 ) => {
3186 assert_eq!(da, db, "Int64 data col {cid}");
3187 assert_eq!(va, vb, "Int64 validity col {cid}");
3188 }
3189 (
3190 NativeColumn::Bytes {
3191 offsets: oa,
3192 values: ua,
3193 validity: va,
3194 },
3195 NativeColumn::Bytes {
3196 offsets: ob,
3197 values: ub,
3198 validity: vb,
3199 },
3200 ) => {
3201 assert_eq!(oa, ob, "Bytes offsets col {cid}");
3202 assert_eq!(ua, ub, "Bytes values col {cid}");
3203 assert_eq!(va, vb, "Bytes validity col {cid}");
3204 }
3205 _ => panic!("type mismatch col {cid}: {a:?} vs {b:?}"),
3206 }
3207 }
3208 }
3209
3210 #[test]
3211 fn page_cache_key_distinguishes_tables() {
3212 let a = page_cache_key(1, 5, 2, 3);
3213 let b = page_cache_key(2, 5, 2, 3);
3214 assert_ne!(a, b, "same run/col/page, different table must differ");
3215 assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 6, 2, 3));
3217 assert_ne!(page_cache_key(1, 5, 2, 3), page_cache_key(1, 5, 2, 4));
3218 }
3219}