1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::collections::{BTreeMap, HashMap, HashSet};
30use std::path::{Path, PathBuf};
31use std::sync::atomic::AtomicBool;
32use std::sync::Arc;
33use zeroize::Zeroizing;
34
35pub const WAL_DIR: &str = "_wal";
36pub const RUNS_DIR: &str = "_runs";
37pub const CACHE_DIR: &str = "_cache";
38pub const META_DIR: &str = "_meta";
39pub const RCACHE_DIR: &str = "_rcache";
40pub const KEYS_FILENAME: &str = "keys";
41pub const SCHEMA_FILENAME: &str = "schema.json";
42
43fn iso_now_bytes() -> Vec<u8> {
46 let secs = std::time::SystemTime::now()
47 .duration_since(std::time::UNIX_EPOCH)
48 .map(|d| d.as_secs() as i64)
49 .unwrap_or(0);
50 let days = secs.div_euclid(86_400);
51 let rem = secs.rem_euclid(86_400);
52 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
53 let (year, month, day) = civil_from_days(days);
54 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
55}
56
57pub(crate) fn unix_nanos_now() -> i64 {
58 std::time::SystemTime::now()
59 .duration_since(std::time::UNIX_EPOCH)
60 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
61 .unwrap_or(0)
62}
63
64fn civil_from_days(z: i64) -> (i64, u32, u32) {
65 let z = z + 719_468;
66 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
67 let doe = z - era * 146_097;
68 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
69 let y = yoe + era * 400;
70 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
71 let mp = (5 * doy + 2) / 153;
72 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
73 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
74 (if m <= 2 { y + 1 } else { y }, m, d)
75}
76
77const DEFAULT_SYNC_BYTE_THRESHOLD: u64 = 0; pub(crate) const PAGE_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; pub(crate) const DECODED_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; const DEFAULT_MUTABLE_RUN_SPILL_BYTES: u64 = 8 * 1024 * 1024;
84
85#[derive(Clone, Copy, Debug)]
100struct AutoIncState {
101 column_id: u16,
102 next: i64,
103 seeded: bool,
104}
105
106type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
107
108fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
111 schema.auto_increment_column().map(|c| AutoIncState {
112 column_id: c.id,
113 next: 0,
114 seeded: false,
115 })
116}
117
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
131pub enum IndexBuildPolicy {
132 #[default]
134 Deferred,
135 Eager,
137}
138
139pub struct Table {
141 dir: PathBuf,
142 table_id: u64,
143 name: String,
147 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
152 read_only: bool,
155 wal: WalSink,
156 memtable: Memtable,
157 mutable_run: MutableRun,
162 mutable_run_spill_bytes: u64,
164 compaction_zstd_level: i32,
167 allocator: RowIdAllocator,
168 epoch: Arc<EpochAuthority>,
169 persisted_epoch: u64,
172 schema: Schema,
173 hot: HotIndex,
174 kek: Option<Arc<Kek>>,
177 column_keys: HashMap<u16, ([u8; 32], u8)>,
181 run_refs: Vec<RunRef>,
182 retiring: Vec<crate::manifest::RetiredRun>,
185 next_run_id: u64,
186 sync_byte_threshold: u64,
187 current_txn_id: u64,
192 bitmap: HashMap<u16, BitmapIndex>,
193 ann: HashMap<u16, AnnIndex>,
194 fm: HashMap<u16, FmIndex>,
195 sparse: HashMap<u16, SparseIndex>,
196 minhash: HashMap<u16, MinHashIndex>,
197 learned_range: HashMap<u16, ColumnLearnedRange>,
200 pk_by_row: HashMap<RowId, Vec<u8>>,
202 pinned: BTreeMap<Epoch, usize>,
205 pub(crate) live_count: u64,
208 reservoir: crate::reservoir::Reservoir,
211 reservoir_complete: bool,
219 had_deletes: bool,
223 agg_cache: HashMap<u64, CachedAgg>,
227 global_idx_epoch: u64,
231 indexes_complete: bool,
236 index_build_policy: IndexBuildPolicy,
238 pk_by_row_complete: bool,
245 flushed_epoch: u64,
248 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
251 snapshots: Arc<crate::retention::SnapshotRegistry>,
254 commit_lock: Arc<parking_lot::Mutex<()>>,
256 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
259 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
269 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
278 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
280 pending_delete_rids: roaring::RoaringBitmap,
283 pending_put_cols: std::collections::HashSet<u16>,
286 pending_rows: Vec<Row>,
292 pending_rows_auto_inc: Vec<bool>,
293 pending_dels: Vec<RowId>,
296 pending_truncate: Option<Epoch>,
300 auto_inc: Option<AutoIncState>,
303 ttl: Option<TtlPolicy>,
306}
307
308const _: () = {
315 const fn assert_sync<T: ?Sized + Sync>() {}
316 assert_sync::<Table>();
317};
318
319enum CachedData {
325 Rows(Arc<Vec<Row>>),
326 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
327}
328
329impl CachedData {
330 fn approx_bytes(&self) -> u64 {
331 match self {
332 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
333 CachedData::Columns(c) => c
334 .iter()
335 .map(|(_, c)| c.approx_bytes())
336 .sum::<u64>()
337 .saturating_add(c.len() as u64 * 16),
338 }
339 }
340}
341
342struct CachedEntry {
346 data: CachedData,
347 footprint: roaring::RoaringBitmap,
348 condition_cols: Vec<u16>,
349}
350
351struct ResultCache {
362 entries: std::collections::HashMap<u64, CachedEntry>,
363 order: std::collections::VecDeque<u64>,
364 bytes: u64,
365 max_bytes: u64,
366 dir: Option<std::path::PathBuf>,
367 #[allow(dead_code)]
368 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
369}
370
371#[derive(serde::Serialize, serde::Deserialize)]
373struct SerializedEntry {
374 condition_cols: Vec<u16>,
375 footprint_bits: Vec<u32>,
376 data: SerializedData,
377}
378
379#[derive(serde::Serialize, serde::Deserialize)]
380enum SerializedData {
381 Rows(Vec<Row>),
382 Columns(Vec<(u16, columnar::NativeColumn)>),
383}
384
385impl SerializedEntry {
386 fn from_entry(entry: &CachedEntry) -> Self {
387 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
388 let data = match &entry.data {
389 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
390 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
391 };
392 Self {
393 condition_cols: entry.condition_cols.clone(),
394 footprint_bits,
395 data,
396 }
397 }
398
399 fn into_entry(self) -> Option<CachedEntry> {
400 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
401 let data = match self.data {
402 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
403 SerializedData::Columns(c) => {
404 if !c.iter().all(|(_, col)| col.validate()) {
407 return None;
408 }
409 CachedData::Columns(Arc::new(c))
410 }
411 };
412 Some(CachedEntry {
413 data,
414 footprint,
415 condition_cols: self.condition_cols,
416 })
417 }
418}
419
420impl ResultCache {
421 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
422
423 fn new() -> Self {
424 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
425 }
426
427 fn with_max_bytes(max_bytes: u64) -> Self {
428 Self {
429 entries: std::collections::HashMap::new(),
430 order: std::collections::VecDeque::new(),
431 bytes: 0,
432 max_bytes,
433 dir: None,
434 cache_dek: None,
435 }
436 }
437
438 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
439 let _ = std::fs::create_dir_all(&dir);
440 self.dir = Some(dir);
441 self
442 }
443
444 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
445 self.cache_dek = dek;
446 self
447 }
448
449 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
450 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
451 }
452
453 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
457 let Some(path) = self.disk_path(key) else {
458 return;
459 };
460 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
461 Ok(s) => s,
462 Err(_) => return,
463 };
464 let on_disk = if let Some(dek) = &self.cache_dek {
466 match self.encrypt_cache(&serialized, dek) {
467 Some(b) => b,
468 None => return,
469 }
470 } else {
471 serialized
472 };
473 let tmp = path.with_extension("tmp");
474 use std::io::Write;
475 let write = || -> std::io::Result<()> {
476 let mut f = std::fs::File::create(&tmp)?;
477 f.write_all(&on_disk)?;
478 f.flush()?;
479 Ok(())
480 };
481 if write().is_err() {
482 let _ = std::fs::remove_file(&tmp);
483 return;
484 }
485 let _ = std::fs::rename(&tmp, &path);
486 }
487
488 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
490 let path = self.disk_path(key)?;
491 let bytes = std::fs::read(&path).ok()?;
492 let plaintext = if let Some(dek) = &self.cache_dek {
493 self.decrypt_cache(&bytes, dek)?
494 } else {
495 bytes
496 };
497 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
498 serialized.into_entry()
499 }
500
501 fn remove_from_disk(&self, key: u64) {
503 if let Some(path) = self.disk_path(key) {
504 let _ = std::fs::remove_file(&path);
505 }
506 }
507
508 #[cfg(feature = "encryption")]
510 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
511 use crate::encryption::Cipher;
512 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
513 let mut nonce = [0u8; 12];
514 crate::encryption::fill_random(&mut nonce);
515 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
516 let mut out = Vec::with_capacity(12 + ct.len());
517 out.extend_from_slice(&nonce);
518 out.extend_from_slice(&ct);
519 Some(out)
520 }
521
522 #[cfg(not(feature = "encryption"))]
523 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
524 None
525 }
526
527 #[cfg(feature = "encryption")]
529 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
530 use crate::encryption::Cipher;
531 if bytes.len() < 28 {
532 return None;
533 }
534 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
535 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
536 let ct = &bytes[12..];
537 cipher.decrypt_page(&nonce, ct).ok()
538 }
539
540 #[cfg(not(feature = "encryption"))]
541 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
542 None
543 }
544
545 fn load_persistent(&mut self) {
548 let Some(dir) = self.dir.as_ref().cloned() else {
549 return;
550 };
551 let entries = match std::fs::read_dir(&dir) {
552 Ok(e) => e,
553 Err(_) => return,
554 };
555 for entry in entries.flatten() {
556 let path = entry.path();
557 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
559 let _ = std::fs::remove_file(&path);
560 continue;
561 }
562 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
563 continue;
564 }
565 let stem = match path.file_stem().and_then(|s| s.to_str()) {
566 Some(s) => s,
567 None => continue,
568 };
569 let key = match u64::from_str_radix(stem, 16) {
570 Ok(k) => k,
571 Err(_) => continue,
572 };
573 let bytes = match std::fs::read(&path) {
574 Ok(b) => b,
575 Err(_) => continue,
576 };
577 let plaintext = if let Some(dek) = &self.cache_dek {
579 match self.decrypt_cache(&bytes, dek) {
580 Some(p) => p,
581 None => {
582 let _ = std::fs::remove_file(&path);
583 continue;
584 }
585 }
586 } else {
587 bytes
588 };
589 match bincode::deserialize::<SerializedEntry>(&plaintext) {
590 Ok(serialized) => {
591 if let Some(entry) = serialized.into_entry() {
592 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
593 self.entries.insert(key, entry);
594 self.order.push_back(key);
595 } else {
596 let _ = std::fs::remove_file(&path);
597 }
598 }
599 Err(_) => {
600 let _ = std::fs::remove_file(&path);
601 }
602 }
603 }
604 self.evict();
605 }
606
607 fn set_max_bytes(&mut self, max_bytes: u64) {
608 self.max_bytes = max_bytes;
609 self.evict();
610 }
611
612 fn touch(&mut self, key: u64) {
614 self.order.retain(|k| *k != key);
615 self.order.push_back(key);
616 }
617
618 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
619 let res = self.entries.get(&key).and_then(|e| match &e.data {
620 CachedData::Rows(r) => Some(r.clone()),
621 CachedData::Columns(_) => None,
622 });
623 if res.is_some() {
624 self.touch(key);
625 return res;
626 }
627 if let Some(entry) = self.load_from_disk(key) {
629 let res = match &entry.data {
630 CachedData::Rows(r) => Some(r.clone()),
631 CachedData::Columns(_) => None,
632 };
633 if res.is_some() {
634 let approx = entry.data.approx_bytes();
635 self.bytes = self.bytes.saturating_add(approx);
636 self.entries.insert(key, entry);
637 self.order.push_back(key);
638 self.evict();
639 return res;
640 }
641 }
642 None
643 }
644
645 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
646 let res = self.entries.get(&key).and_then(|e| match &e.data {
647 CachedData::Columns(c) => Some(c.clone()),
648 CachedData::Rows(_) => None,
649 });
650 if res.is_some() {
651 self.touch(key);
652 return res;
653 }
654 if let Some(entry) = self.load_from_disk(key) {
656 let res = match &entry.data {
657 CachedData::Columns(c) => Some(c.clone()),
658 CachedData::Rows(_) => None,
659 };
660 if res.is_some() {
661 let approx = entry.data.approx_bytes();
662 self.bytes = self.bytes.saturating_add(approx);
663 self.entries.insert(key, entry);
664 self.order.push_back(key);
665 self.evict();
666 return res;
667 }
668 }
669 None
670 }
671
672 fn insert(&mut self, key: u64, entry: CachedEntry) {
673 let approx = entry.data.approx_bytes();
674 if self.entries.remove(&key).is_some() {
675 self.order.retain(|k| *k != key);
676 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
677 }
678 self.store_to_disk(key, &entry);
680 self.bytes = self.bytes.saturating_add(approx);
681 self.entries.insert(key, entry);
682 self.order.push_back(key);
683 self.evict();
684 }
685
686 fn invalidate(
695 &mut self,
696 delete_rids: &roaring::RoaringBitmap,
697 put_cols: &std::collections::HashSet<u16>,
698 ) {
699 if self.entries.is_empty() {
700 return;
701 }
702 let has_deletes = !delete_rids.is_empty();
703 let to_remove: std::collections::HashSet<u64> = self
704 .entries
705 .iter()
706 .filter(|(_, e)| {
707 let delete_hit = if e.footprint.is_empty() {
708 has_deletes
709 } else {
710 e.footprint.intersection_len(delete_rids) > 0
711 };
712 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
713 delete_hit || col_hit
714 })
715 .map(|(&k, _)| k)
716 .collect();
717 for key in &to_remove {
718 if let Some(e) = self.entries.remove(key) {
719 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
720 }
721 self.remove_from_disk(*key);
722 }
723 if !to_remove.is_empty() {
724 self.order.retain(|k| !to_remove.contains(k));
725 }
726 }
727
728 fn clear(&mut self) {
729 if let Some(dir) = &self.dir {
731 if let Ok(entries) = std::fs::read_dir(dir) {
732 for entry in entries.flatten() {
733 let path = entry.path();
734 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
735 let _ = std::fs::remove_file(&path);
736 }
737 }
738 }
739 }
740 self.entries.clear();
741 self.order.clear();
742 self.bytes = 0;
743 }
744
745 fn evict(&mut self) {
746 while self.bytes > self.max_bytes {
747 let Some(k) = self.order.pop_front() else {
748 break;
749 };
750 if let Some(e) = self.entries.remove(&k) {
751 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
752 self.remove_from_disk(k);
756 }
757 }
758 }
759}
760
761type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
768
769fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
770 let _ = kek;
771 #[cfg(feature = "encryption")]
772 {
773 if let Some(k) = kek {
774 return (
775 Some(k.derive_table_wal_key(_table_id)),
776 Some(k.derive_cache_key()),
777 );
778 }
779 }
780 (None, None)
781}
782
783#[cfg(feature = "encryption")]
785fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
786 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
787}
788
789#[cfg(not(feature = "encryption"))]
790fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
791 Box::new(crate::encryption::PlaintextCipher)
792}
793
794fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
795 let Some(kek) = kek else {
796 return HashMap::new();
797 };
798 #[cfg(feature = "encryption")]
799 {
800 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
801 schema
802 .columns
803 .iter()
804 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
805 .map(|c| {
806 let scheme = if schema
807 .indexes
808 .iter()
809 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
810 {
811 SCHEME_OPE_RANGE
812 } else {
813 SCHEME_HMAC_EQ
814 };
815 let key: [u8; 32] = *kek.derive_column_key(c.id);
816 (c.id, (key, scheme))
817 })
818 .collect()
819 }
820 #[cfg(not(feature = "encryption"))]
821 {
822 let _ = (kek, schema);
823 HashMap::new()
824 }
825}
826
827pub(crate) struct SharedCtx {
832 pub epoch: Arc<EpochAuthority>,
833 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
834 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
835 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
836 pub kek: Option<Arc<Kek>>,
837 pub commit_lock: Arc<parking_lot::Mutex<()>>,
843 pub shared: Option<SharedWalCtx>,
847 pub table_name: Option<String>,
850 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
853 pub read_only: bool,
855}
856
857#[derive(Clone)]
863pub(crate) struct SharedWalCtx {
864 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
865 pub group: Arc<GroupCommit>,
866 pub poisoned: Arc<AtomicBool>,
867 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
868 pub change_wake: tokio::sync::broadcast::Sender<()>,
869}
870
871enum WalSink {
874 Private(Wal),
875 Shared(SharedWalCtx),
876}
877
878impl SharedCtx {
879 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
883 let n_shards = if cache_dir.is_some() {
887 1
888 } else {
889 crate::cache::CACHE_SHARDS
890 };
891 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
892 let page_cache = if let Some(d) = cache_dir {
893 Arc::new(crate::cache::Sharded::new(1, || {
894 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
895 }))
896 } else {
897 Arc::new(crate::cache::Sharded::new(n_shards, || {
898 crate::cache::PageCache::new(per_shard)
899 }))
900 };
901 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
902 let decoded_cache = Arc::new(crate::cache::Sharded::new(
903 crate::cache::CACHE_SHARDS,
904 || crate::cache::DecodedPageCache::new(decoded_per_shard),
905 ));
906 Self {
907 epoch: Arc::new(EpochAuthority::new(0)),
908 page_cache,
909 decoded_cache,
910 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
911 kek,
912 commit_lock: Arc::new(parking_lot::Mutex::new(())),
913 shared: None,
914 table_name: None,
915 auth: None,
916 read_only: false,
917 }
918 }
919}
920
921fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
925 use crate::query::Condition;
926 match c {
927 Condition::Pk(_)
929 | Condition::BitmapEq { .. }
930 | Condition::BitmapIn { .. }
931 | Condition::BytesPrefix { .. }
932 | Condition::IsNull { .. }
933 | Condition::IsNotNull { .. } => 0,
934 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
936 1
937 }
938 Condition::FmContains { .. }
940 | Condition::FmContainsAll { .. }
941 | Condition::Ann { .. }
942 | Condition::SparseMatch { .. } => 2,
943 }
944}
945
946impl Table {
947 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
948 let dir = dir.as_ref().to_path_buf();
949 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
950 Self::create_in(&dir, schema, table_id, ctx)
951 }
952
953 #[cfg(feature = "encryption")]
964 pub fn create_encrypted(
965 dir: impl AsRef<Path>,
966 schema: Schema,
967 table_id: u64,
968 passphrase: &str,
969 ) -> Result<Self> {
970 let dir = dir.as_ref();
971 std::fs::create_dir_all(dir.join(META_DIR))?;
972 let salt = crate::encryption::random_salt();
973 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
974 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
975 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
976 Self::create_in(dir, schema, table_id, ctx)
977 }
978
979 #[cfg(feature = "encryption")]
984 pub fn create_with_key(
985 dir: impl AsRef<Path>,
986 schema: Schema,
987 table_id: u64,
988 key: &[u8],
989 ) -> Result<Self> {
990 let dir = dir.as_ref();
991 std::fs::create_dir_all(dir.join(META_DIR))?;
992 let salt = crate::encryption::random_salt();
993 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
994 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
995 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
996 Self::create_in(dir, schema, table_id, ctx)
997 }
998
999 #[cfg(feature = "encryption")]
1001 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1002 let dir = dir.as_ref();
1003 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1004 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1005 MongrelError::NotFound(format!(
1006 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1007 salt_path
1008 ))
1009 })?;
1010 if salt_bytes.len() != crate::encryption::SALT_LEN {
1011 return Err(MongrelError::InvalidArgument(format!(
1012 "salt file is {} bytes, expected {}",
1013 salt_bytes.len(),
1014 crate::encryption::SALT_LEN
1015 )));
1016 }
1017 let mut salt = [0u8; crate::encryption::SALT_LEN];
1018 salt.copy_from_slice(&salt_bytes);
1019 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1020 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1021 Self::open_in(dir, ctx)
1022 }
1023
1024 pub(crate) fn create_in(
1025 dir: impl AsRef<Path>,
1026 schema: Schema,
1027 table_id: u64,
1028 ctx: SharedCtx,
1029 ) -> Result<Self> {
1030 schema.validate_auto_increment()?;
1031 schema.validate_defaults()?;
1032 let dir = dir.as_ref().to_path_buf();
1033 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1034 write_schema(&dir, &schema)?;
1035 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1036 let (wal, current_txn_id) = match ctx.shared.clone() {
1039 Some(s) => (WalSink::Shared(s), 0),
1040 None => {
1041 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1042 let mut w = if let Some(ref dk) = wal_dek {
1043 Wal::create_with_cipher(
1044 dir.join(WAL_DIR).join("seg-000000.wal"),
1045 Epoch(0),
1046 Some(make_cipher(dk)),
1047 0,
1048 )?
1049 } else {
1050 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1051 };
1052 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1053 (WalSink::Private(w), 1)
1054 }
1055 };
1056 let mut manifest = Manifest::new(table_id, schema.schema_id);
1057 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1062 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1063 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1064 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1065 let auto_inc = resolve_auto_inc(&schema);
1066 let rcache_dir = dir.join(RCACHE_DIR);
1067 Ok(Self {
1068 dir,
1069 table_id,
1070 name: ctx.table_name.unwrap_or_default(),
1071 auth: ctx.auth,
1072 read_only: ctx.read_only,
1073 wal,
1074 memtable: Memtable::new(),
1075 mutable_run: MutableRun::new(),
1076 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1077 compaction_zstd_level: 3,
1078 allocator: RowIdAllocator::new(0),
1079 epoch: ctx.epoch,
1080 persisted_epoch: 0,
1081 schema,
1082 hot: HotIndex::new(),
1083 kek: ctx.kek,
1084 column_keys,
1085 run_refs: Vec::new(),
1086 retiring: Vec::new(),
1087 next_run_id: 1,
1088 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1089 current_txn_id,
1090 bitmap,
1091 ann,
1092 fm,
1093 sparse,
1094 minhash,
1095 learned_range: HashMap::new(),
1096 pk_by_row: HashMap::new(),
1097 pinned: BTreeMap::new(),
1098 live_count: 0,
1099 reservoir: crate::reservoir::Reservoir::default(),
1100 reservoir_complete: true,
1101 had_deletes: false,
1102 agg_cache: HashMap::new(),
1103 global_idx_epoch: 0,
1104 indexes_complete: true,
1105 index_build_policy: IndexBuildPolicy::default(),
1106 pk_by_row_complete: false,
1107 flushed_epoch: 0,
1108 page_cache: ctx.page_cache,
1109 decoded_cache: ctx.decoded_cache,
1110 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1111 snapshots: ctx.snapshots,
1112 commit_lock: ctx.commit_lock,
1113 result_cache: Arc::new(parking_lot::Mutex::new(
1114 ResultCache::new()
1115 .with_dir(rcache_dir)
1116 .with_cache_dek(cache_dek.clone()),
1117 )),
1118 pending_delete_rids: roaring::RoaringBitmap::new(),
1119 pending_put_cols: std::collections::HashSet::new(),
1120 pending_rows: Vec::new(),
1121 pending_rows_auto_inc: Vec::new(),
1122 pending_dels: Vec::new(),
1123 pending_truncate: None,
1124 wal_dek,
1125 auto_inc,
1126 ttl: None,
1127 })
1128 }
1129
1130 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1134 let dir = dir.as_ref();
1135 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1136 Self::open_in(dir, ctx)
1137 }
1138
1139 #[cfg(feature = "encryption")]
1142 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1143 let dir = dir.as_ref();
1144 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1145 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1146 MongrelError::NotFound(format!(
1147 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1148 salt_path
1149 ))
1150 })?;
1151 let salt_len = crate::encryption::SALT_LEN;
1152 if salt_bytes.len() != salt_len {
1153 return Err(MongrelError::InvalidArgument(format!(
1154 "encryption salt is {} bytes, expected {salt_len}",
1155 salt_bytes.len()
1156 )));
1157 }
1158 let mut salt = [0u8; 16];
1159 salt.copy_from_slice(&salt_bytes);
1160 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1161 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1162 let t = Self::open_in(dir, ctx)?;
1163 Ok(t)
1164 }
1165
1166 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1167 let dir = dir.as_ref().to_path_buf();
1168 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1169 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1170 let schema: Schema = read_schema(&dir)?;
1171 let replay_epoch = Epoch(manifest.current_epoch);
1172 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1173 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1177 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1178 None => {
1179 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1180 let replayed = match &active {
1182 Some(path) => {
1183 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1184 crate::wal::replay_with_cipher(path, cipher)?
1185 }
1186 None => Vec::new(),
1187 };
1188 let mut w = match &active {
1189 Some(path) => Wal::create_with_cipher(
1190 path,
1191 replay_epoch,
1192 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1193 0,
1194 )?,
1195 None => Wal::create_with_cipher(
1196 dir.join(WAL_DIR).join("seg-000000.wal"),
1197 replay_epoch,
1198 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1199 0,
1200 )?,
1201 };
1202 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1203 (WalSink::Private(w), replayed, 1)
1204 }
1205 };
1206
1207 let mut memtable = Memtable::new();
1208 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1209 let persisted_epoch = manifest.current_epoch;
1210 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1217 s.next = manifest.auto_inc_next;
1218 s.seeded = manifest.auto_inc_next > 0;
1219 s
1220 });
1221
1222 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1229 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1230 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1231 std::collections::BTreeMap::new();
1232 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1233 let mut saw_delete = false;
1234 for record in replayed {
1235 let txn_id = record.txn_id;
1236 match record.op {
1237 Op::Put { rows, .. } => {
1238 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1239 for row in &rows {
1240 allocator.advance_to(row.row_id);
1241 if let Some(ai) = auto_inc.as_mut() {
1242 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1243 if *n + 1 > ai.next {
1244 ai.next = *n + 1;
1245 }
1246 }
1247 }
1248 }
1249 staged_puts.entry(txn_id).or_default().extend(rows);
1250 }
1251 Op::Delete { row_ids, .. } => {
1252 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1253 }
1254 Op::TxnCommit { epoch, .. } => {
1255 let commit_epoch = Epoch(epoch);
1256 if let Some(puts) = staged_puts.remove(&txn_id) {
1257 for row in &puts {
1258 memtable.upsert(row.clone());
1259 }
1260 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1261 }
1262 if let Some(dels) = staged_deletes.remove(&txn_id) {
1263 saw_delete = true;
1264 for rid in dels {
1265 memtable.tombstone(rid, commit_epoch);
1266 replayed_deletes.push((rid, commit_epoch));
1267 }
1268 }
1269 }
1270 Op::TxnAbort => {
1271 staged_puts.remove(&txn_id);
1272 staged_deletes.remove(&txn_id);
1273 }
1274 Op::TruncateTable { .. }
1275 | Op::ExternalTableState { .. }
1276 | Op::Flush { .. }
1277 | Op::Ddl(_)
1278 | Op::BeforeImage { .. }
1279 | Op::CommitTimestamp { .. } => {}
1280 }
1281 }
1282
1283 let rcache_dir = dir.join(RCACHE_DIR);
1284 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1285 let mut db = Self {
1286 dir,
1287 table_id: manifest.table_id,
1288 name: ctx.table_name.unwrap_or_default(),
1289 auth: ctx.auth,
1290 read_only: ctx.read_only,
1291 wal,
1292 memtable,
1293 mutable_run: MutableRun::new(),
1294 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1295 compaction_zstd_level: 3,
1296 allocator,
1297 epoch: ctx.epoch,
1298 persisted_epoch,
1299 schema,
1300 hot: HotIndex::new(),
1301 kek: ctx.kek,
1302 column_keys,
1303 run_refs: manifest.runs.clone(),
1304 retiring: manifest.retiring.clone(),
1305 next_run_id: manifest
1306 .runs
1307 .iter()
1308 .map(|r| r.run_id as u64 + 1)
1309 .max()
1310 .unwrap_or(1),
1311 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1312 current_txn_id,
1313 bitmap: HashMap::new(),
1314 ann: HashMap::new(),
1315 fm: HashMap::new(),
1316 sparse: HashMap::new(),
1317 minhash: HashMap::new(),
1318 learned_range: HashMap::new(),
1319 pk_by_row: HashMap::new(),
1320 pinned: BTreeMap::new(),
1321 live_count: manifest.live_count,
1322 reservoir: crate::reservoir::Reservoir::default(),
1323 reservoir_complete: false,
1324 had_deletes: saw_delete,
1325 agg_cache: HashMap::new(),
1326 global_idx_epoch: manifest.global_idx_epoch,
1327 indexes_complete: true,
1328 index_build_policy: IndexBuildPolicy::default(),
1329 pk_by_row_complete: false,
1330 flushed_epoch: manifest.flushed_epoch,
1331 page_cache: ctx.page_cache,
1332 decoded_cache: ctx.decoded_cache,
1333 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1334 snapshots: ctx.snapshots,
1335 commit_lock: ctx.commit_lock,
1336 result_cache: Arc::new(parking_lot::Mutex::new(
1337 ResultCache::new()
1338 .with_dir(rcache_dir)
1339 .with_cache_dek(cache_dek.clone()),
1340 )),
1341 pending_delete_rids: roaring::RoaringBitmap::new(),
1342 pending_put_cols: std::collections::HashSet::new(),
1343 pending_rows: Vec::new(),
1344 pending_rows_auto_inc: Vec::new(),
1345 pending_dels: Vec::new(),
1346 pending_truncate: None,
1347 wal_dek,
1348 auto_inc,
1349 ttl: manifest.ttl,
1350 };
1351
1352 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1355
1356 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1361 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1362 c.epoch_built == manifest.global_idx_epoch
1363 && manifest.global_idx_epoch > 0
1364 && manifest
1365 .runs
1366 .iter()
1367 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1368 });
1369 if let Some(loaded) = checkpoint {
1370 if checkpoint_valid {
1371 db.hot = loaded.hot;
1372 db.bitmap = loaded.bitmap;
1373 db.ann = loaded.ann;
1374 db.fm = loaded.fm;
1375 db.sparse = loaded.sparse;
1376 db.minhash = loaded.minhash;
1377 db.learned_range = loaded.learned_range;
1378 }
1381 }
1382 if !checkpoint_valid {
1383 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1384 db.bitmap = bitmap;
1385 db.ann = ann;
1386 db.fm = fm;
1387 db.sparse = sparse;
1388 db.minhash = minhash;
1389 db.rebuild_indexes_from_runs()?;
1390 db.build_learned_ranges()?;
1391 }
1392
1393 for (epoch, group) in replayed_puts {
1398 let (losers, winner_pks) = db.partition_pk_winners(&group);
1399 for (key, &row_id) in &winner_pks {
1400 if let Some(old_rid) = db.hot.get(key) {
1401 if old_rid != row_id {
1402 db.tombstone_row(old_rid, epoch, false);
1403 }
1404 }
1405 }
1406 for &loser_rid in &losers {
1407 db.tombstone_row(loser_rid, epoch, false);
1408 }
1409 for (key, row_id) in winner_pks {
1410 db.insert_hot_pk(key, row_id);
1411 }
1412 if db.schema.primary_key().is_none() {
1413 for r in &group {
1414 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1415 }
1416 }
1417 for r in &group {
1418 if !losers.contains(&r.row_id) {
1419 db.index_row(r);
1420 }
1421 }
1422 }
1423 for (rid, epoch) in &replayed_deletes {
1427 db.remove_hot_for_row(*rid, *epoch);
1428 }
1429
1430 db.result_cache.lock().load_persistent();
1437 Ok(db)
1438 }
1439
1440 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1446 if self.reservoir_complete {
1447 return Ok(());
1448 }
1449 self.rebuild_reservoir()?;
1450 self.reservoir_complete = true;
1451 Ok(())
1452 }
1453
1454 fn rebuild_reservoir(&mut self) -> Result<()> {
1457 let snap = self.snapshot();
1458 let rows = self.visible_rows(snap)?;
1459 self.reservoir.reset();
1460 for r in rows {
1461 self.reservoir.offer(r.row_id.0);
1462 }
1463 Ok(())
1464 }
1465
1466 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1467 self.hot = HotIndex::new();
1468 self.pk_by_row.clear();
1469 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1470 self.bitmap = bitmap;
1471 self.ann = ann;
1472 self.fm = fm;
1473 self.sparse = sparse;
1474 self.minhash = minhash;
1475 let snapshot = Epoch(u64::MAX);
1476 let ttl_now = unix_nanos_now();
1477 for rr in self.run_refs.clone() {
1478 let mut reader = self.open_reader(rr.run_id)?;
1479 for row in reader.visible_rows(snapshot)? {
1480 if self.row_expired_at(&row, ttl_now) {
1481 continue;
1482 }
1483 let tok_row = self.tokenized_for_indexes(&row);
1484 index_into(
1485 &self.schema,
1486 &tok_row,
1487 &mut self.hot,
1488 &mut self.bitmap,
1489 &mut self.ann,
1490 &mut self.fm,
1491 &mut self.sparse,
1492 &mut self.minhash,
1493 );
1494 }
1495 }
1496 for row in self.mutable_run.visible_versions(snapshot) {
1497 if row.deleted {
1498 self.remove_hot_for_row(row.row_id, snapshot);
1499 } else if !self.row_expired_at(&row, ttl_now) {
1500 self.index_row(&row);
1501 }
1502 }
1503 for row in self.memtable.visible_versions(snapshot) {
1504 if row.deleted {
1505 self.remove_hot_for_row(row.row_id, snapshot);
1506 } else if !self.row_expired_at(&row, ttl_now) {
1507 self.index_row(&row);
1508 }
1509 }
1510 self.refresh_pk_by_row_from_hot();
1511 Ok(())
1512 }
1513
1514 fn refresh_pk_by_row_from_hot(&mut self) {
1515 self.pk_by_row_complete = true;
1516 if self.schema.primary_key().is_none() {
1517 self.pk_by_row.clear();
1518 return;
1519 }
1520 self.pk_by_row = self
1526 .hot
1527 .entries()
1528 .into_iter()
1529 .map(|(key, row_id)| (row_id, key))
1530 .collect();
1531 }
1532
1533 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1534 if self.schema.primary_key().is_some() {
1535 self.pk_by_row.insert(row_id, key.clone());
1536 }
1537 self.hot.insert(key, row_id);
1538 }
1539
1540 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1544 self.learned_range.clear();
1545 if self.run_refs.len() != 1 {
1546 return Ok(());
1547 }
1548 let cols: Vec<u16> = self
1549 .schema
1550 .indexes
1551 .iter()
1552 .filter(|i| i.kind == IndexKind::LearnedRange)
1553 .map(|i| i.column_id)
1554 .collect();
1555 if cols.is_empty() {
1556 return Ok(());
1557 }
1558 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1559 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1560 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1561 _ => return Ok(()),
1562 };
1563 for cid in cols {
1564 let ty = self
1565 .schema
1566 .columns
1567 .iter()
1568 .find(|c| c.id == cid)
1569 .map(|c| c.ty.clone())
1570 .unwrap_or(TypeId::Int64);
1571 match ty {
1572 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1573 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1574 let pairs: Vec<(i64, u64)> = data
1575 .iter()
1576 .zip(row_ids.iter())
1577 .map(|(v, r)| (*v, *r))
1578 .collect();
1579 self.learned_range
1580 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1581 }
1582 }
1583 TypeId::Float64 => {
1584 if let columnar::NativeColumn::Float64 { data, .. } =
1585 reader.column_native(cid)?
1586 {
1587 let pairs: Vec<(f64, u64)> = data
1588 .iter()
1589 .zip(row_ids.iter())
1590 .map(|(v, r)| (*v, *r))
1591 .collect();
1592 self.learned_range
1593 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1594 }
1595 }
1596 _ => {}
1597 }
1598 }
1599 Ok(())
1600 }
1601
1602 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1609 if self.indexes_complete {
1610 crate::trace::QueryTrace::record(|t| {
1611 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1612 });
1613 return Ok(());
1614 }
1615 crate::trace::QueryTrace::record(|t| {
1616 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1617 });
1618 self.rebuild_indexes_from_runs()?;
1619 self.build_learned_ranges()?;
1620 self.indexes_complete = true;
1621 let epoch = self.current_epoch();
1622 self.checkpoint_indexes(epoch);
1623 Ok(())
1624 }
1625
1626 fn pending_epoch(&self) -> Epoch {
1627 Epoch(self.epoch.visible().0 + 1)
1628 }
1629
1630 fn is_shared(&self) -> bool {
1633 matches!(self.wal, WalSink::Shared(_))
1634 }
1635
1636 fn ensure_txn_id(&mut self) -> u64 {
1640 if self.current_txn_id == 0 {
1641 let id = match &self.wal {
1642 WalSink::Shared(s) => {
1643 let mut g = s.txn_ids.lock();
1644 let v = *g;
1645 *g = g.wrapping_add(1);
1646 v
1647 }
1648 WalSink::Private(_) => 1,
1649 };
1650 self.current_txn_id = id;
1651 }
1652 self.current_txn_id
1653 }
1654
1655 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1658 self.ensure_writable()?;
1659 let txn_id = self.ensure_txn_id();
1660 let table_id = self.table_id;
1661 match &mut self.wal {
1662 WalSink::Private(w) => {
1663 w.append_txn(txn_id, op)?;
1664 }
1665 WalSink::Shared(s) => {
1666 s.wal.lock().append(txn_id, table_id, op)?;
1667 }
1668 }
1669 Ok(())
1670 }
1671
1672 fn ensure_writable(&self) -> Result<()> {
1673 if self.read_only {
1674 Err(MongrelError::ReadOnlyReplica)
1675 } else {
1676 Ok(())
1677 }
1678 }
1679
1680 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1691 match &self.auth {
1692 Some(checker) => checker.check(&self.name, perm),
1693 None => Ok(()),
1694 }
1695 }
1696 pub fn require_select(&self) -> Result<()> {
1701 self.require(crate::auth_state::RequiredPermission::Select)
1702 }
1703 fn require_insert(&self) -> Result<()> {
1704 self.require(crate::auth_state::RequiredPermission::Insert)
1705 }
1706 #[allow(dead_code)]
1710 fn require_update(&self) -> Result<()> {
1711 self.require(crate::auth_state::RequiredPermission::Update)
1712 }
1713 fn require_delete(&self) -> Result<()> {
1714 self.require(crate::auth_state::RequiredPermission::Delete)
1715 }
1716
1717 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1720 self.require_insert()?;
1721 Ok(self.put_returning(columns)?.0)
1722 }
1723
1724 pub fn put_returning(
1729 &mut self,
1730 mut columns: Vec<(u16, Value)>,
1731 ) -> Result<(RowId, Option<i64>)> {
1732 self.require_insert()?;
1733 let assigned = self.fill_auto_inc(&mut columns)?;
1734 self.apply_defaults(&mut columns)?;
1735 self.schema.validate_not_null(&columns)?;
1736 let row_id = if self.schema.clustered {
1741 self.derive_clustered_row_id(&columns)?
1742 } else {
1743 self.allocator.alloc()
1744 };
1745 let epoch = self.pending_epoch();
1746 let mut row = Row::new(row_id, epoch);
1747 for (col_id, val) in columns {
1748 row.columns.insert(col_id, val);
1749 }
1750 self.commit_rows(vec![row], assigned.is_some())?;
1751 Ok((row_id, assigned))
1752 }
1753
1754 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1757 self.require_insert()?;
1758 Ok(self
1759 .put_batch_returning(batch)?
1760 .into_iter()
1761 .map(|(r, _)| r)
1762 .collect())
1763 }
1764
1765 pub fn put_batch_returning(
1768 &mut self,
1769 batch: Vec<Vec<(u16, Value)>>,
1770 ) -> Result<Vec<(RowId, Option<i64>)>> {
1771 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1772 for mut cols in batch {
1773 let assigned = self.fill_auto_inc(&mut cols)?;
1774 self.apply_defaults(&mut cols)?;
1775 filled.push((cols, assigned));
1776 }
1777 for (cols, _) in &filled {
1778 self.schema.validate_not_null(cols)?;
1779 }
1780 let epoch = self.pending_epoch();
1781 let mut rows = Vec::with_capacity(filled.len());
1782 let mut ids = Vec::with_capacity(filled.len());
1783 for (cols, assigned) in filled {
1784 let row_id = if self.schema.clustered {
1785 self.derive_clustered_row_id(&cols)?
1786 } else {
1787 self.allocator.alloc()
1788 };
1789 let mut row = Row::new(row_id, epoch);
1790 for (c, v) in cols {
1791 row.columns.insert(c, v);
1792 }
1793 ids.push((row_id, assigned));
1794 rows.push(row);
1795 }
1796 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1797 self.commit_rows(rows, all_auto_generated)?;
1798 Ok(ids)
1799 }
1800
1801 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1807 self.ensure_writable()?;
1808 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1809 return Ok(None);
1810 };
1811 let pos = columns.iter().position(|(c, _)| *c == cid);
1812 let assigned = match pos {
1813 Some(i) => match &columns[i].1 {
1814 Value::Null => {
1815 let next = self.alloc_auto_inc_value()?;
1816 columns[i].1 = Value::Int64(next);
1817 Some(next)
1818 }
1819 Value::Int64(n) => {
1820 self.advance_auto_inc_past(*n)?;
1821 None
1822 }
1823 other => {
1824 return Err(MongrelError::InvalidArgument(format!(
1825 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1826 other
1827 )))
1828 }
1829 },
1830 None => {
1831 let next = self.alloc_auto_inc_value()?;
1832 columns.push((cid, Value::Int64(next)));
1833 Some(next)
1834 }
1835 };
1836 Ok(assigned)
1837 }
1838
1839 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1845 for col in &self.schema.columns {
1846 let Some(expr) = &col.default_value else {
1847 continue;
1848 };
1849 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
1851 continue;
1852 }
1853 let pos = columns.iter().position(|(c, _)| *c == col.id);
1854 let needs_default = match pos {
1855 None => true,
1856 Some(i) => matches!(columns[i].1, Value::Null),
1857 };
1858 if !needs_default {
1859 continue;
1860 }
1861 let v = match expr {
1862 crate::schema::DefaultExpr::Static(v) => v.clone(),
1863 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
1864 crate::schema::DefaultExpr::Uuid => {
1865 let mut buf = [0u8; 16];
1866 getrandom::getrandom(&mut buf)
1867 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
1868 Value::Uuid(buf)
1869 }
1870 };
1871 match pos {
1872 None => columns.push((col.id, v)),
1873 Some(i) => columns[i].1 = v,
1874 }
1875 }
1876 Ok(())
1877 }
1878
1879 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1881 self.ensure_auto_inc_seeded()?;
1882 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1884 let v = ai.next;
1885 ai.next = ai.next.saturating_add(1);
1886 Ok(v)
1887 }
1888
1889 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1892 self.ensure_auto_inc_seeded()?;
1893 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1894 let floor = used.saturating_add(1).max(1);
1895 if ai.next < floor {
1896 ai.next = floor;
1897 }
1898 Ok(())
1899 }
1900
1901 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1906 let needs_seed = match self.auto_inc {
1907 Some(ai) => !ai.seeded,
1908 None => return Ok(()),
1909 };
1910 if !needs_seed {
1911 return Ok(());
1912 }
1913 if self.seed_empty_auto_inc() {
1914 return Ok(());
1915 }
1916 let cid = self
1917 .auto_inc
1918 .as_ref()
1919 .expect("auto-inc column present")
1920 .column_id;
1921 let max = self.scan_max_int64(cid)?;
1922 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1923 let floor = max.saturating_add(1).max(1);
1924 if ai.next < floor {
1925 ai.next = floor;
1926 }
1927 ai.seeded = true;
1928 Ok(())
1929 }
1930
1931 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1932 if n == 0 || self.auto_inc.is_none() {
1933 return Ok(None);
1934 }
1935 self.ensure_auto_inc_seeded()?;
1936 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1937 let start = ai.next;
1938 ai.next = ai.next.saturating_add(n as i64);
1939 Ok(Some(start))
1940 }
1941
1942 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1946 let mut max: i64 = 0;
1947 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1948 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1949 if *n > max {
1950 max = *n;
1951 }
1952 }
1953 }
1954 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1955 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1956 if *n > max {
1957 max = *n;
1958 }
1959 }
1960 }
1961 for rr in self.run_refs.clone() {
1962 let reader = self.open_reader(rr.run_id)?;
1963 if let Some(stats) = reader.column_page_stats(column_id) {
1964 for s in stats {
1965 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1966 if n > max {
1967 max = n;
1968 }
1969 }
1970 }
1971 } else if reader.has_column(column_id) {
1972 if let columnar::NativeColumn::Int64 { data, validity } =
1973 reader.column_native_shared(column_id)?
1974 {
1975 for (i, n) in data.iter().enumerate() {
1976 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1977 {
1978 max = *n;
1979 }
1980 }
1981 }
1982 }
1983 }
1984 Ok(max)
1985 }
1986
1987 fn seed_empty_auto_inc(&mut self) -> bool {
1988 let Some(ai) = self.auto_inc.as_mut() else {
1989 return false;
1990 };
1991 if ai.seeded || self.live_count != 0 {
1992 return false;
1993 }
1994 if ai.next < 1 {
1995 ai.next = 1;
1996 }
1997 ai.seeded = true;
1998 true
1999 }
2000
2001 fn advance_auto_inc_from_native_columns(
2002 &mut self,
2003 columns: &[(u16, columnar::NativeColumn)],
2004 n: usize,
2005 live_before: u64,
2006 ) -> Result<()> {
2007 let Some(ai) = self.auto_inc.as_mut() else {
2008 return Ok(());
2009 };
2010 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2011 return Ok(());
2012 };
2013 let columnar::NativeColumn::Int64 { data, validity } = col else {
2014 return Err(MongrelError::InvalidArgument(format!(
2015 "AUTO_INCREMENT column {} must be Int64",
2016 ai.column_id
2017 )));
2018 };
2019 let max = if native_int64_strictly_increasing(col, n) {
2020 data.get(n.saturating_sub(1)).copied()
2021 } else {
2022 data.iter()
2023 .take(n)
2024 .enumerate()
2025 .filter_map(|(i, v)| {
2026 if validity.is_empty() || columnar::validity_bit(validity, i) {
2027 Some(*v)
2028 } else {
2029 None
2030 }
2031 })
2032 .max()
2033 };
2034 if let Some(max) = max {
2035 let floor = max.saturating_add(1).max(1);
2036 if ai.next < floor {
2037 ai.next = floor;
2038 }
2039 if ai.seeded || live_before == 0 {
2040 ai.seeded = true;
2041 }
2042 }
2043 Ok(())
2044 }
2045
2046 fn fill_auto_inc_native_columns(
2047 &mut self,
2048 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2049 n: usize,
2050 ) -> Result<()> {
2051 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2052 return Ok(());
2053 };
2054 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2055 if let Some(start) = self.alloc_auto_inc_range(n)? {
2056 columns.push((
2057 cid,
2058 columnar::NativeColumn::Int64 {
2059 data: (start..start.saturating_add(n as i64)).collect(),
2060 validity: vec![0xFF; n.div_ceil(8)],
2061 },
2062 ));
2063 }
2064 return Ok(());
2065 };
2066
2067 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2068 return Err(MongrelError::InvalidArgument(format!(
2069 "AUTO_INCREMENT column {cid} must be Int64"
2070 )));
2071 };
2072 if data.len() < n {
2073 return Err(MongrelError::InvalidArgument(format!(
2074 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2075 data.len()
2076 )));
2077 }
2078 if columnar::all_non_null(validity, n) {
2079 return Ok(());
2080 }
2081 if validity.iter().all(|b| *b == 0) {
2082 if let Some(start) = self.alloc_auto_inc_range(n)? {
2083 for (i, slot) in data.iter_mut().take(n).enumerate() {
2084 *slot = start.saturating_add(i as i64);
2085 }
2086 *validity = vec![0xFF; n.div_ceil(8)];
2087 }
2088 return Ok(());
2089 }
2090
2091 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2092 for (i, slot) in data.iter_mut().enumerate().take(n) {
2093 if columnar::validity_bit(validity, i) {
2094 self.advance_auto_inc_past(*slot)?;
2095 } else {
2096 *slot = self.alloc_auto_inc_value()?;
2097 }
2098 }
2099 *validity = new_validity;
2100 Ok(())
2101 }
2102
2103 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2117 self.ensure_writable()?;
2118 if self.auto_inc.is_none() {
2119 return Ok(None);
2120 }
2121 Ok(Some(self.alloc_auto_inc_value()?))
2122 }
2123
2124 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2130 let payload = bincode::serialize(&rows)?;
2131 self.wal_append_data(Op::Put {
2132 table_id: self.table_id,
2133 rows: payload,
2134 })?;
2135 if self.is_shared() {
2136 self.pending_rows_auto_inc
2137 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2138 self.pending_rows.extend(rows);
2139 } else {
2140 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2141 }
2142 Ok(())
2143 }
2144
2145 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2150 self.apply_put_rows_inner(rows, true)
2151 }
2152
2153 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2154 if check_existing_pk {
2155 self.ensure_indexes_complete()?;
2156 }
2157 if rows.len() == 1 {
2161 let row = rows.into_iter().next().expect("len checked");
2162 return self.apply_put_row_single(row, check_existing_pk);
2163 }
2164 let pk_id = self.schema.primary_key().map(|c| c.id);
2181 let probe = match pk_id {
2182 Some(pid) => {
2183 check_existing_pk
2184 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2185 }
2186 None => false,
2187 };
2188 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2191 for r in rows {
2192 for &cid in r.columns.keys() {
2193 self.pending_put_cols.insert(cid);
2194 }
2195 match pk_id {
2196 Some(pid) if probe || maintain_pk_by_row => {
2197 if let Some(pk_val) = r.columns.get(&pid) {
2198 let key = self.index_lookup_key(pid, pk_val);
2199 if probe {
2200 if let Some(old_rid) = self.hot.get(&key) {
2201 if old_rid != r.row_id {
2202 self.tombstone_row(old_rid, r.committed_epoch, true);
2203 }
2204 }
2205 }
2206 if maintain_pk_by_row {
2207 self.pk_by_row.insert(r.row_id, key);
2208 }
2209 }
2210 }
2211 Some(_) => {}
2212 None => {
2213 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2214 }
2215 }
2216 self.index_row(&r);
2217 self.reservoir.offer(r.row_id.0);
2218 self.memtable.upsert(r);
2219 self.live_count = self.live_count.saturating_add(1);
2222 }
2223 Ok(())
2224 }
2225
2226 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2230 for &cid in row.columns.keys() {
2231 self.pending_put_cols.insert(cid);
2232 }
2233 let epoch = row.committed_epoch;
2234 if let Some(pk_col) = self.schema.primary_key() {
2235 let pk_id = pk_col.id;
2236 if let Some(pk_val) = row.columns.get(&pk_id) {
2237 let maintain_pk_by_row = self.pk_by_row_complete;
2241 if check_existing_pk || maintain_pk_by_row {
2242 let key = self.index_lookup_key(pk_id, pk_val);
2243 if check_existing_pk {
2244 if let Some(old_rid) = self.hot.get(&key) {
2245 if old_rid != row.row_id {
2246 self.tombstone_row(old_rid, epoch, true);
2247 }
2248 }
2249 }
2250 if maintain_pk_by_row {
2251 self.pk_by_row.insert(row.row_id, key);
2252 }
2253 }
2254 }
2255 } else {
2256 self.hot
2257 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2258 }
2259 self.index_row(&row);
2260 self.reservoir.offer(row.row_id.0);
2261 self.memtable.upsert(row);
2262 self.live_count = self.live_count.saturating_add(1);
2263 Ok(())
2264 }
2265
2266 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2269 self.allocator.alloc()
2270 }
2271
2272 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2278 let pk = self.schema.primary_key().ok_or_else(|| {
2279 MongrelError::Schema("clustered table requires a single-column primary key".into())
2280 })?;
2281 let pk_val = columns
2282 .iter()
2283 .find(|(id, _)| *id == pk.id)
2284 .map(|(_, v)| v)
2285 .ok_or_else(|| {
2286 MongrelError::Schema(format!(
2287 "clustered table missing primary key column {} ({})",
2288 pk.id, pk.name
2289 ))
2290 })?;
2291 let key_bytes = pk_val.encode_key();
2292 let mut hash: u64 = 0xcbf29ce484222325;
2294 for &b in &key_bytes {
2295 hash ^= b as u64;
2296 hash = hash.wrapping_mul(0x100000001b3);
2297 }
2298 Ok(RowId(hash.max(1)))
2301 }
2302
2303 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2311 self.ensure_indexes_complete()?;
2312 let n = rows.len();
2313 for r in rows {
2314 for &cid in r.columns.keys() {
2315 self.pending_put_cols.insert(cid);
2316 }
2317 }
2318 let (losers, winner_pks) = self.partition_pk_winners(rows);
2319 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2320 for (key, &row_id) in &winner_pks {
2322 if let Some(old_rid) = self.hot.get(key) {
2323 if old_rid != row_id {
2324 self.tombstone_row(old_rid, epoch, true);
2325 }
2326 }
2327 }
2328 for &loser_rid in &losers {
2331 self.tombstone_row(loser_rid, epoch, false);
2332 }
2333 for (key, row_id) in winner_pks {
2335 self.insert_hot_pk(key, row_id);
2336 }
2337 if self.schema.primary_key().is_none() {
2338 for r in rows {
2339 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2340 }
2341 }
2342 for r in rows {
2343 self.allocator.advance_to(r.row_id);
2344 if !losers.contains(&r.row_id) {
2345 self.index_row(r);
2346 }
2347 }
2348 for r in rows {
2349 if !losers.contains(&r.row_id) {
2350 self.reservoir.offer(r.row_id.0);
2351 }
2352 }
2353 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2354 Ok(())
2355 }
2356
2357 pub(crate) fn recover_apply(
2362 &mut self,
2363 rows: Vec<Row>,
2364 deletes: Vec<(RowId, Epoch)>,
2365 ) -> Result<()> {
2366 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2370 std::collections::BTreeMap::new();
2371 for row in rows {
2372 self.allocator.advance_to(row.row_id);
2373 if let Some(ai) = self.auto_inc.as_mut() {
2378 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2379 if *n + 1 > ai.next {
2380 ai.next = *n + 1;
2381 }
2382 }
2383 }
2384 by_epoch.entry(row.committed_epoch).or_default().push(row);
2385 }
2386 for (epoch, group) in by_epoch {
2387 let (losers, winner_pks) = self.partition_pk_winners(&group);
2388 for (key, &row_id) in &winner_pks {
2390 if let Some(old_rid) = self.hot.get(key) {
2391 if old_rid != row_id {
2392 self.tombstone_row(old_rid, epoch, false);
2393 }
2394 }
2395 }
2396 for (key, row_id) in winner_pks {
2397 self.insert_hot_pk(key, row_id);
2398 }
2399 if self.schema.primary_key().is_none() {
2400 for r in &group {
2401 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2402 }
2403 }
2404 for r in &group {
2405 if !losers.contains(&r.row_id) {
2406 self.memtable.upsert(r.clone());
2407 self.index_row(r);
2408 }
2409 }
2410 }
2411 for (rid, epoch) in deletes {
2412 self.memtable.tombstone(rid, epoch);
2413 self.remove_hot_for_row(rid, epoch);
2414 }
2415 self.reservoir_complete = false;
2418 Ok(())
2419 }
2420
2421 pub(crate) fn flushed_epoch(&self) -> u64 {
2423 self.flushed_epoch
2424 }
2425
2426 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2427 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2428 }
2429
2430 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2432 self.schema.validate_not_null(cells)
2433 }
2434
2435 fn validate_columns_not_null(
2439 &self,
2440 columns: &[(u16, columnar::NativeColumn)],
2441 n: usize,
2442 ) -> Result<()> {
2443 let by_id: HashMap<u16, &columnar::NativeColumn> =
2444 columns.iter().map(|(id, c)| (*id, c)).collect();
2445 for col in &self.schema.columns {
2446 if !col.flags.contains(ColumnFlags::NULLABLE) {
2447 match by_id.get(&col.id) {
2448 None => {
2449 return Err(MongrelError::InvalidArgument(format!(
2450 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2451 col.name, col.id
2452 )));
2453 }
2454 Some(c) => {
2455 if c.null_count(n) != 0 {
2456 return Err(MongrelError::InvalidArgument(format!(
2457 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2458 col.name, col.id
2459 )));
2460 }
2461 }
2462 }
2463 }
2464 if let TypeId::Enum { variants } = &col.ty {
2465 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2466 if by_id.contains_key(&col.id) {
2467 return Err(MongrelError::InvalidArgument(format!(
2468 "column '{}' ({}) enum requires a bytes column",
2469 col.name, col.id
2470 )));
2471 }
2472 continue;
2473 };
2474 for index in 0..n {
2475 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2476 continue;
2477 };
2478 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2479 return Err(MongrelError::InvalidArgument(format!(
2480 "column '{}' ({}) enum value {:?} is not one of {:?}",
2481 col.name,
2482 col.id,
2483 String::from_utf8_lossy(value),
2484 variants
2485 )));
2486 }
2487 }
2488 }
2489 }
2490 Ok(())
2491 }
2492
2493 fn bulk_pk_winner_indices(
2498 &self,
2499 columns: &[(u16, columnar::NativeColumn)],
2500 n: usize,
2501 ) -> Option<Vec<usize>> {
2502 let pk_col = self.schema.primary_key()?;
2503 let pk_id = pk_col.id;
2504 let pk_ty = pk_col.ty.clone();
2505 let by_id: HashMap<u16, &columnar::NativeColumn> =
2506 columns.iter().map(|(id, c)| (*id, c)).collect();
2507 let pk_native = by_id.get(&pk_id)?;
2508 if native_int64_strictly_increasing(pk_native, n) {
2509 return None;
2510 }
2511 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2513 let mut null_pk_rows: Vec<usize> = Vec::new();
2514 for i in 0..n {
2515 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2516 Some(key) => {
2517 last.insert(key, i);
2518 }
2519 None => null_pk_rows.push(i),
2520 }
2521 }
2522 let mut winners: HashSet<usize> = last.values().copied().collect();
2523 for i in null_pk_rows {
2524 winners.insert(i);
2525 }
2526 Some((0..n).filter(|i| winners.contains(i)).collect())
2527 }
2528
2529 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2531 self.require_delete()?;
2532 let epoch = self.pending_epoch();
2533 self.wal_append_data(Op::Delete {
2534 table_id: self.table_id,
2535 row_ids: vec![row_id],
2536 })?;
2537 if self.is_shared() {
2538 self.pending_dels.push(row_id);
2539 } else {
2540 self.apply_delete(row_id, epoch);
2541 }
2542 Ok(())
2543 }
2544
2545 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2546 let pre = self.get(row_id, self.snapshot());
2547 self.delete(row_id)?;
2548 Ok(pre.map(|row| {
2549 let mut columns: Vec<_> = row.columns.into_iter().collect();
2550 columns.sort_by_key(|(id, _)| *id);
2551 OwnedRow { columns }
2552 }))
2553 }
2554
2555 pub fn truncate(&mut self) -> Result<()> {
2557 self.require_delete()?;
2558 let epoch = self.pending_epoch();
2559 self.wal_append_data(Op::TruncateTable {
2560 table_id: self.table_id,
2561 })?;
2562 self.pending_rows.clear();
2563 self.pending_rows_auto_inc.clear();
2564 self.pending_dels.clear();
2565 self.pending_truncate = Some(epoch);
2566 Ok(())
2567 }
2568
2569 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2571 for rr in std::mem::take(&mut self.run_refs) {
2572 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2573 }
2574 for r in std::mem::take(&mut self.retiring) {
2575 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2576 }
2577 self.memtable = Memtable::new();
2578 self.mutable_run = MutableRun::new();
2579 self.hot = HotIndex::new();
2580 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2581 self.bitmap = bitmap;
2582 self.ann = ann;
2583 self.fm = fm;
2584 self.sparse = sparse;
2585 self.minhash = minhash;
2586 self.learned_range.clear();
2587 self.pk_by_row.clear();
2588 self.pk_by_row_complete = false;
2589 self.live_count = 0;
2590 self.reservoir = crate::reservoir::Reservoir::default();
2591 self.reservoir_complete = true;
2592 self.had_deletes = true;
2593 self.agg_cache.clear();
2594 self.global_idx_epoch = 0;
2595 self.indexes_complete = true;
2596 self.pending_delete_rids.clear();
2597 self.pending_put_cols.clear();
2598 self.pending_rows.clear();
2599 self.pending_rows_auto_inc.clear();
2600 self.pending_dels.clear();
2601 self.clear_result_cache();
2602 self.invalidate_index_checkpoint();
2603 Ok(())
2604 }
2605
2606 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2609 self.remove_hot_for_row(row_id, epoch);
2610 self.tombstone_row(row_id, epoch, true);
2611 }
2612
2613 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2617 let tombstone = Row {
2618 row_id,
2619 committed_epoch: epoch,
2620 columns: std::collections::HashMap::new(),
2621 deleted: true,
2622 };
2623 self.memtable.upsert(tombstone);
2624 self.pk_by_row.remove(&row_id);
2625 if adjust_live_count {
2626 self.live_count = self.live_count.saturating_sub(1);
2627 }
2628 self.pending_delete_rids.insert(row_id.0 as u32);
2630 self.had_deletes = true;
2633 self.agg_cache.clear();
2634 }
2635
2636 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2640 let Some(pk_col) = self.schema.primary_key() else {
2641 return;
2642 };
2643 if self.pk_by_row_complete {
2646 if let Some(key) = self.pk_by_row.remove(&row_id) {
2647 if self.hot.get(&key) == Some(row_id) {
2648 self.hot.remove(&key);
2649 }
2650 }
2651 return;
2652 }
2653 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2672 if self.indexes_complete {
2673 let pk_val = self
2674 .memtable
2675 .get_version(row_id, lookup_epoch)
2676 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2677 .or_else(|| {
2678 self.mutable_run
2679 .get_version(row_id, lookup_epoch)
2680 .filter(|(_, r)| !r.deleted)
2681 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2682 })
2683 .or_else(|| {
2684 self.run_refs.iter().find_map(|rr| {
2685 let mut reader = self.open_reader(rr.run_id).ok()?;
2686 let (_, deleted, val) = reader
2687 .get_version_column(row_id, lookup_epoch, pk_col.id)
2688 .ok()??;
2689 if deleted {
2690 return None;
2691 }
2692 val
2693 })
2694 });
2695 if let Some(pk_val) = pk_val {
2696 let key = self.index_lookup_key(pk_col.id, &pk_val);
2697 if self.hot.get(&key) == Some(row_id) {
2698 self.hot.remove(&key);
2699 }
2700 return;
2701 }
2702 }
2703 self.refresh_pk_by_row_from_hot();
2708 if let Some(key) = self.pk_by_row.remove(&row_id) {
2709 if self.hot.get(&key) == Some(row_id) {
2710 self.hot.remove(&key);
2711 }
2712 }
2713 }
2714
2715 fn partition_pk_winners(
2720 &self,
2721 rows: &[Row],
2722 ) -> (
2723 std::collections::HashSet<RowId>,
2724 std::collections::HashMap<Vec<u8>, RowId>,
2725 ) {
2726 let mut losers = std::collections::HashSet::new();
2727 let Some(pk_col) = self.schema.primary_key() else {
2728 return (losers, std::collections::HashMap::new());
2729 };
2730 let pk_id = pk_col.id;
2731 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2732 std::collections::HashMap::new();
2733 for r in rows {
2734 let Some(pk_val) = r.columns.get(&pk_id) else {
2735 continue;
2736 };
2737 let key = self.index_lookup_key(pk_id, pk_val);
2738 if let Some(&old_rid) = winners.get(&key) {
2739 losers.insert(old_rid);
2740 }
2741 winners.insert(key, r.row_id);
2742 }
2743 (losers, winners)
2744 }
2745
2746 fn index_row(&mut self, row: &Row) {
2747 if row.deleted {
2748 return;
2749 }
2750 let any_predicate = self
2758 .schema
2759 .indexes
2760 .iter()
2761 .any(|idx| idx.predicate.is_some());
2762 if any_predicate {
2763 let columns_map: HashMap<u16, &Value> =
2764 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2765 let name_to_id: HashMap<&str, u16> = self
2766 .schema
2767 .columns
2768 .iter()
2769 .map(|c| (c.name.as_str(), c.id))
2770 .collect();
2771 for idx in &self.schema.indexes {
2772 if let Some(pred) = &idx.predicate {
2773 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2774 continue; }
2776 }
2777 index_into_single(
2779 idx,
2780 &self.schema,
2781 row,
2782 &mut self.hot,
2783 &mut self.bitmap,
2784 &mut self.ann,
2785 &mut self.fm,
2786 &mut self.sparse,
2787 &mut self.minhash,
2788 );
2789 }
2790 return;
2791 }
2792 if self.column_keys.is_empty() {
2796 index_into(
2797 &self.schema,
2798 row,
2799 &mut self.hot,
2800 &mut self.bitmap,
2801 &mut self.ann,
2802 &mut self.fm,
2803 &mut self.sparse,
2804 &mut self.minhash,
2805 );
2806 return;
2807 }
2808 let effective_row = self.tokenized_for_indexes(row);
2809 index_into(
2810 &self.schema,
2811 &effective_row,
2812 &mut self.hot,
2813 &mut self.bitmap,
2814 &mut self.ann,
2815 &mut self.fm,
2816 &mut self.sparse,
2817 &mut self.minhash,
2818 );
2819 }
2820
2821 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2827 if self.column_keys.is_empty() {
2828 return row.clone();
2829 }
2830 #[cfg(feature = "encryption")]
2831 {
2832 use crate::encryption::SCHEME_HMAC_EQ;
2833 let mut tok = row.clone();
2834 for (&cid, &(_, scheme)) in &self.column_keys {
2835 if scheme != SCHEME_HMAC_EQ {
2836 continue;
2837 }
2838 if let Some(v) = tok.columns.get(&cid).cloned() {
2839 if let Some(t) = self.tokenize_value(cid, &v) {
2840 tok.columns.insert(cid, t);
2841 }
2842 }
2843 }
2844 tok
2845 }
2846 #[cfg(not(feature = "encryption"))]
2847 {
2848 row.clone()
2849 }
2850 }
2851
2852 pub fn commit(&mut self) -> Result<Epoch> {
2857 self.ensure_writable()?;
2858 if self.is_shared() {
2859 self.commit_shared()
2860 } else {
2861 self.commit_private()
2862 }
2863 }
2864
2865 fn commit_private(&mut self) -> Result<Epoch> {
2867 let commit_lock = Arc::clone(&self.commit_lock);
2871 let _g = commit_lock.lock();
2872 let new_epoch = self.epoch.bump_assigned();
2873 let txn_id = self.current_txn_id;
2874 match &mut self.wal {
2878 WalSink::Private(w) => {
2879 w.append_txn(
2880 txn_id,
2881 Op::TxnCommit {
2882 epoch: new_epoch.0,
2883 added_runs: Vec::new(),
2884 },
2885 )?;
2886 w.sync()?;
2887 }
2888 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2889 }
2890 if let Some(epoch) = self.pending_truncate.take() {
2892 self.apply_truncate(epoch)?;
2893 }
2894 self.invalidate_pending_cache();
2895 self.persist_manifest(new_epoch)?;
2896 self.epoch.publish_in_order(new_epoch);
2900 self.current_txn_id += 1;
2901 Ok(new_epoch)
2902 }
2903
2904 fn commit_shared(&mut self) -> Result<Epoch> {
2910 use std::sync::atomic::Ordering;
2911 let s = match &self.wal {
2912 WalSink::Shared(s) => s.clone(),
2913 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2914 };
2915 if s.poisoned.load(Ordering::Relaxed) {
2916 return Err(MongrelError::Other(
2917 "database poisoned by fsync error".into(),
2918 ));
2919 }
2920 let commit_lock = Arc::clone(&self.commit_lock);
2927 let _g = commit_lock.lock();
2928 let txn_id = self.ensure_txn_id();
2931 let (new_epoch, commit_seq) = {
2932 let mut wal = s.wal.lock();
2933 let new_epoch = self.epoch.bump_assigned();
2934 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2935 (new_epoch, seq)
2936 };
2937 s.group
2938 .await_durable(&s.wal, commit_seq)
2939 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2940
2941 if self.pending_truncate.take().is_some() {
2944 self.apply_truncate(new_epoch)?;
2945 }
2946 let mut rows = std::mem::take(&mut self.pending_rows);
2947 if !rows.is_empty() {
2948 for r in &mut rows {
2949 r.committed_epoch = new_epoch;
2950 }
2951 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2952 let all_auto_generated =
2953 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2954 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2955 } else {
2956 self.pending_rows_auto_inc.clear();
2957 }
2958 let dels = std::mem::take(&mut self.pending_dels);
2959 for rid in dels {
2960 self.apply_delete(rid, new_epoch);
2961 }
2962
2963 self.invalidate_pending_cache();
2964 self.persist_manifest(new_epoch)?;
2965 self.epoch.publish_in_order(new_epoch);
2966 let _ = s.change_wake.send(());
2967 self.current_txn_id = 0;
2969 Ok(new_epoch)
2970 }
2971
2972 pub fn flush(&mut self) -> Result<Epoch> {
2980 self.ensure_indexes_complete()?;
2981 let epoch = self.commit()?;
2982 let rows = self.memtable.drain_sorted();
2983 if !rows.is_empty() {
2984 self.mutable_run.insert_many(rows);
2985 }
2986 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2987 self.spill_mutable_run(epoch)?;
2988 self.mark_flushed(epoch)?;
2992 self.persist_manifest(epoch)?;
2993 self.build_learned_ranges()?;
2994 self.checkpoint_indexes(epoch);
2997 }
2998 Ok(epoch)
3001 }
3002
3003 pub fn force_flush(&mut self) -> Result<Epoch> {
3012 let saved = self.mutable_run_spill_bytes;
3013 self.mutable_run_spill_bytes = 1;
3014 let result = self.flush();
3015 self.mutable_run_spill_bytes = saved;
3016 result
3017 }
3018
3019 pub fn close(&mut self) -> Result<()> {
3026 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3027 self.force_flush()?;
3028 }
3029 Ok(())
3030 }
3031
3032 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3039 let op = Op::Flush {
3040 table_id: self.table_id,
3041 flushed_epoch: epoch.0,
3042 };
3043 match &mut self.wal {
3044 WalSink::Private(w) => {
3045 w.append_system(op)?;
3046 w.sync()?;
3047 }
3048 WalSink::Shared(s) => {
3049 s.wal.lock().append_system(op)?;
3054 }
3055 }
3056 self.flushed_epoch = epoch.0;
3057 if matches!(self.wal, WalSink::Private(_)) {
3058 self.rotate_wal(epoch)?;
3059 }
3060 Ok(())
3061 }
3062
3063 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3067 let rows = self.mutable_run.drain_sorted();
3068 if rows.is_empty() {
3069 return Ok(());
3070 }
3071 let run_id = self.next_run_id;
3072 self.next_run_id += 1;
3073 let path = self.run_path(run_id);
3074 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3075 if let Some(kek) = &self.kek {
3076 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3077 }
3078 let header = writer.write(&path, &rows)?;
3079 self.run_refs.push(RunRef {
3080 run_id: run_id as u128,
3081 level: 0,
3082 epoch_created: epoch.0,
3083 row_count: header.row_count,
3084 });
3085 Ok(())
3086 }
3087
3088 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3092 self.mutable_run_spill_bytes = bytes.max(1);
3093 }
3094
3095 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3099 self.compaction_zstd_level = level;
3100 }
3101
3102 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3106 self.result_cache.lock().set_max_bytes(max_bytes);
3107 }
3108
3109 pub(crate) fn clear_result_cache(&mut self) {
3113 self.result_cache.lock().clear();
3114 }
3115
3116 pub fn mutable_run_len(&self) -> usize {
3118 self.mutable_run.len()
3119 }
3120
3121 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3124 self.mutable_run.drain_sorted()
3125 }
3126
3127 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3132 let epoch = self.commit()?;
3133 let n = batch.len();
3134 if n == 0 {
3135 return Ok(epoch);
3136 }
3137 let live_before = self.live_count;
3138 self.spill_mutable_run(epoch)?;
3142 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3143 && self.indexes_complete
3144 && self.run_refs.is_empty()
3145 && self.memtable.is_empty()
3146 && self.mutable_run.is_empty();
3147 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3153 use rayon::prelude::*;
3154 self.schema
3155 .columns
3156 .par_iter()
3157 .map(|cdef| {
3158 (
3159 cdef.id,
3160 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3161 )
3162 })
3163 .collect::<Vec<_>>()
3164 };
3165 drop(batch);
3166 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3171 self.validate_columns_not_null(&user_columns, n)?;
3172 let winner_idx = self
3173 .bulk_pk_winner_indices(&user_columns, n)
3174 .filter(|idx| idx.len() != n);
3175 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3176 match winner_idx.as_deref() {
3177 Some(idx) => {
3178 let compacted = user_columns
3179 .iter()
3180 .map(|(id, c)| (*id, c.gather(idx)))
3181 .collect();
3182 (compacted, idx.len())
3183 }
3184 None => (std::mem::take(&mut user_columns), n),
3185 };
3186 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3187 let first = self.allocator.alloc_range(write_n as u64).0;
3188 for rid in first..first + write_n as u64 {
3189 self.reservoir.offer(rid);
3190 }
3191 let run_id = self.next_run_id;
3192 self.next_run_id += 1;
3193 let path = self.run_path(run_id);
3194 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3195 .clean(true)
3196 .with_lz4()
3197 .with_native_endian();
3198 if let Some(kek) = &self.kek {
3199 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3200 }
3201 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3202 self.run_refs.push(RunRef {
3203 run_id: run_id as u128,
3204 level: 0,
3205 epoch_created: epoch.0,
3206 row_count: header.row_count,
3207 });
3208 self.live_count = self.live_count.saturating_add(write_n as u64);
3209 if eager_index_build {
3210 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3211 self.index_columns_bulk(&write_columns, &row_ids);
3212 self.indexes_complete = true;
3213 self.build_learned_ranges()?;
3214 } else {
3215 self.indexes_complete = false;
3216 }
3217 self.mark_flushed(epoch)?;
3218 self.persist_manifest(epoch)?;
3219 if eager_index_build {
3220 self.checkpoint_indexes(epoch);
3221 }
3222 self.clear_result_cache();
3223 Ok(epoch)
3224 }
3225
3226 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3229 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3230 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3231 let segment_no = segment
3234 .file_stem()
3235 .and_then(|s| s.to_str())
3236 .and_then(|s| s.strip_prefix("seg-"))
3237 .and_then(|s| s.parse::<u64>().ok())
3238 .unwrap_or(0);
3239 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3240 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3241 wal.sync()?;
3242 self.wal = WalSink::Private(wal);
3243 Ok(())
3244 }
3245
3246 pub(crate) fn invalidate_pending_cache(&mut self) {
3251 self.result_cache
3252 .lock()
3253 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3254 self.pending_delete_rids.clear();
3255 self.pending_put_cols.clear();
3256 }
3257
3258 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3259 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3260 m.current_epoch = epoch.0;
3261 m.next_row_id = self.allocator.current().0;
3262 m.runs = self.run_refs.clone();
3263 m.live_count = self.live_count;
3264 m.global_idx_epoch = self.global_idx_epoch;
3265 m.flushed_epoch = self.flushed_epoch;
3266 m.retiring = self.retiring.clone();
3267 m.auto_inc_next = match self.auto_inc {
3271 Some(ai) if ai.seeded => ai.next,
3272 _ => 0,
3273 };
3274 m.ttl = self.ttl;
3275 let meta_dek = self.manifest_meta_dek();
3276 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3277 Ok(())
3278 }
3279
3280 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3286 if !self.indexes_complete {
3289 return;
3290 }
3291 let snap = global_idx::IndexSnapshot {
3292 hot: &self.hot,
3293 bitmap: &self.bitmap,
3294 ann: &self.ann,
3295 fm: &self.fm,
3296 sparse: &self.sparse,
3297 minhash: &self.minhash,
3298 learned_range: &self.learned_range,
3299 };
3300 let idx_dek = self.idx_dek();
3302 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3303 .is_ok()
3304 {
3305 self.global_idx_epoch = epoch.0;
3306 let _ = self.persist_manifest(epoch);
3307 }
3308 }
3309
3310 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3313 self.global_idx_epoch = 0;
3314 global_idx::remove(&self.dir);
3315 let _ = self.persist_manifest(self.epoch.visible());
3316 }
3317
3318 pub(crate) fn mark_indexes_incomplete(&mut self) {
3319 self.indexes_complete = false;
3320 self.invalidate_index_checkpoint();
3321 }
3322
3323 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3326 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3327 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3328 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3329 best = Some((epoch, row));
3330 }
3331 }
3332 for rr in &self.run_refs {
3333 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3334 continue;
3335 };
3336 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3337 continue;
3338 };
3339 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3340 best = Some((epoch, row));
3341 }
3342 }
3343 let now_nanos = unix_nanos_now();
3344 match best {
3345 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3346 Some((_, r)) => Some(r),
3347 None => None,
3348 }
3349 }
3350
3351 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3355 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3356 let mut fold = |row: Row| {
3357 best.entry(row.row_id.0)
3358 .and_modify(|e| {
3359 if row.committed_epoch > e.0 {
3360 *e = (row.committed_epoch, row.clone());
3361 }
3362 })
3363 .or_insert_with(|| (row.committed_epoch, row));
3364 };
3365 for row in self.memtable.visible_versions(snapshot.epoch) {
3366 fold(row);
3367 }
3368 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3369 fold(row);
3370 }
3371 for rr in &self.run_refs {
3372 let mut reader = self.open_reader(rr.run_id)?;
3373 for row in reader.visible_versions(snapshot.epoch)? {
3374 fold(row);
3375 }
3376 }
3377 let now_nanos = unix_nanos_now();
3378 let mut out: Vec<Row> = best
3379 .into_values()
3380 .filter_map(|(_, r)| {
3381 if r.deleted || self.row_expired_at(&r, now_nanos) {
3382 None
3383 } else {
3384 Some(r)
3385 }
3386 })
3387 .collect();
3388 out.sort_by_key(|r| r.row_id);
3389 Ok(out)
3390 }
3391
3392 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3399 if self.ttl.is_none()
3400 && self.memtable.is_empty()
3401 && self.mutable_run.is_empty()
3402 && self.run_refs.len() == 1
3403 {
3404 let rr = self.run_refs[0].clone();
3405 let mut reader = self.open_reader(rr.run_id)?;
3406 let idxs = reader.visible_indices(snapshot.epoch)?;
3407 let mut cols = Vec::with_capacity(self.schema.columns.len());
3408 for cdef in &self.schema.columns {
3409 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3410 }
3411 return Ok(cols);
3412 }
3413 let rows = self.visible_rows(snapshot)?;
3415 let mut cols: Vec<(u16, Vec<Value>)> = self
3416 .schema
3417 .columns
3418 .iter()
3419 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3420 .collect();
3421 for r in &rows {
3422 for (cid, vec) in cols.iter_mut() {
3423 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3424 }
3425 }
3426 Ok(cols)
3427 }
3428
3429 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3431 let row_id = self.hot.get(key)?;
3432 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3433 Some(row_id)
3434 } else {
3435 None
3436 }
3437 }
3438
3439 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3444 self.require_select()?;
3445 self.ensure_indexes_complete()?;
3446 let snapshot = self.snapshot();
3447 crate::trace::QueryTrace::record(|t| {
3448 t.run_count = self.run_refs.len();
3449 t.memtable_rows = self.memtable.len();
3450 t.mutable_run_rows = self.mutable_run.len();
3451 });
3452 if q.conditions.is_empty() {
3456 crate::trace::QueryTrace::record(|t| {
3457 t.scan_mode = crate::trace::ScanMode::Materialized;
3458 t.row_materialized = true;
3459 });
3460 return self.visible_rows(snapshot);
3461 }
3462 crate::trace::QueryTrace::record(|t| {
3463 t.conditions_pushed = q.conditions.len();
3464 t.scan_mode = crate::trace::ScanMode::Materialized;
3465 t.row_materialized = true;
3466 });
3467 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3474 ordered.sort_by_key(|c| condition_cost_rank(c));
3475 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3476 for c in &ordered {
3477 let s = self.resolve_condition(c, snapshot)?;
3478 let empty = s.is_empty();
3479 sets.push(s);
3480 if empty {
3481 break;
3482 }
3483 }
3484 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3485 self.rows_for_rids(&rids, snapshot)
3486 }
3487
3488 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3493 use std::collections::HashMap;
3494 let mut rows = Vec::with_capacity(rids.len());
3495 let ttl_now = unix_nanos_now();
3496 let tier_size = self.memtable.len() + self.mutable_run.len();
3513 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3514 if rids.len().saturating_mul(24) < tier_size {
3515 for &rid in rids {
3516 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3517 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3518 let newest = match (mem, mrun) {
3519 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3520 (Some((_, mr)), None) => Some(mr),
3521 (None, Some((_, rr))) => Some(rr),
3522 (None, None) => None,
3523 };
3524 if let Some(row) = newest {
3525 overlay.insert(rid, row);
3526 }
3527 }
3528 } else {
3529 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3530 overlay
3531 .entry(row.row_id.0)
3532 .and_modify(|e| {
3533 if row.committed_epoch > e.committed_epoch {
3534 *e = row.clone();
3535 }
3536 })
3537 .or_insert(row);
3538 };
3539 for row in self.memtable.visible_versions(snapshot.epoch) {
3540 fold_newest(row, &mut overlay);
3541 }
3542 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3543 fold_newest(row, &mut overlay);
3544 }
3545 }
3546 if self.run_refs.len() == 1 {
3547 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3548 if rids.len().saturating_mul(24) < reader.row_count() {
3556 for &rid in rids {
3557 if let Some(r) = overlay.get(&rid) {
3558 if !r.deleted {
3559 rows.push(r.clone());
3560 }
3561 continue;
3562 }
3563 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3564 if !row.deleted {
3565 rows.push(row);
3566 }
3567 }
3568 }
3569 rows.retain(|row| !self.row_expired_at(row, ttl_now));
3570 return Ok(rows);
3571 }
3572 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3581 enum Src {
3584 Overlay,
3585 Run,
3586 }
3587 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3588 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3589 for rid in rids {
3590 if overlay.contains_key(rid) {
3591 plan.push(Src::Overlay);
3592 continue;
3593 }
3594 match vis_rids.binary_search(&(*rid as i64)) {
3595 Ok(i) => {
3596 plan.push(Src::Run);
3597 fetch.push(positions[i]);
3598 }
3599 Err(_) => { }
3600 }
3601 }
3602 let fetched = reader.materialize_batch(&fetch)?;
3603 let mut fetched_iter = fetched.into_iter();
3604 for (rid, src) in rids.iter().zip(plan) {
3605 match src {
3606 Src::Overlay => {
3607 if let Some(r) = overlay.get(rid) {
3608 if !r.deleted {
3609 rows.push(r.clone());
3610 }
3611 }
3612 }
3613 Src::Run => {
3614 if let Some(row) = fetched_iter.next() {
3615 if !row.deleted {
3616 rows.push(row);
3617 }
3618 }
3619 }
3620 }
3621 }
3622 rows.retain(|row| !self.row_expired_at(row, ttl_now));
3623 return Ok(rows);
3624 }
3625 let mut readers: Vec<_> = self
3629 .run_refs
3630 .iter()
3631 .map(|rr| self.open_reader(rr.run_id))
3632 .collect::<Result<Vec<_>>>()?;
3633 for rid in rids {
3634 if let Some(r) = overlay.get(rid) {
3635 if !r.deleted {
3636 rows.push(r.clone());
3637 }
3638 continue;
3639 }
3640 let mut best: Option<(Epoch, Row)> = None;
3641 for reader in readers.iter_mut() {
3642 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3643 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3644 best = Some((epoch, row));
3645 }
3646 }
3647 }
3648 if let Some((_, r)) = best {
3649 if !r.deleted {
3650 rows.push(r);
3651 }
3652 }
3653 }
3654 rows.retain(|row| !self.row_expired_at(row, ttl_now));
3655 Ok(rows)
3656 }
3657
3658 pub fn indexes_complete(&self) -> bool {
3668 self.indexes_complete
3669 }
3670
3671 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3673 self.index_build_policy
3674 }
3675
3676 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3680 self.index_build_policy = policy;
3681 }
3682
3683 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3688 if !self.indexes_complete {
3692 return None;
3693 }
3694 let b = self.bitmap.get(&column_id)?;
3695 let result: Vec<Vec<u8>> = b
3696 .keys()
3697 .into_iter()
3698 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3699 .cloned()
3700 .collect();
3701 Some(result)
3702 }
3703
3704 pub fn fk_join_row_ids(
3705 &self,
3706 fk_column_id: u16,
3707 pk_values: &[Vec<u8>],
3708 fk_conditions: &[crate::query::Condition],
3709 snapshot: Snapshot,
3710 ) -> Result<Vec<u64>> {
3711 let Some(b) = self.bitmap.get(&fk_column_id) else {
3712 return Ok(Vec::new());
3713 };
3714 let mut join_set = {
3715 let mut acc = roaring::RoaringBitmap::new();
3716 for v in pk_values {
3717 acc |= b.get(v);
3718 }
3719 RowIdSet::from_roaring(acc)
3720 };
3721 if !fk_conditions.is_empty() {
3722 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3723 sets.push(join_set);
3724 for c in fk_conditions {
3725 sets.push(self.resolve_condition(c, snapshot)?);
3726 }
3727 join_set = RowIdSet::intersect_many(sets);
3728 }
3729 Ok(join_set.into_sorted_vec())
3730 }
3731
3732 pub fn fk_join_count(
3738 &self,
3739 fk_column_id: u16,
3740 pk_values: &[Vec<u8>],
3741 fk_conditions: &[crate::query::Condition],
3742 snapshot: Snapshot,
3743 ) -> Result<u64> {
3744 let Some(b) = self.bitmap.get(&fk_column_id) else {
3745 return Ok(0);
3746 };
3747 let mut acc = roaring::RoaringBitmap::new();
3748 for v in pk_values {
3749 acc |= b.get(v);
3750 }
3751 if fk_conditions.is_empty() {
3752 return Ok(acc.len());
3753 }
3754 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3755 sets.push(RowIdSet::from_roaring(acc));
3756 for c in fk_conditions {
3757 sets.push(self.resolve_condition(c, snapshot)?);
3758 }
3759 Ok(RowIdSet::intersect_many(sets).len() as u64)
3760 }
3761
3762 fn resolve_condition(
3767 &self,
3768 c: &crate::query::Condition,
3769 snapshot: Snapshot,
3770 ) -> Result<RowIdSet> {
3771 use crate::query::Condition;
3772 Ok(match c {
3773 Condition::Pk(key) => {
3774 let lookup = self
3775 .schema
3776 .primary_key()
3777 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3778 .unwrap_or_else(|| key.clone());
3779 self.hot
3780 .get(&lookup)
3781 .map(|r| RowIdSet::one(r.0))
3782 .unwrap_or_else(RowIdSet::empty)
3783 }
3784 Condition::BitmapEq { column_id, value } => {
3785 let lookup = self.index_lookup_key_bytes(*column_id, value);
3786 self.bitmap
3787 .get(column_id)
3788 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3789 .unwrap_or_else(RowIdSet::empty)
3790 }
3791 Condition::BitmapIn { column_id, values } => {
3792 let bm = self.bitmap.get(column_id);
3793 let mut acc = roaring::RoaringBitmap::new();
3794 if let Some(b) = bm {
3795 for v in values {
3796 let lookup = self.index_lookup_key_bytes(*column_id, v);
3797 acc |= b.get(&lookup);
3798 }
3799 }
3800 RowIdSet::from_roaring(acc)
3801 }
3802 Condition::BytesPrefix { column_id, prefix } => {
3803 if let Some(b) = self.bitmap.get(column_id) {
3808 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3809 let mut acc = roaring::RoaringBitmap::new();
3810 for key in b.keys() {
3811 if key.starts_with(&lookup_prefix) {
3812 acc |= b.get(key);
3813 }
3814 }
3815 RowIdSet::from_roaring(acc)
3816 } else {
3817 RowIdSet::empty()
3818 }
3819 }
3820 Condition::FmContains { column_id, pattern } => self
3821 .fm
3822 .get(column_id)
3823 .map(|f| {
3824 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3825 })
3826 .unwrap_or_else(RowIdSet::empty),
3827 Condition::FmContainsAll {
3828 column_id,
3829 patterns,
3830 } => {
3831 if let Some(f) = self.fm.get(column_id) {
3834 let sets: Vec<RowIdSet> = patterns
3835 .iter()
3836 .map(|pat| {
3837 RowIdSet::from_unsorted(
3838 f.locate(pat).into_iter().map(|r| r.0).collect(),
3839 )
3840 })
3841 .collect();
3842 RowIdSet::intersect_many(sets)
3843 } else {
3844 RowIdSet::empty()
3845 }
3846 }
3847 Condition::Ann {
3848 column_id,
3849 query,
3850 k,
3851 } => self
3852 .ann
3853 .get(column_id)
3854 .map(|a| {
3855 RowIdSet::from_unsorted(
3856 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3857 )
3858 })
3859 .unwrap_or_else(RowIdSet::empty),
3860 Condition::SparseMatch {
3861 column_id,
3862 query,
3863 k,
3864 } => self
3865 .sparse
3866 .get(column_id)
3867 .map(|s| {
3868 RowIdSet::from_unsorted(
3869 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3870 )
3871 })
3872 .unwrap_or_else(RowIdSet::empty),
3873 Condition::MinHashSimilar {
3874 column_id,
3875 query,
3876 k,
3877 } => self
3878 .minhash
3879 .get(column_id)
3880 .map(|mh| {
3881 RowIdSet::from_unsorted(
3882 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3883 )
3884 })
3885 .unwrap_or_else(RowIdSet::empty),
3886 Condition::Range { column_id, lo, hi } => {
3887 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3896 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3897 } else if self.run_refs.len() == 1 {
3898 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3899 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3900 } else {
3901 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3902 };
3903 set.remove_many(self.overlay_rid_set(snapshot));
3904 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3905 set
3906 }
3907 Condition::RangeF64 {
3908 column_id,
3909 lo,
3910 lo_inclusive,
3911 hi,
3912 hi_inclusive,
3913 } => {
3914 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3917 RowIdSet::from_unsorted(
3918 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3919 .into_iter()
3920 .collect(),
3921 )
3922 } else if self.run_refs.len() == 1 {
3923 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3924 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3925 } else {
3926 return self.range_scan_f64(
3927 *column_id,
3928 *lo,
3929 *lo_inclusive,
3930 *hi,
3931 *hi_inclusive,
3932 snapshot,
3933 );
3934 };
3935 set.remove_many(self.overlay_rid_set(snapshot));
3936 self.range_scan_overlay_f64(
3937 &mut set,
3938 *column_id,
3939 *lo,
3940 *lo_inclusive,
3941 *hi,
3942 *hi_inclusive,
3943 snapshot,
3944 );
3945 set
3946 }
3947 Condition::IsNull { column_id } => {
3948 let mut set = if self.run_refs.len() == 1 {
3949 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3950 r.null_row_id_set(*column_id, true)?
3951 } else {
3952 return self.null_scan(*column_id, true, snapshot);
3953 };
3954 set.remove_many(self.overlay_rid_set(snapshot));
3955 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3956 set
3957 }
3958 Condition::IsNotNull { column_id } => {
3959 let mut set = if self.run_refs.len() == 1 {
3960 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3961 r.null_row_id_set(*column_id, false)?
3962 } else {
3963 return self.null_scan(*column_id, false, snapshot);
3964 };
3965 set.remove_many(self.overlay_rid_set(snapshot));
3966 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3967 set
3968 }
3969 })
3970 }
3971
3972 fn range_scan_i64(
3980 &self,
3981 column_id: u16,
3982 lo: i64,
3983 hi: i64,
3984 snapshot: Snapshot,
3985 ) -> Result<RowIdSet> {
3986 let mut row_ids = Vec::new();
3987 let overlay_rids = self.overlay_rid_set(snapshot);
3988 for rr in &self.run_refs {
3989 let mut reader = self.open_reader(rr.run_id)?;
3990 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3991 for rid in matched {
3992 if !overlay_rids.contains(&rid) {
3993 row_ids.push(rid);
3994 }
3995 }
3996 }
3997 let mut s = RowIdSet::from_unsorted(row_ids);
3998 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3999 Ok(s)
4000 }
4001
4002 fn range_scan_f64(
4005 &self,
4006 column_id: u16,
4007 lo: f64,
4008 lo_inclusive: bool,
4009 hi: f64,
4010 hi_inclusive: bool,
4011 snapshot: Snapshot,
4012 ) -> Result<RowIdSet> {
4013 let mut row_ids = Vec::new();
4014 let overlay_rids = self.overlay_rid_set(snapshot);
4015 for rr in &self.run_refs {
4016 let mut reader = self.open_reader(rr.run_id)?;
4017 let matched = reader.range_row_ids_visible_f64(
4018 column_id,
4019 lo,
4020 lo_inclusive,
4021 hi,
4022 hi_inclusive,
4023 snapshot.epoch,
4024 )?;
4025 for rid in matched {
4026 if !overlay_rids.contains(&rid) {
4027 row_ids.push(rid);
4028 }
4029 }
4030 }
4031 let mut s = RowIdSet::from_unsorted(row_ids);
4032 self.range_scan_overlay_f64(
4033 &mut s,
4034 column_id,
4035 lo,
4036 lo_inclusive,
4037 hi,
4038 hi_inclusive,
4039 snapshot,
4040 );
4041 Ok(s)
4042 }
4043
4044 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
4046 let mut s = HashSet::new();
4047 for row in self.memtable.visible_versions(snapshot.epoch) {
4048 s.insert(row.row_id.0);
4049 }
4050 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4051 s.insert(row.row_id.0);
4052 }
4053 s
4054 }
4055
4056 fn range_scan_overlay_i64(
4057 &self,
4058 s: &mut RowIdSet,
4059 column_id: u16,
4060 lo: i64,
4061 hi: i64,
4062 snapshot: Snapshot,
4063 ) {
4064 let mut newest: HashMap<u64, &Row> = HashMap::new();
4069 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4070 let memtable = self.memtable.visible_versions(snapshot.epoch);
4071 for r in &mutable {
4072 newest.entry(r.row_id.0).or_insert(r);
4073 }
4074 for r in &memtable {
4075 newest.insert(r.row_id.0, r);
4076 }
4077 for row in newest.values() {
4078 if !row.deleted {
4079 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
4080 if *v >= lo && *v <= hi {
4081 s.insert(row.row_id.0);
4082 }
4083 }
4084 }
4085 }
4086 }
4087
4088 #[allow(clippy::too_many_arguments)]
4089 fn range_scan_overlay_f64(
4090 &self,
4091 s: &mut RowIdSet,
4092 column_id: u16,
4093 lo: f64,
4094 lo_inclusive: bool,
4095 hi: f64,
4096 hi_inclusive: bool,
4097 snapshot: Snapshot,
4098 ) {
4099 let mut newest: HashMap<u64, &Row> = HashMap::new();
4102 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4103 let memtable = self.memtable.visible_versions(snapshot.epoch);
4104 for r in &mutable {
4105 newest.entry(r.row_id.0).or_insert(r);
4106 }
4107 for r in &memtable {
4108 newest.insert(r.row_id.0, r);
4109 }
4110 for row in newest.values() {
4111 if !row.deleted {
4112 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
4113 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
4114 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
4115 if ok_lo && ok_hi {
4116 s.insert(row.row_id.0);
4117 }
4118 }
4119 }
4120 }
4121 }
4122
4123 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
4126 let mut row_ids = Vec::new();
4127 let overlay_rids = self.overlay_rid_set(snapshot);
4128 for rr in &self.run_refs {
4129 let mut reader = self.open_reader(rr.run_id)?;
4130 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
4131 for rid in matched {
4132 if !overlay_rids.contains(&rid) {
4133 row_ids.push(rid);
4134 }
4135 }
4136 }
4137 let mut s = RowIdSet::from_unsorted(row_ids);
4138 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
4139 Ok(s)
4140 }
4141
4142 fn null_scan_overlay(
4146 &self,
4147 s: &mut RowIdSet,
4148 column_id: u16,
4149 want_nulls: bool,
4150 snapshot: Snapshot,
4151 ) {
4152 let mut newest: HashMap<u64, &Row> = HashMap::new();
4153 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4154 let memtable = self.memtable.visible_versions(snapshot.epoch);
4155 for r in &mutable {
4156 newest.entry(r.row_id.0).or_insert(r);
4157 }
4158 for r in &memtable {
4159 newest.insert(r.row_id.0, r);
4160 }
4161 for row in newest.values() {
4162 if row.deleted {
4163 continue;
4164 }
4165 let is_null = !row.columns.contains_key(&column_id)
4166 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
4167 if is_null == want_nulls {
4168 s.insert(row.row_id.0);
4169 }
4170 }
4171 }
4172
4173 pub fn snapshot(&self) -> Snapshot {
4174 Snapshot::at(self.epoch.visible())
4175 }
4176
4177 pub fn pin_snapshot(&mut self) -> Snapshot {
4180 let e = self.epoch.visible();
4181 *self.pinned.entry(e).or_insert(0) += 1;
4182 Snapshot::at(e)
4183 }
4184
4185 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
4187 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
4188 *count -= 1;
4189 if *count == 0 {
4190 self.pinned.remove(&snap.epoch);
4191 }
4192 }
4193 }
4194
4195 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
4205 let local = self.pinned.keys().next().copied();
4206 let global = self.snapshots.min_pinned();
4207 let history = self.snapshots.history_floor(self.current_epoch());
4208 [local, global, history].into_iter().flatten().min()
4209 }
4210
4211 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
4215 self.ensure_writable()?;
4216 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
4217 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
4218 }
4219
4220 pub fn clear_ttl(&mut self) -> Result<()> {
4221 self.ensure_writable()?;
4222 self.apply_ttl_policy_at(None, self.current_epoch())
4223 }
4224
4225 pub fn ttl(&self) -> Option<TtlPolicy> {
4226 self.ttl
4227 }
4228
4229 pub(crate) fn prepare_ttl_policy(
4230 &self,
4231 column_name: &str,
4232 duration_nanos: u64,
4233 ) -> Result<TtlPolicy> {
4234 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
4235 return Err(MongrelError::InvalidArgument(
4236 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
4237 ));
4238 }
4239 let column = self
4240 .schema
4241 .columns
4242 .iter()
4243 .find(|column| column.name == column_name)
4244 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
4245 if column.ty != TypeId::TimestampNanos {
4246 return Err(MongrelError::Schema(format!(
4247 "TTL column {column_name} must be TimestampNanos, is {:?}",
4248 column.ty
4249 )));
4250 }
4251 Ok(TtlPolicy {
4252 column_id: column.id,
4253 duration_nanos,
4254 })
4255 }
4256
4257 pub(crate) fn apply_ttl_policy_at(
4258 &mut self,
4259 policy: Option<TtlPolicy>,
4260 epoch: Epoch,
4261 ) -> Result<()> {
4262 if let Some(policy) = policy {
4263 let column = self
4264 .schema
4265 .columns
4266 .iter()
4267 .find(|column| column.id == policy.column_id)
4268 .ok_or_else(|| {
4269 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
4270 })?;
4271 if column.ty != TypeId::TimestampNanos
4272 || policy.duration_nanos == 0
4273 || policy.duration_nanos > i64::MAX as u64
4274 {
4275 return Err(MongrelError::Schema("invalid TTL policy".into()));
4276 }
4277 }
4278 self.ttl = policy;
4279 self.agg_cache.clear();
4280 self.clear_result_cache();
4281 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
4282 self.persist_manifest(epoch)
4283 }
4284
4285 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
4286 let Some(policy) = self.ttl else {
4287 return false;
4288 };
4289 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
4290 return false;
4291 };
4292 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
4293 }
4294
4295 pub fn current_epoch(&self) -> Epoch {
4296 self.epoch.visible()
4297 }
4298
4299 pub fn memtable_len(&self) -> usize {
4300 self.memtable.len()
4301 }
4302
4303 pub fn count(&self) -> u64 {
4306 if self.ttl.is_none() {
4307 self.live_count
4308 } else {
4309 self.visible_rows(self.snapshot())
4310 .map(|rows| rows.len() as u64)
4311 .unwrap_or(self.live_count)
4312 }
4313 }
4314
4315 pub fn count_conditions(
4319 &mut self,
4320 conditions: &[crate::query::Condition],
4321 snapshot: Snapshot,
4322 ) -> Result<Option<u64>> {
4323 use crate::query::Condition;
4324 if self.ttl.is_some() {
4325 return self
4326 .query(&crate::query::Query {
4327 conditions: conditions.to_vec(),
4328 })
4329 .map(|rows| Some(rows.len() as u64));
4330 }
4331 if conditions.is_empty() {
4332 return Ok(Some(self.live_count));
4333 }
4334 let served = |c: &Condition| {
4335 matches!(
4336 c,
4337 Condition::Pk(_)
4338 | Condition::BitmapEq { .. }
4339 | Condition::BitmapIn { .. }
4340 | Condition::BytesPrefix { .. }
4341 | Condition::FmContains { .. }
4342 | Condition::FmContainsAll { .. }
4343 | Condition::Ann { .. }
4344 | Condition::Range { .. }
4345 | Condition::RangeF64 { .. }
4346 | Condition::SparseMatch { .. }
4347 | Condition::MinHashSimilar { .. }
4348 | Condition::IsNull { .. }
4349 | Condition::IsNotNull { .. }
4350 )
4351 };
4352 if !conditions.iter().all(served) {
4353 return Ok(None);
4354 }
4355 self.ensure_indexes_complete()?;
4356 let mut sets = Vec::with_capacity(conditions.len());
4357 for condition in conditions {
4358 sets.push(self.resolve_condition(condition, snapshot)?);
4359 }
4360 let mut rids = RowIdSet::intersect_many(sets);
4361 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4371 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
4372 }
4373 let count = rids.len() as u64;
4374 crate::trace::QueryTrace::record(|t| {
4375 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
4376 t.survivor_count = Some(count as usize);
4377 t.conditions_pushed = conditions.len();
4378 });
4379 Ok(Some(count))
4380 }
4381
4382 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
4387 let mut out = Vec::new();
4388 for row in self.memtable.visible_versions(snapshot.epoch) {
4389 if row.deleted {
4390 out.push(row.row_id.0);
4391 }
4392 }
4393 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4394 if row.deleted {
4395 out.push(row.row_id.0);
4396 }
4397 }
4398 out
4399 }
4400
4401 pub fn bulk_load_columns(
4410 &mut self,
4411 user_columns: Vec<(u16, columnar::NativeColumn)>,
4412 ) -> Result<Epoch> {
4413 self.bulk_load_columns_with(user_columns, 3, false, true)
4414 }
4415
4416 pub fn bulk_load_fast(
4423 &mut self,
4424 user_columns: Vec<(u16, columnar::NativeColumn)>,
4425 ) -> Result<Epoch> {
4426 self.bulk_load_columns_with(user_columns, -1, true, false)
4427 }
4428
4429 fn bulk_load_columns_with(
4430 &mut self,
4431 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4432 zstd_level: i32,
4433 force_plain: bool,
4434 lz4: bool,
4435 ) -> Result<Epoch> {
4436 let epoch = self.commit()?;
4437 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4438 if n == 0 {
4439 return Ok(epoch);
4440 }
4441 let live_before = self.live_count;
4442 self.spill_mutable_run(epoch)?;
4444 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4445 && self.indexes_complete
4446 && self.run_refs.is_empty()
4447 && self.memtable.is_empty()
4448 && self.mutable_run.is_empty();
4449 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4452 self.validate_columns_not_null(&user_columns, n)?;
4453 let winner_idx = self
4454 .bulk_pk_winner_indices(&user_columns, n)
4455 .filter(|idx| idx.len() != n);
4456 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4457 match winner_idx.as_deref() {
4458 Some(idx) => {
4459 let compacted = user_columns
4460 .iter()
4461 .map(|(id, c)| (*id, c.gather(idx)))
4462 .collect();
4463 (compacted, idx.len())
4464 }
4465 None => (user_columns, n),
4466 };
4467 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4468 let first = self.allocator.alloc_range(write_n as u64).0;
4469 for rid in first..first + write_n as u64 {
4470 self.reservoir.offer(rid);
4471 }
4472 let run_id = self.next_run_id;
4473 self.next_run_id += 1;
4474 let path = self.run_path(run_id);
4475 let mut writer =
4476 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4477 if force_plain {
4478 writer = writer.with_plain();
4479 } else if lz4 {
4480 writer = writer.with_lz4();
4483 } else {
4484 writer = writer.with_zstd_level(zstd_level);
4485 }
4486 if let Some(kek) = &self.kek {
4487 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4488 }
4489 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4490 self.run_refs.push(RunRef {
4491 run_id: run_id as u128,
4492 level: 0,
4493 epoch_created: epoch.0,
4494 row_count: header.row_count,
4495 });
4496 self.live_count = self.live_count.saturating_add(write_n as u64);
4497 if eager_index_build {
4498 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4499 self.index_columns_bulk(&write_columns, &row_ids);
4500 self.indexes_complete = true;
4501 self.build_learned_ranges()?;
4502 } else {
4503 self.indexes_complete = false;
4507 }
4508 self.mark_flushed(epoch)?;
4509 self.persist_manifest(epoch)?;
4510 if eager_index_build {
4511 self.checkpoint_indexes(epoch);
4512 }
4513 self.clear_result_cache();
4514 Ok(epoch)
4515 }
4516
4517 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4535 let n = row_ids.len();
4536 if n == 0 {
4537 return;
4538 }
4539 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4540 columns.iter().map(|(id, c)| (*id, c)).collect();
4541 let ty_of: std::collections::HashMap<u16, TypeId> = self
4542 .schema
4543 .columns
4544 .iter()
4545 .map(|c| (c.id, c.ty.clone()))
4546 .collect();
4547 let pk_id = self.schema.primary_key().map(|c| c.id);
4548
4549 for (i, &rid) in row_ids.iter().enumerate() {
4550 let row_id = RowId(rid);
4551 if let Some(pid) = pk_id {
4552 if let Some(col) = by_id.get(&pid) {
4553 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
4554 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4555 self.insert_hot_pk(key, row_id);
4556 }
4557 }
4558 }
4559 for idef in &self.schema.indexes {
4560 let Some(col) = by_id.get(&idef.column_id) else {
4561 continue;
4562 };
4563 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
4564 match idef.kind {
4565 IndexKind::Bitmap => {
4566 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4567 if let Some(key) =
4568 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4569 {
4570 b.insert(key, row_id);
4571 }
4572 }
4573 }
4574 IndexKind::FmIndex => {
4575 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4576 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4577 f.insert(bytes.to_vec(), row_id);
4578 }
4579 }
4580 }
4581 IndexKind::Sparse => {
4582 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4583 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4584 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4585 s.insert(&terms, row_id);
4586 }
4587 }
4588 }
4589 }
4590 IndexKind::MinHash => {
4591 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4592 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4593 let tokens = crate::index::token_hashes_from_bytes(bytes);
4594 mh.insert(&tokens, row_id);
4595 }
4596 }
4597 }
4598 _ => {}
4599 }
4600 }
4601 }
4602 }
4603
4604 pub fn visible_columns_native(
4609 &self,
4610 snapshot: Snapshot,
4611 projection: Option<&[u16]>,
4612 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4613 let wanted: Vec<u16> = match projection {
4614 Some(p) => p.to_vec(),
4615 None => self.schema.columns.iter().map(|c| c.id).collect(),
4616 };
4617 if self.ttl.is_none()
4618 && self.memtable.is_empty()
4619 && self.mutable_run.is_empty()
4620 && self.run_refs.len() == 1
4621 {
4622 let rr = self.run_refs[0].clone();
4623 let mut reader = self.open_reader(rr.run_id)?;
4624 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4625 let all_visible = idxs.len() == reader.row_count();
4626 if reader.has_mmap() {
4632 use rayon::prelude::*;
4633 let valid: Vec<u16> = wanted
4636 .iter()
4637 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4638 .copied()
4639 .collect();
4640 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4642 .par_iter()
4643 .filter_map(|cid| {
4644 reader
4645 .column_native_shared(*cid)
4646 .ok()
4647 .map(|col| (*cid, col))
4648 })
4649 .collect();
4650 let cols = decoded
4651 .into_iter()
4652 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4653 .collect();
4654 return Ok(cols);
4655 }
4656 let mut cols = Vec::with_capacity(wanted.len());
4657 for cid in &wanted {
4658 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4659 Some(c) => c,
4660 None => continue,
4661 };
4662 let col = reader.column_native(cdef.id)?;
4663 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4664 }
4665 return Ok(cols);
4666 }
4667 let vcols = self.visible_columns(snapshot)?;
4668 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4669 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4670 .into_iter()
4671 .filter(|(id, _)| want_set.contains(id))
4672 .map(|(id, vals)| {
4673 let ty = self
4674 .schema
4675 .columns
4676 .iter()
4677 .find(|c| c.id == id)
4678 .map(|c| c.ty.clone())
4679 .unwrap_or(TypeId::Bytes);
4680 (id, columnar::values_to_native(ty, &vals))
4681 })
4682 .collect();
4683 Ok(out)
4684 }
4685
4686 pub fn run_count(&self) -> usize {
4687 self.run_refs.len()
4688 }
4689
4690 pub fn memtable_is_empty(&self) -> bool {
4692 self.memtable.is_empty()
4693 }
4694
4695 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4699 self.page_cache.stats()
4700 }
4701
4702 pub fn reset_page_cache_stats(&self) {
4704 self.page_cache.reset_stats();
4705 }
4706
4707 pub fn run_ids(&self) -> Vec<u128> {
4710 self.run_refs.iter().map(|r| r.run_id).collect()
4711 }
4712
4713 pub fn single_run_is_clean(&self) -> bool {
4717 if self.ttl.is_some() || self.run_refs.len() != 1 {
4718 return false;
4719 }
4720 self.open_reader(self.run_refs[0].run_id)
4721 .map(|r| r.is_clean())
4722 .unwrap_or(false)
4723 }
4724
4725 fn resolve_footprint(
4732 &self,
4733 conditions: &[crate::query::Condition],
4734 _snapshot: Snapshot,
4735 ) -> roaring::RoaringBitmap {
4736 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4737 return roaring::RoaringBitmap::new();
4738 }
4739 if self.run_refs.is_empty() {
4740 return roaring::RoaringBitmap::new();
4741 }
4742 if self.run_refs.len() == 1 {
4744 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4745 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4746 return rids.to_roaring_lossy();
4747 }
4748 }
4749 }
4750 roaring::RoaringBitmap::new()
4751 }
4752
4753 pub fn query_columns_native_cached(
4764 &mut self,
4765 conditions: &[crate::query::Condition],
4766 projection: Option<&[u16]>,
4767 snapshot: Snapshot,
4768 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4769 if self.ttl.is_some() {
4772 return self.query_columns_native(conditions, projection, snapshot);
4773 }
4774 if conditions.is_empty() {
4775 return self.query_columns_native(conditions, projection, snapshot);
4776 }
4777 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4781 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4782 crate::trace::QueryTrace::record(|t| {
4783 t.result_cache_hit = true;
4784 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4785 });
4786 return Ok(Some((*hit).clone()));
4787 }
4788 let res = self.query_columns_native(conditions, projection, snapshot)?;
4789 if let Some(cols) = &res {
4790 let footprint = self.resolve_footprint(conditions, snapshot);
4791 let condition_cols = crate::query::condition_columns(conditions);
4792 self.result_cache.lock().insert(
4793 key,
4794 CachedEntry {
4795 data: CachedData::Columns(Arc::new(cols.clone())),
4796 footprint,
4797 condition_cols,
4798 },
4799 );
4800 }
4801 Ok(res)
4802 }
4803
4804 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4809 if self.ttl.is_some() {
4810 return self.query(q);
4811 }
4812 if q.conditions.is_empty() {
4813 return self.query(q);
4814 }
4815 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4816 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4817 crate::trace::QueryTrace::record(|t| {
4818 t.result_cache_hit = true;
4819 t.scan_mode = crate::trace::ScanMode::Materialized;
4820 });
4821 return Ok((*hit).clone());
4822 }
4823 let rows = self.query(q)?;
4824 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4825 let condition_cols = crate::query::condition_columns(&q.conditions);
4826 self.result_cache.lock().insert(
4827 key,
4828 CachedEntry {
4829 data: CachedData::Rows(Arc::new(rows.clone())),
4830 footprint,
4831 condition_cols,
4832 },
4833 );
4834 Ok(rows)
4835 }
4836
4837 #[allow(clippy::type_complexity)]
4852 pub fn query_columns_native_traced(
4853 &mut self,
4854 conditions: &[crate::query::Condition],
4855 projection: Option<&[u16]>,
4856 snapshot: Snapshot,
4857 ) -> Result<(
4858 Option<Vec<(u16, columnar::NativeColumn)>>,
4859 crate::trace::QueryTrace,
4860 )> {
4861 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4862 self.query_columns_native(conditions, projection, snapshot)
4863 });
4864 Ok((result?, trace))
4865 }
4866
4867 #[allow(clippy::type_complexity)]
4870 pub fn query_columns_native_cached_traced(
4871 &mut self,
4872 conditions: &[crate::query::Condition],
4873 projection: Option<&[u16]>,
4874 snapshot: Snapshot,
4875 ) -> Result<(
4876 Option<Vec<(u16, columnar::NativeColumn)>>,
4877 crate::trace::QueryTrace,
4878 )> {
4879 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4880 self.query_columns_native_cached(conditions, projection, snapshot)
4881 });
4882 Ok((result?, trace))
4883 }
4884
4885 pub fn native_page_cursor_traced(
4887 &self,
4888 snapshot: Snapshot,
4889 projection: Vec<(u16, TypeId)>,
4890 conditions: &[crate::query::Condition],
4891 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4892 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4893 self.native_page_cursor(snapshot, projection, conditions)
4894 });
4895 Ok((result?, trace))
4896 }
4897
4898 pub fn native_multi_run_cursor_traced(
4900 &self,
4901 snapshot: Snapshot,
4902 projection: Vec<(u16, TypeId)>,
4903 conditions: &[crate::query::Condition],
4904 ) -> Result<(
4905 Option<crate::cursor::MultiRunCursor>,
4906 crate::trace::QueryTrace,
4907 )> {
4908 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4909 self.native_multi_run_cursor(snapshot, projection, conditions)
4910 });
4911 Ok((result?, trace))
4912 }
4913
4914 pub fn count_conditions_traced(
4916 &mut self,
4917 conditions: &[crate::query::Condition],
4918 snapshot: Snapshot,
4919 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4920 let (result, trace) =
4921 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4922 Ok((result?, trace))
4923 }
4924
4925 pub fn query_traced(
4927 &mut self,
4928 q: &crate::query::Query,
4929 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4930 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4931 Ok((result?, trace))
4932 }
4933
4934 pub fn query_columns_native(
4939 &mut self,
4940 conditions: &[crate::query::Condition],
4941 projection: Option<&[u16]>,
4942 snapshot: Snapshot,
4943 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4944 use crate::query::Condition;
4945 if self.ttl.is_some() {
4948 return Ok(None);
4949 }
4950 if conditions.is_empty() {
4951 return Ok(None);
4952 }
4953 self.ensure_indexes_complete()?;
4954
4955 let served = |c: &Condition| {
4960 matches!(
4961 c,
4962 Condition::Pk(_)
4963 | Condition::BitmapEq { .. }
4964 | Condition::BitmapIn { .. }
4965 | Condition::BytesPrefix { .. }
4966 | Condition::FmContains { .. }
4967 | Condition::FmContainsAll { .. }
4968 | Condition::Ann { .. }
4969 | Condition::Range { .. }
4970 | Condition::RangeF64 { .. }
4971 | Condition::SparseMatch { .. }
4972 | Condition::MinHashSimilar { .. }
4973 | Condition::IsNull { .. }
4974 | Condition::IsNotNull { .. }
4975 )
4976 };
4977 if !conditions.iter().all(served) {
4978 return Ok(None);
4979 }
4980 let fast_path =
4981 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4982 crate::trace::QueryTrace::record(|t| {
4983 t.run_count = self.run_refs.len();
4984 t.memtable_rows = self.memtable.len();
4985 t.mutable_run_rows = self.mutable_run.len();
4986 t.conditions_pushed = conditions.len();
4987 t.learned_range_used = conditions.iter().any(|c| match c {
4988 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4989 self.learned_range.contains_key(column_id)
4990 }
4991 _ => false,
4992 });
4993 });
4994 let col_ids: Vec<u16> = projection
4996 .map(|p| p.to_vec())
4997 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4998 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4999 .iter()
5000 .map(|&cid| {
5001 let ty = self
5002 .schema
5003 .columns
5004 .iter()
5005 .find(|c| c.id == cid)
5006 .map(|c| c.ty.clone())
5007 .unwrap_or(TypeId::Bytes);
5008 (cid, ty)
5009 })
5010 .collect();
5011
5012 if fast_path {
5018 let needs_column = conditions.iter().any(|c| match c {
5021 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
5022 Condition::RangeF64 { column_id, .. } => {
5023 !self.learned_range.contains_key(column_id)
5024 }
5025 _ => false,
5026 });
5027 let mut reader_opt: Option<RunReader> = if needs_column {
5028 Some(self.open_reader(self.run_refs[0].run_id)?)
5029 } else {
5030 None
5031 };
5032 let mut sets: Vec<RowIdSet> = Vec::new();
5033 for c in conditions {
5034 let s = match c {
5035 Condition::Range { column_id, lo, hi }
5036 if !self.learned_range.contains_key(column_id) =>
5037 {
5038 if reader_opt.is_none() {
5039 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
5040 }
5041 reader_opt
5042 .as_mut()
5043 .expect("reader opened for range")
5044 .range_row_id_set_i64(*column_id, *lo, *hi)?
5045 }
5046 Condition::RangeF64 {
5047 column_id,
5048 lo,
5049 lo_inclusive,
5050 hi,
5051 hi_inclusive,
5052 } if !self.learned_range.contains_key(column_id) => {
5053 if reader_opt.is_none() {
5054 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
5055 }
5056 reader_opt
5057 .as_mut()
5058 .expect("reader opened for range")
5059 .range_row_id_set_f64(
5060 *column_id,
5061 *lo,
5062 *lo_inclusive,
5063 *hi,
5064 *hi_inclusive,
5065 )?
5066 }
5067 _ => self.resolve_condition(c, snapshot)?,
5068 };
5069 sets.push(s);
5070 }
5071 let candidates = RowIdSet::intersect_many(sets);
5072 crate::trace::QueryTrace::record(|t| {
5073 t.survivor_count = Some(candidates.len());
5074 });
5075 if candidates.is_empty() {
5076 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
5077 .iter()
5078 .map(|&id| {
5079 (
5080 id,
5081 columnar::null_native(
5082 proj_pairs
5083 .iter()
5084 .find(|(c, _)| c == &id)
5085 .map(|(_, t)| t.clone())
5086 .unwrap_or(TypeId::Bytes),
5087 0,
5088 ),
5089 )
5090 })
5091 .collect();
5092 return Ok(Some(cols));
5093 }
5094 let mut reader = match reader_opt.take() {
5095 Some(r) => r,
5096 None => self.open_reader(self.run_refs[0].run_id)?,
5097 };
5098 let candidate_ids = candidates.into_sorted_vec();
5099 let (positions, fast_rid) = if let Some(positions) =
5100 reader.positions_for_row_ids_fast(&candidate_ids)
5101 {
5102 (positions, true)
5103 } else {
5104 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
5105 match col {
5106 columnar::NativeColumn::Int64 { data, .. } => {
5107 let mut p: Vec<usize> = candidate_ids
5108 .iter()
5109 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
5110 .collect();
5111 p.sort_unstable();
5112 (p, false)
5113 }
5114 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
5115 }
5116 };
5117 crate::trace::QueryTrace::record(|t| {
5118 t.scan_mode = crate::trace::ScanMode::NativePushdown;
5119 t.fast_row_id_map = fast_rid;
5120 });
5121 let mut cols = Vec::with_capacity(col_ids.len());
5122 for cid in &col_ids {
5123 let col = reader.column_native(*cid)?;
5124 cols.push((*cid, col.gather(&positions)));
5125 }
5126 return Ok(Some(cols));
5127 }
5128
5129 if !self.run_refs.is_empty() {
5142 use crate::cursor::{drain_cursor_to_columns, Cursor};
5143 let remaining: usize;
5144 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
5145 let c = self
5146 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
5147 .expect("single-run cursor should build when run_refs.len() == 1");
5148 remaining = c.remaining_rows();
5149 Box::new(c)
5150 } else {
5151 let c = self
5152 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
5153 .expect("multi-run cursor should build when run_refs.len() >= 1");
5154 remaining = c.remaining_rows();
5155 Box::new(c)
5156 };
5157 crate::trace::QueryTrace::record(|t| {
5158 if t.survivor_count.is_none() {
5159 t.survivor_count = Some(remaining);
5160 }
5161 });
5162 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
5163 return Ok(Some(cols));
5164 }
5165
5166 crate::trace::QueryTrace::record(|t| {
5171 t.scan_mode = crate::trace::ScanMode::Materialized;
5172 t.row_materialized = true;
5173 });
5174 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5175 for c in conditions {
5176 sets.push(self.resolve_condition(c, snapshot)?);
5177 }
5178 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5179 let rows = self.rows_for_rids(&rids, snapshot)?;
5180 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
5181 for (cid, ty) in &proj_pairs {
5182 let vals: Vec<Value> = rows
5183 .iter()
5184 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5185 .collect();
5186 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
5187 }
5188 Ok(Some(cols))
5189 }
5190
5191 pub fn native_page_cursor(
5206 &self,
5207 snapshot: Snapshot,
5208 projection: Vec<(u16, TypeId)>,
5209 conditions: &[crate::query::Condition],
5210 ) -> Result<Option<NativePageCursor>> {
5211 use crate::cursor::build_page_plans;
5212 if self.ttl.is_some() {
5213 return Ok(None);
5214 }
5215 if !conditions.is_empty() && !self.indexes_complete {
5218 return Ok(None);
5219 }
5220 if self.run_refs.len() != 1 {
5221 return Ok(None);
5222 }
5223 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5224 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
5225
5226 let overlay_rids: HashSet<u64> = {
5229 let mut s = HashSet::new();
5230 for row in self.memtable.visible_versions(snapshot.epoch) {
5231 s.insert(row.row_id.0);
5232 }
5233 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5234 s.insert(row.row_id.0);
5235 }
5236 s
5237 };
5238
5239 let survivors = if conditions.is_empty() {
5243 None
5244 } else {
5245 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
5246 };
5247
5248 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
5255 survivors.clone()
5256 } else if let Some(s) = &survivors {
5257 let mut run_set = s.clone();
5258 run_set.remove_many(overlay_rids.iter().copied());
5259 Some(run_set)
5260 } else {
5261 Some(RowIdSet::from_unsorted(
5262 rids.iter()
5263 .map(|&r| r as u64)
5264 .filter(|r| !overlay_rids.contains(r))
5265 .collect(),
5266 ))
5267 };
5268
5269 let overlay_rows = if overlay_rids.is_empty() {
5270 Vec::new()
5271 } else {
5272 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5273 self.overlay_visible_rows(snapshot, bound)
5274 };
5275
5276 let plans = if positions.is_empty() {
5278 Vec::new()
5279 } else {
5280 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
5281 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
5282 };
5283
5284 let overlay = if overlay_rows.is_empty() {
5286 None
5287 } else {
5288 let filtered =
5289 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5290 if filtered.is_empty() {
5291 None
5292 } else {
5293 Some(self.materialize_overlay(&filtered, &projection))
5294 }
5295 };
5296
5297 let overlay_row_count = overlay
5298 .as_ref()
5299 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5300 .unwrap_or(0);
5301 crate::trace::QueryTrace::record(|t| {
5302 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
5303 t.run_count = self.run_refs.len();
5304 t.memtable_rows = self.memtable.len();
5305 t.mutable_run_rows = self.mutable_run.len();
5306 t.overlay_rows = overlay_row_count;
5307 t.conditions_pushed = conditions.len();
5308 t.pages_decoded = plans
5309 .iter()
5310 .map(|p| p.positions.len())
5311 .sum::<usize>()
5312 .min(1);
5313 });
5314
5315 Ok(Some(NativePageCursor::new_with_overlay(
5316 reader, projection, plans, overlay,
5317 )))
5318 }
5319 #[allow(clippy::type_complexity)]
5329 pub fn native_multi_run_cursor(
5330 &self,
5331 snapshot: Snapshot,
5332 projection: Vec<(u16, TypeId)>,
5333 conditions: &[crate::query::Condition],
5334 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
5335 use crate::cursor::{MultiRunCursor, RunStream};
5336 use crate::sorted_run::SYS_ROW_ID;
5337 use std::collections::{BinaryHeap, HashMap, HashSet};
5338 if self.ttl.is_some() {
5339 return Ok(None);
5340 }
5341 if !conditions.is_empty() && !self.indexes_complete {
5344 return Ok(None);
5345 }
5346 if self.run_refs.is_empty() {
5347 return Ok(None);
5348 }
5349
5350 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
5352 Vec::with_capacity(self.run_refs.len());
5353 for rr in &self.run_refs {
5354 let mut reader = self.open_reader(rr.run_id)?;
5355 let (rids, eps, del) = reader.system_columns_native()?;
5356 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
5357 run_meta.push((reader, rids, eps, del, page_rows));
5358 }
5359
5360 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
5364 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
5365 for i in 0..rids.len() {
5366 let rid = rids[i] as u64;
5367 let e = eps[i] as u64;
5368 if e > snapshot.epoch.0 {
5369 continue;
5370 }
5371 let is_del = del[i] != 0;
5372 best.entry(rid)
5373 .and_modify(|cur| {
5374 if e > cur.0 {
5375 *cur = (e, run_idx, i, is_del);
5376 }
5377 })
5378 .or_insert((e, run_idx, i, is_del));
5379 }
5380 }
5381
5382 let overlay_rids: HashSet<u64> = {
5384 let mut s = HashSet::new();
5385 for row in self.memtable.visible_versions(snapshot.epoch) {
5386 s.insert(row.row_id.0);
5387 }
5388 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5389 s.insert(row.row_id.0);
5390 }
5391 s
5392 };
5393
5394 let survivors: Option<RowIdSet> = if conditions.is_empty() {
5396 None
5397 } else {
5398 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5399 for c in conditions {
5400 sets.push(self.resolve_condition(c, snapshot)?);
5401 }
5402 Some(RowIdSet::intersect_many(sets))
5403 };
5404
5405 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
5409 for (rid, (_, run_idx, pos, deleted)) in &best {
5410 if *deleted {
5411 continue;
5412 }
5413 if overlay_rids.contains(rid) {
5414 continue;
5415 }
5416 if let Some(s) = &survivors {
5417 if !s.contains(*rid) {
5418 continue;
5419 }
5420 }
5421 per_run[*run_idx].push((*rid, *pos));
5422 }
5423 for v in per_run.iter_mut() {
5424 v.sort_unstable_by_key(|&(rid, _)| rid);
5425 }
5426
5427 let mut streams = Vec::with_capacity(run_meta.len());
5429 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
5430 let mut total = 0usize;
5431 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
5432 let mut starts = Vec::with_capacity(page_rows.len());
5433 let mut acc = 0usize;
5434 for &r in &page_rows {
5435 starts.push(acc);
5436 acc += r;
5437 }
5438 let mut survivors_vec: Vec<(u64, usize, usize)> =
5439 Vec::with_capacity(per_run[run_idx].len());
5440 for &(rid, pos) in &per_run[run_idx] {
5441 let page_seq = match starts.partition_point(|&s| s <= pos) {
5442 0 => continue,
5443 p => p - 1,
5444 };
5445 let within = pos - starts[page_seq];
5446 survivors_vec.push((rid, page_seq, within));
5447 }
5448 total += survivors_vec.len();
5449 if let Some(&(rid, _, _)) = survivors_vec.first() {
5450 heap.push(std::cmp::Reverse((rid, run_idx)));
5451 }
5452 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5453 }
5454
5455 let overlay_rows = if overlay_rids.is_empty() {
5457 Vec::new()
5458 } else {
5459 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5460 self.overlay_visible_rows(snapshot, bound)
5461 };
5462 let overlay = if overlay_rows.is_empty() {
5463 None
5464 } else {
5465 let filtered =
5466 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5467 if filtered.is_empty() {
5468 None
5469 } else {
5470 Some(self.materialize_overlay(&filtered, &projection))
5471 }
5472 };
5473
5474 let overlay_row_count = overlay
5475 .as_ref()
5476 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5477 .unwrap_or(0);
5478 crate::trace::QueryTrace::record(|t| {
5479 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5480 t.run_count = self.run_refs.len();
5481 t.memtable_rows = self.memtable.len();
5482 t.mutable_run_rows = self.mutable_run.len();
5483 t.overlay_rows = overlay_row_count;
5484 t.conditions_pushed = conditions.len();
5485 t.survivor_count = Some(total);
5486 });
5487
5488 Ok(Some(MultiRunCursor::new(
5489 streams, projection, heap, total, overlay,
5490 )))
5491 }
5492
5493 fn overlay_materialization_bound<'a>(
5505 conditions: &[crate::query::Condition],
5506 survivors: &'a Option<RowIdSet>,
5507 ) -> Option<&'a RowIdSet> {
5508 use crate::query::Condition;
5509 let has_range = conditions
5510 .iter()
5511 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5512 if has_range {
5513 None
5514 } else {
5515 survivors.as_ref()
5516 }
5517 }
5518
5519 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5531 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5532 let mut fold = |row: Row| {
5533 if let Some(b) = bound {
5534 if !b.contains(row.row_id.0) {
5535 return;
5536 }
5537 }
5538 best.entry(row.row_id.0)
5539 .and_modify(|(be, br)| {
5540 if row.committed_epoch > *be {
5541 *be = row.committed_epoch;
5542 *br = row.clone();
5543 }
5544 })
5545 .or_insert_with(|| (row.committed_epoch, row));
5546 };
5547 for row in self.memtable.visible_versions(snapshot.epoch) {
5548 fold(row);
5549 }
5550 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5551 fold(row);
5552 }
5553 let mut out: Vec<Row> = best
5554 .into_values()
5555 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5556 .collect();
5557 out.sort_by_key(|r| r.row_id);
5558 out
5559 }
5560
5561 fn filter_overlay_rows(
5569 &self,
5570 rows: Vec<Row>,
5571 conditions: &[crate::query::Condition],
5572 survivors: Option<&RowIdSet>,
5573 snapshot: Snapshot,
5574 ) -> Result<Vec<Row>> {
5575 if conditions.is_empty() {
5576 return Ok(rows);
5577 }
5578 use crate::query::Condition;
5579 let all_index_served = !conditions
5583 .iter()
5584 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5585 if all_index_served {
5586 return Ok(rows
5587 .into_iter()
5588 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5589 .collect());
5590 }
5591 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5594 for c in conditions {
5595 let s = match c {
5596 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5597 _ => self.resolve_condition(c, snapshot)?,
5598 };
5599 per_cond_sets.push(s);
5600 }
5601 Ok(rows
5602 .into_iter()
5603 .filter(|row| {
5604 conditions.iter().enumerate().all(|(i, c)| match c {
5605 Condition::Range { column_id, lo, hi } => {
5606 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5607 }
5608 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5609 match row.columns.get(column_id) {
5610 Some(Value::Float64(v)) => {
5611 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5612 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5613 lo_ok && hi_ok
5614 }
5615 _ => false,
5616 }
5617 }
5618 _ => per_cond_sets[i].contains(row.row_id.0),
5619 })
5620 })
5621 .collect())
5622 }
5623
5624 fn materialize_overlay(
5627 &self,
5628 rows: &[Row],
5629 projection: &[(u16, TypeId)],
5630 ) -> Vec<columnar::NativeColumn> {
5631 if projection.is_empty() {
5632 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5633 }
5634 let mut cols = Vec::with_capacity(projection.len());
5635 for (cid, ty) in projection {
5636 let vals: Vec<Value> = rows
5637 .iter()
5638 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5639 .collect();
5640 cols.push(columnar::values_to_native(ty.clone(), &vals));
5641 }
5642 cols
5643 }
5644
5645 fn resolve_survivor_rids(
5650 &self,
5651 conditions: &[crate::query::Condition],
5652 reader: &mut RunReader,
5653 ) -> Result<RowIdSet> {
5654 use crate::query::Condition;
5655 let mut sets: Vec<RowIdSet> = Vec::new();
5656 for c in conditions {
5657 let s: RowIdSet = match c {
5658 Condition::Pk(key) => {
5659 let lookup = self
5660 .schema
5661 .primary_key()
5662 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5663 .unwrap_or_else(|| key.clone());
5664 self.hot
5665 .get(&lookup)
5666 .map(|r| RowIdSet::one(r.0))
5667 .unwrap_or_else(RowIdSet::empty)
5668 }
5669 Condition::BitmapEq { column_id, value } => {
5670 let lookup = self.index_lookup_key_bytes(*column_id, value);
5671 self.bitmap
5672 .get(column_id)
5673 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5674 .unwrap_or_else(RowIdSet::empty)
5675 }
5676 Condition::BitmapIn { column_id, values } => {
5677 let bm = self.bitmap.get(column_id);
5678 let mut acc = roaring::RoaringBitmap::new();
5679 if let Some(b) = bm {
5680 for v in values {
5681 let lookup = self.index_lookup_key_bytes(*column_id, v);
5682 acc |= b.get(&lookup);
5683 }
5684 }
5685 RowIdSet::from_roaring(acc)
5686 }
5687 Condition::BytesPrefix { column_id, prefix } => {
5688 if let Some(b) = self.bitmap.get(column_id) {
5689 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5690 let mut acc = roaring::RoaringBitmap::new();
5691 for key in b.keys() {
5692 if key.starts_with(&lookup_prefix) {
5693 acc |= b.get(key);
5694 }
5695 }
5696 RowIdSet::from_roaring(acc)
5697 } else {
5698 RowIdSet::empty()
5699 }
5700 }
5701 Condition::FmContains { column_id, pattern } => self
5702 .fm
5703 .get(column_id)
5704 .map(|f| {
5705 RowIdSet::from_unsorted(
5706 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5707 )
5708 })
5709 .unwrap_or_else(RowIdSet::empty),
5710 Condition::FmContainsAll {
5711 column_id,
5712 patterns,
5713 } => {
5714 if let Some(f) = self.fm.get(column_id) {
5715 let sets: Vec<RowIdSet> = patterns
5716 .iter()
5717 .map(|pat| {
5718 RowIdSet::from_unsorted(
5719 f.locate(pat).into_iter().map(|r| r.0).collect(),
5720 )
5721 })
5722 .collect();
5723 RowIdSet::intersect_many(sets)
5724 } else {
5725 RowIdSet::empty()
5726 }
5727 }
5728 Condition::Ann {
5729 column_id,
5730 query,
5731 k,
5732 } => self
5733 .ann
5734 .get(column_id)
5735 .map(|a| {
5736 RowIdSet::from_unsorted(
5737 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5738 )
5739 })
5740 .unwrap_or_else(RowIdSet::empty),
5741 Condition::SparseMatch {
5742 column_id,
5743 query,
5744 k,
5745 } => self
5746 .sparse
5747 .get(column_id)
5748 .map(|s| {
5749 RowIdSet::from_unsorted(
5750 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5751 )
5752 })
5753 .unwrap_or_else(RowIdSet::empty),
5754 Condition::MinHashSimilar {
5755 column_id,
5756 query,
5757 k,
5758 } => self
5759 .minhash
5760 .get(column_id)
5761 .map(|mh| {
5762 RowIdSet::from_unsorted(
5763 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5764 )
5765 })
5766 .unwrap_or_else(RowIdSet::empty),
5767 Condition::Range { column_id, lo, hi } => {
5768 if let Some(li) = self.learned_range.get(column_id) {
5769 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5770 } else {
5771 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5772 }
5773 }
5774 Condition::RangeF64 {
5775 column_id,
5776 lo,
5777 lo_inclusive,
5778 hi,
5779 hi_inclusive,
5780 } => {
5781 if let Some(li) = self.learned_range.get(column_id) {
5782 RowIdSet::from_unsorted(
5783 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5784 .into_iter()
5785 .collect(),
5786 )
5787 } else {
5788 reader.range_row_id_set_f64(
5789 *column_id,
5790 *lo,
5791 *lo_inclusive,
5792 *hi,
5793 *hi_inclusive,
5794 )?
5795 }
5796 }
5797 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5798 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5799 };
5800 sets.push(s);
5801 }
5802 Ok(RowIdSet::intersect_many(sets))
5803 }
5804
5805 pub fn scan_cursor(
5826 &self,
5827 snapshot: Snapshot,
5828 projection: Vec<(u16, TypeId)>,
5829 conditions: &[crate::query::Condition],
5830 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5831 if self.ttl.is_some() {
5832 return Ok(None);
5833 }
5834 if !conditions.is_empty() && !self.indexes_complete {
5840 return Ok(None);
5841 }
5842 if self.run_refs.len() == 1 {
5843 Ok(self
5844 .native_page_cursor(snapshot, projection, conditions)?
5845 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5846 } else {
5847 Ok(self
5848 .native_multi_run_cursor(snapshot, projection, conditions)?
5849 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5850 }
5851 }
5852
5853 pub fn aggregate_native(
5867 &self,
5868 snapshot: Snapshot,
5869 column: Option<u16>,
5870 conditions: &[crate::query::Condition],
5871 agg: NativeAgg,
5872 ) -> Result<Option<NativeAggResult>> {
5873 if self.ttl.is_some() {
5874 return Ok(None);
5875 }
5876 if self.run_refs.len() == 1 && conditions.is_empty() {
5878 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5879 return Ok(Some(res));
5880 }
5881 }
5882 if matches!(agg, NativeAgg::Count) && column.is_none() {
5884 return Ok(self
5885 .scan_cursor(snapshot, Vec::new(), conditions)?
5886 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5887 }
5888 let cid = match column {
5891 Some(c) => c,
5892 None => return Ok(None),
5893 };
5894 let ty = self.column_type(cid);
5895 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
5896 else {
5897 return Ok(None);
5898 };
5899 match ty {
5900 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5901 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5902 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5903 }
5904 TypeId::Float64 => {
5905 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5906 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5907 }
5908 _ => Ok(None),
5909 }
5910 }
5911
5912 fn aggregate_from_stats(
5920 &self,
5921 snapshot: Snapshot,
5922 column: Option<u16>,
5923 agg: NativeAgg,
5924 ) -> Result<Option<NativeAggResult>> {
5925 let cid = match (agg, column) {
5926 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5927 _ => return Ok(None), };
5929 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5930 return Ok(None);
5931 };
5932 let Some(cs) = stats.get(&cid) else {
5933 return Ok(None);
5934 };
5935 match agg {
5936 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5938 self.live_count.saturating_sub(cs.null_count),
5939 ))),
5940 NativeAgg::Min | NativeAgg::Max => {
5941 let bound = if agg == NativeAgg::Min {
5942 &cs.min
5943 } else {
5944 &cs.max
5945 };
5946 match bound {
5947 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5948 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5949 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5954 None => Ok(None),
5955 }
5956 }
5957 _ => Ok(None),
5958 }
5959 }
5960
5961 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5970 if self.ttl.is_some() {
5971 return Ok(None);
5972 }
5973 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5974 return Ok(None);
5975 }
5976 self.ensure_indexes_complete()?;
5979 let reader = self.open_reader(self.run_refs[0].run_id)?;
5980 if self.live_count != reader.row_count() as u64 {
5981 return Ok(None);
5982 }
5983 let Some(bm) = self.bitmap.get(&column_id) else {
5984 return Ok(None); };
5986 let mut distinct = bm.value_count() as u64;
5987 if !bm.get(&Value::Null.encode_key()).is_empty() {
5990 distinct = distinct.saturating_sub(1);
5991 }
5992 Ok(Some(distinct))
5993 }
5994
5995 pub fn aggregate_incremental(
6007 &mut self,
6008 cache_key: u64,
6009 conditions: &[crate::query::Condition],
6010 column: Option<u16>,
6011 agg: NativeAgg,
6012 ) -> Result<IncrementalAggResult> {
6013 let snap = self.snapshot();
6014 let cur_wm = self.allocator.current().0;
6015 let cur_epoch = snap.epoch.0;
6016 let incremental_ok = self.ttl.is_none()
6023 && !self.had_deletes
6024 && self.memtable.is_empty()
6025 && self.mutable_run.is_empty();
6026
6027 if incremental_ok {
6030 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
6031 if cached.epoch == cur_epoch {
6032 return Ok(IncrementalAggResult {
6033 state: cached.state,
6034 incremental: true,
6035 delta_rows: 0,
6036 });
6037 }
6038 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
6039 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
6040 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
6041 let index_sets = self.resolve_index_conditions(conditions, snap)?;
6042 let delta_state = agg_state_from_rows(
6043 &delta_rows,
6044 conditions,
6045 &index_sets,
6046 column,
6047 agg,
6048 &self.schema,
6049 )?;
6050 let merged = cached.state.merge(delta_state);
6051 let delta_n = delta_rids.len() as u64;
6052 self.agg_cache.insert(
6053 cache_key,
6054 CachedAgg {
6055 state: merged.clone(),
6056 watermark: cur_wm,
6057 epoch: cur_epoch,
6058 },
6059 );
6060 return Ok(IncrementalAggResult {
6061 state: merged,
6062 incremental: true,
6063 delta_rows: delta_n,
6064 });
6065 }
6066 }
6067 }
6068
6069 let cursor_ok =
6074 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
6075 let state = if cursor_ok && agg != NativeAgg::Avg {
6076 match self.aggregate_native(snap, column, conditions, agg)? {
6077 Some(result) => {
6078 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
6079 }
6080 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
6081 }
6082 } else {
6083 self.agg_state_full_scan(conditions, column, agg, snap)?
6084 };
6085 if incremental_ok {
6087 self.agg_cache.insert(
6088 cache_key,
6089 CachedAgg {
6090 state: state.clone(),
6091 watermark: cur_wm,
6092 epoch: cur_epoch,
6093 },
6094 );
6095 }
6096 Ok(IncrementalAggResult {
6097 state,
6098 incremental: false,
6099 delta_rows: 0,
6100 })
6101 }
6102
6103 fn agg_state_full_scan(
6106 &self,
6107 conditions: &[crate::query::Condition],
6108 column: Option<u16>,
6109 agg: NativeAgg,
6110 snap: Snapshot,
6111 ) -> Result<AggState> {
6112 let rows = self.visible_rows(snap)?;
6113 let index_sets = self.resolve_index_conditions(conditions, snap)?;
6114 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
6115 }
6116
6117 fn resolve_index_conditions(
6120 &self,
6121 conditions: &[crate::query::Condition],
6122 snapshot: Snapshot,
6123 ) -> Result<Vec<RowIdSet>> {
6124 use crate::query::Condition;
6125 let mut sets = Vec::new();
6126 for c in conditions {
6127 if matches!(
6128 c,
6129 Condition::Ann { .. }
6130 | Condition::SparseMatch { .. }
6131 | Condition::MinHashSimilar { .. }
6132 ) {
6133 sets.push(self.resolve_condition(c, snapshot)?);
6134 }
6135 }
6136 Ok(sets)
6137 }
6138
6139 fn column_type(&self, cid: u16) -> TypeId {
6140 self.schema
6141 .columns
6142 .iter()
6143 .find(|c| c.id == cid)
6144 .map(|c| c.ty.clone())
6145 .unwrap_or(TypeId::Bytes)
6146 }
6147
6148 pub fn approx_aggregate(
6157 &mut self,
6158 conditions: &[crate::query::Condition],
6159 column: Option<u16>,
6160 agg: ApproxAgg,
6161 z: f64,
6162 ) -> Result<Option<ApproxResult>> {
6163 use crate::query::Condition;
6164 self.ensure_reservoir_complete()?;
6165 let snapshot = self.snapshot();
6166 let n_pop = self.count();
6167 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
6168 if sample_rids.is_empty() {
6169 return Ok(None);
6170 }
6171 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
6173 let s = live_sample.len();
6174 if s == 0 {
6175 return Ok(None);
6176 }
6177
6178 let mut index_sets: Vec<RowIdSet> = Vec::new();
6181 for c in conditions {
6182 if matches!(
6183 c,
6184 Condition::Ann { .. }
6185 | Condition::SparseMatch { .. }
6186 | Condition::MinHashSimilar { .. }
6187 ) {
6188 index_sets.push(self.resolve_condition(c, snapshot)?);
6189 }
6190 }
6191
6192 let cid = match (agg, column) {
6194 (ApproxAgg::Count, _) => None,
6195 (_, Some(c)) => Some(c),
6196 _ => return Ok(None),
6197 };
6198 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
6199 for r in &live_sample {
6200 if !conditions
6202 .iter()
6203 .all(|c| condition_matches_row(c, r, &self.schema))
6204 {
6205 continue;
6206 }
6207 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
6209 continue;
6210 }
6211 if let Some(cid) = cid {
6212 if let Some(v) = as_f64(r.columns.get(&cid)) {
6213 passing_vals.push(v);
6214 } } else {
6216 passing_vals.push(0.0); }
6218 }
6219 let m = passing_vals.len();
6220
6221 let (point, half) = match agg {
6222 ApproxAgg::Count => {
6223 let p = m as f64 / s as f64;
6225 let point = n_pop as f64 * p;
6226 let var = if s > 1 {
6227 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
6228 * (1.0 - s as f64 / n_pop as f64).max(0.0)
6229 } else {
6230 0.0
6231 };
6232 (point, z * var.sqrt())
6233 }
6234 ApproxAgg::Sum => {
6235 let y: Vec<f64> = live_sample
6237 .iter()
6238 .map(|r| {
6239 let passes_row = conditions
6240 .iter()
6241 .all(|c| condition_matches_row(c, r, &self.schema))
6242 && index_sets.iter().all(|set| set.contains(r.row_id.0));
6243 if passes_row {
6244 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
6245 } else {
6246 0.0
6247 }
6248 })
6249 .collect();
6250 let mean_y = y.iter().sum::<f64>() / s as f64;
6251 let point = n_pop as f64 * mean_y;
6252 let var = if s > 1 {
6253 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
6254 let var_y = ss / (s - 1) as f64;
6255 n_pop as f64 * n_pop as f64 * var_y / s as f64
6256 * (1.0 - s as f64 / n_pop as f64).max(0.0)
6257 } else {
6258 0.0
6259 };
6260 (point, z * var.sqrt())
6261 }
6262 ApproxAgg::Avg => {
6263 if m == 0 {
6264 return Ok(Some(ApproxResult {
6265 point: 0.0,
6266 ci_low: 0.0,
6267 ci_high: 0.0,
6268 n_population: n_pop,
6269 n_sample_live: s,
6270 n_passing: 0,
6271 }));
6272 }
6273 let mean = passing_vals.iter().sum::<f64>() / m as f64;
6274 let half = if m > 1 {
6275 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
6276 let sd = (ss / (m - 1) as f64).sqrt();
6277 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
6278 z * sd / (m as f64).sqrt() * fpc.sqrt()
6279 } else {
6280 0.0
6281 };
6282 (mean, half)
6283 }
6284 };
6285
6286 Ok(Some(ApproxResult {
6287 point,
6288 ci_low: point - half,
6289 ci_high: point + half,
6290 n_population: n_pop,
6291 n_sample_live: s,
6292 n_passing: m,
6293 }))
6294 }
6295
6296 pub fn exact_column_stats(
6304 &self,
6305 _snapshot: Snapshot,
6306 projection: &[u16],
6307 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
6308 if self.ttl.is_some()
6309 || !(self.memtable.is_empty()
6310 && self.mutable_run.is_empty()
6311 && self.run_refs.len() == 1)
6312 {
6313 return Ok(None);
6314 }
6315 let reader = self.open_reader(self.run_refs[0].run_id)?;
6316 if self.live_count != reader.row_count() as u64 {
6317 return Ok(None);
6318 }
6319 let mut out = HashMap::new();
6320 for &cid in projection {
6321 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
6322 Some(c) => c,
6323 None => continue,
6324 };
6325 let Some(stats) = reader.column_page_stats(cid) else {
6327 out.insert(
6328 cid,
6329 ColumnStat {
6330 min: None,
6331 max: None,
6332 null_count: self.live_count,
6333 },
6334 );
6335 continue;
6336 };
6337 let stat = match cdef.ty {
6338 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
6339 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
6340 min: mn.map(Value::Int64),
6341 max: mx.map(Value::Int64),
6342 null_count: n,
6343 })
6344 }
6345 TypeId::Float64 => {
6346 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
6347 min: mn.map(Value::Float64),
6348 max: mx.map(Value::Float64),
6349 null_count: n,
6350 })
6351 }
6352 _ => None,
6353 };
6354 if let Some(s) = stat {
6355 out.insert(cid, s);
6356 }
6357 }
6358 Ok(Some(out))
6359 }
6360
6361 pub fn dir(&self) -> &Path {
6362 &self.dir
6363 }
6364
6365 pub fn schema(&self) -> &Schema {
6366 &self.schema
6367 }
6368
6369 pub(crate) fn set_catalog_name(&mut self, name: String) {
6370 self.name = name;
6371 }
6372
6373 pub(crate) fn prepare_alter_column(
6374 &mut self,
6375 column_name: &str,
6376 change: &AlterColumn,
6377 ) -> Result<ColumnDef> {
6378 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
6379 return Err(MongrelError::InvalidArgument(
6380 "ALTER COLUMN requires committing staged writes first".into(),
6381 ));
6382 }
6383 let old = self
6384 .schema
6385 .columns
6386 .iter()
6387 .find(|c| c.name == column_name)
6388 .cloned()
6389 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6390 let mut next = old.clone();
6391
6392 if let Some(name) = &change.name {
6393 let trimmed = name.trim();
6394 if trimmed.is_empty() {
6395 return Err(MongrelError::InvalidArgument(
6396 "ALTER COLUMN name must not be empty".into(),
6397 ));
6398 }
6399 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
6400 return Err(MongrelError::Schema(format!(
6401 "column {trimmed} already exists"
6402 )));
6403 }
6404 next.name = trimmed.to_string();
6405 }
6406
6407 if let Some(ty) = &change.ty {
6408 next.ty = ty.clone();
6409 }
6410 if let Some(flags) = change.flags {
6411 validate_alter_column_flags(old.flags, flags)?;
6412 next.flags = flags;
6413 }
6414
6415 if let Some(default_change) = &change.default_value {
6416 next.default_value = default_change.clone();
6417 }
6418
6419 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
6420 if old.flags.contains(ColumnFlags::NULLABLE)
6421 && !next.flags.contains(ColumnFlags::NULLABLE)
6422 && self.column_has_nulls(old.id)?
6423 {
6424 return Err(MongrelError::InvalidArgument(format!(
6425 "column '{}' contains NULL values",
6426 old.name
6427 )));
6428 }
6429 Ok(next)
6430 }
6431
6432 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
6433 let idx = self
6434 .schema
6435 .columns
6436 .iter()
6437 .position(|c| c.id == column.id)
6438 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
6439 if self.schema.columns[idx] == column {
6440 return Ok(());
6441 }
6442 self.schema.columns[idx] = column;
6443 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6444 self.schema.validate_auto_increment()?;
6445 self.schema.validate_defaults()?;
6446 self.auto_inc = resolve_auto_inc(&self.schema);
6447 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
6448 write_schema(&self.dir, &self.schema)?;
6449 self.clear_result_cache();
6450 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6451 self.persist_manifest(self.current_epoch())?;
6452 Ok(())
6453 }
6454
6455 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
6456 self.ensure_writable()?;
6457 let column = self.prepare_alter_column(column_name, &change)?;
6458 self.apply_altered_column(column.clone())?;
6459 Ok(column)
6460 }
6461
6462 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
6463 if self.live_count == 0 {
6464 return Ok(false);
6465 }
6466 let snap = self.snapshot();
6467 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6468 Ok(columns
6469 .first()
6470 .map(|(_, col)| col.null_count(col.len()) != 0)
6471 .unwrap_or(true))
6472 }
6473
6474 fn has_stored_versions(&self) -> bool {
6475 !self.memtable.is_empty()
6476 || !self.mutable_run.is_empty()
6477 || self.run_refs.iter().any(|r| r.row_count > 0)
6478 || !self.retiring.is_empty()
6479 }
6480
6481 pub fn add_column(
6486 &mut self,
6487 name: &str,
6488 ty: TypeId,
6489 flags: ColumnFlags,
6490 default_value: Option<crate::schema::DefaultExpr>,
6491 ) -> Result<u16> {
6492 self.add_column_with_id(name, ty, flags, default_value, None)
6493 }
6494
6495 pub fn add_column_with_id(
6496 &mut self,
6497 name: &str,
6498 ty: TypeId,
6499 flags: ColumnFlags,
6500 default_value: Option<crate::schema::DefaultExpr>,
6501 requested_id: Option<u16>,
6502 ) -> Result<u16> {
6503 self.ensure_writable()?;
6504 if self.schema.columns.iter().any(|c| c.name == name) {
6505 return Err(MongrelError::Schema(format!(
6506 "column {name} already exists"
6507 )));
6508 }
6509 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
6510 if self.schema.columns.iter().any(|c| c.id == id) {
6511 return Err(MongrelError::Schema(format!(
6512 "column id {id} already exists"
6513 )));
6514 }
6515 id
6516 } else {
6517 self.schema
6518 .columns
6519 .iter()
6520 .map(|c| c.id)
6521 .max()
6522 .unwrap_or(0)
6523 .checked_add(1)
6524 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
6525 };
6526 self.schema.columns.push(ColumnDef {
6527 id,
6528 name: name.to_string(),
6529 ty,
6530 flags,
6531 default_value,
6532 });
6533 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6534 self.schema.validate_auto_increment()?;
6535 self.schema.validate_defaults()?;
6536 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6537 self.auto_inc = resolve_auto_inc(&self.schema);
6538 }
6539 write_schema(&self.dir, &self.schema)?;
6540 self.clear_result_cache();
6541 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6543 self.persist_manifest(self.current_epoch())?;
6544 Ok(id)
6545 }
6546
6547 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6556 self.ensure_writable()?;
6557 let cid = self
6558 .schema
6559 .columns
6560 .iter()
6561 .find(|c| c.name == column_name)
6562 .map(|c| c.id)
6563 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6564 let ty = self
6565 .schema
6566 .columns
6567 .iter()
6568 .find(|c| c.id == cid)
6569 .map(|c| c.ty.clone())
6570 .unwrap_or(TypeId::Int64);
6571 if !matches!(
6572 ty,
6573 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6574 ) {
6575 return Err(MongrelError::Schema(format!(
6576 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6577 )));
6578 }
6579 if self
6580 .schema
6581 .indexes
6582 .iter()
6583 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6584 {
6585 return Ok(()); }
6587 self.schema.indexes.push(IndexDef {
6588 name: format!("{}_learned_range", column_name),
6589 column_id: cid,
6590 kind: IndexKind::LearnedRange,
6591 predicate: None,
6592 });
6593 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6594 write_schema(&self.dir, &self.schema)?;
6595 self.build_learned_ranges()?;
6596 Ok(())
6597 }
6598
6599 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6602 self.sync_byte_threshold = threshold;
6603 if let WalSink::Private(w) = &mut self.wal {
6604 w.set_sync_byte_threshold(threshold);
6605 }
6606 }
6607
6608 pub fn page_cache_flush(&self) {
6612 self.page_cache.flush_to_disk();
6613 }
6614
6615 pub fn page_cache_len(&self) -> usize {
6617 self.page_cache.len()
6618 }
6619
6620 pub fn decoded_cache_len(&self) -> usize {
6623 self.decoded_cache.len()
6624 }
6625
6626 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6629 self.memtable.drain_sorted()
6630 }
6631
6632 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6633 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6634 }
6635
6636 pub(crate) fn table_dir(&self) -> &Path {
6637 &self.dir
6638 }
6639
6640 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6641 &self.schema
6642 }
6643
6644 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6645 let id = self.next_run_id;
6646 self.next_run_id += 1;
6647 id
6648 }
6649
6650 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6651 self.run_refs.push(run_ref);
6652 }
6653
6654 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6664 self.retiring.push(crate::manifest::RetiredRun {
6665 run_id,
6666 retire_epoch,
6667 });
6668 }
6669
6670 pub(crate) fn reap_retiring(
6674 &mut self,
6675 min_active: Epoch,
6676 backup_pinned: &std::collections::HashSet<u128>,
6677 ) -> Result<usize> {
6678 if self.retiring.is_empty() {
6679 return Ok(0);
6680 }
6681 let mut reaped = 0;
6682 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6683 for r in std::mem::take(&mut self.retiring) {
6689 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
6690 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6691 reaped += 1;
6692 } else {
6693 kept.push(r);
6694 }
6695 }
6696 self.retiring = kept;
6697 if reaped > 0 {
6698 self.persist_manifest(self.current_epoch())?;
6699 }
6700 Ok(reaped)
6701 }
6702
6703 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6704 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6705 return false;
6706 }
6707 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6708 self.run_refs.push(run_ref);
6709 self.indexes_complete = false;
6710 true
6711 }
6712
6713 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6714 self.kek.as_ref()
6715 }
6716
6717 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6718 let mut reader = RunReader::open_with_cache(
6719 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6720 self.schema.clone(),
6721 self.kek.clone(),
6722 Some(self.page_cache.clone()),
6723 Some(self.decoded_cache.clone()),
6724 self.table_id,
6725 Some(&self.verified_runs),
6726 )?;
6727 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6731 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6732 }
6733 Ok(reader)
6734 }
6735
6736 pub(crate) fn run_refs(&self) -> &[RunRef] {
6737 &self.run_refs
6738 }
6739
6740 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
6741 self.retiring.iter().map(|run| run.run_id)
6742 }
6743
6744 pub(crate) fn runs_dir(&self) -> PathBuf {
6745 self.dir.join(RUNS_DIR)
6746 }
6747
6748 pub(crate) fn wal_dir(&self) -> PathBuf {
6749 self.dir.join(WAL_DIR)
6750 }
6751
6752 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6753 self.run_refs = refs;
6754 }
6755
6756 pub(crate) fn next_run_id(&self) -> u64 {
6757 self.next_run_id
6758 }
6759
6760 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6761 self.compaction_zstd_level
6762 }
6763
6764 pub(crate) fn bump_next_run_id(&mut self) {
6765 self.next_run_id += 1;
6766 }
6767
6768 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6769 self.kek.clone()
6770 }
6771
6772 #[cfg(feature = "encryption")]
6776 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6777 self.kek.as_ref().map(|k| k.derive_idx_key())
6778 }
6779
6780 #[cfg(not(feature = "encryption"))]
6781 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6782 None
6783 }
6784
6785 #[cfg(feature = "encryption")]
6789 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6790 self.kek.as_ref().map(|k| *k.derive_meta_key())
6791 }
6792
6793 #[cfg(not(feature = "encryption"))]
6794 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6795 None
6796 }
6797
6798 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6801 self.column_keys
6802 .iter()
6803 .map(|(&id, &(_, scheme))| (id, scheme))
6804 .collect()
6805 }
6806
6807 #[cfg(feature = "encryption")]
6812 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6813 self.tokenize_value_enc(column_id, v)
6814 }
6815
6816 #[cfg(feature = "encryption")]
6817 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6818 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6819 let (key, scheme) = self.column_keys.get(&column_id)?;
6820 let token: Vec<u8> = match (*scheme, v) {
6821 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6822 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6823 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6824 _ => hmac_token(key, &v.encode_key()).to_vec(),
6825 };
6826 Some(Value::Bytes(token))
6827 }
6828
6829 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6831 self.index_lookup_key_bytes(column_id, &v.encode_key())
6832 }
6833
6834 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6837 #[cfg(feature = "encryption")]
6838 {
6839 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6840 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6841 if *scheme == SCHEME_HMAC_EQ {
6842 return hmac_token(key, encoded).to_vec();
6843 }
6844 }
6845 }
6846 let _ = column_id;
6847 encoded.to_vec()
6848 }
6849}
6850
6851fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6852 let columnar::NativeColumn::Int64 { data, validity } = col else {
6853 return false;
6854 };
6855 if data.len() < n || !columnar::all_non_null(validity, n) {
6856 return false;
6857 }
6858 data.iter()
6859 .take(n)
6860 .zip(data.iter().skip(1))
6861 .all(|(a, b)| a < b)
6862}
6863
6864#[derive(Debug, Clone)]
6868pub struct ColumnStat {
6869 pub min: Option<Value>,
6870 pub max: Option<Value>,
6871 pub null_count: u64,
6872}
6873
6874#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6876pub enum NativeAgg {
6877 Count,
6878 Sum,
6879 Min,
6880 Max,
6881 Avg,
6882}
6883
6884#[derive(Debug, Clone, PartialEq)]
6886pub enum NativeAggResult {
6887 Count(u64),
6888 Int(i64),
6889 Float(f64),
6890 Null,
6892}
6893
6894#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6896pub enum ApproxAgg {
6897 Count,
6898 Sum,
6899 Avg,
6900}
6901
6902#[derive(Debug, Clone)]
6906pub struct ApproxResult {
6907 pub point: f64,
6909 pub ci_low: f64,
6911 pub ci_high: f64,
6913 pub n_population: u64,
6915 pub n_sample_live: usize,
6917 pub n_passing: usize,
6919}
6920
6921#[derive(Debug, Clone, PartialEq)]
6926pub enum AggState {
6927 Count(u64),
6929 SumI {
6931 sum: i128,
6932 count: u64,
6933 },
6934 SumF {
6936 sum: f64,
6937 count: u64,
6938 },
6939 AvgI {
6941 sum: i128,
6942 count: u64,
6943 },
6944 AvgF {
6946 sum: f64,
6947 count: u64,
6948 },
6949 MinI(i64),
6951 MaxI(i64),
6952 MinF(f64),
6954 MaxF(f64),
6955 Empty,
6957}
6958
6959impl AggState {
6960 pub fn merge(self, other: AggState) -> AggState {
6962 use AggState::*;
6963 match (self, other) {
6964 (Empty, x) | (x, Empty) => x,
6965 (Count(a), Count(b)) => Count(a + b),
6966 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6967 sum: sa + sb,
6968 count: ca + cb,
6969 },
6970 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6971 sum: sa + sb,
6972 count: ca + cb,
6973 },
6974 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6975 sum: sa + sb,
6976 count: ca + cb,
6977 },
6978 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6979 sum: sa + sb,
6980 count: ca + cb,
6981 },
6982 (MinI(a), MinI(b)) => MinI(a.min(b)),
6983 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6984 (MinF(a), MinF(b)) => MinF(a.min(b)),
6985 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6986 _ => Empty, }
6988 }
6989
6990 pub fn point(&self) -> Option<f64> {
6992 match self {
6993 AggState::Count(n) => Some(*n as f64),
6994 AggState::SumI { sum, .. } => Some(*sum as f64),
6995 AggState::SumF { sum, .. } => Some(*sum),
6996 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6997 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6998 AggState::MinI(n) => Some(*n as f64),
6999 AggState::MaxI(n) => Some(*n as f64),
7000 AggState::MinF(n) => Some(*n),
7001 AggState::MaxF(n) => Some(*n),
7002 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
7003 }
7004 }
7005
7006 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
7010 let is_float = matches!(ty, Some(TypeId::Float64));
7011 match (agg, result) {
7012 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
7013 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
7014 sum: x as i128,
7015 count: 1, },
7017 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
7018 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
7019 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
7020 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
7021 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
7022 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
7023 (NativeAgg::Count, _) => AggState::Empty,
7024 (_, NativeAggResult::Null) => AggState::Empty,
7025 _ => {
7026 let _ = is_float;
7027 AggState::Empty
7028 }
7029 }
7030 }
7031}
7032
7033#[derive(Debug, Clone)]
7036pub struct CachedAgg {
7037 pub state: AggState,
7038 pub watermark: u64,
7039 pub epoch: u64,
7040}
7041
7042#[derive(Debug, Clone)]
7044pub struct IncrementalAggResult {
7045 pub state: AggState,
7047 pub incremental: bool,
7050 pub delta_rows: u64,
7052}
7053
7054fn agg_state_from_rows(
7058 rows: &[Row],
7059 conditions: &[crate::query::Condition],
7060 index_sets: &[RowIdSet],
7061 column: Option<u16>,
7062 agg: NativeAgg,
7063 schema: &Schema,
7064) -> Result<AggState> {
7065 let mut count: u64 = 0;
7066 let mut sum_i: i128 = 0;
7067 let mut sum_f: f64 = 0.0;
7068 let mut mn_i: i64 = i64::MAX;
7069 let mut mx_i: i64 = i64::MIN;
7070 let mut mn_f: f64 = f64::INFINITY;
7071 let mut mx_f: f64 = f64::NEG_INFINITY;
7072 let mut saw_int = false;
7073 let mut saw_float = false;
7074 for r in rows {
7075 if !conditions
7076 .iter()
7077 .all(|c| condition_matches_row(c, r, schema))
7078 {
7079 continue;
7080 }
7081 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
7082 continue;
7083 }
7084 match agg {
7085 NativeAgg::Count => match column {
7086 None => count += 1,
7088 Some(cid) => match r.columns.get(&cid) {
7091 None | Some(Value::Null) => {}
7092 Some(_) => count += 1,
7093 },
7094 },
7095 _ => match column.and_then(|cid| r.columns.get(&cid)) {
7096 Some(Value::Int64(n)) => {
7097 count += 1;
7098 sum_i += *n as i128;
7099 mn_i = mn_i.min(*n);
7100 mx_i = mx_i.max(*n);
7101 saw_int = true;
7102 }
7103 Some(Value::Float64(f)) => {
7104 count += 1;
7105 sum_f += f;
7106 mn_f = mn_f.min(*f);
7107 mx_f = mx_f.max(*f);
7108 saw_float = true;
7109 }
7110 _ => {}
7111 },
7112 }
7113 }
7114 Ok(match agg {
7115 NativeAgg::Count => {
7116 if count == 0 {
7117 AggState::Empty
7118 } else {
7119 AggState::Count(count)
7120 }
7121 }
7122 NativeAgg::Sum => {
7123 if count == 0 {
7124 AggState::Empty
7125 } else if saw_int {
7126 AggState::SumI { sum: sum_i, count }
7127 } else {
7128 AggState::SumF { sum: sum_f, count }
7129 }
7130 }
7131 NativeAgg::Avg => {
7132 if count == 0 {
7133 AggState::Empty
7134 } else if saw_int {
7135 AggState::AvgI { sum: sum_i, count }
7136 } else {
7137 AggState::AvgF { sum: sum_f, count }
7138 }
7139 }
7140 NativeAgg::Min => {
7141 if !saw_int && !saw_float {
7142 AggState::Empty
7143 } else if saw_int {
7144 AggState::MinI(mn_i)
7145 } else {
7146 AggState::MinF(mn_f)
7147 }
7148 }
7149 NativeAgg::Max => {
7150 if !saw_int && !saw_float {
7151 AggState::Empty
7152 } else if saw_int {
7153 AggState::MaxI(mx_i)
7154 } else {
7155 AggState::MaxF(mx_f)
7156 }
7157 }
7158 })
7159}
7160
7161fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
7165 use crate::query::Condition;
7166 match c {
7167 Condition::Pk(key) => match schema.primary_key() {
7168 Some(pk) => row
7169 .columns
7170 .get(&pk.id)
7171 .map(|v| v.encode_key() == *key)
7172 .unwrap_or(false),
7173 None => false,
7174 },
7175 Condition::BitmapEq { column_id, value } => row
7176 .columns
7177 .get(column_id)
7178 .map(|v| v.encode_key() == *value)
7179 .unwrap_or(false),
7180 Condition::BitmapIn { column_id, values } => {
7181 let key = row.columns.get(column_id).map(|v| v.encode_key());
7182 match key {
7183 Some(k) => values.contains(&k),
7184 None => false,
7185 }
7186 }
7187 Condition::BytesPrefix { column_id, prefix } => row
7188 .columns
7189 .get(column_id)
7190 .map(|v| v.encode_key().starts_with(prefix))
7191 .unwrap_or(false),
7192 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
7193 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
7194 _ => false,
7195 },
7196 Condition::RangeF64 {
7197 column_id,
7198 lo,
7199 lo_inclusive,
7200 hi,
7201 hi_inclusive,
7202 } => match row.columns.get(column_id) {
7203 Some(Value::Float64(n)) => {
7204 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
7205 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
7206 lo_ok && hi_ok
7207 }
7208 _ => false,
7209 },
7210 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
7211 Some(Value::Bytes(b)) => {
7212 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
7213 }
7214 _ => false,
7215 },
7216 Condition::FmContainsAll {
7217 column_id,
7218 patterns,
7219 } => match row.columns.get(column_id) {
7220 Some(Value::Bytes(b)) => patterns
7221 .iter()
7222 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
7223 _ => false,
7224 },
7225 Condition::Ann { .. }
7226 | Condition::SparseMatch { .. }
7227 | Condition::MinHashSimilar { .. } => true,
7228 Condition::IsNull { column_id } => {
7229 matches!(row.columns.get(column_id), Some(Value::Null) | None)
7230 }
7231 Condition::IsNotNull { column_id } => {
7232 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
7233 }
7234 }
7235}
7236
7237fn as_f64(v: Option<&Value>) -> Option<f64> {
7239 match v {
7240 Some(Value::Int64(n)) => Some(*n as f64),
7241 Some(Value::Float64(f)) => Some(*f),
7242 _ => None,
7243 }
7244}
7245
7246fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
7250 let mut count: u64 = 0;
7251 let mut sum: i128 = 0;
7252 let mut mn: i64 = i64::MAX;
7253 let mut mx: i64 = i64::MIN;
7254 while let Some(cols) = cursor.next_batch()? {
7255 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
7256 if crate::columnar::all_non_null(validity, data.len()) {
7257 count += data.len() as u64;
7259 sum += data.iter().map(|&v| v as i128).sum::<i128>();
7260 mn = mn.min(*data.iter().min().unwrap_or(&mn));
7261 mx = mx.max(*data.iter().max().unwrap_or(&mx));
7262 } else {
7263 for (i, &v) in data.iter().enumerate() {
7264 if crate::columnar::validity_bit(validity, i) {
7265 count += 1;
7266 sum += v as i128;
7267 mn = mn.min(v);
7268 mx = mx.max(v);
7269 }
7270 }
7271 }
7272 }
7273 }
7274 Ok((count, sum, mn, mx))
7275}
7276
7277fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
7279 let mut count: u64 = 0;
7280 let mut sum: f64 = 0.0;
7281 let mut mn: f64 = f64::INFINITY;
7282 let mut mx: f64 = f64::NEG_INFINITY;
7283 while let Some(cols) = cursor.next_batch()? {
7284 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
7285 if crate::columnar::all_non_null(validity, data.len()) {
7286 count += data.len() as u64;
7287 sum += data.iter().sum::<f64>();
7288 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
7289 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
7290 } else {
7291 for (i, &v) in data.iter().enumerate() {
7292 if crate::columnar::validity_bit(validity, i) {
7293 count += 1;
7294 sum += v;
7295 mn = mn.min(v);
7296 mx = mx.max(v);
7297 }
7298 }
7299 }
7300 }
7301 }
7302 Ok((count, sum, mn, mx))
7303}
7304
7305fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
7306 if count == 0 && !matches!(agg, NativeAgg::Count) {
7307 return NativeAggResult::Null;
7308 }
7309 match agg {
7310 NativeAgg::Count => NativeAggResult::Count(count),
7311 NativeAgg::Sum => match sum.try_into() {
7314 Ok(v) => NativeAggResult::Int(v),
7315 Err(_) => NativeAggResult::Null,
7316 },
7317 NativeAgg::Min => NativeAggResult::Int(mn),
7318 NativeAgg::Max => NativeAggResult::Int(mx),
7319 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
7320 }
7321}
7322
7323fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
7324 if count == 0 && !matches!(agg, NativeAgg::Count) {
7325 return NativeAggResult::Null;
7326 }
7327 match agg {
7328 NativeAgg::Count => NativeAggResult::Count(count),
7329 NativeAgg::Sum => NativeAggResult::Float(sum),
7330 NativeAgg::Min => NativeAggResult::Float(mn),
7331 NativeAgg::Max => NativeAggResult::Float(mx),
7332 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
7333 }
7334}
7335
7336fn agg_int(
7339 stats: &[crate::page::PageStat],
7340 decode: fn(Option<&[u8]>) -> Option<i64>,
7341) -> Option<(Option<i64>, Option<i64>, u64)> {
7342 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
7343 let mut any = false;
7344 for s in stats {
7345 if let Some(v) = decode(s.min.as_deref()) {
7346 mn = mn.min(v);
7347 any = true;
7348 }
7349 if let Some(v) = decode(s.max.as_deref()) {
7350 mx = mx.max(v);
7351 any = true;
7352 }
7353 nulls += s.null_count;
7354 }
7355 any.then_some((Some(mn), Some(mx), nulls))
7356}
7357
7358fn agg_float(
7360 stats: &[crate::page::PageStat],
7361 decode: fn(Option<&[u8]>) -> Option<f64>,
7362) -> Option<(Option<f64>, Option<f64>, u64)> {
7363 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
7364 let mut any = false;
7365 for s in stats {
7366 if let Some(v) = decode(s.min.as_deref()) {
7367 mn = mn.min(v);
7368 any = true;
7369 }
7370 if let Some(v) = decode(s.max.as_deref()) {
7371 mx = mx.max(v);
7372 any = true;
7373 }
7374 nulls += s.null_count;
7375 }
7376 any.then_some((Some(mn), Some(mx), nulls))
7377}
7378
7379type SecondaryIndexes = (
7381 HashMap<u16, BitmapIndex>,
7382 HashMap<u16, AnnIndex>,
7383 HashMap<u16, FmIndex>,
7384 HashMap<u16, SparseIndex>,
7385 HashMap<u16, MinHashIndex>,
7386);
7387
7388fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
7389 let mut bitmap = HashMap::new();
7390 let mut ann = HashMap::new();
7391 let mut fm = HashMap::new();
7392 let mut sparse = HashMap::new();
7393 let mut minhash = HashMap::new();
7394 for idef in &schema.indexes {
7395 match idef.kind {
7396 IndexKind::Bitmap => {
7397 bitmap.insert(idef.column_id, BitmapIndex::new());
7398 }
7399 IndexKind::Ann => {
7400 let dim = schema
7401 .columns
7402 .iter()
7403 .find(|c| c.id == idef.column_id)
7404 .and_then(|c| match c.ty {
7405 TypeId::Embedding { dim } => Some(dim as usize),
7406 _ => None,
7407 })
7408 .unwrap_or(0);
7409 ann.insert(idef.column_id, AnnIndex::new(dim));
7410 }
7411 IndexKind::FmIndex => {
7412 fm.insert(idef.column_id, FmIndex::new());
7413 }
7414 IndexKind::Sparse => {
7415 sparse.insert(idef.column_id, SparseIndex::new());
7416 }
7417 IndexKind::MinHash => {
7418 minhash.insert(idef.column_id, MinHashIndex::new());
7419 }
7420 _ => {}
7421 }
7422 }
7423 (bitmap, ann, fm, sparse, minhash)
7424}
7425
7426const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
7427 | ColumnFlags::AUTO_INCREMENT
7428 | ColumnFlags::ENCRYPTED
7429 | ColumnFlags::ENCRYPTED_INDEXABLE
7430 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
7431
7432fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
7433 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
7434 return Err(MongrelError::Schema(
7435 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
7436 ));
7437 }
7438 Ok(())
7439}
7440
7441fn validate_alter_column_type(
7442 schema: &Schema,
7443 old: &ColumnDef,
7444 next: &ColumnDef,
7445 has_stored_versions: bool,
7446) -> Result<()> {
7447 if old.ty == next.ty {
7448 return Ok(());
7449 }
7450 if schema.indexes.iter().any(|i| i.column_id == old.id) {
7451 return Err(MongrelError::Schema(format!(
7452 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
7453 old.name
7454 )));
7455 }
7456 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
7457 return Ok(());
7458 }
7459 Err(MongrelError::Schema(format!(
7460 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
7461 old.ty, next.ty
7462 )))
7463}
7464
7465fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
7466 matches!(
7467 (old, new),
7468 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
7469 )
7470}
7471
7472fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
7478 let mut prev: Option<i64> = None;
7479 for r in rows {
7480 match r.columns.get(&pk_id) {
7481 Some(Value::Int64(v)) => {
7482 if prev.is_some_and(|p| p >= *v) {
7483 return false;
7484 }
7485 prev = Some(*v);
7486 }
7487 _ => return false,
7488 }
7489 }
7490 true
7491}
7492
7493#[allow(clippy::too_many_arguments)]
7494fn index_into(
7495 schema: &Schema,
7496 row: &Row,
7497 hot: &mut HotIndex,
7498 bitmap: &mut HashMap<u16, BitmapIndex>,
7499 ann: &mut HashMap<u16, AnnIndex>,
7500 fm: &mut HashMap<u16, FmIndex>,
7501 sparse: &mut HashMap<u16, SparseIndex>,
7502 minhash: &mut HashMap<u16, MinHashIndex>,
7503) {
7504 for idef in &schema.indexes {
7505 let Some(val) = row.columns.get(&idef.column_id) else {
7506 continue;
7507 };
7508 match idef.kind {
7509 IndexKind::Bitmap => {
7510 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7511 b.insert(val.encode_key(), row.row_id);
7512 }
7513 }
7514 IndexKind::Ann => {
7515 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7516 a.insert(v, row.row_id);
7517 }
7518 }
7519 IndexKind::FmIndex => {
7520 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7521 f.insert(b.clone(), row.row_id);
7522 }
7523 }
7524 IndexKind::Sparse => {
7525 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7526 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7529 s.insert(&terms, row.row_id);
7530 }
7531 }
7532 }
7533 IndexKind::MinHash => {
7534 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7535 let tokens = crate::index::token_hashes_from_bytes(b);
7538 mh.insert(&tokens, row.row_id);
7539 }
7540 }
7541 _ => {}
7542 }
7543 }
7544 if let Some(pk_col) = schema.primary_key() {
7545 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7546 hot.insert(pk_val.encode_key(), row.row_id);
7547 }
7548 }
7549}
7550
7551#[allow(clippy::too_many_arguments)]
7554fn index_into_single(
7555 idef: &IndexDef,
7556 _schema: &Schema,
7557 row: &Row,
7558 _hot: &mut HotIndex,
7559 bitmap: &mut HashMap<u16, BitmapIndex>,
7560 ann: &mut HashMap<u16, AnnIndex>,
7561 fm: &mut HashMap<u16, FmIndex>,
7562 sparse: &mut HashMap<u16, SparseIndex>,
7563 minhash: &mut HashMap<u16, MinHashIndex>,
7564) {
7565 let Some(val) = row.columns.get(&idef.column_id) else {
7566 return;
7567 };
7568 match idef.kind {
7569 IndexKind::Bitmap => {
7570 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7571 b.insert(val.encode_key(), row.row_id);
7572 }
7573 }
7574 IndexKind::Ann => {
7575 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7576 a.insert(v, row.row_id);
7577 }
7578 }
7579 IndexKind::FmIndex => {
7580 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7581 f.insert(b.clone(), row.row_id);
7582 }
7583 }
7584 IndexKind::Sparse => {
7585 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7586 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7587 s.insert(&terms, row.row_id);
7588 }
7589 }
7590 }
7591 IndexKind::MinHash => {
7592 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7593 let tokens = crate::index::token_hashes_from_bytes(b);
7594 mh.insert(&tokens, row.row_id);
7595 }
7596 }
7597 _ => {}
7598 }
7599}
7600
7601fn eval_partial_predicate(
7607 pred: &str,
7608 columns_map: &HashMap<u16, &Value>,
7609 name_to_id: &HashMap<&str, u16>,
7610) -> bool {
7611 let lower = pred.trim().to_ascii_lowercase();
7612 if let Some(rest) = lower.strip_suffix(" is not null") {
7614 let col_name = rest.trim();
7615 if let Some(col_id) = name_to_id.get(col_name) {
7616 return columns_map
7617 .get(col_id)
7618 .is_some_and(|v| !matches!(v, Value::Null));
7619 }
7620 }
7621 if let Some(rest) = lower.strip_suffix(" is null") {
7623 let col_name = rest.trim();
7624 if let Some(col_id) = name_to_id.get(col_name) {
7625 return columns_map
7626 .get(col_id)
7627 .map_or(true, |v| matches!(v, Value::Null));
7628 }
7629 }
7630 true
7633}
7634
7635#[allow(dead_code)]
7641fn bulk_index_key(
7642 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7643 column_id: u16,
7644 ty: TypeId,
7645 col: &columnar::NativeColumn,
7646 i: usize,
7647) -> Option<Vec<u8>> {
7648 let encoded = columnar::encode_key_native(ty, col, i)?;
7649 #[cfg(feature = "encryption")]
7650 {
7651 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7652 if let Some((key, scheme)) = column_keys.get(&column_id) {
7653 return Some(match (*scheme, col) {
7654 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7655 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7656 ope_token_i64(key, data[i]).to_vec()
7657 }
7658 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7659 ope_token_f64(key, data[i]).to_vec()
7660 }
7661 _ => hmac_token(key, &encoded).to_vec(),
7662 });
7663 }
7664 }
7665 #[cfg(not(feature = "encryption"))]
7666 {
7667 let _ = (column_id, column_keys, col);
7668 }
7669 Some(encoded)
7670}
7671
7672pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7673 let json = serde_json::to_string_pretty(schema)
7674 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7675 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7676 Ok(())
7677}
7678
7679fn read_schema(dir: &Path) -> Result<Schema> {
7680 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7681 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7682}
7683
7684fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7685 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7686}
7687
7688fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7689 let n = list_wal_numbers(wal_dir)?;
7690 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7691}
7692
7693fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7694 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7695}
7696
7697fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7698 let _ = std::fs::create_dir_all(wal_dir);
7699 let mut max_n = None;
7700 for entry in std::fs::read_dir(wal_dir)? {
7701 let entry = entry?;
7702 let fname = entry.file_name();
7703 let Some(s) = fname.to_str() else {
7704 continue;
7705 };
7706 let Some(stripped) = s.strip_prefix("seg-") else {
7707 continue;
7708 };
7709 let Some(stripped) = stripped.strip_suffix(".wal") else {
7710 continue;
7711 };
7712 if let Ok(n) = stripped.parse::<u32>() {
7713 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7714 }
7715 }
7716 Ok(max_n)
7717}