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 data_generation: u64,
175 schema: Schema,
176 hot: HotIndex,
177 kek: Option<Arc<Kek>>,
180 column_keys: HashMap<u16, ([u8; 32], u8)>,
184 run_refs: Vec<RunRef>,
185 retiring: Vec<crate::manifest::RetiredRun>,
188 next_run_id: u64,
189 sync_byte_threshold: u64,
190 current_txn_id: u64,
195 bitmap: HashMap<u16, BitmapIndex>,
196 ann: HashMap<u16, AnnIndex>,
197 fm: HashMap<u16, FmIndex>,
198 sparse: HashMap<u16, SparseIndex>,
199 minhash: HashMap<u16, MinHashIndex>,
200 learned_range: HashMap<u16, ColumnLearnedRange>,
203 pk_by_row: HashMap<RowId, Vec<u8>>,
205 pinned: BTreeMap<Epoch, usize>,
208 pub(crate) live_count: u64,
211 reservoir: crate::reservoir::Reservoir,
214 reservoir_complete: bool,
222 had_deletes: bool,
226 agg_cache: HashMap<u64, CachedAgg>,
230 global_idx_epoch: u64,
234 indexes_complete: bool,
239 index_build_policy: IndexBuildPolicy,
241 pk_by_row_complete: bool,
248 flushed_epoch: u64,
251 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
254 snapshots: Arc<crate::retention::SnapshotRegistry>,
257 commit_lock: Arc<parking_lot::Mutex<()>>,
259 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
262 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
272 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
281 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
283 pending_delete_rids: roaring::RoaringBitmap,
286 pending_put_cols: std::collections::HashSet<u16>,
289 pending_rows: Vec<Row>,
295 pending_rows_auto_inc: Vec<bool>,
296 pending_dels: Vec<RowId>,
299 pending_truncate: Option<Epoch>,
303 auto_inc: Option<AutoIncState>,
306 ttl: Option<TtlPolicy>,
309}
310
311const _: () = {
318 const fn assert_sync<T: ?Sized + Sync>() {}
319 assert_sync::<Table>();
320};
321
322enum CachedData {
328 Rows(Arc<Vec<Row>>),
329 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
330}
331
332impl CachedData {
333 fn approx_bytes(&self) -> u64 {
334 match self {
335 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
336 CachedData::Columns(c) => c
337 .iter()
338 .map(|(_, c)| c.approx_bytes())
339 .sum::<u64>()
340 .saturating_add(c.len() as u64 * 16),
341 }
342 }
343}
344
345struct CachedEntry {
349 data: CachedData,
350 footprint: roaring::RoaringBitmap,
351 condition_cols: Vec<u16>,
352}
353
354struct ResultCache {
365 entries: std::collections::HashMap<u64, CachedEntry>,
366 order: std::collections::VecDeque<u64>,
367 bytes: u64,
368 max_bytes: u64,
369 dir: Option<std::path::PathBuf>,
370 #[allow(dead_code)]
371 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
372}
373
374#[derive(serde::Serialize, serde::Deserialize)]
376struct SerializedEntry {
377 condition_cols: Vec<u16>,
378 footprint_bits: Vec<u32>,
379 data: SerializedData,
380}
381
382#[derive(serde::Serialize, serde::Deserialize)]
383enum SerializedData {
384 Rows(Vec<Row>),
385 Columns(Vec<(u16, columnar::NativeColumn)>),
386}
387
388impl SerializedEntry {
389 fn from_entry(entry: &CachedEntry) -> Self {
390 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
391 let data = match &entry.data {
392 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
393 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
394 };
395 Self {
396 condition_cols: entry.condition_cols.clone(),
397 footprint_bits,
398 data,
399 }
400 }
401
402 fn into_entry(self) -> Option<CachedEntry> {
403 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
404 let data = match self.data {
405 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
406 SerializedData::Columns(c) => {
407 if !c.iter().all(|(_, col)| col.validate()) {
410 return None;
411 }
412 CachedData::Columns(Arc::new(c))
413 }
414 };
415 Some(CachedEntry {
416 data,
417 footprint,
418 condition_cols: self.condition_cols,
419 })
420 }
421}
422
423impl ResultCache {
424 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
425
426 fn new() -> Self {
427 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
428 }
429
430 fn with_max_bytes(max_bytes: u64) -> Self {
431 Self {
432 entries: std::collections::HashMap::new(),
433 order: std::collections::VecDeque::new(),
434 bytes: 0,
435 max_bytes,
436 dir: None,
437 cache_dek: None,
438 }
439 }
440
441 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
442 let _ = std::fs::create_dir_all(&dir);
443 self.dir = Some(dir);
444 self
445 }
446
447 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
448 self.cache_dek = dek;
449 self
450 }
451
452 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
453 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
454 }
455
456 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
460 let Some(path) = self.disk_path(key) else {
461 return;
462 };
463 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
464 Ok(s) => s,
465 Err(_) => return,
466 };
467 let on_disk = if let Some(dek) = &self.cache_dek {
469 match self.encrypt_cache(&serialized, dek) {
470 Some(b) => b,
471 None => return,
472 }
473 } else {
474 serialized
475 };
476 let tmp = path.with_extension("tmp");
477 use std::io::Write;
478 let write = || -> std::io::Result<()> {
479 let mut f = std::fs::File::create(&tmp)?;
480 f.write_all(&on_disk)?;
481 f.flush()?;
482 Ok(())
483 };
484 if write().is_err() {
485 let _ = std::fs::remove_file(&tmp);
486 return;
487 }
488 let _ = std::fs::rename(&tmp, &path);
489 }
490
491 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
493 let path = self.disk_path(key)?;
494 let bytes = std::fs::read(&path).ok()?;
495 let plaintext = if let Some(dek) = &self.cache_dek {
496 self.decrypt_cache(&bytes, dek)?
497 } else {
498 bytes
499 };
500 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
501 serialized.into_entry()
502 }
503
504 fn remove_from_disk(&self, key: u64) {
506 if let Some(path) = self.disk_path(key) {
507 let _ = std::fs::remove_file(&path);
508 }
509 }
510
511 #[cfg(feature = "encryption")]
513 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
514 use crate::encryption::Cipher;
515 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
516 let mut nonce = [0u8; 12];
517 crate::encryption::fill_random(&mut nonce);
518 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
519 let mut out = Vec::with_capacity(12 + ct.len());
520 out.extend_from_slice(&nonce);
521 out.extend_from_slice(&ct);
522 Some(out)
523 }
524
525 #[cfg(not(feature = "encryption"))]
526 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
527 None
528 }
529
530 #[cfg(feature = "encryption")]
532 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
533 use crate::encryption::Cipher;
534 if bytes.len() < 28 {
535 return None;
536 }
537 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
538 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
539 let ct = &bytes[12..];
540 cipher.decrypt_page(&nonce, ct).ok()
541 }
542
543 #[cfg(not(feature = "encryption"))]
544 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
545 None
546 }
547
548 fn load_persistent(&mut self) {
551 let Some(dir) = self.dir.as_ref().cloned() else {
552 return;
553 };
554 let entries = match std::fs::read_dir(&dir) {
555 Ok(e) => e,
556 Err(_) => return,
557 };
558 for entry in entries.flatten() {
559 let path = entry.path();
560 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
562 let _ = std::fs::remove_file(&path);
563 continue;
564 }
565 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
566 continue;
567 }
568 let stem = match path.file_stem().and_then(|s| s.to_str()) {
569 Some(s) => s,
570 None => continue,
571 };
572 let key = match u64::from_str_radix(stem, 16) {
573 Ok(k) => k,
574 Err(_) => continue,
575 };
576 let bytes = match std::fs::read(&path) {
577 Ok(b) => b,
578 Err(_) => continue,
579 };
580 let plaintext = if let Some(dek) = &self.cache_dek {
582 match self.decrypt_cache(&bytes, dek) {
583 Some(p) => p,
584 None => {
585 let _ = std::fs::remove_file(&path);
586 continue;
587 }
588 }
589 } else {
590 bytes
591 };
592 match bincode::deserialize::<SerializedEntry>(&plaintext) {
593 Ok(serialized) => {
594 if let Some(entry) = serialized.into_entry() {
595 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
596 self.entries.insert(key, entry);
597 self.order.push_back(key);
598 } else {
599 let _ = std::fs::remove_file(&path);
600 }
601 }
602 Err(_) => {
603 let _ = std::fs::remove_file(&path);
604 }
605 }
606 }
607 self.evict();
608 }
609
610 fn set_max_bytes(&mut self, max_bytes: u64) {
611 self.max_bytes = max_bytes;
612 self.evict();
613 }
614
615 fn touch(&mut self, key: u64) {
617 self.order.retain(|k| *k != key);
618 self.order.push_back(key);
619 }
620
621 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
622 let res = self.entries.get(&key).and_then(|e| match &e.data {
623 CachedData::Rows(r) => Some(r.clone()),
624 CachedData::Columns(_) => None,
625 });
626 if res.is_some() {
627 self.touch(key);
628 return res;
629 }
630 if let Some(entry) = self.load_from_disk(key) {
632 let res = match &entry.data {
633 CachedData::Rows(r) => Some(r.clone()),
634 CachedData::Columns(_) => None,
635 };
636 if res.is_some() {
637 let approx = entry.data.approx_bytes();
638 self.bytes = self.bytes.saturating_add(approx);
639 self.entries.insert(key, entry);
640 self.order.push_back(key);
641 self.evict();
642 return res;
643 }
644 }
645 None
646 }
647
648 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
649 let res = self.entries.get(&key).and_then(|e| match &e.data {
650 CachedData::Columns(c) => Some(c.clone()),
651 CachedData::Rows(_) => None,
652 });
653 if res.is_some() {
654 self.touch(key);
655 return res;
656 }
657 if let Some(entry) = self.load_from_disk(key) {
659 let res = match &entry.data {
660 CachedData::Columns(c) => Some(c.clone()),
661 CachedData::Rows(_) => None,
662 };
663 if res.is_some() {
664 let approx = entry.data.approx_bytes();
665 self.bytes = self.bytes.saturating_add(approx);
666 self.entries.insert(key, entry);
667 self.order.push_back(key);
668 self.evict();
669 return res;
670 }
671 }
672 None
673 }
674
675 fn insert(&mut self, key: u64, entry: CachedEntry) {
676 let approx = entry.data.approx_bytes();
677 if self.entries.remove(&key).is_some() {
678 self.order.retain(|k| *k != key);
679 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
680 }
681 self.store_to_disk(key, &entry);
683 self.bytes = self.bytes.saturating_add(approx);
684 self.entries.insert(key, entry);
685 self.order.push_back(key);
686 self.evict();
687 }
688
689 fn invalidate(
698 &mut self,
699 delete_rids: &roaring::RoaringBitmap,
700 put_cols: &std::collections::HashSet<u16>,
701 ) {
702 if self.entries.is_empty() {
703 return;
704 }
705 let has_deletes = !delete_rids.is_empty();
706 let to_remove: std::collections::HashSet<u64> = self
707 .entries
708 .iter()
709 .filter(|(_, e)| {
710 let delete_hit = if e.footprint.is_empty() {
711 has_deletes
712 } else {
713 e.footprint.intersection_len(delete_rids) > 0
714 };
715 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
716 delete_hit || col_hit
717 })
718 .map(|(&k, _)| k)
719 .collect();
720 for key in &to_remove {
721 if let Some(e) = self.entries.remove(key) {
722 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
723 }
724 self.remove_from_disk(*key);
725 }
726 if !to_remove.is_empty() {
727 self.order.retain(|k| !to_remove.contains(k));
728 }
729 }
730
731 fn clear(&mut self) {
732 if let Some(dir) = &self.dir {
734 if let Ok(entries) = std::fs::read_dir(dir) {
735 for entry in entries.flatten() {
736 let path = entry.path();
737 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
738 let _ = std::fs::remove_file(&path);
739 }
740 }
741 }
742 }
743 self.entries.clear();
744 self.order.clear();
745 self.bytes = 0;
746 }
747
748 fn evict(&mut self) {
749 while self.bytes > self.max_bytes {
750 let Some(k) = self.order.pop_front() else {
751 break;
752 };
753 if let Some(e) = self.entries.remove(&k) {
754 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
755 self.remove_from_disk(k);
759 }
760 }
761 }
762}
763
764type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
771
772fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
773 let _ = kek;
774 #[cfg(feature = "encryption")]
775 {
776 if let Some(k) = kek {
777 return (
778 Some(k.derive_table_wal_key(_table_id)),
779 Some(k.derive_cache_key()),
780 );
781 }
782 }
783 (None, None)
784}
785
786#[cfg(feature = "encryption")]
788fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
789 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
790}
791
792#[cfg(not(feature = "encryption"))]
793fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
794 Box::new(crate::encryption::PlaintextCipher)
795}
796
797fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
798 let Some(kek) = kek else {
799 return HashMap::new();
800 };
801 #[cfg(feature = "encryption")]
802 {
803 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
804 schema
805 .columns
806 .iter()
807 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
808 .map(|c| {
809 let scheme = if schema
810 .indexes
811 .iter()
812 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
813 {
814 SCHEME_OPE_RANGE
815 } else {
816 SCHEME_HMAC_EQ
817 };
818 let key: [u8; 32] = *kek.derive_column_key(c.id);
819 (c.id, (key, scheme))
820 })
821 .collect()
822 }
823 #[cfg(not(feature = "encryption"))]
824 {
825 let _ = (kek, schema);
826 HashMap::new()
827 }
828}
829
830pub(crate) struct SharedCtx {
835 pub epoch: Arc<EpochAuthority>,
836 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
837 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
838 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
839 pub kek: Option<Arc<Kek>>,
840 pub commit_lock: Arc<parking_lot::Mutex<()>>,
846 pub shared: Option<SharedWalCtx>,
850 pub table_name: Option<String>,
853 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
856 pub read_only: bool,
858}
859
860#[derive(Clone)]
866pub(crate) struct SharedWalCtx {
867 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
868 pub group: Arc<GroupCommit>,
869 pub poisoned: Arc<AtomicBool>,
870 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
871 pub change_wake: tokio::sync::broadcast::Sender<()>,
872}
873
874enum WalSink {
877 Private(Wal),
878 Shared(SharedWalCtx),
879}
880
881impl SharedCtx {
882 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
886 let n_shards = if cache_dir.is_some() {
890 1
891 } else {
892 crate::cache::CACHE_SHARDS
893 };
894 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
895 let page_cache = if let Some(d) = cache_dir {
896 Arc::new(crate::cache::Sharded::new(1, || {
897 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
898 }))
899 } else {
900 Arc::new(crate::cache::Sharded::new(n_shards, || {
901 crate::cache::PageCache::new(per_shard)
902 }))
903 };
904 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
905 let decoded_cache = Arc::new(crate::cache::Sharded::new(
906 crate::cache::CACHE_SHARDS,
907 || crate::cache::DecodedPageCache::new(decoded_per_shard),
908 ));
909 Self {
910 epoch: Arc::new(EpochAuthority::new(0)),
911 page_cache,
912 decoded_cache,
913 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
914 kek,
915 commit_lock: Arc::new(parking_lot::Mutex::new(())),
916 shared: None,
917 table_name: None,
918 auth: None,
919 read_only: false,
920 }
921 }
922}
923
924fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
928 use crate::query::Condition;
929 match c {
930 Condition::Pk(_)
932 | Condition::BitmapEq { .. }
933 | Condition::BitmapIn { .. }
934 | Condition::BytesPrefix { .. }
935 | Condition::IsNull { .. }
936 | Condition::IsNotNull { .. } => 0,
937 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
939 1
940 }
941 Condition::FmContains { .. }
943 | Condition::FmContainsAll { .. }
944 | Condition::Ann { .. }
945 | Condition::SparseMatch { .. } => 2,
946 }
947}
948
949impl Table {
950 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
951 let dir = dir.as_ref().to_path_buf();
952 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
953 Self::create_in(&dir, schema, table_id, ctx)
954 }
955
956 #[cfg(feature = "encryption")]
967 pub fn create_encrypted(
968 dir: impl AsRef<Path>,
969 schema: Schema,
970 table_id: u64,
971 passphrase: &str,
972 ) -> Result<Self> {
973 let dir = dir.as_ref();
974 std::fs::create_dir_all(dir.join(META_DIR))?;
975 let salt = crate::encryption::random_salt();
976 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
977 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
978 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
979 Self::create_in(dir, schema, table_id, ctx)
980 }
981
982 #[cfg(feature = "encryption")]
987 pub fn create_with_key(
988 dir: impl AsRef<Path>,
989 schema: Schema,
990 table_id: u64,
991 key: &[u8],
992 ) -> Result<Self> {
993 let dir = dir.as_ref();
994 std::fs::create_dir_all(dir.join(META_DIR))?;
995 let salt = crate::encryption::random_salt();
996 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
997 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
998 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
999 Self::create_in(dir, schema, table_id, ctx)
1000 }
1001
1002 #[cfg(feature = "encryption")]
1004 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1005 let dir = dir.as_ref();
1006 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1007 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1008 MongrelError::NotFound(format!(
1009 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1010 salt_path
1011 ))
1012 })?;
1013 if salt_bytes.len() != crate::encryption::SALT_LEN {
1014 return Err(MongrelError::InvalidArgument(format!(
1015 "salt file is {} bytes, expected {}",
1016 salt_bytes.len(),
1017 crate::encryption::SALT_LEN
1018 )));
1019 }
1020 let mut salt = [0u8; crate::encryption::SALT_LEN];
1021 salt.copy_from_slice(&salt_bytes);
1022 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1023 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1024 Self::open_in(dir, ctx)
1025 }
1026
1027 pub(crate) fn create_in(
1028 dir: impl AsRef<Path>,
1029 schema: Schema,
1030 table_id: u64,
1031 ctx: SharedCtx,
1032 ) -> Result<Self> {
1033 schema.validate_auto_increment()?;
1034 schema.validate_defaults()?;
1035 schema.validate_ai()?;
1036 for index in &schema.indexes {
1037 index.validate_options()?;
1038 }
1039 let dir = dir.as_ref().to_path_buf();
1040 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1041 write_schema(&dir, &schema)?;
1042 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1043 let (wal, current_txn_id) = match ctx.shared.clone() {
1046 Some(s) => (WalSink::Shared(s), 0),
1047 None => {
1048 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1049 let mut w = if let Some(ref dk) = wal_dek {
1050 Wal::create_with_cipher(
1051 dir.join(WAL_DIR).join("seg-000000.wal"),
1052 Epoch(0),
1053 Some(make_cipher(dk)),
1054 0,
1055 )?
1056 } else {
1057 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1058 };
1059 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1060 (WalSink::Private(w), 1)
1061 }
1062 };
1063 let mut manifest = Manifest::new(table_id, schema.schema_id);
1064 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1069 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1070 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1071 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1072 let auto_inc = resolve_auto_inc(&schema);
1073 let rcache_dir = dir.join(RCACHE_DIR);
1074 Ok(Self {
1075 dir,
1076 table_id,
1077 name: ctx.table_name.unwrap_or_default(),
1078 auth: ctx.auth,
1079 read_only: ctx.read_only,
1080 wal,
1081 memtable: Memtable::new(),
1082 mutable_run: MutableRun::new(),
1083 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1084 compaction_zstd_level: 3,
1085 allocator: RowIdAllocator::new(0),
1086 epoch: ctx.epoch,
1087 persisted_epoch: 0,
1088 data_generation: 0,
1089 schema,
1090 hot: HotIndex::new(),
1091 kek: ctx.kek,
1092 column_keys,
1093 run_refs: Vec::new(),
1094 retiring: Vec::new(),
1095 next_run_id: 1,
1096 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1097 current_txn_id,
1098 bitmap,
1099 ann,
1100 fm,
1101 sparse,
1102 minhash,
1103 learned_range: HashMap::new(),
1104 pk_by_row: HashMap::new(),
1105 pinned: BTreeMap::new(),
1106 live_count: 0,
1107 reservoir: crate::reservoir::Reservoir::default(),
1108 reservoir_complete: true,
1109 had_deletes: false,
1110 agg_cache: HashMap::new(),
1111 global_idx_epoch: 0,
1112 indexes_complete: true,
1113 index_build_policy: IndexBuildPolicy::default(),
1114 pk_by_row_complete: false,
1115 flushed_epoch: 0,
1116 page_cache: ctx.page_cache,
1117 decoded_cache: ctx.decoded_cache,
1118 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1119 snapshots: ctx.snapshots,
1120 commit_lock: ctx.commit_lock,
1121 result_cache: Arc::new(parking_lot::Mutex::new(
1122 ResultCache::new()
1123 .with_dir(rcache_dir)
1124 .with_cache_dek(cache_dek.clone()),
1125 )),
1126 pending_delete_rids: roaring::RoaringBitmap::new(),
1127 pending_put_cols: std::collections::HashSet::new(),
1128 pending_rows: Vec::new(),
1129 pending_rows_auto_inc: Vec::new(),
1130 pending_dels: Vec::new(),
1131 pending_truncate: None,
1132 wal_dek,
1133 auto_inc,
1134 ttl: None,
1135 })
1136 }
1137
1138 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1142 let dir = dir.as_ref();
1143 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1144 Self::open_in(dir, ctx)
1145 }
1146
1147 #[cfg(feature = "encryption")]
1150 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1151 let dir = dir.as_ref();
1152 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1153 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1154 MongrelError::NotFound(format!(
1155 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1156 salt_path
1157 ))
1158 })?;
1159 let salt_len = crate::encryption::SALT_LEN;
1160 if salt_bytes.len() != salt_len {
1161 return Err(MongrelError::InvalidArgument(format!(
1162 "encryption salt is {} bytes, expected {salt_len}",
1163 salt_bytes.len()
1164 )));
1165 }
1166 let mut salt = [0u8; 16];
1167 salt.copy_from_slice(&salt_bytes);
1168 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1169 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1170 let t = Self::open_in(dir, ctx)?;
1171 Ok(t)
1172 }
1173
1174 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1175 let dir = dir.as_ref().to_path_buf();
1176 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1177 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1178 let schema: Schema = read_schema(&dir)?;
1179 schema.validate_ai()?;
1180 for index in &schema.indexes {
1181 index.validate_options()?;
1182 }
1183 let replay_epoch = Epoch(manifest.current_epoch);
1184 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1185 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1189 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1190 None => {
1191 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1192 let replayed = match &active {
1194 Some(path) => {
1195 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1196 crate::wal::replay_with_cipher(path, cipher)?
1197 }
1198 None => Vec::new(),
1199 };
1200 let mut w = match &active {
1201 Some(path) => Wal::create_with_cipher(
1202 path,
1203 replay_epoch,
1204 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1205 0,
1206 )?,
1207 None => Wal::create_with_cipher(
1208 dir.join(WAL_DIR).join("seg-000000.wal"),
1209 replay_epoch,
1210 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1211 0,
1212 )?,
1213 };
1214 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1215 (WalSink::Private(w), replayed, 1)
1216 }
1217 };
1218
1219 let mut memtable = Memtable::new();
1220 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1221 let persisted_epoch = manifest.current_epoch;
1222 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1229 s.next = manifest.auto_inc_next;
1230 s.seeded = manifest.auto_inc_next > 0;
1231 s
1232 });
1233
1234 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1241 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1242 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1243 std::collections::BTreeMap::new();
1244 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1245 let mut saw_delete = false;
1246 for record in replayed {
1247 let txn_id = record.txn_id;
1248 match record.op {
1249 Op::Put { rows, .. } => {
1250 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1251 for row in &rows {
1252 allocator.advance_to(row.row_id);
1253 if let Some(ai) = auto_inc.as_mut() {
1254 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1255 if *n + 1 > ai.next {
1256 ai.next = *n + 1;
1257 }
1258 }
1259 }
1260 }
1261 staged_puts.entry(txn_id).or_default().extend(rows);
1262 }
1263 Op::Delete { row_ids, .. } => {
1264 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1265 }
1266 Op::TxnCommit { epoch, .. } => {
1267 let commit_epoch = Epoch(epoch);
1268 if let Some(puts) = staged_puts.remove(&txn_id) {
1269 if commit_epoch.0 > manifest.flushed_epoch {
1270 for row in &puts {
1271 memtable.upsert(row.clone());
1272 }
1273 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1274 }
1275 }
1276 if let Some(dels) = staged_deletes.remove(&txn_id) {
1277 saw_delete = true;
1278 if commit_epoch.0 > manifest.flushed_epoch {
1279 for rid in dels {
1280 memtable.tombstone(rid, commit_epoch);
1281 replayed_deletes.push((rid, commit_epoch));
1282 }
1283 }
1284 }
1285 }
1286 Op::TxnAbort => {
1287 staged_puts.remove(&txn_id);
1288 staged_deletes.remove(&txn_id);
1289 }
1290 Op::TruncateTable { .. }
1291 | Op::ExternalTableState { .. }
1292 | Op::Flush { .. }
1293 | Op::Ddl(_)
1294 | Op::BeforeImage { .. }
1295 | Op::CommitTimestamp { .. } => {}
1296 }
1297 }
1298
1299 let rcache_dir = dir.join(RCACHE_DIR);
1300 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1301 let mut db = Self {
1302 dir,
1303 table_id: manifest.table_id,
1304 name: ctx.table_name.unwrap_or_default(),
1305 auth: ctx.auth,
1306 read_only: ctx.read_only,
1307 wal,
1308 memtable,
1309 mutable_run: MutableRun::new(),
1310 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1311 compaction_zstd_level: 3,
1312 allocator,
1313 epoch: ctx.epoch,
1314 persisted_epoch,
1315 data_generation: persisted_epoch,
1316 schema,
1317 hot: HotIndex::new(),
1318 kek: ctx.kek,
1319 column_keys,
1320 run_refs: manifest.runs.clone(),
1321 retiring: manifest.retiring.clone(),
1322 next_run_id: manifest
1323 .runs
1324 .iter()
1325 .map(|r| r.run_id as u64 + 1)
1326 .max()
1327 .unwrap_or(1),
1328 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1329 current_txn_id,
1330 bitmap: HashMap::new(),
1331 ann: HashMap::new(),
1332 fm: HashMap::new(),
1333 sparse: HashMap::new(),
1334 minhash: HashMap::new(),
1335 learned_range: HashMap::new(),
1336 pk_by_row: HashMap::new(),
1337 pinned: BTreeMap::new(),
1338 live_count: manifest.live_count,
1339 reservoir: crate::reservoir::Reservoir::default(),
1340 reservoir_complete: false,
1341 had_deletes: saw_delete
1342 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1343 != manifest.live_count,
1344 agg_cache: HashMap::new(),
1345 global_idx_epoch: manifest.global_idx_epoch,
1346 indexes_complete: true,
1347 index_build_policy: IndexBuildPolicy::default(),
1348 pk_by_row_complete: false,
1349 flushed_epoch: manifest.flushed_epoch,
1350 page_cache: ctx.page_cache,
1351 decoded_cache: ctx.decoded_cache,
1352 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1353 snapshots: ctx.snapshots,
1354 commit_lock: ctx.commit_lock,
1355 result_cache: Arc::new(parking_lot::Mutex::new(
1356 ResultCache::new()
1357 .with_dir(rcache_dir)
1358 .with_cache_dek(cache_dek.clone()),
1359 )),
1360 pending_delete_rids: roaring::RoaringBitmap::new(),
1361 pending_put_cols: std::collections::HashSet::new(),
1362 pending_rows: Vec::new(),
1363 pending_rows_auto_inc: Vec::new(),
1364 pending_dels: Vec::new(),
1365 pending_truncate: None,
1366 wal_dek,
1367 auto_inc,
1368 ttl: manifest.ttl,
1369 };
1370
1371 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1374
1375 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1380 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1381 c.epoch_built == manifest.global_idx_epoch
1382 && manifest.global_idx_epoch > 0
1383 && manifest
1384 .runs
1385 .iter()
1386 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1387 });
1388 if let Some(loaded) = checkpoint {
1389 if checkpoint_valid {
1390 db.hot = loaded.hot;
1391 db.bitmap = loaded.bitmap;
1392 db.ann = loaded.ann;
1393 db.fm = loaded.fm;
1394 db.sparse = loaded.sparse;
1395 db.minhash = loaded.minhash;
1396 db.learned_range = loaded.learned_range;
1397 }
1400 }
1401 if !checkpoint_valid {
1402 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1403 db.bitmap = bitmap;
1404 db.ann = ann;
1405 db.fm = fm;
1406 db.sparse = sparse;
1407 db.minhash = minhash;
1408 db.rebuild_indexes_from_runs()?;
1409 db.build_learned_ranges()?;
1410 }
1411
1412 for (epoch, group) in replayed_puts {
1417 let (losers, winner_pks) = db.partition_pk_winners(&group);
1418 for (key, &row_id) in &winner_pks {
1419 if let Some(old_rid) = db.hot.get(key) {
1420 if old_rid != row_id {
1421 db.tombstone_row(old_rid, epoch, false);
1422 }
1423 }
1424 }
1425 for &loser_rid in &losers {
1426 db.tombstone_row(loser_rid, epoch, false);
1427 }
1428 for (key, row_id) in winner_pks {
1429 db.insert_hot_pk(key, row_id);
1430 }
1431 if db.schema.primary_key().is_none() {
1432 for r in &group {
1433 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1434 }
1435 }
1436 for r in &group {
1437 if !losers.contains(&r.row_id) {
1438 db.index_row(r);
1439 }
1440 }
1441 }
1442 for (rid, epoch) in &replayed_deletes {
1446 db.remove_hot_for_row(*rid, *epoch);
1447 }
1448
1449 db.result_cache.lock().load_persistent();
1456 Ok(db)
1457 }
1458
1459 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1465 if self.reservoir_complete {
1466 return Ok(());
1467 }
1468 self.rebuild_reservoir()?;
1469 self.reservoir_complete = true;
1470 Ok(())
1471 }
1472
1473 fn rebuild_reservoir(&mut self) -> Result<()> {
1476 let snap = self.snapshot();
1477 let rows = self.visible_rows(snap)?;
1478 self.reservoir.reset();
1479 for r in rows {
1480 self.reservoir.offer(r.row_id.0);
1481 }
1482 Ok(())
1483 }
1484
1485 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1486 self.hot = HotIndex::new();
1487 self.pk_by_row.clear();
1488 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1489 self.bitmap = bitmap;
1490 self.ann = ann;
1491 self.fm = fm;
1492 self.sparse = sparse;
1493 self.minhash = minhash;
1494 let snapshot = Epoch(u64::MAX);
1495 let ttl_now = unix_nanos_now();
1496 for rr in self.run_refs.clone() {
1497 let mut reader = self.open_reader(rr.run_id)?;
1498 for row in reader.visible_rows(snapshot)? {
1499 if self.row_expired_at(&row, ttl_now) {
1500 continue;
1501 }
1502 let tok_row = self.tokenized_for_indexes(&row);
1503 index_into(
1504 &self.schema,
1505 &tok_row,
1506 &mut self.hot,
1507 &mut self.bitmap,
1508 &mut self.ann,
1509 &mut self.fm,
1510 &mut self.sparse,
1511 &mut self.minhash,
1512 );
1513 }
1514 }
1515 for row in self.mutable_run.visible_versions(snapshot) {
1516 if row.deleted {
1517 self.remove_hot_for_row(row.row_id, snapshot);
1518 } else if !self.row_expired_at(&row, ttl_now) {
1519 self.index_row(&row);
1520 }
1521 }
1522 for row in self.memtable.visible_versions(snapshot) {
1523 if row.deleted {
1524 self.remove_hot_for_row(row.row_id, snapshot);
1525 } else if !self.row_expired_at(&row, ttl_now) {
1526 self.index_row(&row);
1527 }
1528 }
1529 self.refresh_pk_by_row_from_hot();
1530 Ok(())
1531 }
1532
1533 fn refresh_pk_by_row_from_hot(&mut self) {
1534 self.pk_by_row_complete = true;
1535 if self.schema.primary_key().is_none() {
1536 self.pk_by_row.clear();
1537 return;
1538 }
1539 self.pk_by_row = self
1545 .hot
1546 .entries()
1547 .into_iter()
1548 .map(|(key, row_id)| (row_id, key))
1549 .collect();
1550 }
1551
1552 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1553 if self.schema.primary_key().is_some() {
1554 self.pk_by_row.insert(row_id, key.clone());
1555 }
1556 self.hot.insert(key, row_id);
1557 }
1558
1559 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1563 self.learned_range.clear();
1564 if self.run_refs.len() != 1 {
1565 return Ok(());
1566 }
1567 let cols: Vec<(u16, usize)> = self
1568 .schema
1569 .indexes
1570 .iter()
1571 .filter(|i| i.kind == IndexKind::LearnedRange)
1572 .map(|i| {
1573 (
1574 i.column_id,
1575 i.options
1576 .learned_range
1577 .as_ref()
1578 .map(|options| options.epsilon)
1579 .unwrap_or(16),
1580 )
1581 })
1582 .collect();
1583 if cols.is_empty() {
1584 return Ok(());
1585 }
1586 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1587 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1588 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1589 _ => return Ok(()),
1590 };
1591 for (cid, epsilon) in cols {
1592 let ty = self
1593 .schema
1594 .columns
1595 .iter()
1596 .find(|c| c.id == cid)
1597 .map(|c| c.ty.clone())
1598 .unwrap_or(TypeId::Int64);
1599 match ty {
1600 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1601 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1602 let pairs: Vec<(i64, u64)> = data
1603 .iter()
1604 .zip(row_ids.iter())
1605 .map(|(v, r)| (*v, *r))
1606 .collect();
1607 self.learned_range.insert(
1608 cid,
1609 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
1610 );
1611 }
1612 }
1613 TypeId::Float64 => {
1614 if let columnar::NativeColumn::Float64 { data, .. } =
1615 reader.column_native(cid)?
1616 {
1617 let pairs: Vec<(f64, u64)> = data
1618 .iter()
1619 .zip(row_ids.iter())
1620 .map(|(v, r)| (*v, *r))
1621 .collect();
1622 self.learned_range.insert(
1623 cid,
1624 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
1625 );
1626 }
1627 }
1628 _ => {}
1629 }
1630 }
1631 Ok(())
1632 }
1633
1634 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1641 if self.indexes_complete {
1642 crate::trace::QueryTrace::record(|t| {
1643 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1644 });
1645 return Ok(());
1646 }
1647 crate::trace::QueryTrace::record(|t| {
1648 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1649 });
1650 self.rebuild_indexes_from_runs()?;
1651 self.build_learned_ranges()?;
1652 self.indexes_complete = true;
1653 let epoch = self.current_epoch();
1654 self.checkpoint_indexes(epoch);
1655 Ok(())
1656 }
1657
1658 fn pending_epoch(&self) -> Epoch {
1659 Epoch(self.epoch.visible().0 + 1)
1660 }
1661
1662 fn is_shared(&self) -> bool {
1665 matches!(self.wal, WalSink::Shared(_))
1666 }
1667
1668 fn ensure_txn_id(&mut self) -> u64 {
1672 if self.current_txn_id == 0 {
1673 let id = match &self.wal {
1674 WalSink::Shared(s) => {
1675 let mut g = s.txn_ids.lock();
1676 let v = *g;
1677 *g = g.wrapping_add(1);
1678 v
1679 }
1680 WalSink::Private(_) => 1,
1681 };
1682 self.current_txn_id = id;
1683 }
1684 self.current_txn_id
1685 }
1686
1687 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1690 self.ensure_writable()?;
1691 let txn_id = self.ensure_txn_id();
1692 let table_id = self.table_id;
1693 match &mut self.wal {
1694 WalSink::Private(w) => {
1695 w.append_txn(txn_id, op)?;
1696 }
1697 WalSink::Shared(s) => {
1698 s.wal.lock().append(txn_id, table_id, op)?;
1699 }
1700 }
1701 Ok(())
1702 }
1703
1704 fn ensure_writable(&self) -> Result<()> {
1705 if self.read_only {
1706 Err(MongrelError::ReadOnlyReplica)
1707 } else {
1708 Ok(())
1709 }
1710 }
1711
1712 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1723 match &self.auth {
1724 Some(checker) => checker.check(&self.name, perm),
1725 None => Ok(()),
1726 }
1727 }
1728 pub fn require_select(&self) -> Result<()> {
1733 self.require(crate::auth_state::RequiredPermission::Select)
1734 }
1735 fn require_insert(&self) -> Result<()> {
1736 self.require(crate::auth_state::RequiredPermission::Insert)
1737 }
1738 #[allow(dead_code)]
1742 fn require_update(&self) -> Result<()> {
1743 self.require(crate::auth_state::RequiredPermission::Update)
1744 }
1745 fn require_delete(&self) -> Result<()> {
1746 self.require(crate::auth_state::RequiredPermission::Delete)
1747 }
1748
1749 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1752 self.require_insert()?;
1753 Ok(self.put_returning(columns)?.0)
1754 }
1755
1756 pub fn put_returning(
1761 &mut self,
1762 mut columns: Vec<(u16, Value)>,
1763 ) -> Result<(RowId, Option<i64>)> {
1764 self.require_insert()?;
1765 let assigned = self.fill_auto_inc(&mut columns)?;
1766 self.apply_defaults(&mut columns)?;
1767 self.schema.validate_values(&columns)?;
1768 let row_id = if self.schema.clustered {
1773 self.derive_clustered_row_id(&columns)?
1774 } else {
1775 self.allocator.alloc()
1776 };
1777 let epoch = self.pending_epoch();
1778 let mut row = Row::new(row_id, epoch);
1779 for (col_id, val) in columns {
1780 row.columns.insert(col_id, val);
1781 }
1782 self.commit_rows(vec![row], assigned.is_some())?;
1783 Ok((row_id, assigned))
1784 }
1785
1786 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1789 self.require_insert()?;
1790 Ok(self
1791 .put_batch_returning(batch)?
1792 .into_iter()
1793 .map(|(r, _)| r)
1794 .collect())
1795 }
1796
1797 pub fn put_batch_returning(
1800 &mut self,
1801 batch: Vec<Vec<(u16, Value)>>,
1802 ) -> Result<Vec<(RowId, Option<i64>)>> {
1803 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1804 for mut cols in batch {
1805 let assigned = self.fill_auto_inc(&mut cols)?;
1806 self.apply_defaults(&mut cols)?;
1807 filled.push((cols, assigned));
1808 }
1809 for (cols, _) in &filled {
1810 self.schema.validate_values(cols)?;
1811 }
1812 let epoch = self.pending_epoch();
1813 let mut rows = Vec::with_capacity(filled.len());
1814 let mut ids = Vec::with_capacity(filled.len());
1815 for (cols, assigned) in filled {
1816 let row_id = if self.schema.clustered {
1817 self.derive_clustered_row_id(&cols)?
1818 } else {
1819 self.allocator.alloc()
1820 };
1821 let mut row = Row::new(row_id, epoch);
1822 for (c, v) in cols {
1823 row.columns.insert(c, v);
1824 }
1825 ids.push((row_id, assigned));
1826 rows.push(row);
1827 }
1828 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1829 self.commit_rows(rows, all_auto_generated)?;
1830 Ok(ids)
1831 }
1832
1833 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1839 self.ensure_writable()?;
1840 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1841 return Ok(None);
1842 };
1843 let pos = columns.iter().position(|(c, _)| *c == cid);
1844 let assigned = match pos {
1845 Some(i) => match &columns[i].1 {
1846 Value::Null => {
1847 let next = self.alloc_auto_inc_value()?;
1848 columns[i].1 = Value::Int64(next);
1849 Some(next)
1850 }
1851 Value::Int64(n) => {
1852 self.advance_auto_inc_past(*n)?;
1853 None
1854 }
1855 other => {
1856 return Err(MongrelError::InvalidArgument(format!(
1857 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1858 other
1859 )))
1860 }
1861 },
1862 None => {
1863 let next = self.alloc_auto_inc_value()?;
1864 columns.push((cid, Value::Int64(next)));
1865 Some(next)
1866 }
1867 };
1868 Ok(assigned)
1869 }
1870
1871 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1877 for col in &self.schema.columns {
1878 let Some(expr) = &col.default_value else {
1879 continue;
1880 };
1881 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
1883 continue;
1884 }
1885 let pos = columns.iter().position(|(c, _)| *c == col.id);
1886 let needs_default = match pos {
1887 None => true,
1888 Some(i) => matches!(columns[i].1, Value::Null),
1889 };
1890 if !needs_default {
1891 continue;
1892 }
1893 let v = match expr {
1894 crate::schema::DefaultExpr::Static(v) => v.clone(),
1895 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
1896 crate::schema::DefaultExpr::Uuid => {
1897 let mut buf = [0u8; 16];
1898 getrandom::getrandom(&mut buf)
1899 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
1900 Value::Uuid(buf)
1901 }
1902 };
1903 match pos {
1904 None => columns.push((col.id, v)),
1905 Some(i) => columns[i].1 = v,
1906 }
1907 }
1908 Ok(())
1909 }
1910
1911 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1913 self.ensure_auto_inc_seeded()?;
1914 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1916 let v = ai.next;
1917 ai.next = ai.next.saturating_add(1);
1918 Ok(v)
1919 }
1920
1921 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1924 self.ensure_auto_inc_seeded()?;
1925 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1926 let floor = used.saturating_add(1).max(1);
1927 if ai.next < floor {
1928 ai.next = floor;
1929 }
1930 Ok(())
1931 }
1932
1933 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1938 let needs_seed = match self.auto_inc {
1939 Some(ai) => !ai.seeded,
1940 None => return Ok(()),
1941 };
1942 if !needs_seed {
1943 return Ok(());
1944 }
1945 if self.seed_empty_auto_inc() {
1946 return Ok(());
1947 }
1948 let cid = self
1949 .auto_inc
1950 .as_ref()
1951 .expect("auto-inc column present")
1952 .column_id;
1953 let max = self.scan_max_int64(cid)?;
1954 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1955 let floor = max.saturating_add(1).max(1);
1956 if ai.next < floor {
1957 ai.next = floor;
1958 }
1959 ai.seeded = true;
1960 Ok(())
1961 }
1962
1963 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1964 if n == 0 || self.auto_inc.is_none() {
1965 return Ok(None);
1966 }
1967 self.ensure_auto_inc_seeded()?;
1968 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1969 let start = ai.next;
1970 ai.next = ai.next.saturating_add(n as i64);
1971 Ok(Some(start))
1972 }
1973
1974 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1978 let mut max: i64 = 0;
1979 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1980 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1981 if *n > max {
1982 max = *n;
1983 }
1984 }
1985 }
1986 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1987 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1988 if *n > max {
1989 max = *n;
1990 }
1991 }
1992 }
1993 for rr in self.run_refs.clone() {
1994 let reader = self.open_reader(rr.run_id)?;
1995 if let Some(stats) = reader.column_page_stats(column_id) {
1996 for s in stats {
1997 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1998 if n > max {
1999 max = n;
2000 }
2001 }
2002 }
2003 } else if reader.has_column(column_id) {
2004 if let columnar::NativeColumn::Int64 { data, validity } =
2005 reader.column_native_shared(column_id)?
2006 {
2007 for (i, n) in data.iter().enumerate() {
2008 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2009 {
2010 max = *n;
2011 }
2012 }
2013 }
2014 }
2015 }
2016 Ok(max)
2017 }
2018
2019 fn seed_empty_auto_inc(&mut self) -> bool {
2020 let Some(ai) = self.auto_inc.as_mut() else {
2021 return false;
2022 };
2023 if ai.seeded || self.live_count != 0 {
2024 return false;
2025 }
2026 if ai.next < 1 {
2027 ai.next = 1;
2028 }
2029 ai.seeded = true;
2030 true
2031 }
2032
2033 fn advance_auto_inc_from_native_columns(
2034 &mut self,
2035 columns: &[(u16, columnar::NativeColumn)],
2036 n: usize,
2037 live_before: u64,
2038 ) -> Result<()> {
2039 let Some(ai) = self.auto_inc.as_mut() else {
2040 return Ok(());
2041 };
2042 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2043 return Ok(());
2044 };
2045 let columnar::NativeColumn::Int64 { data, validity } = col else {
2046 return Err(MongrelError::InvalidArgument(format!(
2047 "AUTO_INCREMENT column {} must be Int64",
2048 ai.column_id
2049 )));
2050 };
2051 let max = if native_int64_strictly_increasing(col, n) {
2052 data.get(n.saturating_sub(1)).copied()
2053 } else {
2054 data.iter()
2055 .take(n)
2056 .enumerate()
2057 .filter_map(|(i, v)| {
2058 if validity.is_empty() || columnar::validity_bit(validity, i) {
2059 Some(*v)
2060 } else {
2061 None
2062 }
2063 })
2064 .max()
2065 };
2066 if let Some(max) = max {
2067 let floor = max.saturating_add(1).max(1);
2068 if ai.next < floor {
2069 ai.next = floor;
2070 }
2071 if ai.seeded || live_before == 0 {
2072 ai.seeded = true;
2073 }
2074 }
2075 Ok(())
2076 }
2077
2078 fn fill_auto_inc_native_columns(
2079 &mut self,
2080 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2081 n: usize,
2082 ) -> Result<()> {
2083 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2084 return Ok(());
2085 };
2086 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2087 if let Some(start) = self.alloc_auto_inc_range(n)? {
2088 columns.push((
2089 cid,
2090 columnar::NativeColumn::Int64 {
2091 data: (start..start.saturating_add(n as i64)).collect(),
2092 validity: vec![0xFF; n.div_ceil(8)],
2093 },
2094 ));
2095 }
2096 return Ok(());
2097 };
2098
2099 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2100 return Err(MongrelError::InvalidArgument(format!(
2101 "AUTO_INCREMENT column {cid} must be Int64"
2102 )));
2103 };
2104 if data.len() < n {
2105 return Err(MongrelError::InvalidArgument(format!(
2106 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2107 data.len()
2108 )));
2109 }
2110 if columnar::all_non_null(validity, n) {
2111 return Ok(());
2112 }
2113 if validity.iter().all(|b| *b == 0) {
2114 if let Some(start) = self.alloc_auto_inc_range(n)? {
2115 for (i, slot) in data.iter_mut().take(n).enumerate() {
2116 *slot = start.saturating_add(i as i64);
2117 }
2118 *validity = vec![0xFF; n.div_ceil(8)];
2119 }
2120 return Ok(());
2121 }
2122
2123 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2124 for (i, slot) in data.iter_mut().enumerate().take(n) {
2125 if columnar::validity_bit(validity, i) {
2126 self.advance_auto_inc_past(*slot)?;
2127 } else {
2128 *slot = self.alloc_auto_inc_value()?;
2129 }
2130 }
2131 *validity = new_validity;
2132 Ok(())
2133 }
2134
2135 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2149 self.ensure_writable()?;
2150 if self.auto_inc.is_none() {
2151 return Ok(None);
2152 }
2153 Ok(Some(self.alloc_auto_inc_value()?))
2154 }
2155
2156 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2162 let payload = bincode::serialize(&rows)?;
2163 self.wal_append_data(Op::Put {
2164 table_id: self.table_id,
2165 rows: payload,
2166 })?;
2167 if self.is_shared() {
2168 self.pending_rows_auto_inc
2169 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2170 self.pending_rows.extend(rows);
2171 } else {
2172 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2173 }
2174 Ok(())
2175 }
2176
2177 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2182 self.apply_put_rows_inner(rows, true)
2183 }
2184
2185 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2186 if check_existing_pk {
2187 self.ensure_indexes_complete()?;
2188 }
2189 if rows.len() == 1 {
2193 let row = rows.into_iter().next().expect("len checked");
2194 return self.apply_put_row_single(row, check_existing_pk);
2195 }
2196 let pk_id = self.schema.primary_key().map(|c| c.id);
2213 let probe = match pk_id {
2214 Some(pid) => {
2215 check_existing_pk
2216 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2217 }
2218 None => false,
2219 };
2220 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2223 for r in rows {
2224 for &cid in r.columns.keys() {
2225 self.pending_put_cols.insert(cid);
2226 }
2227 match pk_id {
2228 Some(pid) if probe || maintain_pk_by_row => {
2229 if let Some(pk_val) = r.columns.get(&pid) {
2230 let key = self.index_lookup_key(pid, pk_val);
2231 if probe {
2232 if let Some(old_rid) = self.hot.get(&key) {
2233 if old_rid != r.row_id {
2234 self.tombstone_row(old_rid, r.committed_epoch, true);
2235 }
2236 }
2237 }
2238 if maintain_pk_by_row {
2239 self.pk_by_row.insert(r.row_id, key);
2240 }
2241 }
2242 }
2243 Some(_) => {}
2244 None => {
2245 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2246 }
2247 }
2248 self.index_row(&r);
2249 self.reservoir.offer(r.row_id.0);
2250 self.memtable.upsert(r);
2251 self.live_count = self.live_count.saturating_add(1);
2254 }
2255 self.data_generation = self.data_generation.wrapping_add(1);
2256 Ok(())
2257 }
2258
2259 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2263 for &cid in row.columns.keys() {
2264 self.pending_put_cols.insert(cid);
2265 }
2266 let epoch = row.committed_epoch;
2267 if let Some(pk_col) = self.schema.primary_key() {
2268 let pk_id = pk_col.id;
2269 if let Some(pk_val) = row.columns.get(&pk_id) {
2270 let maintain_pk_by_row = self.pk_by_row_complete;
2274 if check_existing_pk || maintain_pk_by_row {
2275 let key = self.index_lookup_key(pk_id, pk_val);
2276 if check_existing_pk {
2277 if let Some(old_rid) = self.hot.get(&key) {
2278 if old_rid != row.row_id {
2279 self.tombstone_row(old_rid, epoch, true);
2280 }
2281 }
2282 }
2283 if maintain_pk_by_row {
2284 self.pk_by_row.insert(row.row_id, key);
2285 }
2286 }
2287 }
2288 } else {
2289 self.hot
2290 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2291 }
2292 self.index_row(&row);
2293 self.reservoir.offer(row.row_id.0);
2294 self.memtable.upsert(row);
2295 self.live_count = self.live_count.saturating_add(1);
2296 self.data_generation = self.data_generation.wrapping_add(1);
2297 Ok(())
2298 }
2299
2300 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2303 self.allocator.alloc()
2304 }
2305
2306 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2312 let pk = self.schema.primary_key().ok_or_else(|| {
2313 MongrelError::Schema("clustered table requires a single-column primary key".into())
2314 })?;
2315 let pk_val = columns
2316 .iter()
2317 .find(|(id, _)| *id == pk.id)
2318 .map(|(_, v)| v)
2319 .ok_or_else(|| {
2320 MongrelError::Schema(format!(
2321 "clustered table missing primary key column {} ({})",
2322 pk.id, pk.name
2323 ))
2324 })?;
2325 let key_bytes = pk_val.encode_key();
2326 let mut hash: u64 = 0xcbf29ce484222325;
2328 for &b in &key_bytes {
2329 hash ^= b as u64;
2330 hash = hash.wrapping_mul(0x100000001b3);
2331 }
2332 Ok(RowId(hash.max(1)))
2335 }
2336
2337 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2345 self.ensure_indexes_complete()?;
2346 let n = rows.len();
2347 for r in rows {
2348 for &cid in r.columns.keys() {
2349 self.pending_put_cols.insert(cid);
2350 }
2351 }
2352 let (losers, winner_pks) = self.partition_pk_winners(rows);
2353 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2354 for (key, &row_id) in &winner_pks {
2356 if let Some(old_rid) = self.hot.get(key) {
2357 if old_rid != row_id {
2358 self.tombstone_row(old_rid, epoch, true);
2359 }
2360 }
2361 }
2362 for &loser_rid in &losers {
2365 self.tombstone_row(loser_rid, epoch, false);
2366 }
2367 for (key, row_id) in winner_pks {
2369 self.insert_hot_pk(key, row_id);
2370 }
2371 if self.schema.primary_key().is_none() {
2372 for r in rows {
2373 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2374 }
2375 }
2376 for r in rows {
2377 self.allocator.advance_to(r.row_id);
2378 if !losers.contains(&r.row_id) {
2379 self.index_row(r);
2380 }
2381 }
2382 for r in rows {
2383 if !losers.contains(&r.row_id) {
2384 self.reservoir.offer(r.row_id.0);
2385 }
2386 }
2387 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2388 self.data_generation = self.data_generation.wrapping_add(1);
2389 Ok(())
2390 }
2391
2392 pub(crate) fn recover_apply(
2397 &mut self,
2398 rows: Vec<Row>,
2399 deletes: Vec<(RowId, Epoch)>,
2400 ) -> Result<()> {
2401 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2405 std::collections::BTreeMap::new();
2406 for row in rows {
2407 self.allocator.advance_to(row.row_id);
2408 if let Some(ai) = self.auto_inc.as_mut() {
2413 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2414 if *n + 1 > ai.next {
2415 ai.next = *n + 1;
2416 }
2417 }
2418 }
2419 by_epoch.entry(row.committed_epoch).or_default().push(row);
2420 }
2421 for (epoch, group) in by_epoch {
2422 let (losers, winner_pks) = self.partition_pk_winners(&group);
2423 for (key, &row_id) in &winner_pks {
2425 if let Some(old_rid) = self.hot.get(key) {
2426 if old_rid != row_id {
2427 self.tombstone_row(old_rid, epoch, false);
2428 }
2429 }
2430 }
2431 for (key, row_id) in winner_pks {
2432 self.insert_hot_pk(key, row_id);
2433 }
2434 if self.schema.primary_key().is_none() {
2435 for r in &group {
2436 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2437 }
2438 }
2439 for r in &group {
2440 if !losers.contains(&r.row_id) {
2441 self.memtable.upsert(r.clone());
2442 self.index_row(r);
2443 }
2444 }
2445 }
2446 for (rid, epoch) in deletes {
2447 self.memtable.tombstone(rid, epoch);
2448 self.remove_hot_for_row(rid, epoch);
2449 }
2450 self.reservoir_complete = false;
2453 Ok(())
2454 }
2455
2456 pub(crate) fn flushed_epoch(&self) -> u64 {
2458 self.flushed_epoch
2459 }
2460
2461 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2462 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2463 }
2464
2465 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2467 self.schema.validate_values(cells)
2468 }
2469
2470 fn validate_columns_not_null(
2474 &self,
2475 columns: &[(u16, columnar::NativeColumn)],
2476 n: usize,
2477 ) -> Result<()> {
2478 let by_id: HashMap<u16, &columnar::NativeColumn> =
2479 columns.iter().map(|(id, c)| (*id, c)).collect();
2480 for col in &self.schema.columns {
2481 if !col.flags.contains(ColumnFlags::NULLABLE) {
2482 match by_id.get(&col.id) {
2483 None => {
2484 return Err(MongrelError::InvalidArgument(format!(
2485 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2486 col.name, col.id
2487 )));
2488 }
2489 Some(c) => {
2490 if c.null_count(n) != 0 {
2491 return Err(MongrelError::InvalidArgument(format!(
2492 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2493 col.name, col.id
2494 )));
2495 }
2496 }
2497 }
2498 }
2499 if let TypeId::Enum { variants } = &col.ty {
2500 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2501 if by_id.contains_key(&col.id) {
2502 return Err(MongrelError::InvalidArgument(format!(
2503 "column '{}' ({}) enum requires a bytes column",
2504 col.name, col.id
2505 )));
2506 }
2507 continue;
2508 };
2509 for index in 0..n {
2510 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2511 continue;
2512 };
2513 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2514 return Err(MongrelError::InvalidArgument(format!(
2515 "column '{}' ({}) enum value {:?} is not one of {:?}",
2516 col.name,
2517 col.id,
2518 String::from_utf8_lossy(value),
2519 variants
2520 )));
2521 }
2522 }
2523 }
2524 }
2525 Ok(())
2526 }
2527
2528 fn bulk_pk_winner_indices(
2533 &self,
2534 columns: &[(u16, columnar::NativeColumn)],
2535 n: usize,
2536 ) -> Option<Vec<usize>> {
2537 let pk_col = self.schema.primary_key()?;
2538 let pk_id = pk_col.id;
2539 let pk_ty = pk_col.ty.clone();
2540 let by_id: HashMap<u16, &columnar::NativeColumn> =
2541 columns.iter().map(|(id, c)| (*id, c)).collect();
2542 let pk_native = by_id.get(&pk_id)?;
2543 if native_int64_strictly_increasing(pk_native, n) {
2544 return None;
2545 }
2546 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2548 let mut null_pk_rows: Vec<usize> = Vec::new();
2549 for i in 0..n {
2550 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2551 Some(key) => {
2552 last.insert(key, i);
2553 }
2554 None => null_pk_rows.push(i),
2555 }
2556 }
2557 let mut winners: HashSet<usize> = last.values().copied().collect();
2558 for i in null_pk_rows {
2559 winners.insert(i);
2560 }
2561 Some((0..n).filter(|i| winners.contains(i)).collect())
2562 }
2563
2564 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2566 self.require_delete()?;
2567 let epoch = self.pending_epoch();
2568 self.wal_append_data(Op::Delete {
2569 table_id: self.table_id,
2570 row_ids: vec![row_id],
2571 })?;
2572 if self.is_shared() {
2573 self.pending_dels.push(row_id);
2574 } else {
2575 self.apply_delete(row_id, epoch);
2576 }
2577 Ok(())
2578 }
2579
2580 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2581 let pre = self.get(row_id, self.snapshot());
2582 self.delete(row_id)?;
2583 Ok(pre.map(|row| {
2584 let mut columns: Vec<_> = row.columns.into_iter().collect();
2585 columns.sort_by_key(|(id, _)| *id);
2586 OwnedRow { columns }
2587 }))
2588 }
2589
2590 pub fn truncate(&mut self) -> Result<()> {
2592 self.require_delete()?;
2593 let epoch = self.pending_epoch();
2594 self.wal_append_data(Op::TruncateTable {
2595 table_id: self.table_id,
2596 })?;
2597 self.pending_rows.clear();
2598 self.pending_rows_auto_inc.clear();
2599 self.pending_dels.clear();
2600 self.pending_truncate = Some(epoch);
2601 Ok(())
2602 }
2603
2604 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2606 for rr in std::mem::take(&mut self.run_refs) {
2607 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2608 }
2609 for r in std::mem::take(&mut self.retiring) {
2610 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2611 }
2612 self.memtable = Memtable::new();
2613 self.mutable_run = MutableRun::new();
2614 self.hot = HotIndex::new();
2615 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2616 self.bitmap = bitmap;
2617 self.ann = ann;
2618 self.fm = fm;
2619 self.sparse = sparse;
2620 self.minhash = minhash;
2621 self.learned_range.clear();
2622 self.pk_by_row.clear();
2623 self.pk_by_row_complete = false;
2624 self.live_count = 0;
2625 self.reservoir = crate::reservoir::Reservoir::default();
2626 self.reservoir_complete = true;
2627 self.had_deletes = true;
2628 self.agg_cache.clear();
2629 self.global_idx_epoch = 0;
2630 self.indexes_complete = true;
2631 self.pending_delete_rids.clear();
2632 self.pending_put_cols.clear();
2633 self.pending_rows.clear();
2634 self.pending_rows_auto_inc.clear();
2635 self.pending_dels.clear();
2636 self.clear_result_cache();
2637 self.invalidate_index_checkpoint();
2638 self.data_generation = self.data_generation.wrapping_add(1);
2639 Ok(())
2640 }
2641
2642 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2645 self.remove_hot_for_row(row_id, epoch);
2646 self.tombstone_row(row_id, epoch, true);
2647 self.data_generation = self.data_generation.wrapping_add(1);
2648 }
2649
2650 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2654 let tombstone = Row {
2655 row_id,
2656 committed_epoch: epoch,
2657 columns: std::collections::HashMap::new(),
2658 deleted: true,
2659 };
2660 self.memtable.upsert(tombstone);
2661 self.pk_by_row.remove(&row_id);
2662 if adjust_live_count {
2663 self.live_count = self.live_count.saturating_sub(1);
2664 }
2665 self.pending_delete_rids.insert(row_id.0 as u32);
2667 self.had_deletes = true;
2670 self.agg_cache.clear();
2671 }
2672
2673 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2677 let Some(pk_col) = self.schema.primary_key() else {
2678 return;
2679 };
2680 if self.pk_by_row_complete {
2683 if let Some(key) = self.pk_by_row.remove(&row_id) {
2684 if self.hot.get(&key) == Some(row_id) {
2685 self.hot.remove(&key);
2686 }
2687 }
2688 return;
2689 }
2690 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2709 if self.indexes_complete {
2710 let pk_val = self
2711 .memtable
2712 .get_version(row_id, lookup_epoch)
2713 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2714 .or_else(|| {
2715 self.mutable_run
2716 .get_version(row_id, lookup_epoch)
2717 .filter(|(_, r)| !r.deleted)
2718 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2719 })
2720 .or_else(|| {
2721 self.run_refs.iter().find_map(|rr| {
2722 let mut reader = self.open_reader(rr.run_id).ok()?;
2723 let (_, deleted, val) = reader
2724 .get_version_column(row_id, lookup_epoch, pk_col.id)
2725 .ok()??;
2726 if deleted {
2727 return None;
2728 }
2729 val
2730 })
2731 });
2732 if let Some(pk_val) = pk_val {
2733 let key = self.index_lookup_key(pk_col.id, &pk_val);
2734 if self.hot.get(&key) == Some(row_id) {
2735 self.hot.remove(&key);
2736 }
2737 return;
2738 }
2739 }
2740 self.refresh_pk_by_row_from_hot();
2745 if let Some(key) = self.pk_by_row.remove(&row_id) {
2746 if self.hot.get(&key) == Some(row_id) {
2747 self.hot.remove(&key);
2748 }
2749 }
2750 }
2751
2752 fn partition_pk_winners(
2757 &self,
2758 rows: &[Row],
2759 ) -> (
2760 std::collections::HashSet<RowId>,
2761 std::collections::HashMap<Vec<u8>, RowId>,
2762 ) {
2763 let mut losers = std::collections::HashSet::new();
2764 let Some(pk_col) = self.schema.primary_key() else {
2765 return (losers, std::collections::HashMap::new());
2766 };
2767 let pk_id = pk_col.id;
2768 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2769 std::collections::HashMap::new();
2770 for r in rows {
2771 let Some(pk_val) = r.columns.get(&pk_id) else {
2772 continue;
2773 };
2774 let key = self.index_lookup_key(pk_id, pk_val);
2775 if let Some(&old_rid) = winners.get(&key) {
2776 losers.insert(old_rid);
2777 }
2778 winners.insert(key, r.row_id);
2779 }
2780 (losers, winners)
2781 }
2782
2783 fn index_row(&mut self, row: &Row) {
2784 if row.deleted {
2785 return;
2786 }
2787 let any_predicate = self
2795 .schema
2796 .indexes
2797 .iter()
2798 .any(|idx| idx.predicate.is_some());
2799 if any_predicate {
2800 let columns_map: HashMap<u16, &Value> =
2801 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2802 let name_to_id: HashMap<&str, u16> = self
2803 .schema
2804 .columns
2805 .iter()
2806 .map(|c| (c.name.as_str(), c.id))
2807 .collect();
2808 for idx in &self.schema.indexes {
2809 if let Some(pred) = &idx.predicate {
2810 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2811 continue; }
2813 }
2814 index_into_single(
2816 idx,
2817 &self.schema,
2818 row,
2819 &mut self.hot,
2820 &mut self.bitmap,
2821 &mut self.ann,
2822 &mut self.fm,
2823 &mut self.sparse,
2824 &mut self.minhash,
2825 );
2826 }
2827 return;
2828 }
2829 if self.column_keys.is_empty() {
2833 index_into(
2834 &self.schema,
2835 row,
2836 &mut self.hot,
2837 &mut self.bitmap,
2838 &mut self.ann,
2839 &mut self.fm,
2840 &mut self.sparse,
2841 &mut self.minhash,
2842 );
2843 return;
2844 }
2845 let effective_row = self.tokenized_for_indexes(row);
2846 index_into(
2847 &self.schema,
2848 &effective_row,
2849 &mut self.hot,
2850 &mut self.bitmap,
2851 &mut self.ann,
2852 &mut self.fm,
2853 &mut self.sparse,
2854 &mut self.minhash,
2855 );
2856 }
2857
2858 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2864 if self.column_keys.is_empty() {
2865 return row.clone();
2866 }
2867 #[cfg(feature = "encryption")]
2868 {
2869 use crate::encryption::SCHEME_HMAC_EQ;
2870 let mut tok = row.clone();
2871 for (&cid, &(_, scheme)) in &self.column_keys {
2872 if scheme != SCHEME_HMAC_EQ {
2873 continue;
2874 }
2875 if let Some(v) = tok.columns.get(&cid).cloned() {
2876 if let Some(t) = self.tokenize_value(cid, &v) {
2877 tok.columns.insert(cid, t);
2878 }
2879 }
2880 }
2881 tok
2882 }
2883 #[cfg(not(feature = "encryption"))]
2884 {
2885 row.clone()
2886 }
2887 }
2888
2889 pub fn commit(&mut self) -> Result<Epoch> {
2894 self.ensure_writable()?;
2895 if self.is_shared() {
2896 self.commit_shared()
2897 } else {
2898 self.commit_private()
2899 }
2900 }
2901
2902 fn commit_private(&mut self) -> Result<Epoch> {
2904 let commit_lock = Arc::clone(&self.commit_lock);
2908 let _g = commit_lock.lock();
2909 let new_epoch = self.epoch.bump_assigned();
2910 let txn_id = self.current_txn_id;
2911 match &mut self.wal {
2915 WalSink::Private(w) => {
2916 w.append_txn(
2917 txn_id,
2918 Op::TxnCommit {
2919 epoch: new_epoch.0,
2920 added_runs: Vec::new(),
2921 },
2922 )?;
2923 w.sync()?;
2924 }
2925 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2926 }
2927 if let Some(epoch) = self.pending_truncate.take() {
2929 self.apply_truncate(epoch)?;
2930 }
2931 self.invalidate_pending_cache();
2932 self.persist_manifest(new_epoch)?;
2933 self.epoch.publish_in_order(new_epoch);
2937 self.current_txn_id += 1;
2938 self.data_generation = self.data_generation.wrapping_add(1);
2939 Ok(new_epoch)
2940 }
2941
2942 fn commit_shared(&mut self) -> Result<Epoch> {
2948 use std::sync::atomic::Ordering;
2949 let s = match &self.wal {
2950 WalSink::Shared(s) => s.clone(),
2951 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2952 };
2953 if s.poisoned.load(Ordering::Relaxed) {
2954 return Err(MongrelError::Other(
2955 "database poisoned by fsync error".into(),
2956 ));
2957 }
2958 let commit_lock = Arc::clone(&self.commit_lock);
2965 let _g = commit_lock.lock();
2966 let txn_id = self.ensure_txn_id();
2969 let (new_epoch, commit_seq) = {
2970 let mut wal = s.wal.lock();
2971 let new_epoch = self.epoch.bump_assigned();
2972 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2973 (new_epoch, seq)
2974 };
2975 s.group
2976 .await_durable(&s.wal, commit_seq)
2977 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2978
2979 if self.pending_truncate.take().is_some() {
2982 self.apply_truncate(new_epoch)?;
2983 }
2984 let mut rows = std::mem::take(&mut self.pending_rows);
2985 if !rows.is_empty() {
2986 for r in &mut rows {
2987 r.committed_epoch = new_epoch;
2988 }
2989 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2990 let all_auto_generated =
2991 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2992 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2993 } else {
2994 self.pending_rows_auto_inc.clear();
2995 }
2996 let dels = std::mem::take(&mut self.pending_dels);
2997 for rid in dels {
2998 self.apply_delete(rid, new_epoch);
2999 }
3000
3001 self.invalidate_pending_cache();
3002 self.persist_manifest(new_epoch)?;
3003 self.epoch.publish_in_order(new_epoch);
3004 let _ = s.change_wake.send(());
3005 self.current_txn_id = 0;
3007 self.data_generation = self.data_generation.wrapping_add(1);
3008 Ok(new_epoch)
3009 }
3010
3011 pub fn flush(&mut self) -> Result<Epoch> {
3019 self.ensure_indexes_complete()?;
3020 let epoch = self.commit()?;
3021 let rows = self.memtable.drain_sorted();
3022 if !rows.is_empty() {
3023 self.mutable_run.insert_many(rows);
3024 }
3025 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3026 self.spill_mutable_run(epoch)?;
3027 self.mark_flushed(epoch)?;
3031 self.persist_manifest(epoch)?;
3032 self.build_learned_ranges()?;
3033 self.checkpoint_indexes(epoch);
3036 }
3037 Ok(epoch)
3040 }
3041
3042 pub fn force_flush(&mut self) -> Result<Epoch> {
3051 let saved = self.mutable_run_spill_bytes;
3052 self.mutable_run_spill_bytes = 1;
3053 let result = self.flush();
3054 self.mutable_run_spill_bytes = saved;
3055 result
3056 }
3057
3058 pub fn close(&mut self) -> Result<()> {
3065 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3066 self.force_flush()?;
3067 }
3068 Ok(())
3069 }
3070
3071 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3078 let op = Op::Flush {
3079 table_id: self.table_id,
3080 flushed_epoch: epoch.0,
3081 };
3082 match &mut self.wal {
3083 WalSink::Private(w) => {
3084 w.append_system(op)?;
3085 w.sync()?;
3086 }
3087 WalSink::Shared(s) => {
3088 s.wal.lock().append_system(op)?;
3093 }
3094 }
3095 self.flushed_epoch = epoch.0;
3096 if matches!(self.wal, WalSink::Private(_)) {
3097 self.rotate_wal(epoch)?;
3098 }
3099 Ok(())
3100 }
3101
3102 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3106 let rows = self.mutable_run.drain_sorted();
3107 if rows.is_empty() {
3108 return Ok(());
3109 }
3110 let run_id = self.next_run_id;
3111 self.next_run_id += 1;
3112 let path = self.run_path(run_id);
3113 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3114 if let Some(kek) = &self.kek {
3115 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3116 }
3117 let header = writer.write(&path, &rows)?;
3118 self.run_refs.push(RunRef {
3119 run_id: run_id as u128,
3120 level: 0,
3121 epoch_created: epoch.0,
3122 row_count: header.row_count,
3123 });
3124 Ok(())
3125 }
3126
3127 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3131 self.mutable_run_spill_bytes = bytes.max(1);
3132 }
3133
3134 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3138 self.compaction_zstd_level = level;
3139 }
3140
3141 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3145 self.result_cache.lock().set_max_bytes(max_bytes);
3146 }
3147
3148 pub(crate) fn clear_result_cache(&mut self) {
3152 self.result_cache.lock().clear();
3153 }
3154
3155 pub fn mutable_run_len(&self) -> usize {
3157 self.mutable_run.len()
3158 }
3159
3160 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3163 self.mutable_run.drain_sorted()
3164 }
3165
3166 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3171 let epoch = self.commit()?;
3172 let n = batch.len();
3173 if n == 0 {
3174 return Ok(epoch);
3175 }
3176 for row in &batch {
3177 self.schema.validate_values(row)?;
3178 }
3179 let live_before = self.live_count;
3180 self.spill_mutable_run(epoch)?;
3184 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3185 && self.indexes_complete
3186 && self.run_refs.is_empty()
3187 && self.memtable.is_empty()
3188 && self.mutable_run.is_empty();
3189 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3195 use rayon::prelude::*;
3196 self.schema
3197 .columns
3198 .par_iter()
3199 .map(|cdef| {
3200 (
3201 cdef.id,
3202 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3203 )
3204 })
3205 .collect::<Vec<_>>()
3206 };
3207 drop(batch);
3208 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3213 self.validate_columns_not_null(&user_columns, n)?;
3214 let winner_idx = self
3215 .bulk_pk_winner_indices(&user_columns, n)
3216 .filter(|idx| idx.len() != n);
3217 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3218 match winner_idx.as_deref() {
3219 Some(idx) => {
3220 let compacted = user_columns
3221 .iter()
3222 .map(|(id, c)| (*id, c.gather(idx)))
3223 .collect();
3224 (compacted, idx.len())
3225 }
3226 None => (std::mem::take(&mut user_columns), n),
3227 };
3228 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3229 let first = self.allocator.alloc_range(write_n as u64).0;
3230 for rid in first..first + write_n as u64 {
3231 self.reservoir.offer(rid);
3232 }
3233 let run_id = self.next_run_id;
3234 self.next_run_id += 1;
3235 let path = self.run_path(run_id);
3236 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3237 .clean(true)
3238 .with_lz4()
3239 .with_native_endian();
3240 if let Some(kek) = &self.kek {
3241 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3242 }
3243 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3244 self.run_refs.push(RunRef {
3245 run_id: run_id as u128,
3246 level: 0,
3247 epoch_created: epoch.0,
3248 row_count: header.row_count,
3249 });
3250 self.live_count = self.live_count.saturating_add(write_n as u64);
3251 if eager_index_build {
3252 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3253 self.index_columns_bulk(&write_columns, &row_ids);
3254 self.indexes_complete = true;
3255 self.build_learned_ranges()?;
3256 } else {
3257 self.indexes_complete = false;
3258 }
3259 self.mark_flushed(epoch)?;
3260 self.persist_manifest(epoch)?;
3261 if eager_index_build {
3262 self.checkpoint_indexes(epoch);
3263 }
3264 self.clear_result_cache();
3265 Ok(epoch)
3266 }
3267
3268 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3271 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3272 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3273 let segment_no = segment
3276 .file_stem()
3277 .and_then(|s| s.to_str())
3278 .and_then(|s| s.strip_prefix("seg-"))
3279 .and_then(|s| s.parse::<u64>().ok())
3280 .unwrap_or(0);
3281 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3282 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3283 wal.sync()?;
3284 self.wal = WalSink::Private(wal);
3285 Ok(())
3286 }
3287
3288 pub(crate) fn invalidate_pending_cache(&mut self) {
3293 self.result_cache
3294 .lock()
3295 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3296 self.pending_delete_rids.clear();
3297 self.pending_put_cols.clear();
3298 }
3299
3300 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3301 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3302 m.current_epoch = epoch.0;
3303 m.next_row_id = self.allocator.current().0;
3304 m.runs = self.run_refs.clone();
3305 m.live_count = self.live_count;
3306 m.global_idx_epoch = self.global_idx_epoch;
3307 m.flushed_epoch = self.flushed_epoch;
3308 m.retiring = self.retiring.clone();
3309 m.auto_inc_next = match self.auto_inc {
3313 Some(ai) if ai.seeded => ai.next,
3314 _ => 0,
3315 };
3316 m.ttl = self.ttl;
3317 let meta_dek = self.manifest_meta_dek();
3318 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3319 Ok(())
3320 }
3321
3322 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3328 if !self.indexes_complete {
3331 return;
3332 }
3333 let snap = global_idx::IndexSnapshot {
3334 hot: &self.hot,
3335 bitmap: &self.bitmap,
3336 ann: &self.ann,
3337 fm: &self.fm,
3338 sparse: &self.sparse,
3339 minhash: &self.minhash,
3340 learned_range: &self.learned_range,
3341 };
3342 let idx_dek = self.idx_dek();
3344 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3345 .is_ok()
3346 {
3347 self.global_idx_epoch = epoch.0;
3348 let _ = self.persist_manifest(epoch);
3349 }
3350 }
3351
3352 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3355 self.global_idx_epoch = 0;
3356 global_idx::remove(&self.dir);
3357 let _ = self.persist_manifest(self.epoch.visible());
3358 }
3359
3360 pub(crate) fn mark_indexes_incomplete(&mut self) {
3361 self.indexes_complete = false;
3362 self.invalidate_index_checkpoint();
3363 }
3364
3365 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3368 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3369 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3370 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3371 best = Some((epoch, row));
3372 }
3373 }
3374 for rr in &self.run_refs {
3375 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3376 continue;
3377 };
3378 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3379 continue;
3380 };
3381 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3382 best = Some((epoch, row));
3383 }
3384 }
3385 let now_nanos = unix_nanos_now();
3386 match best {
3387 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3388 Some((_, r)) => Some(r),
3389 None => None,
3390 }
3391 }
3392
3393 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3397 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3398 let mut fold = |row: Row| {
3399 best.entry(row.row_id.0)
3400 .and_modify(|e| {
3401 if row.committed_epoch > e.0 {
3402 *e = (row.committed_epoch, row.clone());
3403 }
3404 })
3405 .or_insert_with(|| (row.committed_epoch, row));
3406 };
3407 for row in self.memtable.visible_versions(snapshot.epoch) {
3408 fold(row);
3409 }
3410 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3411 fold(row);
3412 }
3413 for rr in &self.run_refs {
3414 let mut reader = self.open_reader(rr.run_id)?;
3415 for row in reader.visible_versions(snapshot.epoch)? {
3416 fold(row);
3417 }
3418 }
3419 let now_nanos = unix_nanos_now();
3420 let mut out: Vec<Row> = best
3421 .into_values()
3422 .filter_map(|(_, r)| {
3423 if r.deleted || self.row_expired_at(&r, now_nanos) {
3424 None
3425 } else {
3426 Some(r)
3427 }
3428 })
3429 .collect();
3430 out.sort_by_key(|r| r.row_id);
3431 Ok(out)
3432 }
3433
3434 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3441 if self.ttl.is_none()
3442 && self.memtable.is_empty()
3443 && self.mutable_run.is_empty()
3444 && self.run_refs.len() == 1
3445 {
3446 let rr = self.run_refs[0].clone();
3447 let mut reader = self.open_reader(rr.run_id)?;
3448 let idxs = reader.visible_indices(snapshot.epoch)?;
3449 let mut cols = Vec::with_capacity(self.schema.columns.len());
3450 for cdef in &self.schema.columns {
3451 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3452 }
3453 return Ok(cols);
3454 }
3455 let rows = self.visible_rows(snapshot)?;
3457 let mut cols: Vec<(u16, Vec<Value>)> = self
3458 .schema
3459 .columns
3460 .iter()
3461 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3462 .collect();
3463 for r in &rows {
3464 for (cid, vec) in cols.iter_mut() {
3465 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3466 }
3467 }
3468 Ok(cols)
3469 }
3470
3471 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3473 let row_id = self.hot.get(key)?;
3474 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3475 Some(row_id)
3476 } else {
3477 None
3478 }
3479 }
3480
3481 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3486 self.query_at_with_allowed(q, self.snapshot(), None)
3487 }
3488
3489 pub fn query_at_with_allowed(
3492 &mut self,
3493 q: &crate::query::Query,
3494 snapshot: Snapshot,
3495 allowed: Option<&std::collections::HashSet<RowId>>,
3496 ) -> Result<Vec<Row>> {
3497 self.require_select()?;
3498 self.ensure_indexes_complete()?;
3499 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3500 return Err(MongrelError::InvalidArgument(format!(
3501 "query exceeds {} conditions",
3502 crate::query::MAX_HARD_CONDITIONS
3503 )));
3504 }
3505 if let Some(limit) = q.limit {
3506 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
3507 return Err(MongrelError::InvalidArgument(format!(
3508 "query limit must be between 1 and {}",
3509 crate::query::MAX_FINAL_LIMIT
3510 )));
3511 }
3512 }
3513 self.query_conditions_at(&q.conditions, snapshot, allowed, q.limit, q.offset)
3514 }
3515
3516 #[doc(hidden)]
3519 pub fn query_all_at(
3520 &mut self,
3521 conditions: &[crate::query::Condition],
3522 snapshot: Snapshot,
3523 ) -> Result<Vec<Row>> {
3524 self.require_select()?;
3525 self.ensure_indexes_complete()?;
3526 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3527 return Err(MongrelError::InvalidArgument(format!(
3528 "query exceeds {} conditions",
3529 crate::query::MAX_HARD_CONDITIONS
3530 )));
3531 }
3532 self.query_conditions_at(conditions, snapshot, None, None, 0)
3533 }
3534
3535 fn query_conditions_at(
3536 &self,
3537 conditions: &[crate::query::Condition],
3538 snapshot: Snapshot,
3539 allowed: Option<&std::collections::HashSet<RowId>>,
3540 limit: Option<usize>,
3541 offset: usize,
3542 ) -> Result<Vec<Row>> {
3543 crate::trace::QueryTrace::record(|t| {
3544 t.run_count = self.run_refs.len();
3545 t.memtable_rows = self.memtable.len();
3546 t.mutable_run_rows = self.mutable_run.len();
3547 });
3548 if conditions.is_empty() {
3552 crate::trace::QueryTrace::record(|t| {
3553 t.scan_mode = crate::trace::ScanMode::Materialized;
3554 t.row_materialized = true;
3555 });
3556 let mut rows = self.visible_rows(snapshot)?;
3557 if let Some(allowed) = allowed {
3558 rows.retain(|row| allowed.contains(&row.row_id));
3559 }
3560 rows.drain(..offset.min(rows.len()));
3561 if let Some(limit) = limit {
3562 rows.truncate(limit);
3563 }
3564 return Ok(rows);
3565 }
3566 crate::trace::QueryTrace::record(|t| {
3567 t.conditions_pushed = conditions.len();
3568 t.scan_mode = crate::trace::ScanMode::Materialized;
3569 t.row_materialized = true;
3570 });
3571 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
3578 ordered.sort_by_key(|c| condition_cost_rank(c));
3579 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3580 for c in &ordered {
3581 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
3582 let empty = s.is_empty();
3583 sets.push(s);
3584 if empty {
3585 break;
3586 }
3587 }
3588 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3589 if let Some(allowed) = allowed {
3590 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
3591 }
3592 rids.drain(..offset.min(rids.len()));
3593 if let Some(limit) = limit {
3594 rids.truncate(limit);
3595 }
3596 self.rows_for_rids(&rids, snapshot)
3597 }
3598
3599 pub fn retrieve(
3601 &mut self,
3602 retriever: &crate::query::Retriever,
3603 ) -> Result<Vec<crate::query::RetrieverHit>> {
3604 self.retrieve_with_allowed(retriever, None)
3605 }
3606
3607 pub fn retrieve_at(
3608 &mut self,
3609 retriever: &crate::query::Retriever,
3610 snapshot: Snapshot,
3611 allowed: Option<&std::collections::HashSet<RowId>>,
3612 ) -> Result<Vec<crate::query::RetrieverHit>> {
3613 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
3614 }
3615
3616 pub fn retrieve_with_allowed(
3619 &mut self,
3620 retriever: &crate::query::Retriever,
3621 allowed: Option<&std::collections::HashSet<RowId>>,
3622 ) -> Result<Vec<crate::query::RetrieverHit>> {
3623 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
3624 }
3625
3626 pub fn retrieve_at_with_allowed(
3627 &mut self,
3628 retriever: &crate::query::Retriever,
3629 snapshot: Snapshot,
3630 allowed: Option<&std::collections::HashSet<RowId>>,
3631 ) -> Result<Vec<crate::query::RetrieverHit>> {
3632 self.require_select()?;
3633 self.ensure_indexes_complete()?;
3634 self.validate_retriever(retriever)?;
3635 self.retrieve_filtered(retriever, snapshot, None, allowed)
3636 }
3637
3638 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3639 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
3640 let (column_id, k) = match retriever {
3641 Retriever::Ann {
3642 column_id,
3643 query,
3644 k,
3645 } => {
3646 let index = self.ann.get(column_id).ok_or_else(|| {
3647 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3648 })?;
3649 if query.len() != index.dim() {
3650 return Err(MongrelError::InvalidArgument(format!(
3651 "ANN query dimension must be {}, got {}",
3652 index.dim(),
3653 query.len()
3654 )));
3655 }
3656 if query.iter().any(|value| !value.is_finite()) {
3657 return Err(MongrelError::InvalidArgument(
3658 "ANN query values must be finite".into(),
3659 ));
3660 }
3661 (*column_id, *k)
3662 }
3663 Retriever::Sparse {
3664 column_id,
3665 query,
3666 k,
3667 } => {
3668 if !self.sparse.contains_key(column_id) {
3669 return Err(MongrelError::InvalidArgument(format!(
3670 "column {column_id} has no Sparse index"
3671 )));
3672 }
3673 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3674 return Err(MongrelError::InvalidArgument(
3675 "Sparse query must be non-empty with finite weights".into(),
3676 ));
3677 }
3678 if query.len() > MAX_SPARSE_TERMS {
3679 return Err(MongrelError::InvalidArgument(format!(
3680 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
3681 )));
3682 }
3683 (*column_id, *k)
3684 }
3685 Retriever::MinHash {
3686 column_id,
3687 members,
3688 k,
3689 } => {
3690 if !self.minhash.contains_key(column_id) {
3691 return Err(MongrelError::InvalidArgument(format!(
3692 "column {column_id} has no MinHash index"
3693 )));
3694 }
3695 if members.is_empty() {
3696 return Err(MongrelError::InvalidArgument(
3697 "MinHash members must not be empty".into(),
3698 ));
3699 }
3700 if members.len() > MAX_SET_MEMBERS {
3701 return Err(MongrelError::InvalidArgument(format!(
3702 "MinHash query exceeds {MAX_SET_MEMBERS} members"
3703 )));
3704 }
3705 (*column_id, *k)
3706 }
3707 };
3708 if k == 0 {
3709 return Err(MongrelError::InvalidArgument(
3710 "retriever k must be > 0".into(),
3711 ));
3712 }
3713 if k > MAX_RETRIEVER_K {
3714 return Err(MongrelError::InvalidArgument(format!(
3715 "retriever k exceeds {MAX_RETRIEVER_K}"
3716 )));
3717 }
3718 debug_assert!(self
3719 .schema
3720 .columns
3721 .iter()
3722 .any(|column| column.id == column_id));
3723 Ok(())
3724 }
3725
3726 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3727 use crate::query::Condition;
3728 match condition {
3729 Condition::Ann {
3730 column_id,
3731 query,
3732 k,
3733 } => self.validate_retriever(&crate::query::Retriever::Ann {
3734 column_id: *column_id,
3735 query: query.clone(),
3736 k: *k,
3737 }),
3738 Condition::SparseMatch {
3739 column_id,
3740 query,
3741 k,
3742 } => self.validate_retriever(&crate::query::Retriever::Sparse {
3743 column_id: *column_id,
3744 query: query.clone(),
3745 k: *k,
3746 }),
3747 Condition::MinHashSimilar {
3748 column_id,
3749 query,
3750 k,
3751 } => {
3752 if !self.minhash.contains_key(column_id) {
3753 return Err(MongrelError::InvalidArgument(format!(
3754 "column {column_id} has no MinHash index"
3755 )));
3756 }
3757 if query.is_empty() || *k == 0 {
3758 return Err(MongrelError::InvalidArgument(
3759 "MinHash query must be non-empty and k must be > 0".into(),
3760 ));
3761 }
3762 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
3763 {
3764 return Err(MongrelError::InvalidArgument(format!(
3765 "MinHash query must have <= {} members and k <= {}",
3766 crate::query::MAX_SET_MEMBERS,
3767 crate::query::MAX_RETRIEVER_K
3768 )));
3769 }
3770 Ok(())
3771 }
3772 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
3773 Err(MongrelError::InvalidArgument(format!(
3774 "bitmap IN exceeds {} values",
3775 crate::query::MAX_SET_MEMBERS
3776 )))
3777 }
3778 Condition::FmContainsAll { patterns, .. }
3779 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
3780 {
3781 Err(MongrelError::InvalidArgument(format!(
3782 "FM query exceeds {} patterns",
3783 crate::query::MAX_HARD_CONDITIONS
3784 )))
3785 }
3786 _ => Ok(()),
3787 }
3788 }
3789
3790 fn retrieve_filtered(
3791 &self,
3792 retriever: &crate::query::Retriever,
3793 snapshot: Snapshot,
3794 hard_filter: Option<&RowIdSet>,
3795 allowed: Option<&std::collections::HashSet<RowId>>,
3796 ) -> Result<Vec<crate::query::RetrieverHit>> {
3797 use crate::query::{Retriever, RetrieverHit, RetrieverScore, SetMember};
3798 let started = std::time::Instant::now();
3799 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
3800 Retriever::Ann {
3801 column_id,
3802 query,
3803 k,
3804 } => {
3805 let Some(index) = self.ann.get(column_id) else {
3806 return Ok(Vec::new());
3807 };
3808 let cap = index.len();
3809 if cap == 0 {
3810 return Ok(Vec::new());
3811 }
3812 let mut breadth = (*k).max(1).min(cap);
3813 let mut eligibility = std::collections::HashMap::new();
3814 let mut filtered = loop {
3815 let mut seen = std::collections::HashSet::new();
3816 let raw = index.search(query, breadth)?;
3817 let unchecked: Vec<_> = raw
3818 .iter()
3819 .map(|(row_id, _)| *row_id)
3820 .filter(|row_id| !eligibility.contains_key(row_id))
3821 .filter(|row_id| {
3822 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3823 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3824 })
3825 .collect();
3826 let eligible = self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3827 for row_id in unchecked {
3828 eligibility.insert(row_id, eligible.contains(&row_id));
3829 }
3830 let filtered: Vec<_> = raw
3831 .into_iter()
3832 .filter(|(row_id, _)| {
3833 seen.insert(*row_id)
3834 && eligibility.get(row_id).copied().unwrap_or(false)
3835 })
3836 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
3837 .collect();
3838 if filtered.len() >= *k || breadth >= cap {
3839 break filtered;
3840 }
3841 breadth = breadth.saturating_mul(2).min(cap);
3842 };
3843 filtered.truncate(*k);
3844 filtered
3845 }
3846 Retriever::Sparse {
3847 column_id,
3848 query,
3849 k,
3850 } => self
3851 .sparse
3852 .get(column_id)
3853 .map(|index| -> Result<Vec<_>> {
3854 let mut breadth = (*k).max(1);
3855 let mut eligibility = std::collections::HashMap::new();
3856 loop {
3857 let raw = index.search(query, breadth);
3858 let unchecked: Vec<_> = raw
3859 .iter()
3860 .map(|(row_id, _)| *row_id)
3861 .filter(|row_id| !eligibility.contains_key(row_id))
3862 .filter(|row_id| {
3863 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3864 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3865 })
3866 .collect();
3867 let eligible =
3868 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3869 for row_id in unchecked {
3870 eligibility.insert(row_id, eligible.contains(&row_id));
3871 }
3872 let filtered: Vec<_> = raw
3873 .iter()
3874 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3875 .take(*k)
3876 .map(|(row_id, score)| {
3877 (*row_id, RetrieverScore::SparseDotProduct(*score))
3878 })
3879 .collect();
3880 if filtered.len() >= *k || raw.len() < breadth {
3881 break Ok(filtered);
3882 }
3883 let next = breadth.saturating_mul(2);
3884 if next == breadth {
3885 break Ok(filtered);
3886 }
3887 breadth = next;
3888 }
3889 })
3890 .transpose()?
3891 .unwrap_or_default(),
3892 Retriever::MinHash {
3893 column_id,
3894 members,
3895 k,
3896 } => self
3897 .minhash
3898 .get(column_id)
3899 .map(|index| -> Result<Vec<_>> {
3900 let hashes: Vec<_> = members.iter().map(SetMember::hash_v1).collect();
3901 let mut breadth = (*k).max(1);
3902 let mut eligibility = std::collections::HashMap::new();
3903 loop {
3904 let raw = index.search(&hashes, breadth);
3905 let unchecked: Vec<_> = raw
3906 .iter()
3907 .map(|(row_id, _)| *row_id)
3908 .filter(|row_id| !eligibility.contains_key(row_id))
3909 .filter(|row_id| {
3910 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3911 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3912 })
3913 .collect();
3914 let eligible =
3915 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3916 for row_id in unchecked {
3917 eligibility.insert(row_id, eligible.contains(&row_id));
3918 }
3919 let filtered: Vec<_> = raw
3920 .iter()
3921 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3922 .take(*k)
3923 .map(|(row_id, score)| {
3924 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
3925 })
3926 .collect();
3927 if filtered.len() >= *k || raw.len() < breadth {
3928 break Ok(filtered);
3929 }
3930 let next = breadth.saturating_mul(2);
3931 if next == breadth {
3932 break Ok(filtered);
3933 }
3934 breadth = next;
3935 }
3936 })
3937 .transpose()?
3938 .unwrap_or_default(),
3939 };
3940 let elapsed = started.elapsed().as_nanos() as u64;
3941 crate::trace::QueryTrace::record(|trace| {
3942 match retriever {
3943 Retriever::Ann { .. } => {
3944 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
3945 }
3946 Retriever::Sparse { .. } => {
3947 trace.sparse_candidate_nanos =
3948 trace.sparse_candidate_nanos.saturating_add(elapsed)
3949 }
3950 Retriever::MinHash { .. } => {
3951 trace.minhash_candidate_nanos =
3952 trace.minhash_candidate_nanos.saturating_add(elapsed)
3953 }
3954 }
3955 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
3956 });
3957 Ok(scored
3958 .into_iter()
3959 .enumerate()
3960 .map(|(rank, (row_id, score))| RetrieverHit {
3961 row_id,
3962 rank: rank + 1,
3963 score,
3964 })
3965 .collect())
3966 }
3967
3968 fn eligible_candidate_ids(
3969 &self,
3970 candidates: &[RowId],
3971 _column_id: u16,
3972 snapshot: Snapshot,
3973 ) -> Result<std::collections::HashSet<RowId>> {
3974 if !self.had_deletes
3975 && self.ttl.is_none()
3976 && self.pending_put_cols.is_empty()
3977 && snapshot.epoch == self.snapshot().epoch
3978 {
3979 return Ok(candidates.iter().copied().collect());
3980 }
3981 let mut readers: Vec<_> = self
3982 .run_refs
3983 .iter()
3984 .map(|run| self.open_reader(run.run_id))
3985 .collect::<Result<_>>()?;
3986 let now = unix_nanos_now();
3987 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
3988 for &row_id in candidates {
3989 let mem = self.memtable.get_version(row_id, snapshot.epoch);
3990 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
3991 let overlay = match (mem, mutable) {
3992 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
3993 (Some(value), None) | (None, Some(value)) => Some(value),
3994 (None, None) => None,
3995 };
3996 if let Some((_, row)) = overlay {
3997 if !row.deleted && !self.row_expired_at(&row, now) {
3998 eligible.insert(row_id);
3999 }
4000 continue;
4001 }
4002 let mut best: Option<(Epoch, bool, usize)> = None;
4003 for (index, reader) in readers.iter_mut().enumerate() {
4004 if let Some((epoch, deleted)) =
4005 reader.get_version_visibility(row_id, snapshot.epoch)?
4006 {
4007 if best
4008 .as_ref()
4009 .map(|(best_epoch, ..)| epoch > *best_epoch)
4010 .unwrap_or(true)
4011 {
4012 best = Some((epoch, deleted, index));
4013 }
4014 }
4015 }
4016 let Some((_, false, reader_index)) = best else {
4017 continue;
4018 };
4019 if let Some(ttl) = self.ttl {
4020 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4021 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4022 {
4023 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4024 continue;
4025 }
4026 }
4027 }
4028 eligible.insert(row_id);
4029 }
4030 Ok(eligible)
4031 }
4032
4033 pub fn search(
4035 &mut self,
4036 request: &crate::query::SearchRequest,
4037 ) -> Result<Vec<crate::query::SearchHit>> {
4038 self.search_with_allowed(request, None)
4039 }
4040
4041 pub fn search_at(
4042 &mut self,
4043 request: &crate::query::SearchRequest,
4044 snapshot: Snapshot,
4045 authorized: Option<&std::collections::HashSet<RowId>>,
4046 ) -> Result<Vec<crate::query::SearchHit>> {
4047 self.search_at_with_allowed(request, snapshot, authorized)
4048 }
4049
4050 pub fn search_with_allowed(
4051 &mut self,
4052 request: &crate::query::SearchRequest,
4053 authorized: Option<&std::collections::HashSet<RowId>>,
4054 ) -> Result<Vec<crate::query::SearchHit>> {
4055 self.search_at_with_allowed(request, self.snapshot(), authorized)
4056 }
4057
4058 pub fn search_at_with_allowed(
4059 &mut self,
4060 request: &crate::query::SearchRequest,
4061 snapshot: Snapshot,
4062 authorized: Option<&std::collections::HashSet<RowId>>,
4063 ) -> Result<Vec<crate::query::SearchHit>> {
4064 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
4065 }
4066
4067 pub fn search_at_with_allowed_and_context(
4068 &mut self,
4069 request: &crate::query::SearchRequest,
4070 snapshot: Snapshot,
4071 authorized: Option<&std::collections::HashSet<RowId>>,
4072 context: Option<&crate::query::AiExecutionContext>,
4073 ) -> Result<Vec<crate::query::SearchHit>> {
4074 use crate::query::{
4075 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
4076 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
4077 };
4078 let total_started = std::time::Instant::now();
4079 self.require_select()?;
4080 self.ensure_indexes_complete()?;
4081 if request.limit == 0 {
4082 return Err(MongrelError::InvalidArgument(
4083 "search limit must be > 0".into(),
4084 ));
4085 }
4086 if request.limit > MAX_FINAL_LIMIT {
4087 return Err(MongrelError::InvalidArgument(format!(
4088 "search limit exceeds {MAX_FINAL_LIMIT}"
4089 )));
4090 }
4091 if request.retrievers.is_empty() {
4092 return Err(MongrelError::InvalidArgument(
4093 "search requires at least one retriever".into(),
4094 ));
4095 }
4096 if request.retrievers.len() > MAX_RETRIEVERS {
4097 return Err(MongrelError::InvalidArgument(format!(
4098 "search exceeds {MAX_RETRIEVERS} retrievers"
4099 )));
4100 }
4101 if request.must.len() > MAX_HARD_CONDITIONS {
4102 return Err(MongrelError::InvalidArgument(format!(
4103 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
4104 )));
4105 }
4106 if request.must.iter().any(|condition| {
4107 matches!(
4108 condition,
4109 Condition::Ann { .. }
4110 | Condition::SparseMatch { .. }
4111 | Condition::MinHashSimilar { .. }
4112 )
4113 }) {
4114 return Err(MongrelError::InvalidArgument(
4115 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
4116 .into(),
4117 ));
4118 }
4119 let mut names = std::collections::HashSet::new();
4120 for named in &request.retrievers {
4121 if named.name.is_empty() || !names.insert(named.name.as_str()) {
4122 return Err(MongrelError::InvalidArgument(
4123 "retriever names must be non-empty and unique".into(),
4124 ));
4125 }
4126 if !named.weight.is_finite()
4127 || named.weight < 0.0
4128 || named.weight > MAX_RETRIEVER_WEIGHT
4129 {
4130 return Err(MongrelError::InvalidArgument(format!(
4131 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4132 )));
4133 }
4134 self.validate_retriever(&named.retriever)?;
4135 }
4136 let projection = request
4137 .projection
4138 .clone()
4139 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
4140 if projection.len() > MAX_PROJECTION_COLUMNS {
4141 return Err(MongrelError::InvalidArgument(format!(
4142 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
4143 )));
4144 }
4145 for column_id in &projection {
4146 if !self
4147 .schema
4148 .columns
4149 .iter()
4150 .any(|column| column.id == *column_id)
4151 {
4152 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
4153 }
4154 }
4155
4156 let hard_filter_started = std::time::Instant::now();
4157 let hard_filter = if request.must.is_empty() {
4158 None
4159 } else {
4160 let mut sets = Vec::with_capacity(request.must.len());
4161 for condition in &request.must {
4162 if let Some(context) = context {
4163 context.checkpoint()?;
4164 }
4165 sets.push(self.resolve_condition(condition, snapshot)?);
4166 }
4167 Some(RowIdSet::intersect_many(sets))
4168 };
4169 crate::trace::QueryTrace::record(|trace| {
4170 trace.hard_filter_nanos = trace
4171 .hard_filter_nanos
4172 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
4173 });
4174 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
4175 return Ok(Vec::new());
4176 }
4177
4178 let constant = match request.fusion {
4179 Fusion::ReciprocalRank { constant } => constant,
4180 };
4181 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
4182 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
4183 let mut fusion_nanos = 0u64;
4184 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
4185 std::collections::HashMap::new();
4186 for named in retrievers {
4187 if let Some(context) = context {
4188 context.checkpoint()?;
4189 }
4190 let hits = self.retrieve_filtered(
4191 &named.retriever,
4192 snapshot,
4193 hard_filter.as_ref(),
4194 authorized,
4195 )?;
4196 let fusion_started = std::time::Instant::now();
4197 for hit in hits {
4198 if let Some(context) = context {
4199 context.consume(1)?;
4200 }
4201 let contribution = named.weight / (constant as f64 + hit.rank as f64);
4202 if !contribution.is_finite() {
4203 return Err(MongrelError::InvalidArgument(
4204 "retriever contribution must be finite".into(),
4205 ));
4206 }
4207 if !fused.contains_key(&hit.row_id)
4208 && fused.len() >= crate::query::MAX_FUSED_CANDIDATES
4209 {
4210 return Err(MongrelError::WorkBudgetExceeded);
4211 }
4212 let entry = fused.entry(hit.row_id).or_default();
4213 entry.0 += contribution;
4214 if !entry.0.is_finite() {
4215 return Err(MongrelError::InvalidArgument(
4216 "fused score must be finite".into(),
4217 ));
4218 }
4219 entry.1.push(ComponentScore {
4220 retriever_name: named.name.clone(),
4221 rank: hit.rank,
4222 raw_score: hit.score,
4223 contribution,
4224 });
4225 }
4226 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
4227 }
4228 let union_size = fused.len();
4229 let mut ranked: Vec<_> = fused.into_iter().collect();
4230 let sort_started = std::time::Instant::now();
4231 let order =
4232 |(a_row, (a_score, _)): &(RowId, (f64, Vec<ComponentScore>)),
4233 (b_row, (b_score, _)): &(RowId, (f64, Vec<ComponentScore>))| {
4234 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
4235 };
4236 if ranked.len() > request.limit {
4237 let (_, _, _) = ranked.select_nth_unstable_by(request.limit, order);
4238 ranked.truncate(request.limit);
4239 }
4240 ranked.sort_by(order);
4241 fusion_nanos = fusion_nanos.saturating_add(sort_started.elapsed().as_nanos() as u64);
4242 crate::trace::QueryTrace::record(|trace| {
4243 trace.union_size = union_size;
4244 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
4245 });
4246
4247 let projection_started = std::time::Instant::now();
4248 let row_ids: Vec<_> = ranked.iter().map(|(row_id, _)| row_id.0).collect();
4249 if let Some(context) = context {
4250 context.consume(row_ids.len().saturating_mul(projection.len().max(1)))?;
4251 }
4252 let sentinel = projection
4253 .first()
4254 .copied()
4255 .or_else(|| self.schema.columns.first().map(|column| column.id));
4256 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
4257 std::collections::HashMap::new();
4258 if let Some(column_id) = sentinel {
4259 for (row_id, value) in self.values_for_rids_batch(&row_ids, column_id, snapshot)? {
4260 cells.entry(row_id).or_default().insert(column_id, value);
4261 }
4262 }
4263 for &column_id in &projection {
4264 if Some(column_id) == sentinel {
4265 continue;
4266 }
4267 for (row_id, value) in self.values_for_rids_batch(&row_ids, column_id, snapshot)? {
4268 cells.entry(row_id).or_default().insert(column_id, value);
4269 }
4270 }
4271
4272 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
4273 for (row_id, (fused_score, mut components)) in ranked {
4274 let Some(row_cells) = cells.remove(&row_id) else {
4275 continue;
4276 };
4277 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4278 out.push(SearchHit {
4279 row_id,
4280 cells: projection
4281 .iter()
4282 .filter_map(|column_id| {
4283 row_cells
4284 .get(column_id)
4285 .cloned()
4286 .map(|value| (*column_id, value))
4287 })
4288 .collect(),
4289 components,
4290 fused_score,
4291 });
4292 if out.len() == request.limit {
4293 break;
4294 }
4295 }
4296 crate::trace::QueryTrace::record(|trace| {
4297 trace.projection_nanos = trace
4298 .projection_nanos
4299 .saturating_add(projection_started.elapsed().as_nanos() as u64);
4300 trace.total_nanos = trace
4301 .total_nanos
4302 .saturating_add(total_started.elapsed().as_nanos() as u64);
4303 });
4304 Ok(out)
4305 }
4306
4307 pub fn set_similarity(
4310 &mut self,
4311 request: &crate::query::SetSimilarityRequest,
4312 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4313 self.set_similarity_with_allowed(request, None)
4314 }
4315
4316 pub fn set_similarity_at(
4317 &mut self,
4318 request: &crate::query::SetSimilarityRequest,
4319 snapshot: Snapshot,
4320 allowed: Option<&std::collections::HashSet<RowId>>,
4321 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4322 self.set_similarity_explained_at(request, snapshot, allowed)
4323 .map(|(hits, _)| hits)
4324 }
4325
4326 pub fn ann_rerank(
4328 &mut self,
4329 request: &crate::query::AnnRerankRequest,
4330 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4331 self.ann_rerank_with_allowed(request, None)
4332 }
4333
4334 pub fn ann_rerank_with_allowed(
4335 &mut self,
4336 request: &crate::query::AnnRerankRequest,
4337 allowed: Option<&std::collections::HashSet<RowId>>,
4338 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4339 self.ann_rerank_at(request, self.snapshot(), allowed)
4340 }
4341
4342 pub fn ann_rerank_at(
4343 &mut self,
4344 request: &crate::query::AnnRerankRequest,
4345 snapshot: Snapshot,
4346 allowed: Option<&std::collections::HashSet<RowId>>,
4347 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4348 use crate::query::{
4349 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
4350 };
4351 if request.candidate_k == 0 || request.limit == 0 {
4352 return Err(MongrelError::InvalidArgument(
4353 "candidate_k and limit must be > 0".into(),
4354 ));
4355 }
4356 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
4357 return Err(MongrelError::InvalidArgument(format!(
4358 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
4359 )));
4360 }
4361 let hits = self.retrieve_at_with_allowed(
4362 &Retriever::Ann {
4363 column_id: request.column_id,
4364 query: request.query.clone(),
4365 k: request.candidate_k,
4366 },
4367 snapshot,
4368 allowed,
4369 )?;
4370 let distances: std::collections::HashMap<_, _> = hits
4371 .iter()
4372 .filter_map(|hit| match hit.score {
4373 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
4374 _ => None,
4375 })
4376 .collect();
4377 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4378 let gather_started = std::time::Instant::now();
4379 let values = self.values_for_rids_batch(&row_ids, request.column_id, snapshot)?;
4380 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
4381 let score_started = std::time::Instant::now();
4382 let query_norm = request
4383 .query
4384 .iter()
4385 .map(|value| f64::from(*value).powi(2))
4386 .sum::<f64>()
4387 .sqrt();
4388 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
4389 for (row_id, value) in values {
4390 let Value::Embedding(vector) = value else {
4391 continue;
4392 };
4393 let dot = request
4394 .query
4395 .iter()
4396 .zip(&vector)
4397 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4398 .sum::<f64>();
4399 let exact_score = match request.metric {
4400 VectorMetric::DotProduct => dot,
4401 VectorMetric::Cosine => {
4402 let norm = vector
4403 .iter()
4404 .map(|value| f64::from(*value).powi(2))
4405 .sum::<f64>()
4406 .sqrt();
4407 if query_norm == 0.0 || norm == 0.0 {
4408 0.0
4409 } else {
4410 dot / (query_norm * norm)
4411 }
4412 }
4413 VectorMetric::Euclidean => request
4414 .query
4415 .iter()
4416 .zip(&vector)
4417 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
4418 .sum::<f64>()
4419 .sqrt(),
4420 };
4421 let exact_score = exact_score as f32;
4422 if !exact_score.is_finite() {
4423 return Err(MongrelError::InvalidArgument(
4424 "exact ANN score must be finite".into(),
4425 ));
4426 }
4427 reranked.push(AnnRerankHit {
4428 row_id,
4429 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
4430 exact_score,
4431 });
4432 }
4433 reranked.sort_by(|left, right| {
4434 let score = match request.metric {
4435 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
4436 VectorMetric::Cosine | VectorMetric::DotProduct => {
4437 right.exact_score.total_cmp(&left.exact_score)
4438 }
4439 };
4440 score.then_with(|| left.row_id.cmp(&right.row_id))
4441 });
4442 reranked.truncate(request.limit);
4443 crate::trace::QueryTrace::record(|trace| {
4444 trace.exact_vector_gather_nanos =
4445 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
4446 trace.exact_vector_score_nanos = trace
4447 .exact_vector_score_nanos
4448 .saturating_add(score_started.elapsed().as_nanos() as u64);
4449 });
4450 Ok(reranked)
4451 }
4452
4453 pub fn set_similarity_with_allowed(
4454 &mut self,
4455 request: &crate::query::SetSimilarityRequest,
4456 allowed: Option<&std::collections::HashSet<RowId>>,
4457 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4458 self.set_similarity_explained_at(request, self.snapshot(), allowed)
4459 .map(|(hits, _)| hits)
4460 }
4461
4462 pub fn set_similarity_explained(
4463 &mut self,
4464 request: &crate::query::SetSimilarityRequest,
4465 ) -> Result<(
4466 Vec<crate::query::SetSimilarityHit>,
4467 crate::query::SetSimilarityTrace,
4468 )> {
4469 self.set_similarity_explained_at(request, self.snapshot(), None)
4470 }
4471
4472 fn set_similarity_explained_at(
4473 &mut self,
4474 request: &crate::query::SetSimilarityRequest,
4475 snapshot: Snapshot,
4476 allowed: Option<&std::collections::HashSet<RowId>>,
4477 ) -> Result<(
4478 Vec<crate::query::SetSimilarityHit>,
4479 crate::query::SetSimilarityTrace,
4480 )> {
4481 use crate::query::{
4482 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
4483 MAX_SET_MEMBERS,
4484 };
4485 let mut trace = crate::query::SetSimilarityTrace::default();
4486 if request.members.is_empty() {
4487 return Ok((Vec::new(), trace));
4488 }
4489 if request.candidate_k == 0 || request.limit == 0 {
4490 return Err(MongrelError::InvalidArgument(
4491 "candidate_k and limit must be > 0".into(),
4492 ));
4493 }
4494 if request.candidate_k > MAX_RETRIEVER_K
4495 || request.limit > MAX_FINAL_LIMIT
4496 || request.members.len() > MAX_SET_MEMBERS
4497 {
4498 return Err(MongrelError::InvalidArgument(format!(
4499 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
4500 )));
4501 }
4502 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
4503 return Err(MongrelError::InvalidArgument(
4504 "min_jaccard must be finite and between 0 and 1".into(),
4505 ));
4506 }
4507 let started = std::time::Instant::now();
4508 let hits = self.retrieve_at_with_allowed(
4509 &Retriever::MinHash {
4510 column_id: request.column_id,
4511 members: request.members.clone(),
4512 k: request.candidate_k,
4513 },
4514 snapshot,
4515 allowed,
4516 )?;
4517 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
4518 trace.candidate_count = hits.len();
4519 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4520 let started = std::time::Instant::now();
4521 let values = self.values_for_rids_batch(&row_ids, request.column_id, snapshot)?;
4522 trace.gather_us = started.elapsed().as_micros() as u64;
4523 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
4524 let estimates: std::collections::HashMap<_, _> = hits
4525 .into_iter()
4526 .filter_map(|hit| match hit.score {
4527 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
4528 _ => None,
4529 })
4530 .collect();
4531 let started = std::time::Instant::now();
4532 let parsed: Vec<_> = values
4533 .into_iter()
4534 .filter_map(|(row_id, value)| {
4535 let Value::Bytes(bytes) = value else {
4536 return None;
4537 };
4538 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
4539 return None;
4540 };
4541 let stored = members
4542 .into_iter()
4543 .filter_map(|member| match member {
4544 serde_json::Value::String(value) => {
4545 Some(crate::query::SetMember::String(value))
4546 }
4547 serde_json::Value::Number(value) => {
4548 Some(crate::query::SetMember::Number(value))
4549 }
4550 serde_json::Value::Bool(value) => {
4551 Some(crate::query::SetMember::Boolean(value))
4552 }
4553 _ => None,
4554 })
4555 .collect::<std::collections::HashSet<_>>();
4556 Some((row_id, stored))
4557 })
4558 .collect();
4559 trace.parse_us = started.elapsed().as_micros() as u64;
4560 trace.verified_count = parsed.len();
4561 let started = std::time::Instant::now();
4562 let mut exact = Vec::new();
4563 for (row_id, stored) in parsed {
4564 let union = query.union(&stored).count();
4565 let score = if union == 0 {
4566 1.0
4567 } else {
4568 query.intersection(&stored).count() as f32 / union as f32
4569 };
4570 if score >= request.min_jaccard {
4571 exact.push(SetSimilarityHit {
4572 row_id,
4573 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
4574 exact_jaccard: score,
4575 });
4576 }
4577 }
4578 exact.sort_by(|a, b| {
4579 b.exact_jaccard
4580 .total_cmp(&a.exact_jaccard)
4581 .then_with(|| a.row_id.cmp(&b.row_id))
4582 });
4583 exact.truncate(request.limit);
4584 trace.score_us = started.elapsed().as_micros() as u64;
4585 crate::trace::QueryTrace::record(|query_trace| {
4586 query_trace.exact_set_gather_nanos = query_trace
4587 .exact_set_gather_nanos
4588 .saturating_add(trace.gather_us.saturating_mul(1_000));
4589 query_trace.exact_set_parse_nanos = query_trace
4590 .exact_set_parse_nanos
4591 .saturating_add(trace.parse_us.saturating_mul(1_000));
4592 query_trace.exact_set_score_nanos = query_trace
4593 .exact_set_score_nanos
4594 .saturating_add(trace.score_us.saturating_mul(1_000));
4595 });
4596 Ok((exact, trace))
4597 }
4598
4599 fn values_for_rids_batch(
4601 &self,
4602 row_ids: &[u64],
4603 column_id: u16,
4604 snapshot: Snapshot,
4605 ) -> Result<Vec<(RowId, Value)>> {
4606 if self.ttl.is_none()
4607 && self.memtable.is_empty()
4608 && self.mutable_run.is_empty()
4609 && self.run_refs.len() == 1
4610 {
4611 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4612 let (positions, visible_row_ids) =
4613 reader.visible_positions_with_rids(snapshot.epoch)?;
4614 let requested: Vec<(RowId, usize)> = row_ids
4615 .iter()
4616 .filter_map(|raw| {
4617 visible_row_ids
4618 .binary_search(&(*raw as i64))
4619 .ok()
4620 .map(|index| (RowId(*raw), positions[index]))
4621 })
4622 .collect();
4623 let values = reader.gather_column(
4624 column_id,
4625 &requested
4626 .iter()
4627 .map(|(_, position)| *position)
4628 .collect::<Vec<_>>(),
4629 )?;
4630 return Ok(requested
4631 .into_iter()
4632 .zip(values)
4633 .map(|((row_id, _), value)| (row_id, value))
4634 .collect());
4635 }
4636 self.values_for_rids(row_ids, column_id, snapshot)
4637 }
4638
4639 fn values_for_rids(
4641 &self,
4642 row_ids: &[u64],
4643 column_id: u16,
4644 snapshot: Snapshot,
4645 ) -> Result<Vec<(RowId, Value)>> {
4646 let mut readers: Vec<_> = self
4647 .run_refs
4648 .iter()
4649 .map(|run| self.open_reader(run.run_id))
4650 .collect::<Result<_>>()?;
4651 let now = unix_nanos_now();
4652 let mut out = Vec::with_capacity(row_ids.len());
4653 for &raw_row_id in row_ids {
4654 let row_id = RowId(raw_row_id);
4655 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4656 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4657 let overlay = match (mem, mutable) {
4658 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
4659 (a_epoch, a)
4660 } else {
4661 (b_epoch, b)
4662 }),
4663 (Some(value), None) | (None, Some(value)) => Some(value),
4664 (None, None) => None,
4665 };
4666 if let Some((_, row)) = overlay {
4667 if !row.deleted && !self.row_expired_at(&row, now) {
4668 if let Some(value) = row.columns.get(&column_id) {
4669 out.push((row_id, value.clone()));
4670 }
4671 }
4672 continue;
4673 }
4674
4675 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
4676 for (index, reader) in readers.iter_mut().enumerate() {
4677 if let Some((epoch, deleted, value)) =
4678 reader.get_version_column(row_id, snapshot.epoch, column_id)?
4679 {
4680 if best
4681 .as_ref()
4682 .map(|(best_epoch, ..)| epoch > *best_epoch)
4683 .unwrap_or(true)
4684 {
4685 best = Some((epoch, deleted, value, index));
4686 }
4687 }
4688 }
4689 let Some((_, false, Some(value), reader_index)) = best else {
4690 continue;
4691 };
4692 if let Some(ttl) = self.ttl {
4693 if ttl.column_id != column_id {
4694 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4695 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4696 {
4697 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4698 continue;
4699 }
4700 }
4701 } else if let Value::Int64(timestamp) = value {
4702 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4703 continue;
4704 }
4705 }
4706 }
4707 out.push((row_id, value));
4708 }
4709 Ok(out)
4710 }
4711
4712 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
4717 use std::collections::HashMap;
4718 let mut rows = Vec::with_capacity(rids.len());
4719 let ttl_now = unix_nanos_now();
4720 let tier_size = self.memtable.len() + self.mutable_run.len();
4737 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
4738 if rids.len().saturating_mul(24) < tier_size {
4739 for &rid in rids {
4740 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
4741 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
4742 let newest = match (mem, mrun) {
4743 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
4744 (Some((_, mr)), None) => Some(mr),
4745 (None, Some((_, rr))) => Some(rr),
4746 (None, None) => None,
4747 };
4748 if let Some(row) = newest {
4749 overlay.insert(rid, row);
4750 }
4751 }
4752 } else {
4753 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
4754 overlay
4755 .entry(row.row_id.0)
4756 .and_modify(|e| {
4757 if row.committed_epoch > e.committed_epoch {
4758 *e = row.clone();
4759 }
4760 })
4761 .or_insert(row);
4762 };
4763 for row in self.memtable.visible_versions(snapshot.epoch) {
4764 fold_newest(row, &mut overlay);
4765 }
4766 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4767 fold_newest(row, &mut overlay);
4768 }
4769 }
4770 if self.run_refs.len() == 1 {
4771 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4772 if rids.len().saturating_mul(24) < reader.row_count() {
4780 for &rid in rids {
4781 if let Some(r) = overlay.get(&rid) {
4782 if !r.deleted {
4783 rows.push(r.clone());
4784 }
4785 continue;
4786 }
4787 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
4788 if !row.deleted {
4789 rows.push(row);
4790 }
4791 }
4792 }
4793 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4794 return Ok(rows);
4795 }
4796 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4805 enum Src {
4808 Overlay,
4809 Run,
4810 }
4811 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
4812 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
4813 for rid in rids {
4814 if overlay.contains_key(rid) {
4815 plan.push(Src::Overlay);
4816 continue;
4817 }
4818 match vis_rids.binary_search(&(*rid as i64)) {
4819 Ok(i) => {
4820 plan.push(Src::Run);
4821 fetch.push(positions[i]);
4822 }
4823 Err(_) => { }
4824 }
4825 }
4826 let fetched = reader.materialize_batch(&fetch)?;
4827 let mut fetched_iter = fetched.into_iter();
4828 for (rid, src) in rids.iter().zip(plan) {
4829 match src {
4830 Src::Overlay => {
4831 if let Some(r) = overlay.get(rid) {
4832 if !r.deleted {
4833 rows.push(r.clone());
4834 }
4835 }
4836 }
4837 Src::Run => {
4838 if let Some(row) = fetched_iter.next() {
4839 if !row.deleted {
4840 rows.push(row);
4841 }
4842 }
4843 }
4844 }
4845 }
4846 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4847 return Ok(rows);
4848 }
4849 let mut readers: Vec<_> = self
4853 .run_refs
4854 .iter()
4855 .map(|rr| self.open_reader(rr.run_id))
4856 .collect::<Result<Vec<_>>>()?;
4857 for rid in rids {
4858 if let Some(r) = overlay.get(rid) {
4859 if !r.deleted {
4860 rows.push(r.clone());
4861 }
4862 continue;
4863 }
4864 let mut best: Option<(Epoch, Row)> = None;
4865 for reader in readers.iter_mut() {
4866 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
4867 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4868 best = Some((epoch, row));
4869 }
4870 }
4871 }
4872 if let Some((_, r)) = best {
4873 if !r.deleted {
4874 rows.push(r);
4875 }
4876 }
4877 }
4878 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4879 Ok(rows)
4880 }
4881
4882 pub fn indexes_complete(&self) -> bool {
4892 self.indexes_complete
4893 }
4894
4895 pub fn index_build_policy(&self) -> IndexBuildPolicy {
4897 self.index_build_policy
4898 }
4899
4900 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
4904 self.index_build_policy = policy;
4905 }
4906
4907 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
4912 if !self.indexes_complete {
4916 return None;
4917 }
4918 let b = self.bitmap.get(&column_id)?;
4919 let result: Vec<Vec<u8>> = b
4920 .keys()
4921 .into_iter()
4922 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
4923 .cloned()
4924 .collect();
4925 Some(result)
4926 }
4927
4928 pub fn fk_join_row_ids(
4929 &self,
4930 fk_column_id: u16,
4931 pk_values: &[Vec<u8>],
4932 fk_conditions: &[crate::query::Condition],
4933 snapshot: Snapshot,
4934 ) -> Result<Vec<u64>> {
4935 let Some(b) = self.bitmap.get(&fk_column_id) else {
4936 return Ok(Vec::new());
4937 };
4938 let mut join_set = {
4939 let mut acc = roaring::RoaringBitmap::new();
4940 for v in pk_values {
4941 acc |= b.get(v);
4942 }
4943 RowIdSet::from_roaring(acc)
4944 };
4945 if !fk_conditions.is_empty() {
4946 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4947 sets.push(join_set);
4948 for c in fk_conditions {
4949 sets.push(self.resolve_condition(c, snapshot)?);
4950 }
4951 join_set = RowIdSet::intersect_many(sets);
4952 }
4953 Ok(join_set.into_sorted_vec())
4954 }
4955
4956 pub fn fk_join_count(
4962 &self,
4963 fk_column_id: u16,
4964 pk_values: &[Vec<u8>],
4965 fk_conditions: &[crate::query::Condition],
4966 snapshot: Snapshot,
4967 ) -> Result<u64> {
4968 let Some(b) = self.bitmap.get(&fk_column_id) else {
4969 return Ok(0);
4970 };
4971 let mut acc = roaring::RoaringBitmap::new();
4972 for v in pk_values {
4973 acc |= b.get(v);
4974 }
4975 if fk_conditions.is_empty() {
4976 return Ok(acc.len());
4977 }
4978 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4979 sets.push(RowIdSet::from_roaring(acc));
4980 for c in fk_conditions {
4981 sets.push(self.resolve_condition(c, snapshot)?);
4982 }
4983 Ok(RowIdSet::intersect_many(sets).len() as u64)
4984 }
4985
4986 fn resolve_condition(
4991 &self,
4992 c: &crate::query::Condition,
4993 snapshot: Snapshot,
4994 ) -> Result<RowIdSet> {
4995 self.resolve_condition_with_allowed(c, snapshot, None)
4996 }
4997
4998 fn resolve_condition_with_allowed(
4999 &self,
5000 c: &crate::query::Condition,
5001 snapshot: Snapshot,
5002 allowed: Option<&std::collections::HashSet<RowId>>,
5003 ) -> Result<RowIdSet> {
5004 use crate::query::Condition;
5005 self.validate_condition(c)?;
5006 Ok(match c {
5007 Condition::Pk(key) => {
5008 let lookup = self
5009 .schema
5010 .primary_key()
5011 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5012 .unwrap_or_else(|| key.clone());
5013 self.hot
5014 .get(&lookup)
5015 .map(|r| RowIdSet::one(r.0))
5016 .unwrap_or_else(RowIdSet::empty)
5017 }
5018 Condition::BitmapEq { column_id, value } => {
5019 let lookup = self.index_lookup_key_bytes(*column_id, value);
5020 self.bitmap
5021 .get(column_id)
5022 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5023 .unwrap_or_else(RowIdSet::empty)
5024 }
5025 Condition::BitmapIn { column_id, values } => {
5026 let bm = self.bitmap.get(column_id);
5027 let mut acc = roaring::RoaringBitmap::new();
5028 if let Some(b) = bm {
5029 for v in values {
5030 let lookup = self.index_lookup_key_bytes(*column_id, v);
5031 acc |= b.get(&lookup);
5032 }
5033 }
5034 RowIdSet::from_roaring(acc)
5035 }
5036 Condition::BytesPrefix { column_id, prefix } => {
5037 if let Some(b) = self.bitmap.get(column_id) {
5042 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5043 let mut acc = roaring::RoaringBitmap::new();
5044 for key in b.keys() {
5045 if key.starts_with(&lookup_prefix) {
5046 acc |= b.get(key);
5047 }
5048 }
5049 RowIdSet::from_roaring(acc)
5050 } else {
5051 RowIdSet::empty()
5052 }
5053 }
5054 Condition::FmContains { column_id, pattern } => self
5055 .fm
5056 .get(column_id)
5057 .map(|f| {
5058 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
5059 })
5060 .unwrap_or_else(RowIdSet::empty),
5061 Condition::FmContainsAll {
5062 column_id,
5063 patterns,
5064 } => {
5065 if let Some(f) = self.fm.get(column_id) {
5068 let sets: Vec<RowIdSet> = patterns
5069 .iter()
5070 .map(|pat| {
5071 RowIdSet::from_unsorted(
5072 f.locate(pat).into_iter().map(|r| r.0).collect(),
5073 )
5074 })
5075 .collect();
5076 RowIdSet::intersect_many(sets)
5077 } else {
5078 RowIdSet::empty()
5079 }
5080 }
5081 Condition::Ann {
5082 column_id,
5083 query,
5084 k,
5085 } => RowIdSet::from_unsorted(
5086 self.retrieve_filtered(
5087 &crate::query::Retriever::Ann {
5088 column_id: *column_id,
5089 query: query.clone(),
5090 k: *k,
5091 },
5092 snapshot,
5093 None,
5094 allowed,
5095 )?
5096 .into_iter()
5097 .map(|hit| hit.row_id.0)
5098 .collect(),
5099 ),
5100 Condition::SparseMatch {
5101 column_id,
5102 query,
5103 k,
5104 } => RowIdSet::from_unsorted(
5105 self.retrieve_filtered(
5106 &crate::query::Retriever::Sparse {
5107 column_id: *column_id,
5108 query: query.clone(),
5109 k: *k,
5110 },
5111 snapshot,
5112 None,
5113 allowed,
5114 )?
5115 .into_iter()
5116 .map(|hit| hit.row_id.0)
5117 .collect(),
5118 ),
5119 Condition::MinHashSimilar {
5120 column_id,
5121 query,
5122 k,
5123 } => match self.minhash.get(column_id) {
5124 Some(index) => {
5125 let candidates = index.candidate_row_ids(query);
5126 let eligible =
5127 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
5128 RowIdSet::from_unsorted(
5129 index
5130 .search_filtered(query, *k, |row_id| {
5131 eligible.contains(&row_id)
5132 && allowed.map_or(true, |allowed| allowed.contains(&row_id))
5133 })
5134 .into_iter()
5135 .map(|(row_id, _)| row_id.0)
5136 .collect(),
5137 )
5138 }
5139 None => RowIdSet::empty(),
5140 },
5141 Condition::Range { column_id, lo, hi } => {
5142 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5151 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5152 } else if self.run_refs.len() == 1 {
5153 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5154 r.range_row_id_set_i64(*column_id, *lo, *hi)?
5155 } else {
5156 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
5157 };
5158 set.remove_many(self.overlay_rid_set(snapshot));
5159 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
5160 set
5161 }
5162 Condition::RangeF64 {
5163 column_id,
5164 lo,
5165 lo_inclusive,
5166 hi,
5167 hi_inclusive,
5168 } => {
5169 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5172 RowIdSet::from_unsorted(
5173 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5174 .into_iter()
5175 .collect(),
5176 )
5177 } else if self.run_refs.len() == 1 {
5178 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5179 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
5180 } else {
5181 return self.range_scan_f64(
5182 *column_id,
5183 *lo,
5184 *lo_inclusive,
5185 *hi,
5186 *hi_inclusive,
5187 snapshot,
5188 );
5189 };
5190 set.remove_many(self.overlay_rid_set(snapshot));
5191 self.range_scan_overlay_f64(
5192 &mut set,
5193 *column_id,
5194 *lo,
5195 *lo_inclusive,
5196 *hi,
5197 *hi_inclusive,
5198 snapshot,
5199 );
5200 set
5201 }
5202 Condition::IsNull { column_id } => {
5203 let mut set = if self.run_refs.len() == 1 {
5204 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5205 r.null_row_id_set(*column_id, true)?
5206 } else {
5207 return self.null_scan(*column_id, true, snapshot);
5208 };
5209 set.remove_many(self.overlay_rid_set(snapshot));
5210 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
5211 set
5212 }
5213 Condition::IsNotNull { column_id } => {
5214 let mut set = if self.run_refs.len() == 1 {
5215 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5216 r.null_row_id_set(*column_id, false)?
5217 } else {
5218 return self.null_scan(*column_id, false, snapshot);
5219 };
5220 set.remove_many(self.overlay_rid_set(snapshot));
5221 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
5222 set
5223 }
5224 })
5225 }
5226
5227 fn range_scan_i64(
5235 &self,
5236 column_id: u16,
5237 lo: i64,
5238 hi: i64,
5239 snapshot: Snapshot,
5240 ) -> Result<RowIdSet> {
5241 let mut row_ids = Vec::new();
5242 let overlay_rids = self.overlay_rid_set(snapshot);
5243 for rr in &self.run_refs {
5244 let mut reader = self.open_reader(rr.run_id)?;
5245 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
5246 for rid in matched {
5247 if !overlay_rids.contains(&rid) {
5248 row_ids.push(rid);
5249 }
5250 }
5251 }
5252 let mut s = RowIdSet::from_unsorted(row_ids);
5253 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
5254 Ok(s)
5255 }
5256
5257 fn range_scan_f64(
5260 &self,
5261 column_id: u16,
5262 lo: f64,
5263 lo_inclusive: bool,
5264 hi: f64,
5265 hi_inclusive: bool,
5266 snapshot: Snapshot,
5267 ) -> Result<RowIdSet> {
5268 let mut row_ids = Vec::new();
5269 let overlay_rids = self.overlay_rid_set(snapshot);
5270 for rr in &self.run_refs {
5271 let mut reader = self.open_reader(rr.run_id)?;
5272 let matched = reader.range_row_ids_visible_f64(
5273 column_id,
5274 lo,
5275 lo_inclusive,
5276 hi,
5277 hi_inclusive,
5278 snapshot.epoch,
5279 )?;
5280 for rid in matched {
5281 if !overlay_rids.contains(&rid) {
5282 row_ids.push(rid);
5283 }
5284 }
5285 }
5286 let mut s = RowIdSet::from_unsorted(row_ids);
5287 self.range_scan_overlay_f64(
5288 &mut s,
5289 column_id,
5290 lo,
5291 lo_inclusive,
5292 hi,
5293 hi_inclusive,
5294 snapshot,
5295 );
5296 Ok(s)
5297 }
5298
5299 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
5301 let mut s = HashSet::new();
5302 for row in self.memtable.visible_versions(snapshot.epoch) {
5303 s.insert(row.row_id.0);
5304 }
5305 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5306 s.insert(row.row_id.0);
5307 }
5308 s
5309 }
5310
5311 fn range_scan_overlay_i64(
5312 &self,
5313 s: &mut RowIdSet,
5314 column_id: u16,
5315 lo: i64,
5316 hi: i64,
5317 snapshot: Snapshot,
5318 ) {
5319 let mut newest: HashMap<u64, &Row> = HashMap::new();
5324 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5325 let memtable = self.memtable.visible_versions(snapshot.epoch);
5326 for r in &mutable {
5327 newest.entry(r.row_id.0).or_insert(r);
5328 }
5329 for r in &memtable {
5330 newest.insert(r.row_id.0, r);
5331 }
5332 for row in newest.values() {
5333 if !row.deleted {
5334 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
5335 if *v >= lo && *v <= hi {
5336 s.insert(row.row_id.0);
5337 }
5338 }
5339 }
5340 }
5341 }
5342
5343 #[allow(clippy::too_many_arguments)]
5344 fn range_scan_overlay_f64(
5345 &self,
5346 s: &mut RowIdSet,
5347 column_id: u16,
5348 lo: f64,
5349 lo_inclusive: bool,
5350 hi: f64,
5351 hi_inclusive: bool,
5352 snapshot: Snapshot,
5353 ) {
5354 let mut newest: HashMap<u64, &Row> = HashMap::new();
5357 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5358 let memtable = self.memtable.visible_versions(snapshot.epoch);
5359 for r in &mutable {
5360 newest.entry(r.row_id.0).or_insert(r);
5361 }
5362 for r in &memtable {
5363 newest.insert(r.row_id.0, r);
5364 }
5365 for row in newest.values() {
5366 if !row.deleted {
5367 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
5368 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
5369 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
5370 if ok_lo && ok_hi {
5371 s.insert(row.row_id.0);
5372 }
5373 }
5374 }
5375 }
5376 }
5377
5378 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
5381 let mut row_ids = Vec::new();
5382 let overlay_rids = self.overlay_rid_set(snapshot);
5383 for rr in &self.run_refs {
5384 let mut reader = self.open_reader(rr.run_id)?;
5385 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
5386 for rid in matched {
5387 if !overlay_rids.contains(&rid) {
5388 row_ids.push(rid);
5389 }
5390 }
5391 }
5392 let mut s = RowIdSet::from_unsorted(row_ids);
5393 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
5394 Ok(s)
5395 }
5396
5397 fn null_scan_overlay(
5401 &self,
5402 s: &mut RowIdSet,
5403 column_id: u16,
5404 want_nulls: bool,
5405 snapshot: Snapshot,
5406 ) {
5407 let mut newest: HashMap<u64, &Row> = HashMap::new();
5408 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5409 let memtable = self.memtable.visible_versions(snapshot.epoch);
5410 for r in &mutable {
5411 newest.entry(r.row_id.0).or_insert(r);
5412 }
5413 for r in &memtable {
5414 newest.insert(r.row_id.0, r);
5415 }
5416 for row in newest.values() {
5417 if row.deleted {
5418 continue;
5419 }
5420 let is_null = !row.columns.contains_key(&column_id)
5421 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
5422 if is_null == want_nulls {
5423 s.insert(row.row_id.0);
5424 }
5425 }
5426 }
5427
5428 pub fn snapshot(&self) -> Snapshot {
5429 Snapshot::at(self.epoch.visible())
5430 }
5431
5432 pub fn data_generation(&self) -> u64 {
5434 self.data_generation
5435 }
5436
5437 pub fn pin_snapshot(&mut self) -> Snapshot {
5440 let e = self.epoch.visible();
5441 *self.pinned.entry(e).or_insert(0) += 1;
5442 Snapshot::at(e)
5443 }
5444
5445 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
5447 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
5448 *count -= 1;
5449 if *count == 0 {
5450 self.pinned.remove(&snap.epoch);
5451 }
5452 }
5453 }
5454
5455 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
5465 let local = self.pinned.keys().next().copied();
5466 let global = self.snapshots.min_pinned();
5467 let history = self.snapshots.history_floor(self.current_epoch());
5468 [local, global, history].into_iter().flatten().min()
5469 }
5470
5471 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
5475 self.ensure_writable()?;
5476 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
5477 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
5478 }
5479
5480 pub fn clear_ttl(&mut self) -> Result<()> {
5481 self.ensure_writable()?;
5482 self.apply_ttl_policy_at(None, self.current_epoch())
5483 }
5484
5485 pub fn ttl(&self) -> Option<TtlPolicy> {
5486 self.ttl
5487 }
5488
5489 pub(crate) fn prepare_ttl_policy(
5490 &self,
5491 column_name: &str,
5492 duration_nanos: u64,
5493 ) -> Result<TtlPolicy> {
5494 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
5495 return Err(MongrelError::InvalidArgument(
5496 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
5497 ));
5498 }
5499 let column = self
5500 .schema
5501 .columns
5502 .iter()
5503 .find(|column| column.name == column_name)
5504 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
5505 if column.ty != TypeId::TimestampNanos {
5506 return Err(MongrelError::Schema(format!(
5507 "TTL column {column_name} must be TimestampNanos, is {:?}",
5508 column.ty
5509 )));
5510 }
5511 Ok(TtlPolicy {
5512 column_id: column.id,
5513 duration_nanos,
5514 })
5515 }
5516
5517 pub(crate) fn apply_ttl_policy_at(
5518 &mut self,
5519 policy: Option<TtlPolicy>,
5520 epoch: Epoch,
5521 ) -> Result<()> {
5522 if let Some(policy) = policy {
5523 let column = self
5524 .schema
5525 .columns
5526 .iter()
5527 .find(|column| column.id == policy.column_id)
5528 .ok_or_else(|| {
5529 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
5530 })?;
5531 if column.ty != TypeId::TimestampNanos
5532 || policy.duration_nanos == 0
5533 || policy.duration_nanos > i64::MAX as u64
5534 {
5535 return Err(MongrelError::Schema("invalid TTL policy".into()));
5536 }
5537 }
5538 self.ttl = policy;
5539 self.agg_cache.clear();
5540 self.clear_result_cache();
5541 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5542 self.persist_manifest(epoch)
5543 }
5544
5545 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
5546 let Some(policy) = self.ttl else {
5547 return false;
5548 };
5549 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
5550 return false;
5551 };
5552 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
5553 }
5554
5555 pub fn current_epoch(&self) -> Epoch {
5556 self.epoch.visible()
5557 }
5558
5559 pub fn memtable_len(&self) -> usize {
5560 self.memtable.len()
5561 }
5562
5563 pub fn count(&self) -> u64 {
5566 if self.ttl.is_none()
5567 && self.pending_put_cols.is_empty()
5568 && self.pending_delete_rids.is_empty()
5569 && self.pending_rows.is_empty()
5570 && self.pending_dels.is_empty()
5571 && self.pending_truncate.is_none()
5572 {
5573 self.live_count
5574 } else {
5575 self.visible_rows(self.snapshot())
5576 .map(|rows| rows.len() as u64)
5577 .unwrap_or(self.live_count)
5578 }
5579 }
5580
5581 pub fn count_conditions(
5585 &mut self,
5586 conditions: &[crate::query::Condition],
5587 snapshot: Snapshot,
5588 ) -> Result<Option<u64>> {
5589 use crate::query::Condition;
5590 if self.ttl.is_some() {
5591 if conditions.is_empty() {
5592 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
5593 }
5594 let mut sets = Vec::with_capacity(conditions.len());
5595 for condition in conditions {
5596 sets.push(self.resolve_condition(condition, snapshot)?);
5597 }
5598 let survivors = RowIdSet::intersect_many(sets);
5599 let rows = self.visible_rows(snapshot)?;
5600 return Ok(Some(
5601 rows.into_iter()
5602 .filter(|row| survivors.contains(row.row_id.0))
5603 .count() as u64,
5604 ));
5605 }
5606 if conditions.is_empty() {
5607 return Ok(Some(self.count()));
5608 }
5609 let served = |c: &Condition| {
5610 matches!(
5611 c,
5612 Condition::Pk(_)
5613 | Condition::BitmapEq { .. }
5614 | Condition::BitmapIn { .. }
5615 | Condition::BytesPrefix { .. }
5616 | Condition::FmContains { .. }
5617 | Condition::FmContainsAll { .. }
5618 | Condition::Ann { .. }
5619 | Condition::Range { .. }
5620 | Condition::RangeF64 { .. }
5621 | Condition::SparseMatch { .. }
5622 | Condition::MinHashSimilar { .. }
5623 | Condition::IsNull { .. }
5624 | Condition::IsNotNull { .. }
5625 )
5626 };
5627 if !conditions.iter().all(served) {
5628 return Ok(None);
5629 }
5630 self.ensure_indexes_complete()?;
5631 if !self.pending_put_cols.is_empty()
5632 || !self.pending_delete_rids.is_empty()
5633 || !self.pending_rows.is_empty()
5634 || !self.pending_dels.is_empty()
5635 || self.pending_truncate.is_some()
5636 {
5637 let mut sets = Vec::with_capacity(conditions.len());
5638 for condition in conditions {
5639 sets.push(self.resolve_condition(condition, snapshot)?);
5640 }
5641 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5642 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
5643 }
5644 let mut sets = Vec::with_capacity(conditions.len());
5645 for condition in conditions {
5646 sets.push(self.resolve_condition(condition, snapshot)?);
5647 }
5648 let mut rids = RowIdSet::intersect_many(sets);
5649 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
5659 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
5660 }
5661 let count = rids.len() as u64;
5662 crate::trace::QueryTrace::record(|t| {
5663 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
5664 t.survivor_count = Some(count as usize);
5665 t.conditions_pushed = conditions.len();
5666 });
5667 Ok(Some(count))
5668 }
5669
5670 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
5675 let mut out = Vec::new();
5676 for row in self.memtable.visible_versions(snapshot.epoch) {
5677 if row.deleted {
5678 out.push(row.row_id.0);
5679 }
5680 }
5681 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5682 if row.deleted {
5683 out.push(row.row_id.0);
5684 }
5685 }
5686 out
5687 }
5688
5689 pub fn bulk_load_columns(
5698 &mut self,
5699 user_columns: Vec<(u16, columnar::NativeColumn)>,
5700 ) -> Result<Epoch> {
5701 self.bulk_load_columns_with(user_columns, 3, false, true)
5702 }
5703
5704 pub fn bulk_load_fast(
5711 &mut self,
5712 user_columns: Vec<(u16, columnar::NativeColumn)>,
5713 ) -> Result<Epoch> {
5714 self.bulk_load_columns_with(user_columns, -1, true, false)
5715 }
5716
5717 fn bulk_load_columns_with(
5718 &mut self,
5719 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
5720 zstd_level: i32,
5721 force_plain: bool,
5722 lz4: bool,
5723 ) -> Result<Epoch> {
5724 let epoch = self.commit()?;
5725 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
5726 if n == 0 {
5727 return Ok(epoch);
5728 }
5729 let live_before = self.live_count;
5730 self.spill_mutable_run(epoch)?;
5732 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
5733 && self.indexes_complete
5734 && self.run_refs.is_empty()
5735 && self.memtable.is_empty()
5736 && self.mutable_run.is_empty();
5737 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
5740 self.validate_columns_not_null(&user_columns, n)?;
5741 let winner_idx = self
5742 .bulk_pk_winner_indices(&user_columns, n)
5743 .filter(|idx| idx.len() != n);
5744 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
5745 match winner_idx.as_deref() {
5746 Some(idx) => {
5747 let compacted = user_columns
5748 .iter()
5749 .map(|(id, c)| (*id, c.gather(idx)))
5750 .collect();
5751 (compacted, idx.len())
5752 }
5753 None => (user_columns, n),
5754 };
5755 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
5756 let first = self.allocator.alloc_range(write_n as u64).0;
5757 for rid in first..first + write_n as u64 {
5758 self.reservoir.offer(rid);
5759 }
5760 let run_id = self.next_run_id;
5761 self.next_run_id += 1;
5762 let path = self.run_path(run_id);
5763 let mut writer =
5764 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
5765 if force_plain {
5766 writer = writer.with_plain();
5767 } else if lz4 {
5768 writer = writer.with_lz4();
5771 } else {
5772 writer = writer.with_zstd_level(zstd_level);
5773 }
5774 if let Some(kek) = &self.kek {
5775 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5776 }
5777 let header = writer.write_native(&path, &write_columns, write_n, first)?;
5778 self.run_refs.push(RunRef {
5779 run_id: run_id as u128,
5780 level: 0,
5781 epoch_created: epoch.0,
5782 row_count: header.row_count,
5783 });
5784 self.live_count = self.live_count.saturating_add(write_n as u64);
5785 if eager_index_build {
5786 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
5787 self.index_columns_bulk(&write_columns, &row_ids);
5788 self.indexes_complete = true;
5789 self.build_learned_ranges()?;
5790 } else {
5791 self.indexes_complete = false;
5795 }
5796 self.mark_flushed(epoch)?;
5797 self.persist_manifest(epoch)?;
5798 if eager_index_build {
5799 self.checkpoint_indexes(epoch);
5800 }
5801 self.clear_result_cache();
5802 self.data_generation = self.data_generation.wrapping_add(1);
5803 Ok(epoch)
5804 }
5805
5806 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
5824 let n = row_ids.len();
5825 if n == 0 {
5826 return;
5827 }
5828 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
5829 columns.iter().map(|(id, c)| (*id, c)).collect();
5830 let ty_of: std::collections::HashMap<u16, TypeId> = self
5831 .schema
5832 .columns
5833 .iter()
5834 .map(|c| (c.id, c.ty.clone()))
5835 .collect();
5836 let pk_id = self.schema.primary_key().map(|c| c.id);
5837
5838 for (i, &rid) in row_ids.iter().enumerate() {
5839 let row_id = RowId(rid);
5840 if let Some(pid) = pk_id {
5841 if let Some(col) = by_id.get(&pid) {
5842 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
5843 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
5844 self.insert_hot_pk(key, row_id);
5845 }
5846 }
5847 }
5848 for idef in &self.schema.indexes {
5849 let Some(col) = by_id.get(&idef.column_id) else {
5850 continue;
5851 };
5852 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
5853 match idef.kind {
5854 IndexKind::Bitmap => {
5855 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5856 if let Some(key) =
5857 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
5858 {
5859 b.insert(key, row_id);
5860 }
5861 }
5862 }
5863 IndexKind::FmIndex => {
5864 if let Some(f) = self.fm.get_mut(&idef.column_id) {
5865 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5866 f.insert(bytes.to_vec(), row_id);
5867 }
5868 }
5869 }
5870 IndexKind::Sparse => {
5871 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
5872 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5873 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
5874 s.insert(&terms, row_id);
5875 }
5876 }
5877 }
5878 }
5879 IndexKind::MinHash => {
5880 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
5881 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5882 let tokens = crate::index::token_hashes_from_bytes(bytes);
5883 mh.insert(&tokens, row_id);
5884 }
5885 }
5886 }
5887 _ => {}
5888 }
5889 }
5890 }
5891 }
5892
5893 pub fn visible_columns_native(
5898 &self,
5899 snapshot: Snapshot,
5900 projection: Option<&[u16]>,
5901 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
5902 let wanted: Vec<u16> = match projection {
5903 Some(p) => p.to_vec(),
5904 None => self.schema.columns.iter().map(|c| c.id).collect(),
5905 };
5906 if self.ttl.is_none()
5907 && self.memtable.is_empty()
5908 && self.mutable_run.is_empty()
5909 && self.run_refs.len() == 1
5910 {
5911 let rr = self.run_refs[0].clone();
5912 let mut reader = self.open_reader(rr.run_id)?;
5913 let idxs = reader.visible_indices_native(snapshot.epoch)?;
5914 let all_visible = idxs.len() == reader.row_count();
5915 if reader.has_mmap() {
5921 use rayon::prelude::*;
5922 let valid: Vec<u16> = wanted
5925 .iter()
5926 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
5927 .copied()
5928 .collect();
5929 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
5931 .par_iter()
5932 .filter_map(|cid| {
5933 reader
5934 .column_native_shared(*cid)
5935 .ok()
5936 .map(|col| (*cid, col))
5937 })
5938 .collect();
5939 let cols = decoded
5940 .into_iter()
5941 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
5942 .collect();
5943 return Ok(cols);
5944 }
5945 let mut cols = Vec::with_capacity(wanted.len());
5946 for cid in &wanted {
5947 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
5948 Some(c) => c,
5949 None => continue,
5950 };
5951 let col = reader.column_native(cdef.id)?;
5952 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
5953 }
5954 return Ok(cols);
5955 }
5956 let vcols = self.visible_columns(snapshot)?;
5957 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
5958 let out: Vec<(u16, columnar::NativeColumn)> = vcols
5959 .into_iter()
5960 .filter(|(id, _)| want_set.contains(id))
5961 .map(|(id, vals)| {
5962 let ty = self
5963 .schema
5964 .columns
5965 .iter()
5966 .find(|c| c.id == id)
5967 .map(|c| c.ty.clone())
5968 .unwrap_or(TypeId::Bytes);
5969 (id, columnar::values_to_native(ty, &vals))
5970 })
5971 .collect();
5972 Ok(out)
5973 }
5974
5975 pub fn run_count(&self) -> usize {
5976 self.run_refs.len()
5977 }
5978
5979 pub fn memtable_is_empty(&self) -> bool {
5981 self.memtable.is_empty()
5982 }
5983
5984 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
5988 self.page_cache.stats()
5989 }
5990
5991 pub fn reset_page_cache_stats(&self) {
5993 self.page_cache.reset_stats();
5994 }
5995
5996 pub fn run_ids(&self) -> Vec<u128> {
5999 self.run_refs.iter().map(|r| r.run_id).collect()
6000 }
6001
6002 pub fn single_run_is_clean(&self) -> bool {
6006 if self.ttl.is_some() || self.run_refs.len() != 1 {
6007 return false;
6008 }
6009 self.open_reader(self.run_refs[0].run_id)
6010 .map(|r| r.is_clean())
6011 .unwrap_or(false)
6012 }
6013
6014 fn resolve_footprint(
6021 &self,
6022 conditions: &[crate::query::Condition],
6023 snapshot: Snapshot,
6024 ) -> roaring::RoaringBitmap {
6025 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6026 return roaring::RoaringBitmap::new();
6027 }
6028 if self.run_refs.is_empty() {
6029 return roaring::RoaringBitmap::new();
6030 }
6031 if self.run_refs.len() == 1 {
6033 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
6034 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
6035 return rids.to_roaring_lossy();
6036 }
6037 }
6038 }
6039 roaring::RoaringBitmap::new()
6040 }
6041
6042 pub fn query_columns_native_cached(
6053 &mut self,
6054 conditions: &[crate::query::Condition],
6055 projection: Option<&[u16]>,
6056 snapshot: Snapshot,
6057 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
6058 if self.ttl.is_some() {
6061 return self.query_columns_native(conditions, projection, snapshot);
6062 }
6063 if conditions.is_empty() {
6064 return self.query_columns_native(conditions, projection, snapshot);
6065 }
6066 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
6070 if let Some(hit) = self.result_cache.lock().get_columns(key) {
6071 crate::trace::QueryTrace::record(|t| {
6072 t.result_cache_hit = true;
6073 t.scan_mode = crate::trace::ScanMode::NativePushdown;
6074 });
6075 return Ok(Some((*hit).clone()));
6076 }
6077 let res = self.query_columns_native(conditions, projection, snapshot)?;
6078 if let Some(cols) = &res {
6079 let footprint = self.resolve_footprint(conditions, snapshot);
6080 let condition_cols = crate::query::condition_columns(conditions);
6081 self.result_cache.lock().insert(
6082 key,
6083 CachedEntry {
6084 data: CachedData::Columns(Arc::new(cols.clone())),
6085 footprint,
6086 condition_cols,
6087 },
6088 );
6089 }
6090 Ok(res)
6091 }
6092
6093 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
6098 if self.ttl.is_some() {
6099 return self.query(q);
6100 }
6101 if q.conditions.is_empty() {
6102 return self.query(q);
6103 }
6104 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
6105 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
6106 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
6107 if let Some(hit) = self.result_cache.lock().get_rows(key) {
6108 crate::trace::QueryTrace::record(|t| {
6109 t.result_cache_hit = true;
6110 t.scan_mode = crate::trace::ScanMode::Materialized;
6111 });
6112 return Ok((*hit).clone());
6113 }
6114 let rows = self.query(q)?;
6115 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
6116 let condition_cols = crate::query::condition_columns(&q.conditions);
6117 self.result_cache.lock().insert(
6118 key,
6119 CachedEntry {
6120 data: CachedData::Rows(Arc::new(rows.clone())),
6121 footprint,
6122 condition_cols,
6123 },
6124 );
6125 Ok(rows)
6126 }
6127
6128 #[allow(clippy::type_complexity)]
6143 pub fn query_columns_native_traced(
6144 &mut self,
6145 conditions: &[crate::query::Condition],
6146 projection: Option<&[u16]>,
6147 snapshot: Snapshot,
6148 ) -> Result<(
6149 Option<Vec<(u16, columnar::NativeColumn)>>,
6150 crate::trace::QueryTrace,
6151 )> {
6152 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6153 self.query_columns_native(conditions, projection, snapshot)
6154 });
6155 Ok((result?, trace))
6156 }
6157
6158 #[allow(clippy::type_complexity)]
6161 pub fn query_columns_native_cached_traced(
6162 &mut self,
6163 conditions: &[crate::query::Condition],
6164 projection: Option<&[u16]>,
6165 snapshot: Snapshot,
6166 ) -> Result<(
6167 Option<Vec<(u16, columnar::NativeColumn)>>,
6168 crate::trace::QueryTrace,
6169 )> {
6170 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6171 self.query_columns_native_cached(conditions, projection, snapshot)
6172 });
6173 Ok((result?, trace))
6174 }
6175
6176 pub fn native_page_cursor_traced(
6178 &self,
6179 snapshot: Snapshot,
6180 projection: Vec<(u16, TypeId)>,
6181 conditions: &[crate::query::Condition],
6182 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
6183 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6184 self.native_page_cursor(snapshot, projection, conditions)
6185 });
6186 Ok((result?, trace))
6187 }
6188
6189 pub fn native_multi_run_cursor_traced(
6191 &self,
6192 snapshot: Snapshot,
6193 projection: Vec<(u16, TypeId)>,
6194 conditions: &[crate::query::Condition],
6195 ) -> Result<(
6196 Option<crate::cursor::MultiRunCursor>,
6197 crate::trace::QueryTrace,
6198 )> {
6199 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6200 self.native_multi_run_cursor(snapshot, projection, conditions)
6201 });
6202 Ok((result?, trace))
6203 }
6204
6205 pub fn count_conditions_traced(
6207 &mut self,
6208 conditions: &[crate::query::Condition],
6209 snapshot: Snapshot,
6210 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
6211 let (result, trace) =
6212 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
6213 Ok((result?, trace))
6214 }
6215
6216 pub fn query_traced(
6218 &mut self,
6219 q: &crate::query::Query,
6220 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
6221 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
6222 Ok((result?, trace))
6223 }
6224
6225 pub fn query_columns_native(
6230 &mut self,
6231 conditions: &[crate::query::Condition],
6232 projection: Option<&[u16]>,
6233 snapshot: Snapshot,
6234 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
6235 use crate::query::Condition;
6236 if self.ttl.is_some() {
6239 return Ok(None);
6240 }
6241 if conditions.is_empty() {
6242 return Ok(None);
6243 }
6244 self.ensure_indexes_complete()?;
6245
6246 let served = |c: &Condition| {
6251 matches!(
6252 c,
6253 Condition::Pk(_)
6254 | Condition::BitmapEq { .. }
6255 | Condition::BitmapIn { .. }
6256 | Condition::BytesPrefix { .. }
6257 | Condition::FmContains { .. }
6258 | Condition::FmContainsAll { .. }
6259 | Condition::Ann { .. }
6260 | Condition::Range { .. }
6261 | Condition::RangeF64 { .. }
6262 | Condition::SparseMatch { .. }
6263 | Condition::MinHashSimilar { .. }
6264 | Condition::IsNull { .. }
6265 | Condition::IsNotNull { .. }
6266 )
6267 };
6268 if !conditions.iter().all(served) {
6269 return Ok(None);
6270 }
6271 let fast_path =
6272 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
6273 crate::trace::QueryTrace::record(|t| {
6274 t.run_count = self.run_refs.len();
6275 t.memtable_rows = self.memtable.len();
6276 t.mutable_run_rows = self.mutable_run.len();
6277 t.conditions_pushed = conditions.len();
6278 t.learned_range_used = conditions.iter().any(|c| match c {
6279 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
6280 self.learned_range.contains_key(column_id)
6281 }
6282 _ => false,
6283 });
6284 });
6285 let col_ids: Vec<u16> = projection
6287 .map(|p| p.to_vec())
6288 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
6289 let proj_pairs: Vec<(u16, TypeId)> = col_ids
6290 .iter()
6291 .map(|&cid| {
6292 let ty = self
6293 .schema
6294 .columns
6295 .iter()
6296 .find(|c| c.id == cid)
6297 .map(|c| c.ty.clone())
6298 .unwrap_or(TypeId::Bytes);
6299 (cid, ty)
6300 })
6301 .collect();
6302
6303 if fast_path {
6309 let needs_column = conditions.iter().any(|c| match c {
6312 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
6313 Condition::RangeF64 { column_id, .. } => {
6314 !self.learned_range.contains_key(column_id)
6315 }
6316 _ => false,
6317 });
6318 let mut reader_opt: Option<RunReader> = if needs_column {
6319 Some(self.open_reader(self.run_refs[0].run_id)?)
6320 } else {
6321 None
6322 };
6323 let mut sets: Vec<RowIdSet> = Vec::new();
6324 for c in conditions {
6325 let s = match c {
6326 Condition::Range { column_id, lo, hi }
6327 if !self.learned_range.contains_key(column_id) =>
6328 {
6329 if reader_opt.is_none() {
6330 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
6331 }
6332 reader_opt
6333 .as_mut()
6334 .expect("reader opened for range")
6335 .range_row_id_set_i64(*column_id, *lo, *hi)?
6336 }
6337 Condition::RangeF64 {
6338 column_id,
6339 lo,
6340 lo_inclusive,
6341 hi,
6342 hi_inclusive,
6343 } if !self.learned_range.contains_key(column_id) => {
6344 if reader_opt.is_none() {
6345 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
6346 }
6347 reader_opt
6348 .as_mut()
6349 .expect("reader opened for range")
6350 .range_row_id_set_f64(
6351 *column_id,
6352 *lo,
6353 *lo_inclusive,
6354 *hi,
6355 *hi_inclusive,
6356 )?
6357 }
6358 _ => self.resolve_condition(c, snapshot)?,
6359 };
6360 sets.push(s);
6361 }
6362 let candidates = RowIdSet::intersect_many(sets);
6363 crate::trace::QueryTrace::record(|t| {
6364 t.survivor_count = Some(candidates.len());
6365 });
6366 if candidates.is_empty() {
6367 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
6368 .iter()
6369 .map(|&id| {
6370 (
6371 id,
6372 columnar::null_native(
6373 proj_pairs
6374 .iter()
6375 .find(|(c, _)| c == &id)
6376 .map(|(_, t)| t.clone())
6377 .unwrap_or(TypeId::Bytes),
6378 0,
6379 ),
6380 )
6381 })
6382 .collect();
6383 return Ok(Some(cols));
6384 }
6385 let mut reader = match reader_opt.take() {
6386 Some(r) => r,
6387 None => self.open_reader(self.run_refs[0].run_id)?,
6388 };
6389 let candidate_ids = candidates.into_sorted_vec();
6390 let (positions, fast_rid) = if let Some(positions) =
6391 reader.positions_for_row_ids_fast(&candidate_ids)
6392 {
6393 (positions, true)
6394 } else {
6395 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
6396 match col {
6397 columnar::NativeColumn::Int64 { data, .. } => {
6398 let mut p: Vec<usize> = candidate_ids
6399 .iter()
6400 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
6401 .collect();
6402 p.sort_unstable();
6403 (p, false)
6404 }
6405 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
6406 }
6407 };
6408 crate::trace::QueryTrace::record(|t| {
6409 t.scan_mode = crate::trace::ScanMode::NativePushdown;
6410 t.fast_row_id_map = fast_rid;
6411 });
6412 let mut cols = Vec::with_capacity(col_ids.len());
6413 for cid in &col_ids {
6414 let col = reader.column_native(*cid)?;
6415 cols.push((*cid, col.gather(&positions)));
6416 }
6417 return Ok(Some(cols));
6418 }
6419
6420 if !self.run_refs.is_empty() {
6433 use crate::cursor::{drain_cursor_to_columns, Cursor};
6434 let remaining: usize;
6435 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
6436 let c = self
6437 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
6438 .expect("single-run cursor should build when run_refs.len() == 1");
6439 remaining = c.remaining_rows();
6440 Box::new(c)
6441 } else {
6442 let c = self
6443 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
6444 .expect("multi-run cursor should build when run_refs.len() >= 1");
6445 remaining = c.remaining_rows();
6446 Box::new(c)
6447 };
6448 crate::trace::QueryTrace::record(|t| {
6449 if t.survivor_count.is_none() {
6450 t.survivor_count = Some(remaining);
6451 }
6452 });
6453 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
6454 return Ok(Some(cols));
6455 }
6456
6457 crate::trace::QueryTrace::record(|t| {
6462 t.scan_mode = crate::trace::ScanMode::Materialized;
6463 t.row_materialized = true;
6464 });
6465 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6466 for c in conditions {
6467 sets.push(self.resolve_condition(c, snapshot)?);
6468 }
6469 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6470 let rows = self.rows_for_rids(&rids, snapshot)?;
6471 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
6472 for (cid, ty) in &proj_pairs {
6473 let vals: Vec<Value> = rows
6474 .iter()
6475 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6476 .collect();
6477 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
6478 }
6479 Ok(Some(cols))
6480 }
6481
6482 pub fn native_page_cursor(
6497 &self,
6498 snapshot: Snapshot,
6499 projection: Vec<(u16, TypeId)>,
6500 conditions: &[crate::query::Condition],
6501 ) -> Result<Option<NativePageCursor>> {
6502 use crate::cursor::build_page_plans;
6503 if self.ttl.is_some() {
6504 return Ok(None);
6505 }
6506 if !conditions.is_empty() && !self.indexes_complete {
6509 return Ok(None);
6510 }
6511 if self.run_refs.len() != 1 {
6512 return Ok(None);
6513 }
6514 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6515 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6516
6517 let overlay_rids: HashSet<u64> = {
6520 let mut s = HashSet::new();
6521 for row in self.memtable.visible_versions(snapshot.epoch) {
6522 s.insert(row.row_id.0);
6523 }
6524 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6525 s.insert(row.row_id.0);
6526 }
6527 s
6528 };
6529
6530 let survivors = if conditions.is_empty() {
6534 None
6535 } else {
6536 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
6537 };
6538
6539 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
6546 survivors.clone()
6547 } else if let Some(s) = &survivors {
6548 let mut run_set = s.clone();
6549 run_set.remove_many(overlay_rids.iter().copied());
6550 Some(run_set)
6551 } else {
6552 Some(RowIdSet::from_unsorted(
6553 rids.iter()
6554 .map(|&r| r as u64)
6555 .filter(|r| !overlay_rids.contains(r))
6556 .collect(),
6557 ))
6558 };
6559
6560 let overlay_rows = if overlay_rids.is_empty() {
6561 Vec::new()
6562 } else {
6563 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6564 self.overlay_visible_rows(snapshot, bound)
6565 };
6566
6567 let plans = if positions.is_empty() {
6569 Vec::new()
6570 } else {
6571 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
6572 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
6573 };
6574
6575 let overlay = if overlay_rows.is_empty() {
6577 None
6578 } else {
6579 let filtered =
6580 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6581 if filtered.is_empty() {
6582 None
6583 } else {
6584 Some(self.materialize_overlay(&filtered, &projection))
6585 }
6586 };
6587
6588 let overlay_row_count = overlay
6589 .as_ref()
6590 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6591 .unwrap_or(0);
6592 crate::trace::QueryTrace::record(|t| {
6593 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
6594 t.run_count = self.run_refs.len();
6595 t.memtable_rows = self.memtable.len();
6596 t.mutable_run_rows = self.mutable_run.len();
6597 t.overlay_rows = overlay_row_count;
6598 t.conditions_pushed = conditions.len();
6599 t.pages_decoded = plans
6600 .iter()
6601 .map(|p| p.positions.len())
6602 .sum::<usize>()
6603 .min(1);
6604 });
6605
6606 Ok(Some(NativePageCursor::new_with_overlay(
6607 reader, projection, plans, overlay,
6608 )))
6609 }
6610 #[allow(clippy::type_complexity)]
6620 pub fn native_multi_run_cursor(
6621 &self,
6622 snapshot: Snapshot,
6623 projection: Vec<(u16, TypeId)>,
6624 conditions: &[crate::query::Condition],
6625 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
6626 use crate::cursor::{MultiRunCursor, RunStream};
6627 use crate::sorted_run::SYS_ROW_ID;
6628 use std::collections::{BinaryHeap, HashMap, HashSet};
6629 if self.ttl.is_some() {
6630 return Ok(None);
6631 }
6632 if !conditions.is_empty() && !self.indexes_complete {
6635 return Ok(None);
6636 }
6637 if self.run_refs.is_empty() {
6638 return Ok(None);
6639 }
6640
6641 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
6643 Vec::with_capacity(self.run_refs.len());
6644 for rr in &self.run_refs {
6645 let mut reader = self.open_reader(rr.run_id)?;
6646 let (rids, eps, del) = reader.system_columns_native()?;
6647 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
6648 run_meta.push((reader, rids, eps, del, page_rows));
6649 }
6650
6651 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
6655 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
6656 for i in 0..rids.len() {
6657 let rid = rids[i] as u64;
6658 let e = eps[i] as u64;
6659 if e > snapshot.epoch.0 {
6660 continue;
6661 }
6662 let is_del = del[i] != 0;
6663 best.entry(rid)
6664 .and_modify(|cur| {
6665 if e > cur.0 {
6666 *cur = (e, run_idx, i, is_del);
6667 }
6668 })
6669 .or_insert((e, run_idx, i, is_del));
6670 }
6671 }
6672
6673 let overlay_rids: HashSet<u64> = {
6675 let mut s = HashSet::new();
6676 for row in self.memtable.visible_versions(snapshot.epoch) {
6677 s.insert(row.row_id.0);
6678 }
6679 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6680 s.insert(row.row_id.0);
6681 }
6682 s
6683 };
6684
6685 let survivors: Option<RowIdSet> = if conditions.is_empty() {
6687 None
6688 } else {
6689 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6690 for c in conditions {
6691 sets.push(self.resolve_condition(c, snapshot)?);
6692 }
6693 Some(RowIdSet::intersect_many(sets))
6694 };
6695
6696 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
6700 for (rid, (_, run_idx, pos, deleted)) in &best {
6701 if *deleted {
6702 continue;
6703 }
6704 if overlay_rids.contains(rid) {
6705 continue;
6706 }
6707 if let Some(s) = &survivors {
6708 if !s.contains(*rid) {
6709 continue;
6710 }
6711 }
6712 per_run[*run_idx].push((*rid, *pos));
6713 }
6714 for v in per_run.iter_mut() {
6715 v.sort_unstable_by_key(|&(rid, _)| rid);
6716 }
6717
6718 let mut streams = Vec::with_capacity(run_meta.len());
6720 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
6721 let mut total = 0usize;
6722 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
6723 let mut starts = Vec::with_capacity(page_rows.len());
6724 let mut acc = 0usize;
6725 for &r in &page_rows {
6726 starts.push(acc);
6727 acc += r;
6728 }
6729 let mut survivors_vec: Vec<(u64, usize, usize)> =
6730 Vec::with_capacity(per_run[run_idx].len());
6731 for &(rid, pos) in &per_run[run_idx] {
6732 let page_seq = match starts.partition_point(|&s| s <= pos) {
6733 0 => continue,
6734 p => p - 1,
6735 };
6736 let within = pos - starts[page_seq];
6737 survivors_vec.push((rid, page_seq, within));
6738 }
6739 total += survivors_vec.len();
6740 if let Some(&(rid, _, _)) = survivors_vec.first() {
6741 heap.push(std::cmp::Reverse((rid, run_idx)));
6742 }
6743 streams.push(RunStream::new(reader, survivors_vec, page_rows));
6744 }
6745
6746 let overlay_rows = if overlay_rids.is_empty() {
6748 Vec::new()
6749 } else {
6750 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6751 self.overlay_visible_rows(snapshot, bound)
6752 };
6753 let overlay = if overlay_rows.is_empty() {
6754 None
6755 } else {
6756 let filtered =
6757 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6758 if filtered.is_empty() {
6759 None
6760 } else {
6761 Some(self.materialize_overlay(&filtered, &projection))
6762 }
6763 };
6764
6765 let overlay_row_count = overlay
6766 .as_ref()
6767 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6768 .unwrap_or(0);
6769 crate::trace::QueryTrace::record(|t| {
6770 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
6771 t.run_count = self.run_refs.len();
6772 t.memtable_rows = self.memtable.len();
6773 t.mutable_run_rows = self.mutable_run.len();
6774 t.overlay_rows = overlay_row_count;
6775 t.conditions_pushed = conditions.len();
6776 t.survivor_count = Some(total);
6777 });
6778
6779 Ok(Some(MultiRunCursor::new(
6780 streams, projection, heap, total, overlay,
6781 )))
6782 }
6783
6784 fn overlay_materialization_bound<'a>(
6796 conditions: &[crate::query::Condition],
6797 survivors: &'a Option<RowIdSet>,
6798 ) -> Option<&'a RowIdSet> {
6799 use crate::query::Condition;
6800 let has_range = conditions
6801 .iter()
6802 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6803 if has_range {
6804 None
6805 } else {
6806 survivors.as_ref()
6807 }
6808 }
6809
6810 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
6822 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
6823 let mut fold = |row: Row| {
6824 if let Some(b) = bound {
6825 if !b.contains(row.row_id.0) {
6826 return;
6827 }
6828 }
6829 best.entry(row.row_id.0)
6830 .and_modify(|(be, br)| {
6831 if row.committed_epoch > *be {
6832 *be = row.committed_epoch;
6833 *br = row.clone();
6834 }
6835 })
6836 .or_insert_with(|| (row.committed_epoch, row));
6837 };
6838 for row in self.memtable.visible_versions(snapshot.epoch) {
6839 fold(row);
6840 }
6841 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6842 fold(row);
6843 }
6844 let mut out: Vec<Row> = best
6845 .into_values()
6846 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
6847 .collect();
6848 out.sort_by_key(|r| r.row_id);
6849 out
6850 }
6851
6852 fn filter_overlay_rows(
6860 &self,
6861 rows: Vec<Row>,
6862 conditions: &[crate::query::Condition],
6863 survivors: Option<&RowIdSet>,
6864 snapshot: Snapshot,
6865 ) -> Result<Vec<Row>> {
6866 if conditions.is_empty() {
6867 return Ok(rows);
6868 }
6869 use crate::query::Condition;
6870 let all_index_served = !conditions
6874 .iter()
6875 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6876 if all_index_served {
6877 return Ok(rows
6878 .into_iter()
6879 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
6880 .collect());
6881 }
6882 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6885 for c in conditions {
6886 let s = match c {
6887 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
6888 _ => self.resolve_condition(c, snapshot)?,
6889 };
6890 per_cond_sets.push(s);
6891 }
6892 Ok(rows
6893 .into_iter()
6894 .filter(|row| {
6895 conditions.iter().enumerate().all(|(i, c)| match c {
6896 Condition::Range { column_id, lo, hi } => {
6897 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
6898 }
6899 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
6900 match row.columns.get(column_id) {
6901 Some(Value::Float64(v)) => {
6902 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
6903 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
6904 lo_ok && hi_ok
6905 }
6906 _ => false,
6907 }
6908 }
6909 _ => per_cond_sets[i].contains(row.row_id.0),
6910 })
6911 })
6912 .collect())
6913 }
6914
6915 fn materialize_overlay(
6918 &self,
6919 rows: &[Row],
6920 projection: &[(u16, TypeId)],
6921 ) -> Vec<columnar::NativeColumn> {
6922 if projection.is_empty() {
6923 return vec![columnar::null_native(TypeId::Int64, rows.len())];
6924 }
6925 let mut cols = Vec::with_capacity(projection.len());
6926 for (cid, ty) in projection {
6927 let vals: Vec<Value> = rows
6928 .iter()
6929 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6930 .collect();
6931 cols.push(columnar::values_to_native(ty.clone(), &vals));
6932 }
6933 cols
6934 }
6935
6936 fn resolve_survivor_rids(
6941 &self,
6942 conditions: &[crate::query::Condition],
6943 reader: &mut RunReader,
6944 snapshot: Snapshot,
6945 ) -> Result<RowIdSet> {
6946 use crate::query::Condition;
6947 let mut sets: Vec<RowIdSet> = Vec::new();
6948 for c in conditions {
6949 self.validate_condition(c)?;
6950 let s: RowIdSet = match c {
6951 Condition::Pk(key) => {
6952 let lookup = self
6953 .schema
6954 .primary_key()
6955 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
6956 .unwrap_or_else(|| key.clone());
6957 self.hot
6958 .get(&lookup)
6959 .map(|r| RowIdSet::one(r.0))
6960 .unwrap_or_else(RowIdSet::empty)
6961 }
6962 Condition::BitmapEq { column_id, value } => {
6963 let lookup = self.index_lookup_key_bytes(*column_id, value);
6964 self.bitmap
6965 .get(column_id)
6966 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
6967 .unwrap_or_else(RowIdSet::empty)
6968 }
6969 Condition::BitmapIn { column_id, values } => {
6970 let bm = self.bitmap.get(column_id);
6971 let mut acc = roaring::RoaringBitmap::new();
6972 if let Some(b) = bm {
6973 for v in values {
6974 let lookup = self.index_lookup_key_bytes(*column_id, v);
6975 acc |= b.get(&lookup);
6976 }
6977 }
6978 RowIdSet::from_roaring(acc)
6979 }
6980 Condition::BytesPrefix { column_id, prefix } => {
6981 if let Some(b) = self.bitmap.get(column_id) {
6982 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
6983 let mut acc = roaring::RoaringBitmap::new();
6984 for key in b.keys() {
6985 if key.starts_with(&lookup_prefix) {
6986 acc |= b.get(key);
6987 }
6988 }
6989 RowIdSet::from_roaring(acc)
6990 } else {
6991 RowIdSet::empty()
6992 }
6993 }
6994 Condition::FmContains { column_id, pattern } => self
6995 .fm
6996 .get(column_id)
6997 .map(|f| {
6998 RowIdSet::from_unsorted(
6999 f.locate(pattern).into_iter().map(|r| r.0).collect(),
7000 )
7001 })
7002 .unwrap_or_else(RowIdSet::empty),
7003 Condition::FmContainsAll {
7004 column_id,
7005 patterns,
7006 } => {
7007 if let Some(f) = self.fm.get(column_id) {
7008 let sets: Vec<RowIdSet> = patterns
7009 .iter()
7010 .map(|pat| {
7011 RowIdSet::from_unsorted(
7012 f.locate(pat).into_iter().map(|r| r.0).collect(),
7013 )
7014 })
7015 .collect();
7016 RowIdSet::intersect_many(sets)
7017 } else {
7018 RowIdSet::empty()
7019 }
7020 }
7021 Condition::Ann {
7022 column_id,
7023 query,
7024 k,
7025 } => RowIdSet::from_unsorted(
7026 self.retrieve_filtered(
7027 &crate::query::Retriever::Ann {
7028 column_id: *column_id,
7029 query: query.clone(),
7030 k: *k,
7031 },
7032 snapshot,
7033 None,
7034 None,
7035 )?
7036 .into_iter()
7037 .map(|hit| hit.row_id.0)
7038 .collect(),
7039 ),
7040 Condition::SparseMatch {
7041 column_id,
7042 query,
7043 k,
7044 } => RowIdSet::from_unsorted(
7045 self.retrieve_filtered(
7046 &crate::query::Retriever::Sparse {
7047 column_id: *column_id,
7048 query: query.clone(),
7049 k: *k,
7050 },
7051 snapshot,
7052 None,
7053 None,
7054 )?
7055 .into_iter()
7056 .map(|hit| hit.row_id.0)
7057 .collect(),
7058 ),
7059 Condition::MinHashSimilar {
7060 column_id,
7061 query,
7062 k,
7063 } => match self.minhash.get(column_id) {
7064 Some(index) => {
7065 let candidates = index.candidate_row_ids(query);
7066 let eligible =
7067 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
7068 RowIdSet::from_unsorted(
7069 index
7070 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
7071 .into_iter()
7072 .map(|(row_id, _)| row_id.0)
7073 .collect(),
7074 )
7075 }
7076 None => RowIdSet::empty(),
7077 },
7078 Condition::Range { column_id, lo, hi } => {
7079 if let Some(li) = self.learned_range.get(column_id) {
7080 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7081 } else {
7082 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
7083 }
7084 }
7085 Condition::RangeF64 {
7086 column_id,
7087 lo,
7088 lo_inclusive,
7089 hi,
7090 hi_inclusive,
7091 } => {
7092 if let Some(li) = self.learned_range.get(column_id) {
7093 RowIdSet::from_unsorted(
7094 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7095 .into_iter()
7096 .collect(),
7097 )
7098 } else {
7099 reader.range_row_id_set_f64(
7100 *column_id,
7101 *lo,
7102 *lo_inclusive,
7103 *hi,
7104 *hi_inclusive,
7105 )?
7106 }
7107 }
7108 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
7109 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
7110 };
7111 sets.push(s);
7112 }
7113 Ok(RowIdSet::intersect_many(sets))
7114 }
7115
7116 pub fn scan_cursor(
7137 &self,
7138 snapshot: Snapshot,
7139 projection: Vec<(u16, TypeId)>,
7140 conditions: &[crate::query::Condition],
7141 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
7142 if self.ttl.is_some() {
7143 return Ok(None);
7144 }
7145 if !conditions.is_empty() && !self.indexes_complete {
7151 return Ok(None);
7152 }
7153 if self.run_refs.len() == 1 {
7154 Ok(self
7155 .native_page_cursor(snapshot, projection, conditions)?
7156 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7157 } else {
7158 Ok(self
7159 .native_multi_run_cursor(snapshot, projection, conditions)?
7160 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7161 }
7162 }
7163
7164 pub fn aggregate_native(
7178 &self,
7179 snapshot: Snapshot,
7180 column: Option<u16>,
7181 conditions: &[crate::query::Condition],
7182 agg: NativeAgg,
7183 ) -> Result<Option<NativeAggResult>> {
7184 if self.ttl.is_some() {
7185 return Ok(None);
7186 }
7187 if self.run_refs.len() == 1 && conditions.is_empty() {
7189 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
7190 return Ok(Some(res));
7191 }
7192 }
7193 if matches!(agg, NativeAgg::Count) && column.is_none() {
7195 return Ok(self
7196 .scan_cursor(snapshot, Vec::new(), conditions)?
7197 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
7198 }
7199 let cid = match column {
7202 Some(c) => c,
7203 None => return Ok(None),
7204 };
7205 let ty = self.column_type(cid);
7206 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
7207 else {
7208 return Ok(None);
7209 };
7210 match ty {
7211 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
7212 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
7213 Ok(Some(pack_int(agg, count, sum, mn, mx)))
7214 }
7215 TypeId::Float64 => {
7216 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
7217 Ok(Some(pack_float(agg, count, sum, mn, mx)))
7218 }
7219 _ => Ok(None),
7220 }
7221 }
7222
7223 fn aggregate_from_stats(
7231 &self,
7232 snapshot: Snapshot,
7233 column: Option<u16>,
7234 agg: NativeAgg,
7235 ) -> Result<Option<NativeAggResult>> {
7236 let cid = match (agg, column) {
7237 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
7238 _ => return Ok(None), };
7240 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
7241 return Ok(None);
7242 };
7243 let Some(cs) = stats.get(&cid) else {
7244 return Ok(None);
7245 };
7246 match agg {
7247 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
7249 self.live_count.saturating_sub(cs.null_count),
7250 ))),
7251 NativeAgg::Min | NativeAgg::Max => {
7252 let bound = if agg == NativeAgg::Min {
7253 &cs.min
7254 } else {
7255 &cs.max
7256 };
7257 match bound {
7258 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
7259 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
7260 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
7265 None => Ok(None),
7266 }
7267 }
7268 _ => Ok(None),
7269 }
7270 }
7271
7272 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
7281 if self.ttl.is_some() {
7282 return Ok(None);
7283 }
7284 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
7285 return Ok(None);
7286 }
7287 self.ensure_indexes_complete()?;
7290 let reader = self.open_reader(self.run_refs[0].run_id)?;
7291 if self.live_count != reader.row_count() as u64 {
7292 return Ok(None);
7293 }
7294 let Some(bm) = self.bitmap.get(&column_id) else {
7295 return Ok(None); };
7297 let mut distinct = bm.value_count() as u64;
7298 if !bm.get(&Value::Null.encode_key()).is_empty() {
7301 distinct = distinct.saturating_sub(1);
7302 }
7303 Ok(Some(distinct))
7304 }
7305
7306 pub fn aggregate_incremental(
7318 &mut self,
7319 cache_key: u64,
7320 conditions: &[crate::query::Condition],
7321 column: Option<u16>,
7322 agg: NativeAgg,
7323 ) -> Result<IncrementalAggResult> {
7324 let snap = self.snapshot();
7325 let cur_wm = self.allocator.current().0;
7326 let cur_epoch = snap.epoch.0;
7327 let incremental_ok = self.ttl.is_none()
7334 && !self.had_deletes
7335 && self.memtable.is_empty()
7336 && self.mutable_run.is_empty();
7337
7338 if incremental_ok {
7341 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
7342 if cached.epoch == cur_epoch {
7343 return Ok(IncrementalAggResult {
7344 state: cached.state,
7345 incremental: true,
7346 delta_rows: 0,
7347 });
7348 }
7349 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
7350 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
7351 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
7352 let index_sets = self.resolve_index_conditions(conditions, snap)?;
7353 let delta_state = agg_state_from_rows(
7354 &delta_rows,
7355 conditions,
7356 &index_sets,
7357 column,
7358 agg,
7359 &self.schema,
7360 )?;
7361 let merged = cached.state.merge(delta_state);
7362 let delta_n = delta_rids.len() as u64;
7363 self.agg_cache.insert(
7364 cache_key,
7365 CachedAgg {
7366 state: merged.clone(),
7367 watermark: cur_wm,
7368 epoch: cur_epoch,
7369 },
7370 );
7371 return Ok(IncrementalAggResult {
7372 state: merged,
7373 incremental: true,
7374 delta_rows: delta_n,
7375 });
7376 }
7377 }
7378 }
7379
7380 let cursor_ok =
7385 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
7386 let state = if cursor_ok && agg != NativeAgg::Avg {
7387 match self.aggregate_native(snap, column, conditions, agg)? {
7388 Some(result) => {
7389 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
7390 }
7391 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
7392 }
7393 } else {
7394 self.agg_state_full_scan(conditions, column, agg, snap)?
7395 };
7396 if incremental_ok {
7398 self.agg_cache.insert(
7399 cache_key,
7400 CachedAgg {
7401 state: state.clone(),
7402 watermark: cur_wm,
7403 epoch: cur_epoch,
7404 },
7405 );
7406 }
7407 Ok(IncrementalAggResult {
7408 state,
7409 incremental: false,
7410 delta_rows: 0,
7411 })
7412 }
7413
7414 fn agg_state_full_scan(
7417 &self,
7418 conditions: &[crate::query::Condition],
7419 column: Option<u16>,
7420 agg: NativeAgg,
7421 snap: Snapshot,
7422 ) -> Result<AggState> {
7423 let rows = self.visible_rows(snap)?;
7424 let index_sets = self.resolve_index_conditions(conditions, snap)?;
7425 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
7426 }
7427
7428 fn resolve_index_conditions(
7431 &self,
7432 conditions: &[crate::query::Condition],
7433 snapshot: Snapshot,
7434 ) -> Result<Vec<RowIdSet>> {
7435 use crate::query::Condition;
7436 let mut sets = Vec::new();
7437 for c in conditions {
7438 if matches!(
7439 c,
7440 Condition::Ann { .. }
7441 | Condition::SparseMatch { .. }
7442 | Condition::MinHashSimilar { .. }
7443 ) {
7444 sets.push(self.resolve_condition(c, snapshot)?);
7445 }
7446 }
7447 Ok(sets)
7448 }
7449
7450 fn column_type(&self, cid: u16) -> TypeId {
7451 self.schema
7452 .columns
7453 .iter()
7454 .find(|c| c.id == cid)
7455 .map(|c| c.ty.clone())
7456 .unwrap_or(TypeId::Bytes)
7457 }
7458
7459 pub fn approx_aggregate(
7468 &mut self,
7469 conditions: &[crate::query::Condition],
7470 column: Option<u16>,
7471 agg: ApproxAgg,
7472 z: f64,
7473 ) -> Result<Option<ApproxResult>> {
7474 use crate::query::Condition;
7475 self.ensure_reservoir_complete()?;
7476 let snapshot = self.snapshot();
7477 let n_pop = self.count();
7478 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
7479 if sample_rids.is_empty() {
7480 return Ok(None);
7481 }
7482 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
7484 let s = live_sample.len();
7485 if s == 0 {
7486 return Ok(None);
7487 }
7488
7489 let mut index_sets: Vec<RowIdSet> = Vec::new();
7492 for c in conditions {
7493 if matches!(
7494 c,
7495 Condition::Ann { .. }
7496 | Condition::SparseMatch { .. }
7497 | Condition::MinHashSimilar { .. }
7498 ) {
7499 index_sets.push(self.resolve_condition(c, snapshot)?);
7500 }
7501 }
7502
7503 let cid = match (agg, column) {
7505 (ApproxAgg::Count, _) => None,
7506 (_, Some(c)) => Some(c),
7507 _ => return Ok(None),
7508 };
7509 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
7510 for r in &live_sample {
7511 if !conditions
7513 .iter()
7514 .all(|c| condition_matches_row(c, r, &self.schema))
7515 {
7516 continue;
7517 }
7518 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
7520 continue;
7521 }
7522 if let Some(cid) = cid {
7523 if let Some(v) = as_f64(r.columns.get(&cid)) {
7524 passing_vals.push(v);
7525 } } else {
7527 passing_vals.push(0.0); }
7529 }
7530 let m = passing_vals.len();
7531
7532 let (point, half) = match agg {
7533 ApproxAgg::Count => {
7534 let p = m as f64 / s as f64;
7536 let point = n_pop as f64 * p;
7537 let var = if s > 1 {
7538 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
7539 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7540 } else {
7541 0.0
7542 };
7543 (point, z * var.sqrt())
7544 }
7545 ApproxAgg::Sum => {
7546 let y: Vec<f64> = live_sample
7548 .iter()
7549 .map(|r| {
7550 let passes_row = conditions
7551 .iter()
7552 .all(|c| condition_matches_row(c, r, &self.schema))
7553 && index_sets.iter().all(|set| set.contains(r.row_id.0));
7554 if passes_row {
7555 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
7556 } else {
7557 0.0
7558 }
7559 })
7560 .collect();
7561 let mean_y = y.iter().sum::<f64>() / s as f64;
7562 let point = n_pop as f64 * mean_y;
7563 let var = if s > 1 {
7564 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
7565 let var_y = ss / (s - 1) as f64;
7566 n_pop as f64 * n_pop as f64 * var_y / s as f64
7567 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7568 } else {
7569 0.0
7570 };
7571 (point, z * var.sqrt())
7572 }
7573 ApproxAgg::Avg => {
7574 if m == 0 {
7575 return Ok(Some(ApproxResult {
7576 point: 0.0,
7577 ci_low: 0.0,
7578 ci_high: 0.0,
7579 n_population: n_pop,
7580 n_sample_live: s,
7581 n_passing: 0,
7582 }));
7583 }
7584 let mean = passing_vals.iter().sum::<f64>() / m as f64;
7585 let half = if m > 1 {
7586 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
7587 let sd = (ss / (m - 1) as f64).sqrt();
7588 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
7589 z * sd / (m as f64).sqrt() * fpc.sqrt()
7590 } else {
7591 0.0
7592 };
7593 (mean, half)
7594 }
7595 };
7596
7597 Ok(Some(ApproxResult {
7598 point,
7599 ci_low: point - half,
7600 ci_high: point + half,
7601 n_population: n_pop,
7602 n_sample_live: s,
7603 n_passing: m,
7604 }))
7605 }
7606
7607 pub fn exact_column_stats(
7615 &self,
7616 _snapshot: Snapshot,
7617 projection: &[u16],
7618 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
7619 if self.ttl.is_some()
7620 || !(self.memtable.is_empty()
7621 && self.mutable_run.is_empty()
7622 && self.run_refs.len() == 1)
7623 {
7624 return Ok(None);
7625 }
7626 let reader = self.open_reader(self.run_refs[0].run_id)?;
7627 if self.live_count != reader.row_count() as u64 {
7628 return Ok(None);
7629 }
7630 let mut out = HashMap::new();
7631 for &cid in projection {
7632 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
7633 Some(c) => c,
7634 None => continue,
7635 };
7636 let Some(stats) = reader.column_page_stats(cid) else {
7638 out.insert(
7639 cid,
7640 ColumnStat {
7641 min: None,
7642 max: None,
7643 null_count: self.live_count,
7644 },
7645 );
7646 continue;
7647 };
7648 let stat = match cdef.ty {
7649 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
7650 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
7651 min: mn.map(Value::Int64),
7652 max: mx.map(Value::Int64),
7653 null_count: n,
7654 })
7655 }
7656 TypeId::Float64 => {
7657 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
7658 min: mn.map(Value::Float64),
7659 max: mx.map(Value::Float64),
7660 null_count: n,
7661 })
7662 }
7663 _ => None,
7664 };
7665 if let Some(s) = stat {
7666 out.insert(cid, s);
7667 }
7668 }
7669 Ok(Some(out))
7670 }
7671
7672 pub fn dir(&self) -> &Path {
7673 &self.dir
7674 }
7675
7676 pub fn schema(&self) -> &Schema {
7677 &self.schema
7678 }
7679
7680 pub(crate) fn set_catalog_name(&mut self, name: String) {
7681 self.name = name;
7682 }
7683
7684 pub(crate) fn prepare_alter_column(
7685 &mut self,
7686 column_name: &str,
7687 change: &AlterColumn,
7688 ) -> Result<ColumnDef> {
7689 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
7690 return Err(MongrelError::InvalidArgument(
7691 "ALTER COLUMN requires committing staged writes first".into(),
7692 ));
7693 }
7694 let old = self
7695 .schema
7696 .columns
7697 .iter()
7698 .find(|c| c.name == column_name)
7699 .cloned()
7700 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7701 let mut next = old.clone();
7702
7703 if let Some(name) = &change.name {
7704 let trimmed = name.trim();
7705 if trimmed.is_empty() {
7706 return Err(MongrelError::InvalidArgument(
7707 "ALTER COLUMN name must not be empty".into(),
7708 ));
7709 }
7710 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
7711 return Err(MongrelError::Schema(format!(
7712 "column {trimmed} already exists"
7713 )));
7714 }
7715 next.name = trimmed.to_string();
7716 }
7717
7718 if let Some(ty) = &change.ty {
7719 next.ty = ty.clone();
7720 }
7721 if let Some(flags) = change.flags {
7722 validate_alter_column_flags(old.flags, flags)?;
7723 next.flags = flags;
7724 }
7725
7726 if let Some(default_change) = &change.default_value {
7727 next.default_value = default_change.clone();
7728 }
7729
7730 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
7731 if old.flags.contains(ColumnFlags::NULLABLE)
7732 && !next.flags.contains(ColumnFlags::NULLABLE)
7733 && self.column_has_nulls(old.id)?
7734 {
7735 return Err(MongrelError::InvalidArgument(format!(
7736 "column '{}' contains NULL values",
7737 old.name
7738 )));
7739 }
7740 Ok(next)
7741 }
7742
7743 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
7744 let idx = self
7745 .schema
7746 .columns
7747 .iter()
7748 .position(|c| c.id == column.id)
7749 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
7750 if self.schema.columns[idx] == column {
7751 return Ok(());
7752 }
7753 self.schema.columns[idx] = column;
7754 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7755 self.schema.validate_auto_increment()?;
7756 self.schema.validate_defaults()?;
7757 self.auto_inc = resolve_auto_inc(&self.schema);
7758 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
7759 write_schema(&self.dir, &self.schema)?;
7760 self.clear_result_cache();
7761 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7762 self.persist_manifest(self.current_epoch())?;
7763 Ok(())
7764 }
7765
7766 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
7767 self.ensure_writable()?;
7768 let column = self.prepare_alter_column(column_name, &change)?;
7769 self.apply_altered_column(column.clone())?;
7770 Ok(column)
7771 }
7772
7773 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
7774 if self.live_count == 0 {
7775 return Ok(false);
7776 }
7777 let snap = self.snapshot();
7778 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
7779 Ok(columns
7780 .first()
7781 .map(|(_, col)| col.null_count(col.len()) != 0)
7782 .unwrap_or(true))
7783 }
7784
7785 fn has_stored_versions(&self) -> bool {
7786 !self.memtable.is_empty()
7787 || !self.mutable_run.is_empty()
7788 || self.run_refs.iter().any(|r| r.row_count > 0)
7789 || !self.retiring.is_empty()
7790 }
7791
7792 pub fn add_column(
7797 &mut self,
7798 name: &str,
7799 ty: TypeId,
7800 flags: ColumnFlags,
7801 default_value: Option<crate::schema::DefaultExpr>,
7802 ) -> Result<u16> {
7803 self.add_column_with_id(name, ty, flags, default_value, None)
7804 }
7805
7806 pub fn add_column_with_id(
7807 &mut self,
7808 name: &str,
7809 ty: TypeId,
7810 flags: ColumnFlags,
7811 default_value: Option<crate::schema::DefaultExpr>,
7812 requested_id: Option<u16>,
7813 ) -> Result<u16> {
7814 self.ensure_writable()?;
7815 if self.schema.columns.iter().any(|c| c.name == name) {
7816 return Err(MongrelError::Schema(format!(
7817 "column {name} already exists"
7818 )));
7819 }
7820 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
7821 if self.schema.columns.iter().any(|c| c.id == id) {
7822 return Err(MongrelError::Schema(format!(
7823 "column id {id} already exists"
7824 )));
7825 }
7826 id
7827 } else {
7828 self.schema
7829 .columns
7830 .iter()
7831 .map(|c| c.id)
7832 .max()
7833 .unwrap_or(0)
7834 .checked_add(1)
7835 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
7836 };
7837 self.schema.columns.push(ColumnDef {
7838 id,
7839 name: name.to_string(),
7840 ty,
7841 flags,
7842 default_value,
7843 });
7844 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7845 self.schema.validate_auto_increment()?;
7846 self.schema.validate_defaults()?;
7847 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
7848 self.auto_inc = resolve_auto_inc(&self.schema);
7849 }
7850 write_schema(&self.dir, &self.schema)?;
7851 self.clear_result_cache();
7852 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7854 self.persist_manifest(self.current_epoch())?;
7855 Ok(id)
7856 }
7857
7858 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
7867 self.ensure_writable()?;
7868 let cid = self
7869 .schema
7870 .columns
7871 .iter()
7872 .find(|c| c.name == column_name)
7873 .map(|c| c.id)
7874 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7875 let ty = self
7876 .schema
7877 .columns
7878 .iter()
7879 .find(|c| c.id == cid)
7880 .map(|c| c.ty.clone())
7881 .unwrap_or(TypeId::Int64);
7882 if !matches!(
7883 ty,
7884 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
7885 ) {
7886 return Err(MongrelError::Schema(format!(
7887 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
7888 )));
7889 }
7890 if self
7891 .schema
7892 .indexes
7893 .iter()
7894 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
7895 {
7896 return Ok(()); }
7898 self.schema.indexes.push(IndexDef {
7899 name: format!("{}_learned_range", column_name),
7900 column_id: cid,
7901 kind: IndexKind::LearnedRange,
7902 predicate: None,
7903 options: Default::default(),
7904 });
7905 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7906 write_schema(&self.dir, &self.schema)?;
7907 self.build_learned_ranges()?;
7908 Ok(())
7909 }
7910
7911 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
7914 self.sync_byte_threshold = threshold;
7915 if let WalSink::Private(w) = &mut self.wal {
7916 w.set_sync_byte_threshold(threshold);
7917 }
7918 }
7919
7920 pub fn page_cache_flush(&self) {
7924 self.page_cache.flush_to_disk();
7925 }
7926
7927 pub fn page_cache_len(&self) -> usize {
7929 self.page_cache.len()
7930 }
7931
7932 pub fn decoded_cache_len(&self) -> usize {
7935 self.decoded_cache.len()
7936 }
7937
7938 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
7941 self.memtable.drain_sorted()
7942 }
7943
7944 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
7945 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
7946 }
7947
7948 pub(crate) fn table_dir(&self) -> &Path {
7949 &self.dir
7950 }
7951
7952 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
7953 &self.schema
7954 }
7955
7956 pub(crate) fn alloc_run_id(&mut self) -> u64 {
7957 let id = self.next_run_id;
7958 self.next_run_id += 1;
7959 id
7960 }
7961
7962 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
7963 self.run_refs.push(run_ref);
7964 }
7965
7966 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
7976 self.retiring.push(crate::manifest::RetiredRun {
7977 run_id,
7978 retire_epoch,
7979 });
7980 }
7981
7982 pub(crate) fn reap_retiring(
7986 &mut self,
7987 min_active: Epoch,
7988 backup_pinned: &std::collections::HashSet<u128>,
7989 ) -> Result<usize> {
7990 if self.retiring.is_empty() {
7991 return Ok(0);
7992 }
7993 let mut reaped = 0;
7994 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
7995 for r in std::mem::take(&mut self.retiring) {
8001 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
8002 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
8003 reaped += 1;
8004 } else {
8005 kept.push(r);
8006 }
8007 }
8008 self.retiring = kept;
8009 if reaped > 0 {
8010 self.persist_manifest(self.current_epoch())?;
8011 }
8012 Ok(reaped)
8013 }
8014
8015 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
8016 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
8017 return false;
8018 }
8019 self.live_count = self.live_count.saturating_add(run_ref.row_count);
8020 self.run_refs.push(run_ref);
8021 self.indexes_complete = false;
8022 true
8023 }
8024
8025 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
8026 self.kek.as_ref()
8027 }
8028
8029 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
8030 let mut reader = RunReader::open_with_cache(
8031 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
8032 self.schema.clone(),
8033 self.kek.clone(),
8034 Some(self.page_cache.clone()),
8035 Some(self.decoded_cache.clone()),
8036 self.table_id,
8037 Some(&self.verified_runs),
8038 )?;
8039 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
8043 reader.set_uniform_epoch(Epoch(rr.epoch_created));
8044 }
8045 Ok(reader)
8046 }
8047
8048 pub(crate) fn run_refs(&self) -> &[RunRef] {
8049 &self.run_refs
8050 }
8051
8052 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
8053 self.retiring.iter().map(|run| run.run_id)
8054 }
8055
8056 pub(crate) fn runs_dir(&self) -> PathBuf {
8057 self.dir.join(RUNS_DIR)
8058 }
8059
8060 pub(crate) fn wal_dir(&self) -> PathBuf {
8061 self.dir.join(WAL_DIR)
8062 }
8063
8064 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
8065 self.run_refs = refs;
8066 }
8067
8068 pub(crate) fn next_run_id(&self) -> u64 {
8069 self.next_run_id
8070 }
8071
8072 pub(crate) fn compaction_zstd_level(&self) -> i32 {
8073 self.compaction_zstd_level
8074 }
8075
8076 pub(crate) fn bump_next_run_id(&mut self) {
8077 self.next_run_id += 1;
8078 }
8079
8080 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
8081 self.kek.clone()
8082 }
8083
8084 #[cfg(feature = "encryption")]
8088 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8089 self.kek.as_ref().map(|k| k.derive_idx_key())
8090 }
8091
8092 #[cfg(not(feature = "encryption"))]
8093 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8094 None
8095 }
8096
8097 #[cfg(feature = "encryption")]
8101 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8102 self.kek.as_ref().map(|k| *k.derive_meta_key())
8103 }
8104
8105 #[cfg(not(feature = "encryption"))]
8106 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8107 None
8108 }
8109
8110 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
8113 self.column_keys
8114 .iter()
8115 .map(|(&id, &(_, scheme))| (id, scheme))
8116 .collect()
8117 }
8118
8119 #[cfg(feature = "encryption")]
8124 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
8125 self.tokenize_value_enc(column_id, v)
8126 }
8127
8128 #[cfg(feature = "encryption")]
8129 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
8130 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8131 let (key, scheme) = self.column_keys.get(&column_id)?;
8132 let token: Vec<u8> = match (*scheme, v) {
8133 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
8134 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
8135 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
8136 _ => hmac_token(key, &v.encode_key()).to_vec(),
8137 };
8138 Some(Value::Bytes(token))
8139 }
8140
8141 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
8143 self.index_lookup_key_bytes(column_id, &v.encode_key())
8144 }
8145
8146 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
8149 #[cfg(feature = "encryption")]
8150 {
8151 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
8152 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
8153 if *scheme == SCHEME_HMAC_EQ {
8154 return hmac_token(key, encoded).to_vec();
8155 }
8156 }
8157 }
8158 let _ = column_id;
8159 encoded.to_vec()
8160 }
8161}
8162
8163fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
8164 let columnar::NativeColumn::Int64 { data, validity } = col else {
8165 return false;
8166 };
8167 if data.len() < n || !columnar::all_non_null(validity, n) {
8168 return false;
8169 }
8170 data.iter()
8171 .take(n)
8172 .zip(data.iter().skip(1))
8173 .all(|(a, b)| a < b)
8174}
8175
8176#[derive(Debug, Clone)]
8180pub struct ColumnStat {
8181 pub min: Option<Value>,
8182 pub max: Option<Value>,
8183 pub null_count: u64,
8184}
8185
8186#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8188pub enum NativeAgg {
8189 Count,
8190 Sum,
8191 Min,
8192 Max,
8193 Avg,
8194}
8195
8196#[derive(Debug, Clone, PartialEq)]
8198pub enum NativeAggResult {
8199 Count(u64),
8200 Int(i64),
8201 Float(f64),
8202 Null,
8204}
8205
8206#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8208pub enum ApproxAgg {
8209 Count,
8210 Sum,
8211 Avg,
8212}
8213
8214#[derive(Debug, Clone)]
8218pub struct ApproxResult {
8219 pub point: f64,
8221 pub ci_low: f64,
8223 pub ci_high: f64,
8225 pub n_population: u64,
8227 pub n_sample_live: usize,
8229 pub n_passing: usize,
8231}
8232
8233#[derive(Debug, Clone, PartialEq)]
8238pub enum AggState {
8239 Count(u64),
8241 SumI {
8243 sum: i128,
8244 count: u64,
8245 },
8246 SumF {
8248 sum: f64,
8249 count: u64,
8250 },
8251 AvgI {
8253 sum: i128,
8254 count: u64,
8255 },
8256 AvgF {
8258 sum: f64,
8259 count: u64,
8260 },
8261 MinI(i64),
8263 MaxI(i64),
8264 MinF(f64),
8266 MaxF(f64),
8267 Empty,
8269}
8270
8271impl AggState {
8272 pub fn merge(self, other: AggState) -> AggState {
8274 use AggState::*;
8275 match (self, other) {
8276 (Empty, x) | (x, Empty) => x,
8277 (Count(a), Count(b)) => Count(a + b),
8278 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
8279 sum: sa + sb,
8280 count: ca + cb,
8281 },
8282 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
8283 sum: sa + sb,
8284 count: ca + cb,
8285 },
8286 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
8287 sum: sa + sb,
8288 count: ca + cb,
8289 },
8290 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
8291 sum: sa + sb,
8292 count: ca + cb,
8293 },
8294 (MinI(a), MinI(b)) => MinI(a.min(b)),
8295 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
8296 (MinF(a), MinF(b)) => MinF(a.min(b)),
8297 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
8298 _ => Empty, }
8300 }
8301
8302 pub fn point(&self) -> Option<f64> {
8304 match self {
8305 AggState::Count(n) => Some(*n as f64),
8306 AggState::SumI { sum, .. } => Some(*sum as f64),
8307 AggState::SumF { sum, .. } => Some(*sum),
8308 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
8309 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
8310 AggState::MinI(n) => Some(*n as f64),
8311 AggState::MaxI(n) => Some(*n as f64),
8312 AggState::MinF(n) => Some(*n),
8313 AggState::MaxF(n) => Some(*n),
8314 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
8315 }
8316 }
8317
8318 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
8322 let is_float = matches!(ty, Some(TypeId::Float64));
8323 match (agg, result) {
8324 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
8325 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
8326 sum: x as i128,
8327 count: 1, },
8329 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
8330 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
8331 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
8332 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
8333 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
8334 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
8335 (NativeAgg::Count, _) => AggState::Empty,
8336 (_, NativeAggResult::Null) => AggState::Empty,
8337 _ => {
8338 let _ = is_float;
8339 AggState::Empty
8340 }
8341 }
8342 }
8343}
8344
8345#[derive(Debug, Clone)]
8348pub struct CachedAgg {
8349 pub state: AggState,
8350 pub watermark: u64,
8351 pub epoch: u64,
8352}
8353
8354#[derive(Debug, Clone)]
8356pub struct IncrementalAggResult {
8357 pub state: AggState,
8359 pub incremental: bool,
8362 pub delta_rows: u64,
8364}
8365
8366fn agg_state_from_rows(
8370 rows: &[Row],
8371 conditions: &[crate::query::Condition],
8372 index_sets: &[RowIdSet],
8373 column: Option<u16>,
8374 agg: NativeAgg,
8375 schema: &Schema,
8376) -> Result<AggState> {
8377 let mut count: u64 = 0;
8378 let mut sum_i: i128 = 0;
8379 let mut sum_f: f64 = 0.0;
8380 let mut mn_i: i64 = i64::MAX;
8381 let mut mx_i: i64 = i64::MIN;
8382 let mut mn_f: f64 = f64::INFINITY;
8383 let mut mx_f: f64 = f64::NEG_INFINITY;
8384 let mut saw_int = false;
8385 let mut saw_float = false;
8386 for r in rows {
8387 if !conditions
8388 .iter()
8389 .all(|c| condition_matches_row(c, r, schema))
8390 {
8391 continue;
8392 }
8393 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
8394 continue;
8395 }
8396 match agg {
8397 NativeAgg::Count => match column {
8398 None => count += 1,
8400 Some(cid) => match r.columns.get(&cid) {
8403 None | Some(Value::Null) => {}
8404 Some(_) => count += 1,
8405 },
8406 },
8407 _ => match column.and_then(|cid| r.columns.get(&cid)) {
8408 Some(Value::Int64(n)) => {
8409 count += 1;
8410 sum_i += *n as i128;
8411 mn_i = mn_i.min(*n);
8412 mx_i = mx_i.max(*n);
8413 saw_int = true;
8414 }
8415 Some(Value::Float64(f)) => {
8416 count += 1;
8417 sum_f += f;
8418 mn_f = mn_f.min(*f);
8419 mx_f = mx_f.max(*f);
8420 saw_float = true;
8421 }
8422 _ => {}
8423 },
8424 }
8425 }
8426 Ok(match agg {
8427 NativeAgg::Count => {
8428 if count == 0 {
8429 AggState::Empty
8430 } else {
8431 AggState::Count(count)
8432 }
8433 }
8434 NativeAgg::Sum => {
8435 if count == 0 {
8436 AggState::Empty
8437 } else if saw_int {
8438 AggState::SumI { sum: sum_i, count }
8439 } else {
8440 AggState::SumF { sum: sum_f, count }
8441 }
8442 }
8443 NativeAgg::Avg => {
8444 if count == 0 {
8445 AggState::Empty
8446 } else if saw_int {
8447 AggState::AvgI { sum: sum_i, count }
8448 } else {
8449 AggState::AvgF { sum: sum_f, count }
8450 }
8451 }
8452 NativeAgg::Min => {
8453 if !saw_int && !saw_float {
8454 AggState::Empty
8455 } else if saw_int {
8456 AggState::MinI(mn_i)
8457 } else {
8458 AggState::MinF(mn_f)
8459 }
8460 }
8461 NativeAgg::Max => {
8462 if !saw_int && !saw_float {
8463 AggState::Empty
8464 } else if saw_int {
8465 AggState::MaxI(mx_i)
8466 } else {
8467 AggState::MaxF(mx_f)
8468 }
8469 }
8470 })
8471}
8472
8473fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
8477 use crate::query::Condition;
8478 match c {
8479 Condition::Pk(key) => match schema.primary_key() {
8480 Some(pk) => row
8481 .columns
8482 .get(&pk.id)
8483 .map(|v| v.encode_key() == *key)
8484 .unwrap_or(false),
8485 None => false,
8486 },
8487 Condition::BitmapEq { column_id, value } => row
8488 .columns
8489 .get(column_id)
8490 .map(|v| v.encode_key() == *value)
8491 .unwrap_or(false),
8492 Condition::BitmapIn { column_id, values } => {
8493 let key = row.columns.get(column_id).map(|v| v.encode_key());
8494 match key {
8495 Some(k) => values.contains(&k),
8496 None => false,
8497 }
8498 }
8499 Condition::BytesPrefix { column_id, prefix } => row
8500 .columns
8501 .get(column_id)
8502 .map(|v| v.encode_key().starts_with(prefix))
8503 .unwrap_or(false),
8504 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
8505 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
8506 _ => false,
8507 },
8508 Condition::RangeF64 {
8509 column_id,
8510 lo,
8511 lo_inclusive,
8512 hi,
8513 hi_inclusive,
8514 } => match row.columns.get(column_id) {
8515 Some(Value::Float64(n)) => {
8516 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
8517 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
8518 lo_ok && hi_ok
8519 }
8520 _ => false,
8521 },
8522 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
8523 Some(Value::Bytes(b)) => {
8524 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
8525 }
8526 _ => false,
8527 },
8528 Condition::FmContainsAll {
8529 column_id,
8530 patterns,
8531 } => match row.columns.get(column_id) {
8532 Some(Value::Bytes(b)) => patterns
8533 .iter()
8534 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
8535 _ => false,
8536 },
8537 Condition::Ann { .. }
8538 | Condition::SparseMatch { .. }
8539 | Condition::MinHashSimilar { .. } => true,
8540 Condition::IsNull { column_id } => {
8541 matches!(row.columns.get(column_id), Some(Value::Null) | None)
8542 }
8543 Condition::IsNotNull { column_id } => {
8544 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
8545 }
8546 }
8547}
8548
8549fn as_f64(v: Option<&Value>) -> Option<f64> {
8551 match v {
8552 Some(Value::Int64(n)) => Some(*n as f64),
8553 Some(Value::Float64(f)) => Some(*f),
8554 _ => None,
8555 }
8556}
8557
8558fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
8562 let mut count: u64 = 0;
8563 let mut sum: i128 = 0;
8564 let mut mn: i64 = i64::MAX;
8565 let mut mx: i64 = i64::MIN;
8566 while let Some(cols) = cursor.next_batch()? {
8567 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
8568 if crate::columnar::all_non_null(validity, data.len()) {
8569 count += data.len() as u64;
8571 sum += data.iter().map(|&v| v as i128).sum::<i128>();
8572 mn = mn.min(*data.iter().min().unwrap_or(&mn));
8573 mx = mx.max(*data.iter().max().unwrap_or(&mx));
8574 } else {
8575 for (i, &v) in data.iter().enumerate() {
8576 if crate::columnar::validity_bit(validity, i) {
8577 count += 1;
8578 sum += v as i128;
8579 mn = mn.min(v);
8580 mx = mx.max(v);
8581 }
8582 }
8583 }
8584 }
8585 }
8586 Ok((count, sum, mn, mx))
8587}
8588
8589fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
8591 let mut count: u64 = 0;
8592 let mut sum: f64 = 0.0;
8593 let mut mn: f64 = f64::INFINITY;
8594 let mut mx: f64 = f64::NEG_INFINITY;
8595 while let Some(cols) = cursor.next_batch()? {
8596 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
8597 if crate::columnar::all_non_null(validity, data.len()) {
8598 count += data.len() as u64;
8599 sum += data.iter().sum::<f64>();
8600 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
8601 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
8602 } else {
8603 for (i, &v) in data.iter().enumerate() {
8604 if crate::columnar::validity_bit(validity, i) {
8605 count += 1;
8606 sum += v;
8607 mn = mn.min(v);
8608 mx = mx.max(v);
8609 }
8610 }
8611 }
8612 }
8613 }
8614 Ok((count, sum, mn, mx))
8615}
8616
8617fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
8618 if count == 0 && !matches!(agg, NativeAgg::Count) {
8619 return NativeAggResult::Null;
8620 }
8621 match agg {
8622 NativeAgg::Count => NativeAggResult::Count(count),
8623 NativeAgg::Sum => match sum.try_into() {
8626 Ok(v) => NativeAggResult::Int(v),
8627 Err(_) => NativeAggResult::Null,
8628 },
8629 NativeAgg::Min => NativeAggResult::Int(mn),
8630 NativeAgg::Max => NativeAggResult::Int(mx),
8631 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
8632 }
8633}
8634
8635fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
8636 if count == 0 && !matches!(agg, NativeAgg::Count) {
8637 return NativeAggResult::Null;
8638 }
8639 match agg {
8640 NativeAgg::Count => NativeAggResult::Count(count),
8641 NativeAgg::Sum => NativeAggResult::Float(sum),
8642 NativeAgg::Min => NativeAggResult::Float(mn),
8643 NativeAgg::Max => NativeAggResult::Float(mx),
8644 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
8645 }
8646}
8647
8648fn agg_int(
8651 stats: &[crate::page::PageStat],
8652 decode: fn(Option<&[u8]>) -> Option<i64>,
8653) -> Option<(Option<i64>, Option<i64>, u64)> {
8654 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
8655 let mut any = false;
8656 for s in stats {
8657 if let Some(v) = decode(s.min.as_deref()) {
8658 mn = mn.min(v);
8659 any = true;
8660 }
8661 if let Some(v) = decode(s.max.as_deref()) {
8662 mx = mx.max(v);
8663 any = true;
8664 }
8665 nulls += s.null_count;
8666 }
8667 any.then_some((Some(mn), Some(mx), nulls))
8668}
8669
8670fn agg_float(
8672 stats: &[crate::page::PageStat],
8673 decode: fn(Option<&[u8]>) -> Option<f64>,
8674) -> Option<(Option<f64>, Option<f64>, u64)> {
8675 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
8676 let mut any = false;
8677 for s in stats {
8678 if let Some(v) = decode(s.min.as_deref()) {
8679 mn = mn.min(v);
8680 any = true;
8681 }
8682 if let Some(v) = decode(s.max.as_deref()) {
8683 mx = mx.max(v);
8684 any = true;
8685 }
8686 nulls += s.null_count;
8687 }
8688 any.then_some((Some(mn), Some(mx), nulls))
8689}
8690
8691type SecondaryIndexes = (
8693 HashMap<u16, BitmapIndex>,
8694 HashMap<u16, AnnIndex>,
8695 HashMap<u16, FmIndex>,
8696 HashMap<u16, SparseIndex>,
8697 HashMap<u16, MinHashIndex>,
8698);
8699
8700fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
8701 let mut bitmap = HashMap::new();
8702 let mut ann = HashMap::new();
8703 let mut fm = HashMap::new();
8704 let mut sparse = HashMap::new();
8705 let mut minhash = HashMap::new();
8706 for idef in &schema.indexes {
8707 match idef.kind {
8708 IndexKind::Bitmap => {
8709 bitmap.insert(idef.column_id, BitmapIndex::new());
8710 }
8711 IndexKind::Ann => {
8712 let dim = schema
8713 .columns
8714 .iter()
8715 .find(|c| c.id == idef.column_id)
8716 .and_then(|c| match c.ty {
8717 TypeId::Embedding { dim } => Some(dim as usize),
8718 _ => None,
8719 })
8720 .unwrap_or(0);
8721 let options = idef.options.ann.clone().unwrap_or_default();
8722 ann.insert(
8723 idef.column_id,
8724 AnnIndex::with_options(
8725 dim,
8726 options.m,
8727 options.ef_construction,
8728 options.ef_search,
8729 ),
8730 );
8731 }
8732 IndexKind::FmIndex => {
8733 fm.insert(idef.column_id, FmIndex::new());
8734 }
8735 IndexKind::Sparse => {
8736 sparse.insert(idef.column_id, SparseIndex::new());
8737 }
8738 IndexKind::MinHash => {
8739 let options = idef.options.minhash.clone().unwrap_or_default();
8740 minhash.insert(
8741 idef.column_id,
8742 MinHashIndex::with_options(options.permutations, options.bands),
8743 );
8744 }
8745 _ => {}
8746 }
8747 }
8748 (bitmap, ann, fm, sparse, minhash)
8749}
8750
8751const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
8752 | ColumnFlags::AUTO_INCREMENT
8753 | ColumnFlags::ENCRYPTED
8754 | ColumnFlags::ENCRYPTED_INDEXABLE
8755 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
8756
8757fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
8758 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
8759 return Err(MongrelError::Schema(
8760 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
8761 ));
8762 }
8763 Ok(())
8764}
8765
8766fn validate_alter_column_type(
8767 schema: &Schema,
8768 old: &ColumnDef,
8769 next: &ColumnDef,
8770 has_stored_versions: bool,
8771) -> Result<()> {
8772 if old.ty == next.ty {
8773 return Ok(());
8774 }
8775 if schema.indexes.iter().any(|i| i.column_id == old.id) {
8776 return Err(MongrelError::Schema(format!(
8777 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
8778 old.name
8779 )));
8780 }
8781 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
8782 return Ok(());
8783 }
8784 Err(MongrelError::Schema(format!(
8785 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
8786 old.ty, next.ty
8787 )))
8788}
8789
8790fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
8791 matches!(
8792 (old, new),
8793 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
8794 )
8795}
8796
8797fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
8803 let mut prev: Option<i64> = None;
8804 for r in rows {
8805 match r.columns.get(&pk_id) {
8806 Some(Value::Int64(v)) => {
8807 if prev.is_some_and(|p| p >= *v) {
8808 return false;
8809 }
8810 prev = Some(*v);
8811 }
8812 _ => return false,
8813 }
8814 }
8815 true
8816}
8817
8818#[allow(clippy::too_many_arguments)]
8819fn index_into(
8820 schema: &Schema,
8821 row: &Row,
8822 hot: &mut HotIndex,
8823 bitmap: &mut HashMap<u16, BitmapIndex>,
8824 ann: &mut HashMap<u16, AnnIndex>,
8825 fm: &mut HashMap<u16, FmIndex>,
8826 sparse: &mut HashMap<u16, SparseIndex>,
8827 minhash: &mut HashMap<u16, MinHashIndex>,
8828) {
8829 for idef in &schema.indexes {
8830 let Some(val) = row.columns.get(&idef.column_id) else {
8831 continue;
8832 };
8833 match idef.kind {
8834 IndexKind::Bitmap => {
8835 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8836 b.insert(val.encode_key(), row.row_id);
8837 }
8838 }
8839 IndexKind::Ann => {
8840 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8841 a.insert_validated(v, row.row_id);
8842 }
8843 }
8844 IndexKind::FmIndex => {
8845 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8846 f.insert(b.clone(), row.row_id);
8847 }
8848 }
8849 IndexKind::Sparse => {
8850 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8851 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8854 s.insert(&terms, row.row_id);
8855 }
8856 }
8857 }
8858 IndexKind::MinHash => {
8859 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8860 let tokens = crate::index::token_hashes_from_bytes(b);
8863 mh.insert(&tokens, row.row_id);
8864 }
8865 }
8866 _ => {}
8867 }
8868 }
8869 if let Some(pk_col) = schema.primary_key() {
8870 if let Some(pk_val) = row.columns.get(&pk_col.id) {
8871 hot.insert(pk_val.encode_key(), row.row_id);
8872 }
8873 }
8874}
8875
8876#[allow(clippy::too_many_arguments)]
8879fn index_into_single(
8880 idef: &IndexDef,
8881 _schema: &Schema,
8882 row: &Row,
8883 _hot: &mut HotIndex,
8884 bitmap: &mut HashMap<u16, BitmapIndex>,
8885 ann: &mut HashMap<u16, AnnIndex>,
8886 fm: &mut HashMap<u16, FmIndex>,
8887 sparse: &mut HashMap<u16, SparseIndex>,
8888 minhash: &mut HashMap<u16, MinHashIndex>,
8889) {
8890 let Some(val) = row.columns.get(&idef.column_id) else {
8891 return;
8892 };
8893 match idef.kind {
8894 IndexKind::Bitmap => {
8895 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8896 b.insert(val.encode_key(), row.row_id);
8897 }
8898 }
8899 IndexKind::Ann => {
8900 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8901 a.insert_validated(v, row.row_id);
8902 }
8903 }
8904 IndexKind::FmIndex => {
8905 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8906 f.insert(b.clone(), row.row_id);
8907 }
8908 }
8909 IndexKind::Sparse => {
8910 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8911 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8912 s.insert(&terms, row.row_id);
8913 }
8914 }
8915 }
8916 IndexKind::MinHash => {
8917 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8918 let tokens = crate::index::token_hashes_from_bytes(b);
8919 mh.insert(&tokens, row.row_id);
8920 }
8921 }
8922 _ => {}
8923 }
8924}
8925
8926fn eval_partial_predicate(
8932 pred: &str,
8933 columns_map: &HashMap<u16, &Value>,
8934 name_to_id: &HashMap<&str, u16>,
8935) -> bool {
8936 let lower = pred.trim().to_ascii_lowercase();
8937 if let Some(rest) = lower.strip_suffix(" is not null") {
8939 let col_name = rest.trim();
8940 if let Some(col_id) = name_to_id.get(col_name) {
8941 return columns_map
8942 .get(col_id)
8943 .is_some_and(|v| !matches!(v, Value::Null));
8944 }
8945 }
8946 if let Some(rest) = lower.strip_suffix(" is null") {
8948 let col_name = rest.trim();
8949 if let Some(col_id) = name_to_id.get(col_name) {
8950 return columns_map
8951 .get(col_id)
8952 .map_or(true, |v| matches!(v, Value::Null));
8953 }
8954 }
8955 true
8958}
8959
8960#[allow(dead_code)]
8966fn bulk_index_key(
8967 column_keys: &HashMap<u16, ([u8; 32], u8)>,
8968 column_id: u16,
8969 ty: TypeId,
8970 col: &columnar::NativeColumn,
8971 i: usize,
8972) -> Option<Vec<u8>> {
8973 let encoded = columnar::encode_key_native(ty, col, i)?;
8974 #[cfg(feature = "encryption")]
8975 {
8976 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8977 if let Some((key, scheme)) = column_keys.get(&column_id) {
8978 return Some(match (*scheme, col) {
8979 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
8980 (_, columnar::NativeColumn::Int64 { data, .. }) => {
8981 ope_token_i64(key, data[i]).to_vec()
8982 }
8983 (_, columnar::NativeColumn::Float64 { data, .. }) => {
8984 ope_token_f64(key, data[i]).to_vec()
8985 }
8986 _ => hmac_token(key, &encoded).to_vec(),
8987 });
8988 }
8989 }
8990 #[cfg(not(feature = "encryption"))]
8991 {
8992 let _ = (column_id, column_keys, col);
8993 }
8994 Some(encoded)
8995}
8996
8997pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
8998 let json = serde_json::to_string_pretty(schema)
8999 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
9000 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
9001 Ok(())
9002}
9003
9004fn read_schema(dir: &Path) -> Result<Schema> {
9005 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
9006 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
9007}
9008
9009fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
9010 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
9011}
9012
9013fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
9014 let n = list_wal_numbers(wal_dir)?;
9015 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
9016}
9017
9018fn next_wal_number(wal_dir: &Path) -> Result<u32> {
9019 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
9020}
9021
9022fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
9023 let _ = std::fs::create_dir_all(wal_dir);
9024 let mut max_n = None;
9025 for entry in std::fs::read_dir(wal_dir)? {
9026 let entry = entry?;
9027 let fname = entry.file_name();
9028 let Some(s) = fname.to_str() else {
9029 continue;
9030 };
9031 let Some(stripped) = s.strip_prefix("seg-") else {
9032 continue;
9033 };
9034 let Some(stripped) = stripped.strip_suffix(".wal") else {
9035 continue;
9036 };
9037 if let Ok(n) = stripped.parse::<u32>() {
9038 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
9039 }
9040 }
9041 Ok(max_n)
9042}