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)
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)
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 ) -> Result<Vec<Row>> {
3542 crate::trace::QueryTrace::record(|t| {
3543 t.run_count = self.run_refs.len();
3544 t.memtable_rows = self.memtable.len();
3545 t.mutable_run_rows = self.mutable_run.len();
3546 });
3547 if conditions.is_empty() {
3551 crate::trace::QueryTrace::record(|t| {
3552 t.scan_mode = crate::trace::ScanMode::Materialized;
3553 t.row_materialized = true;
3554 });
3555 let mut rows = self.visible_rows(snapshot)?;
3556 if let Some(allowed) = allowed {
3557 rows.retain(|row| allowed.contains(&row.row_id));
3558 }
3559 if let Some(limit) = limit {
3560 rows.truncate(limit);
3561 }
3562 return Ok(rows);
3563 }
3564 crate::trace::QueryTrace::record(|t| {
3565 t.conditions_pushed = conditions.len();
3566 t.scan_mode = crate::trace::ScanMode::Materialized;
3567 t.row_materialized = true;
3568 });
3569 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
3576 ordered.sort_by_key(|c| condition_cost_rank(c));
3577 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3578 for c in &ordered {
3579 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
3580 let empty = s.is_empty();
3581 sets.push(s);
3582 if empty {
3583 break;
3584 }
3585 }
3586 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3587 if let Some(allowed) = allowed {
3588 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
3589 }
3590 if let Some(limit) = limit {
3591 rids.truncate(limit);
3592 }
3593 self.rows_for_rids(&rids, snapshot)
3594 }
3595
3596 pub fn retrieve(
3598 &mut self,
3599 retriever: &crate::query::Retriever,
3600 ) -> Result<Vec<crate::query::RetrieverHit>> {
3601 self.retrieve_with_allowed(retriever, None)
3602 }
3603
3604 pub fn retrieve_at(
3605 &mut self,
3606 retriever: &crate::query::Retriever,
3607 snapshot: Snapshot,
3608 allowed: Option<&std::collections::HashSet<RowId>>,
3609 ) -> Result<Vec<crate::query::RetrieverHit>> {
3610 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
3611 }
3612
3613 pub fn retrieve_with_allowed(
3616 &mut self,
3617 retriever: &crate::query::Retriever,
3618 allowed: Option<&std::collections::HashSet<RowId>>,
3619 ) -> Result<Vec<crate::query::RetrieverHit>> {
3620 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
3621 }
3622
3623 pub fn retrieve_at_with_allowed(
3624 &mut self,
3625 retriever: &crate::query::Retriever,
3626 snapshot: Snapshot,
3627 allowed: Option<&std::collections::HashSet<RowId>>,
3628 ) -> Result<Vec<crate::query::RetrieverHit>> {
3629 self.require_select()?;
3630 self.ensure_indexes_complete()?;
3631 self.validate_retriever(retriever)?;
3632 self.retrieve_filtered(retriever, snapshot, None, allowed)
3633 }
3634
3635 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3636 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
3637 let (column_id, k) = match retriever {
3638 Retriever::Ann {
3639 column_id,
3640 query,
3641 k,
3642 } => {
3643 let index = self.ann.get(column_id).ok_or_else(|| {
3644 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3645 })?;
3646 if query.len() != index.dim() {
3647 return Err(MongrelError::InvalidArgument(format!(
3648 "ANN query dimension must be {}, got {}",
3649 index.dim(),
3650 query.len()
3651 )));
3652 }
3653 if query.iter().any(|value| !value.is_finite()) {
3654 return Err(MongrelError::InvalidArgument(
3655 "ANN query values must be finite".into(),
3656 ));
3657 }
3658 (*column_id, *k)
3659 }
3660 Retriever::Sparse {
3661 column_id,
3662 query,
3663 k,
3664 } => {
3665 if !self.sparse.contains_key(column_id) {
3666 return Err(MongrelError::InvalidArgument(format!(
3667 "column {column_id} has no Sparse index"
3668 )));
3669 }
3670 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3671 return Err(MongrelError::InvalidArgument(
3672 "Sparse query must be non-empty with finite weights".into(),
3673 ));
3674 }
3675 if query.len() > MAX_SPARSE_TERMS {
3676 return Err(MongrelError::InvalidArgument(format!(
3677 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
3678 )));
3679 }
3680 (*column_id, *k)
3681 }
3682 Retriever::MinHash {
3683 column_id,
3684 members,
3685 k,
3686 } => {
3687 if !self.minhash.contains_key(column_id) {
3688 return Err(MongrelError::InvalidArgument(format!(
3689 "column {column_id} has no MinHash index"
3690 )));
3691 }
3692 if members.is_empty() {
3693 return Err(MongrelError::InvalidArgument(
3694 "MinHash members must not be empty".into(),
3695 ));
3696 }
3697 if members.len() > MAX_SET_MEMBERS {
3698 return Err(MongrelError::InvalidArgument(format!(
3699 "MinHash query exceeds {MAX_SET_MEMBERS} members"
3700 )));
3701 }
3702 (*column_id, *k)
3703 }
3704 };
3705 if k == 0 {
3706 return Err(MongrelError::InvalidArgument(
3707 "retriever k must be > 0".into(),
3708 ));
3709 }
3710 if k > MAX_RETRIEVER_K {
3711 return Err(MongrelError::InvalidArgument(format!(
3712 "retriever k exceeds {MAX_RETRIEVER_K}"
3713 )));
3714 }
3715 debug_assert!(self
3716 .schema
3717 .columns
3718 .iter()
3719 .any(|column| column.id == column_id));
3720 Ok(())
3721 }
3722
3723 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3724 use crate::query::Condition;
3725 match condition {
3726 Condition::Ann {
3727 column_id,
3728 query,
3729 k,
3730 } => self.validate_retriever(&crate::query::Retriever::Ann {
3731 column_id: *column_id,
3732 query: query.clone(),
3733 k: *k,
3734 }),
3735 Condition::SparseMatch {
3736 column_id,
3737 query,
3738 k,
3739 } => self.validate_retriever(&crate::query::Retriever::Sparse {
3740 column_id: *column_id,
3741 query: query.clone(),
3742 k: *k,
3743 }),
3744 Condition::MinHashSimilar {
3745 column_id,
3746 query,
3747 k,
3748 } => {
3749 if !self.minhash.contains_key(column_id) {
3750 return Err(MongrelError::InvalidArgument(format!(
3751 "column {column_id} has no MinHash index"
3752 )));
3753 }
3754 if query.is_empty() || *k == 0 {
3755 return Err(MongrelError::InvalidArgument(
3756 "MinHash query must be non-empty and k must be > 0".into(),
3757 ));
3758 }
3759 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
3760 {
3761 return Err(MongrelError::InvalidArgument(format!(
3762 "MinHash query must have <= {} members and k <= {}",
3763 crate::query::MAX_SET_MEMBERS,
3764 crate::query::MAX_RETRIEVER_K
3765 )));
3766 }
3767 Ok(())
3768 }
3769 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
3770 Err(MongrelError::InvalidArgument(format!(
3771 "bitmap IN exceeds {} values",
3772 crate::query::MAX_SET_MEMBERS
3773 )))
3774 }
3775 Condition::FmContainsAll { patterns, .. }
3776 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
3777 {
3778 Err(MongrelError::InvalidArgument(format!(
3779 "FM query exceeds {} patterns",
3780 crate::query::MAX_HARD_CONDITIONS
3781 )))
3782 }
3783 _ => Ok(()),
3784 }
3785 }
3786
3787 fn retrieve_filtered(
3788 &self,
3789 retriever: &crate::query::Retriever,
3790 snapshot: Snapshot,
3791 hard_filter: Option<&RowIdSet>,
3792 allowed: Option<&std::collections::HashSet<RowId>>,
3793 ) -> Result<Vec<crate::query::RetrieverHit>> {
3794 use crate::query::{Retriever, RetrieverHit, RetrieverScore, SetMember};
3795 let started = std::time::Instant::now();
3796 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
3797 Retriever::Ann {
3798 column_id,
3799 query,
3800 k,
3801 } => {
3802 let Some(index) = self.ann.get(column_id) else {
3803 return Ok(Vec::new());
3804 };
3805 let cap = index.len();
3806 if cap == 0 {
3807 return Ok(Vec::new());
3808 }
3809 let mut breadth = (*k).max(1).min(cap);
3810 let mut eligibility = std::collections::HashMap::new();
3811 let mut filtered = loop {
3812 let mut seen = std::collections::HashSet::new();
3813 let raw = index.search(query, breadth)?;
3814 let unchecked: Vec<_> = raw
3815 .iter()
3816 .map(|(row_id, _)| *row_id)
3817 .filter(|row_id| !eligibility.contains_key(row_id))
3818 .filter(|row_id| {
3819 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3820 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3821 })
3822 .collect();
3823 let eligible = self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3824 for row_id in unchecked {
3825 eligibility.insert(row_id, eligible.contains(&row_id));
3826 }
3827 let filtered: Vec<_> = raw
3828 .into_iter()
3829 .filter(|(row_id, _)| {
3830 seen.insert(*row_id)
3831 && eligibility.get(row_id).copied().unwrap_or(false)
3832 })
3833 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
3834 .collect();
3835 if filtered.len() >= *k || breadth >= cap {
3836 break filtered;
3837 }
3838 breadth = breadth.saturating_mul(2).min(cap);
3839 };
3840 filtered.truncate(*k);
3841 filtered
3842 }
3843 Retriever::Sparse {
3844 column_id,
3845 query,
3846 k,
3847 } => self
3848 .sparse
3849 .get(column_id)
3850 .map(|index| -> Result<Vec<_>> {
3851 let mut breadth = (*k).max(1);
3852 let mut eligibility = std::collections::HashMap::new();
3853 loop {
3854 let raw = index.search(query, breadth);
3855 let unchecked: Vec<_> = raw
3856 .iter()
3857 .map(|(row_id, _)| *row_id)
3858 .filter(|row_id| !eligibility.contains_key(row_id))
3859 .filter(|row_id| {
3860 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3861 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3862 })
3863 .collect();
3864 let eligible =
3865 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3866 for row_id in unchecked {
3867 eligibility.insert(row_id, eligible.contains(&row_id));
3868 }
3869 let filtered: Vec<_> = raw
3870 .iter()
3871 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3872 .take(*k)
3873 .map(|(row_id, score)| {
3874 (*row_id, RetrieverScore::SparseDotProduct(*score))
3875 })
3876 .collect();
3877 if filtered.len() >= *k || raw.len() < breadth {
3878 break Ok(filtered);
3879 }
3880 let next = breadth.saturating_mul(2);
3881 if next == breadth {
3882 break Ok(filtered);
3883 }
3884 breadth = next;
3885 }
3886 })
3887 .transpose()?
3888 .unwrap_or_default(),
3889 Retriever::MinHash {
3890 column_id,
3891 members,
3892 k,
3893 } => self
3894 .minhash
3895 .get(column_id)
3896 .map(|index| -> Result<Vec<_>> {
3897 let hashes: Vec<_> = members.iter().map(SetMember::hash_v1).collect();
3898 let mut breadth = (*k).max(1);
3899 let mut eligibility = std::collections::HashMap::new();
3900 loop {
3901 let raw = index.search(&hashes, breadth);
3902 let unchecked: Vec<_> = raw
3903 .iter()
3904 .map(|(row_id, _)| *row_id)
3905 .filter(|row_id| !eligibility.contains_key(row_id))
3906 .filter(|row_id| {
3907 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3908 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3909 })
3910 .collect();
3911 let eligible =
3912 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3913 for row_id in unchecked {
3914 eligibility.insert(row_id, eligible.contains(&row_id));
3915 }
3916 let filtered: Vec<_> = raw
3917 .iter()
3918 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3919 .take(*k)
3920 .map(|(row_id, score)| {
3921 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
3922 })
3923 .collect();
3924 if filtered.len() >= *k || raw.len() < breadth {
3925 break Ok(filtered);
3926 }
3927 let next = breadth.saturating_mul(2);
3928 if next == breadth {
3929 break Ok(filtered);
3930 }
3931 breadth = next;
3932 }
3933 })
3934 .transpose()?
3935 .unwrap_or_default(),
3936 };
3937 let elapsed = started.elapsed().as_nanos() as u64;
3938 crate::trace::QueryTrace::record(|trace| {
3939 match retriever {
3940 Retriever::Ann { .. } => {
3941 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
3942 }
3943 Retriever::Sparse { .. } => {
3944 trace.sparse_candidate_nanos =
3945 trace.sparse_candidate_nanos.saturating_add(elapsed)
3946 }
3947 Retriever::MinHash { .. } => {
3948 trace.minhash_candidate_nanos =
3949 trace.minhash_candidate_nanos.saturating_add(elapsed)
3950 }
3951 }
3952 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
3953 });
3954 Ok(scored
3955 .into_iter()
3956 .enumerate()
3957 .map(|(rank, (row_id, score))| RetrieverHit {
3958 row_id,
3959 rank: rank + 1,
3960 score,
3961 })
3962 .collect())
3963 }
3964
3965 fn eligible_candidate_ids(
3966 &self,
3967 candidates: &[RowId],
3968 _column_id: u16,
3969 snapshot: Snapshot,
3970 ) -> Result<std::collections::HashSet<RowId>> {
3971 if !self.had_deletes
3972 && self.ttl.is_none()
3973 && self.pending_put_cols.is_empty()
3974 && snapshot.epoch == self.snapshot().epoch
3975 {
3976 return Ok(candidates.iter().copied().collect());
3977 }
3978 let mut readers: Vec<_> = self
3979 .run_refs
3980 .iter()
3981 .map(|run| self.open_reader(run.run_id))
3982 .collect::<Result<_>>()?;
3983 let now = unix_nanos_now();
3984 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
3985 for &row_id in candidates {
3986 let mem = self.memtable.get_version(row_id, snapshot.epoch);
3987 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
3988 let overlay = match (mem, mutable) {
3989 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
3990 (Some(value), None) | (None, Some(value)) => Some(value),
3991 (None, None) => None,
3992 };
3993 if let Some((_, row)) = overlay {
3994 if !row.deleted && !self.row_expired_at(&row, now) {
3995 eligible.insert(row_id);
3996 }
3997 continue;
3998 }
3999 let mut best: Option<(Epoch, bool, usize)> = None;
4000 for (index, reader) in readers.iter_mut().enumerate() {
4001 if let Some((epoch, deleted)) =
4002 reader.get_version_visibility(row_id, snapshot.epoch)?
4003 {
4004 if best
4005 .as_ref()
4006 .map(|(best_epoch, ..)| epoch > *best_epoch)
4007 .unwrap_or(true)
4008 {
4009 best = Some((epoch, deleted, index));
4010 }
4011 }
4012 }
4013 let Some((_, false, reader_index)) = best else {
4014 continue;
4015 };
4016 if let Some(ttl) = self.ttl {
4017 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4018 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4019 {
4020 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4021 continue;
4022 }
4023 }
4024 }
4025 eligible.insert(row_id);
4026 }
4027 Ok(eligible)
4028 }
4029
4030 pub fn search(
4032 &mut self,
4033 request: &crate::query::SearchRequest,
4034 ) -> Result<Vec<crate::query::SearchHit>> {
4035 self.search_with_allowed(request, None)
4036 }
4037
4038 pub fn search_at(
4039 &mut self,
4040 request: &crate::query::SearchRequest,
4041 snapshot: Snapshot,
4042 authorized: Option<&std::collections::HashSet<RowId>>,
4043 ) -> Result<Vec<crate::query::SearchHit>> {
4044 self.search_at_with_allowed(request, snapshot, authorized)
4045 }
4046
4047 pub fn search_with_allowed(
4048 &mut self,
4049 request: &crate::query::SearchRequest,
4050 authorized: Option<&std::collections::HashSet<RowId>>,
4051 ) -> Result<Vec<crate::query::SearchHit>> {
4052 self.search_at_with_allowed(request, self.snapshot(), authorized)
4053 }
4054
4055 pub fn search_at_with_allowed(
4056 &mut self,
4057 request: &crate::query::SearchRequest,
4058 snapshot: Snapshot,
4059 authorized: Option<&std::collections::HashSet<RowId>>,
4060 ) -> Result<Vec<crate::query::SearchHit>> {
4061 use crate::query::{
4062 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
4063 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
4064 };
4065 let total_started = std::time::Instant::now();
4066 self.require_select()?;
4067 self.ensure_indexes_complete()?;
4068 if request.limit == 0 {
4069 return Err(MongrelError::InvalidArgument(
4070 "search limit must be > 0".into(),
4071 ));
4072 }
4073 if request.limit > MAX_FINAL_LIMIT {
4074 return Err(MongrelError::InvalidArgument(format!(
4075 "search limit exceeds {MAX_FINAL_LIMIT}"
4076 )));
4077 }
4078 if request.retrievers.is_empty() {
4079 return Err(MongrelError::InvalidArgument(
4080 "search requires at least one retriever".into(),
4081 ));
4082 }
4083 if request.retrievers.len() > MAX_RETRIEVERS {
4084 return Err(MongrelError::InvalidArgument(format!(
4085 "search exceeds {MAX_RETRIEVERS} retrievers"
4086 )));
4087 }
4088 if request.must.len() > MAX_HARD_CONDITIONS {
4089 return Err(MongrelError::InvalidArgument(format!(
4090 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
4091 )));
4092 }
4093 if request.must.iter().any(|condition| {
4094 matches!(
4095 condition,
4096 Condition::Ann { .. }
4097 | Condition::SparseMatch { .. }
4098 | Condition::MinHashSimilar { .. }
4099 )
4100 }) {
4101 return Err(MongrelError::InvalidArgument(
4102 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
4103 .into(),
4104 ));
4105 }
4106 let mut names = std::collections::HashSet::new();
4107 for named in &request.retrievers {
4108 if named.name.is_empty() || !names.insert(named.name.as_str()) {
4109 return Err(MongrelError::InvalidArgument(
4110 "retriever names must be non-empty and unique".into(),
4111 ));
4112 }
4113 if !named.weight.is_finite()
4114 || named.weight < 0.0
4115 || named.weight > MAX_RETRIEVER_WEIGHT
4116 {
4117 return Err(MongrelError::InvalidArgument(format!(
4118 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4119 )));
4120 }
4121 self.validate_retriever(&named.retriever)?;
4122 }
4123 let projection = request
4124 .projection
4125 .clone()
4126 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
4127 if projection.len() > MAX_PROJECTION_COLUMNS {
4128 return Err(MongrelError::InvalidArgument(format!(
4129 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
4130 )));
4131 }
4132 for column_id in &projection {
4133 if !self
4134 .schema
4135 .columns
4136 .iter()
4137 .any(|column| column.id == *column_id)
4138 {
4139 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
4140 }
4141 }
4142
4143 let hard_filter_started = std::time::Instant::now();
4144 let hard_filter = if request.must.is_empty() {
4145 None
4146 } else {
4147 let mut sets = Vec::with_capacity(request.must.len());
4148 for condition in &request.must {
4149 sets.push(self.resolve_condition(condition, snapshot)?);
4150 }
4151 Some(RowIdSet::intersect_many(sets))
4152 };
4153 crate::trace::QueryTrace::record(|trace| {
4154 trace.hard_filter_nanos = trace
4155 .hard_filter_nanos
4156 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
4157 });
4158 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
4159 return Ok(Vec::new());
4160 }
4161
4162 let constant = match request.fusion {
4163 Fusion::ReciprocalRank { constant } => constant,
4164 };
4165 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
4166 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
4167 let mut fusion_nanos = 0u64;
4168 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
4169 std::collections::HashMap::new();
4170 for named in retrievers {
4171 let hits = self.retrieve_filtered(
4172 &named.retriever,
4173 snapshot,
4174 hard_filter.as_ref(),
4175 authorized,
4176 )?;
4177 let fusion_started = std::time::Instant::now();
4178 for hit in hits {
4179 let contribution = named.weight / (constant as f64 + hit.rank as f64);
4180 if !contribution.is_finite() {
4181 return Err(MongrelError::InvalidArgument(
4182 "retriever contribution must be finite".into(),
4183 ));
4184 }
4185 let entry = fused.entry(hit.row_id).or_default();
4186 entry.0 += contribution;
4187 if !entry.0.is_finite() {
4188 return Err(MongrelError::InvalidArgument(
4189 "fused score must be finite".into(),
4190 ));
4191 }
4192 entry.1.push(ComponentScore {
4193 retriever_name: named.name.clone(),
4194 rank: hit.rank,
4195 raw_score: hit.score,
4196 contribution,
4197 });
4198 }
4199 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
4200 }
4201 let union_size = fused.len();
4202 let mut ranked: Vec<_> = fused.into_iter().collect();
4203 let sort_started = std::time::Instant::now();
4204 ranked.sort_by(|(a_row, (a_score, _)), (b_row, (b_score, _))| {
4205 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
4206 });
4207 fusion_nanos = fusion_nanos.saturating_add(sort_started.elapsed().as_nanos() as u64);
4208 crate::trace::QueryTrace::record(|trace| {
4209 trace.union_size = union_size;
4210 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
4211 });
4212
4213 let projection_started = std::time::Instant::now();
4214 let row_ids: Vec<_> = ranked.iter().map(|(row_id, _)| row_id.0).collect();
4215 let sentinel = projection
4216 .first()
4217 .copied()
4218 .or_else(|| self.schema.columns.first().map(|column| column.id));
4219 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
4220 std::collections::HashMap::new();
4221 if let Some(column_id) = sentinel {
4222 for (row_id, value) in self.values_for_rids_batch(&row_ids, column_id, snapshot)? {
4223 cells.entry(row_id).or_default().insert(column_id, value);
4224 }
4225 }
4226 for &column_id in &projection {
4227 if Some(column_id) == sentinel {
4228 continue;
4229 }
4230 for (row_id, value) in self.values_for_rids_batch(&row_ids, column_id, snapshot)? {
4231 cells.entry(row_id).or_default().insert(column_id, value);
4232 }
4233 }
4234
4235 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
4236 for (row_id, (fused_score, mut components)) in ranked {
4237 let Some(row_cells) = cells.remove(&row_id) else {
4238 continue;
4239 };
4240 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4241 out.push(SearchHit {
4242 row_id,
4243 cells: projection
4244 .iter()
4245 .filter_map(|column_id| {
4246 row_cells
4247 .get(column_id)
4248 .cloned()
4249 .map(|value| (*column_id, value))
4250 })
4251 .collect(),
4252 components,
4253 fused_score,
4254 });
4255 if out.len() == request.limit {
4256 break;
4257 }
4258 }
4259 crate::trace::QueryTrace::record(|trace| {
4260 trace.projection_nanos = trace
4261 .projection_nanos
4262 .saturating_add(projection_started.elapsed().as_nanos() as u64);
4263 trace.total_nanos = trace
4264 .total_nanos
4265 .saturating_add(total_started.elapsed().as_nanos() as u64);
4266 });
4267 Ok(out)
4268 }
4269
4270 pub fn set_similarity(
4273 &mut self,
4274 request: &crate::query::SetSimilarityRequest,
4275 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4276 self.set_similarity_with_allowed(request, None)
4277 }
4278
4279 pub fn set_similarity_at(
4280 &mut self,
4281 request: &crate::query::SetSimilarityRequest,
4282 snapshot: Snapshot,
4283 allowed: Option<&std::collections::HashSet<RowId>>,
4284 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4285 self.set_similarity_explained_at(request, snapshot, allowed)
4286 .map(|(hits, _)| hits)
4287 }
4288
4289 pub fn ann_rerank(
4291 &mut self,
4292 request: &crate::query::AnnRerankRequest,
4293 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4294 self.ann_rerank_with_allowed(request, None)
4295 }
4296
4297 pub fn ann_rerank_with_allowed(
4298 &mut self,
4299 request: &crate::query::AnnRerankRequest,
4300 allowed: Option<&std::collections::HashSet<RowId>>,
4301 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4302 self.ann_rerank_at(request, self.snapshot(), allowed)
4303 }
4304
4305 pub fn ann_rerank_at(
4306 &mut self,
4307 request: &crate::query::AnnRerankRequest,
4308 snapshot: Snapshot,
4309 allowed: Option<&std::collections::HashSet<RowId>>,
4310 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4311 use crate::query::{
4312 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
4313 };
4314 if request.candidate_k == 0 || request.limit == 0 {
4315 return Err(MongrelError::InvalidArgument(
4316 "candidate_k and limit must be > 0".into(),
4317 ));
4318 }
4319 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
4320 return Err(MongrelError::InvalidArgument(format!(
4321 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
4322 )));
4323 }
4324 let hits = self.retrieve_at_with_allowed(
4325 &Retriever::Ann {
4326 column_id: request.column_id,
4327 query: request.query.clone(),
4328 k: request.candidate_k,
4329 },
4330 snapshot,
4331 allowed,
4332 )?;
4333 let distances: std::collections::HashMap<_, _> = hits
4334 .iter()
4335 .filter_map(|hit| match hit.score {
4336 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
4337 _ => None,
4338 })
4339 .collect();
4340 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4341 let gather_started = std::time::Instant::now();
4342 let values = self.values_for_rids_batch(&row_ids, request.column_id, snapshot)?;
4343 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
4344 let score_started = std::time::Instant::now();
4345 let query_norm = request
4346 .query
4347 .iter()
4348 .map(|value| f64::from(*value).powi(2))
4349 .sum::<f64>()
4350 .sqrt();
4351 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
4352 for (row_id, value) in values {
4353 let Value::Embedding(vector) = value else {
4354 continue;
4355 };
4356 let dot = request
4357 .query
4358 .iter()
4359 .zip(&vector)
4360 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4361 .sum::<f64>();
4362 let exact_score = match request.metric {
4363 VectorMetric::DotProduct => dot,
4364 VectorMetric::Cosine => {
4365 let norm = vector
4366 .iter()
4367 .map(|value| f64::from(*value).powi(2))
4368 .sum::<f64>()
4369 .sqrt();
4370 if query_norm == 0.0 || norm == 0.0 {
4371 0.0
4372 } else {
4373 dot / (query_norm * norm)
4374 }
4375 }
4376 VectorMetric::Euclidean => request
4377 .query
4378 .iter()
4379 .zip(&vector)
4380 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
4381 .sum::<f64>()
4382 .sqrt(),
4383 };
4384 let exact_score = exact_score as f32;
4385 if !exact_score.is_finite() {
4386 return Err(MongrelError::InvalidArgument(
4387 "exact ANN score must be finite".into(),
4388 ));
4389 }
4390 reranked.push(AnnRerankHit {
4391 row_id,
4392 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
4393 exact_score,
4394 });
4395 }
4396 reranked.sort_by(|left, right| {
4397 let score = match request.metric {
4398 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
4399 VectorMetric::Cosine | VectorMetric::DotProduct => {
4400 right.exact_score.total_cmp(&left.exact_score)
4401 }
4402 };
4403 score.then_with(|| left.row_id.cmp(&right.row_id))
4404 });
4405 reranked.truncate(request.limit);
4406 crate::trace::QueryTrace::record(|trace| {
4407 trace.exact_vector_gather_nanos =
4408 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
4409 trace.exact_vector_score_nanos = trace
4410 .exact_vector_score_nanos
4411 .saturating_add(score_started.elapsed().as_nanos() as u64);
4412 });
4413 Ok(reranked)
4414 }
4415
4416 pub fn set_similarity_with_allowed(
4417 &mut self,
4418 request: &crate::query::SetSimilarityRequest,
4419 allowed: Option<&std::collections::HashSet<RowId>>,
4420 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4421 self.set_similarity_explained_at(request, self.snapshot(), allowed)
4422 .map(|(hits, _)| hits)
4423 }
4424
4425 pub fn set_similarity_explained(
4426 &mut self,
4427 request: &crate::query::SetSimilarityRequest,
4428 ) -> Result<(
4429 Vec<crate::query::SetSimilarityHit>,
4430 crate::query::SetSimilarityTrace,
4431 )> {
4432 self.set_similarity_explained_at(request, self.snapshot(), None)
4433 }
4434
4435 fn set_similarity_explained_at(
4436 &mut self,
4437 request: &crate::query::SetSimilarityRequest,
4438 snapshot: Snapshot,
4439 allowed: Option<&std::collections::HashSet<RowId>>,
4440 ) -> Result<(
4441 Vec<crate::query::SetSimilarityHit>,
4442 crate::query::SetSimilarityTrace,
4443 )> {
4444 use crate::query::{
4445 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
4446 MAX_SET_MEMBERS,
4447 };
4448 let mut trace = crate::query::SetSimilarityTrace::default();
4449 if request.members.is_empty() {
4450 return Ok((Vec::new(), trace));
4451 }
4452 if request.candidate_k == 0 || request.limit == 0 {
4453 return Err(MongrelError::InvalidArgument(
4454 "candidate_k and limit must be > 0".into(),
4455 ));
4456 }
4457 if request.candidate_k > MAX_RETRIEVER_K
4458 || request.limit > MAX_FINAL_LIMIT
4459 || request.members.len() > MAX_SET_MEMBERS
4460 {
4461 return Err(MongrelError::InvalidArgument(format!(
4462 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
4463 )));
4464 }
4465 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
4466 return Err(MongrelError::InvalidArgument(
4467 "min_jaccard must be finite and between 0 and 1".into(),
4468 ));
4469 }
4470 let started = std::time::Instant::now();
4471 let hits = self.retrieve_at_with_allowed(
4472 &Retriever::MinHash {
4473 column_id: request.column_id,
4474 members: request.members.clone(),
4475 k: request.candidate_k,
4476 },
4477 snapshot,
4478 allowed,
4479 )?;
4480 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
4481 trace.candidate_count = hits.len();
4482 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4483 let started = std::time::Instant::now();
4484 let values = self.values_for_rids_batch(&row_ids, request.column_id, snapshot)?;
4485 trace.gather_us = started.elapsed().as_micros() as u64;
4486 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
4487 let estimates: std::collections::HashMap<_, _> = hits
4488 .into_iter()
4489 .filter_map(|hit| match hit.score {
4490 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
4491 _ => None,
4492 })
4493 .collect();
4494 let started = std::time::Instant::now();
4495 let parsed: Vec<_> = values
4496 .into_iter()
4497 .filter_map(|(row_id, value)| {
4498 let Value::Bytes(bytes) = value else {
4499 return None;
4500 };
4501 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
4502 return None;
4503 };
4504 let stored = members
4505 .into_iter()
4506 .filter_map(|member| match member {
4507 serde_json::Value::String(value) => {
4508 Some(crate::query::SetMember::String(value))
4509 }
4510 serde_json::Value::Number(value) => {
4511 Some(crate::query::SetMember::Number(value))
4512 }
4513 serde_json::Value::Bool(value) => {
4514 Some(crate::query::SetMember::Boolean(value))
4515 }
4516 _ => None,
4517 })
4518 .collect::<std::collections::HashSet<_>>();
4519 Some((row_id, stored))
4520 })
4521 .collect();
4522 trace.parse_us = started.elapsed().as_micros() as u64;
4523 trace.verified_count = parsed.len();
4524 let started = std::time::Instant::now();
4525 let mut exact = Vec::new();
4526 for (row_id, stored) in parsed {
4527 let union = query.union(&stored).count();
4528 let score = if union == 0 {
4529 1.0
4530 } else {
4531 query.intersection(&stored).count() as f32 / union as f32
4532 };
4533 if score >= request.min_jaccard {
4534 exact.push(SetSimilarityHit {
4535 row_id,
4536 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
4537 exact_jaccard: score,
4538 });
4539 }
4540 }
4541 exact.sort_by(|a, b| {
4542 b.exact_jaccard
4543 .total_cmp(&a.exact_jaccard)
4544 .then_with(|| a.row_id.cmp(&b.row_id))
4545 });
4546 exact.truncate(request.limit);
4547 trace.score_us = started.elapsed().as_micros() as u64;
4548 crate::trace::QueryTrace::record(|query_trace| {
4549 query_trace.exact_set_gather_nanos = query_trace
4550 .exact_set_gather_nanos
4551 .saturating_add(trace.gather_us.saturating_mul(1_000));
4552 query_trace.exact_set_parse_nanos = query_trace
4553 .exact_set_parse_nanos
4554 .saturating_add(trace.parse_us.saturating_mul(1_000));
4555 query_trace.exact_set_score_nanos = query_trace
4556 .exact_set_score_nanos
4557 .saturating_add(trace.score_us.saturating_mul(1_000));
4558 });
4559 Ok((exact, trace))
4560 }
4561
4562 fn values_for_rids_batch(
4564 &self,
4565 row_ids: &[u64],
4566 column_id: u16,
4567 snapshot: Snapshot,
4568 ) -> Result<Vec<(RowId, Value)>> {
4569 if self.ttl.is_none()
4570 && self.memtable.is_empty()
4571 && self.mutable_run.is_empty()
4572 && self.run_refs.len() == 1
4573 {
4574 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4575 let (positions, visible_row_ids) =
4576 reader.visible_positions_with_rids(snapshot.epoch)?;
4577 let requested: Vec<(RowId, usize)> = row_ids
4578 .iter()
4579 .filter_map(|raw| {
4580 visible_row_ids
4581 .binary_search(&(*raw as i64))
4582 .ok()
4583 .map(|index| (RowId(*raw), positions[index]))
4584 })
4585 .collect();
4586 let values = reader.gather_column(
4587 column_id,
4588 &requested
4589 .iter()
4590 .map(|(_, position)| *position)
4591 .collect::<Vec<_>>(),
4592 )?;
4593 return Ok(requested
4594 .into_iter()
4595 .zip(values)
4596 .map(|((row_id, _), value)| (row_id, value))
4597 .collect());
4598 }
4599 self.values_for_rids(row_ids, column_id, snapshot)
4600 }
4601
4602 fn values_for_rids(
4604 &self,
4605 row_ids: &[u64],
4606 column_id: u16,
4607 snapshot: Snapshot,
4608 ) -> Result<Vec<(RowId, Value)>> {
4609 let mut readers: Vec<_> = self
4610 .run_refs
4611 .iter()
4612 .map(|run| self.open_reader(run.run_id))
4613 .collect::<Result<_>>()?;
4614 let now = unix_nanos_now();
4615 let mut out = Vec::with_capacity(row_ids.len());
4616 for &raw_row_id in row_ids {
4617 let row_id = RowId(raw_row_id);
4618 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4619 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4620 let overlay = match (mem, mutable) {
4621 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
4622 (a_epoch, a)
4623 } else {
4624 (b_epoch, b)
4625 }),
4626 (Some(value), None) | (None, Some(value)) => Some(value),
4627 (None, None) => None,
4628 };
4629 if let Some((_, row)) = overlay {
4630 if !row.deleted && !self.row_expired_at(&row, now) {
4631 if let Some(value) = row.columns.get(&column_id) {
4632 out.push((row_id, value.clone()));
4633 }
4634 }
4635 continue;
4636 }
4637
4638 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
4639 for (index, reader) in readers.iter_mut().enumerate() {
4640 if let Some((epoch, deleted, value)) =
4641 reader.get_version_column(row_id, snapshot.epoch, column_id)?
4642 {
4643 if best
4644 .as_ref()
4645 .map(|(best_epoch, ..)| epoch > *best_epoch)
4646 .unwrap_or(true)
4647 {
4648 best = Some((epoch, deleted, value, index));
4649 }
4650 }
4651 }
4652 let Some((_, false, Some(value), reader_index)) = best else {
4653 continue;
4654 };
4655 if let Some(ttl) = self.ttl {
4656 if ttl.column_id != column_id {
4657 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4658 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4659 {
4660 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4661 continue;
4662 }
4663 }
4664 } else if let Value::Int64(timestamp) = value {
4665 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4666 continue;
4667 }
4668 }
4669 }
4670 out.push((row_id, value));
4671 }
4672 Ok(out)
4673 }
4674
4675 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
4680 use std::collections::HashMap;
4681 let mut rows = Vec::with_capacity(rids.len());
4682 let ttl_now = unix_nanos_now();
4683 let tier_size = self.memtable.len() + self.mutable_run.len();
4700 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
4701 if rids.len().saturating_mul(24) < tier_size {
4702 for &rid in rids {
4703 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
4704 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
4705 let newest = match (mem, mrun) {
4706 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
4707 (Some((_, mr)), None) => Some(mr),
4708 (None, Some((_, rr))) => Some(rr),
4709 (None, None) => None,
4710 };
4711 if let Some(row) = newest {
4712 overlay.insert(rid, row);
4713 }
4714 }
4715 } else {
4716 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
4717 overlay
4718 .entry(row.row_id.0)
4719 .and_modify(|e| {
4720 if row.committed_epoch > e.committed_epoch {
4721 *e = row.clone();
4722 }
4723 })
4724 .or_insert(row);
4725 };
4726 for row in self.memtable.visible_versions(snapshot.epoch) {
4727 fold_newest(row, &mut overlay);
4728 }
4729 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4730 fold_newest(row, &mut overlay);
4731 }
4732 }
4733 if self.run_refs.len() == 1 {
4734 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4735 if rids.len().saturating_mul(24) < reader.row_count() {
4743 for &rid in rids {
4744 if let Some(r) = overlay.get(&rid) {
4745 if !r.deleted {
4746 rows.push(r.clone());
4747 }
4748 continue;
4749 }
4750 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
4751 if !row.deleted {
4752 rows.push(row);
4753 }
4754 }
4755 }
4756 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4757 return Ok(rows);
4758 }
4759 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4768 enum Src {
4771 Overlay,
4772 Run,
4773 }
4774 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
4775 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
4776 for rid in rids {
4777 if overlay.contains_key(rid) {
4778 plan.push(Src::Overlay);
4779 continue;
4780 }
4781 match vis_rids.binary_search(&(*rid as i64)) {
4782 Ok(i) => {
4783 plan.push(Src::Run);
4784 fetch.push(positions[i]);
4785 }
4786 Err(_) => { }
4787 }
4788 }
4789 let fetched = reader.materialize_batch(&fetch)?;
4790 let mut fetched_iter = fetched.into_iter();
4791 for (rid, src) in rids.iter().zip(plan) {
4792 match src {
4793 Src::Overlay => {
4794 if let Some(r) = overlay.get(rid) {
4795 if !r.deleted {
4796 rows.push(r.clone());
4797 }
4798 }
4799 }
4800 Src::Run => {
4801 if let Some(row) = fetched_iter.next() {
4802 if !row.deleted {
4803 rows.push(row);
4804 }
4805 }
4806 }
4807 }
4808 }
4809 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4810 return Ok(rows);
4811 }
4812 let mut readers: Vec<_> = self
4816 .run_refs
4817 .iter()
4818 .map(|rr| self.open_reader(rr.run_id))
4819 .collect::<Result<Vec<_>>>()?;
4820 for rid in rids {
4821 if let Some(r) = overlay.get(rid) {
4822 if !r.deleted {
4823 rows.push(r.clone());
4824 }
4825 continue;
4826 }
4827 let mut best: Option<(Epoch, Row)> = None;
4828 for reader in readers.iter_mut() {
4829 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
4830 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4831 best = Some((epoch, row));
4832 }
4833 }
4834 }
4835 if let Some((_, r)) = best {
4836 if !r.deleted {
4837 rows.push(r);
4838 }
4839 }
4840 }
4841 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4842 Ok(rows)
4843 }
4844
4845 pub fn indexes_complete(&self) -> bool {
4855 self.indexes_complete
4856 }
4857
4858 pub fn index_build_policy(&self) -> IndexBuildPolicy {
4860 self.index_build_policy
4861 }
4862
4863 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
4867 self.index_build_policy = policy;
4868 }
4869
4870 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
4875 if !self.indexes_complete {
4879 return None;
4880 }
4881 let b = self.bitmap.get(&column_id)?;
4882 let result: Vec<Vec<u8>> = b
4883 .keys()
4884 .into_iter()
4885 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
4886 .cloned()
4887 .collect();
4888 Some(result)
4889 }
4890
4891 pub fn fk_join_row_ids(
4892 &self,
4893 fk_column_id: u16,
4894 pk_values: &[Vec<u8>],
4895 fk_conditions: &[crate::query::Condition],
4896 snapshot: Snapshot,
4897 ) -> Result<Vec<u64>> {
4898 let Some(b) = self.bitmap.get(&fk_column_id) else {
4899 return Ok(Vec::new());
4900 };
4901 let mut join_set = {
4902 let mut acc = roaring::RoaringBitmap::new();
4903 for v in pk_values {
4904 acc |= b.get(v);
4905 }
4906 RowIdSet::from_roaring(acc)
4907 };
4908 if !fk_conditions.is_empty() {
4909 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4910 sets.push(join_set);
4911 for c in fk_conditions {
4912 sets.push(self.resolve_condition(c, snapshot)?);
4913 }
4914 join_set = RowIdSet::intersect_many(sets);
4915 }
4916 Ok(join_set.into_sorted_vec())
4917 }
4918
4919 pub fn fk_join_count(
4925 &self,
4926 fk_column_id: u16,
4927 pk_values: &[Vec<u8>],
4928 fk_conditions: &[crate::query::Condition],
4929 snapshot: Snapshot,
4930 ) -> Result<u64> {
4931 let Some(b) = self.bitmap.get(&fk_column_id) else {
4932 return Ok(0);
4933 };
4934 let mut acc = roaring::RoaringBitmap::new();
4935 for v in pk_values {
4936 acc |= b.get(v);
4937 }
4938 if fk_conditions.is_empty() {
4939 return Ok(acc.len());
4940 }
4941 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4942 sets.push(RowIdSet::from_roaring(acc));
4943 for c in fk_conditions {
4944 sets.push(self.resolve_condition(c, snapshot)?);
4945 }
4946 Ok(RowIdSet::intersect_many(sets).len() as u64)
4947 }
4948
4949 fn resolve_condition(
4954 &self,
4955 c: &crate::query::Condition,
4956 snapshot: Snapshot,
4957 ) -> Result<RowIdSet> {
4958 self.resolve_condition_with_allowed(c, snapshot, None)
4959 }
4960
4961 fn resolve_condition_with_allowed(
4962 &self,
4963 c: &crate::query::Condition,
4964 snapshot: Snapshot,
4965 allowed: Option<&std::collections::HashSet<RowId>>,
4966 ) -> Result<RowIdSet> {
4967 use crate::query::Condition;
4968 self.validate_condition(c)?;
4969 Ok(match c {
4970 Condition::Pk(key) => {
4971 let lookup = self
4972 .schema
4973 .primary_key()
4974 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
4975 .unwrap_or_else(|| key.clone());
4976 self.hot
4977 .get(&lookup)
4978 .map(|r| RowIdSet::one(r.0))
4979 .unwrap_or_else(RowIdSet::empty)
4980 }
4981 Condition::BitmapEq { column_id, value } => {
4982 let lookup = self.index_lookup_key_bytes(*column_id, value);
4983 self.bitmap
4984 .get(column_id)
4985 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
4986 .unwrap_or_else(RowIdSet::empty)
4987 }
4988 Condition::BitmapIn { column_id, values } => {
4989 let bm = self.bitmap.get(column_id);
4990 let mut acc = roaring::RoaringBitmap::new();
4991 if let Some(b) = bm {
4992 for v in values {
4993 let lookup = self.index_lookup_key_bytes(*column_id, v);
4994 acc |= b.get(&lookup);
4995 }
4996 }
4997 RowIdSet::from_roaring(acc)
4998 }
4999 Condition::BytesPrefix { column_id, prefix } => {
5000 if let Some(b) = self.bitmap.get(column_id) {
5005 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5006 let mut acc = roaring::RoaringBitmap::new();
5007 for key in b.keys() {
5008 if key.starts_with(&lookup_prefix) {
5009 acc |= b.get(key);
5010 }
5011 }
5012 RowIdSet::from_roaring(acc)
5013 } else {
5014 RowIdSet::empty()
5015 }
5016 }
5017 Condition::FmContains { column_id, pattern } => self
5018 .fm
5019 .get(column_id)
5020 .map(|f| {
5021 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
5022 })
5023 .unwrap_or_else(RowIdSet::empty),
5024 Condition::FmContainsAll {
5025 column_id,
5026 patterns,
5027 } => {
5028 if let Some(f) = self.fm.get(column_id) {
5031 let sets: Vec<RowIdSet> = patterns
5032 .iter()
5033 .map(|pat| {
5034 RowIdSet::from_unsorted(
5035 f.locate(pat).into_iter().map(|r| r.0).collect(),
5036 )
5037 })
5038 .collect();
5039 RowIdSet::intersect_many(sets)
5040 } else {
5041 RowIdSet::empty()
5042 }
5043 }
5044 Condition::Ann {
5045 column_id,
5046 query,
5047 k,
5048 } => RowIdSet::from_unsorted(
5049 self.retrieve_filtered(
5050 &crate::query::Retriever::Ann {
5051 column_id: *column_id,
5052 query: query.clone(),
5053 k: *k,
5054 },
5055 snapshot,
5056 None,
5057 allowed,
5058 )?
5059 .into_iter()
5060 .map(|hit| hit.row_id.0)
5061 .collect(),
5062 ),
5063 Condition::SparseMatch {
5064 column_id,
5065 query,
5066 k,
5067 } => RowIdSet::from_unsorted(
5068 self.retrieve_filtered(
5069 &crate::query::Retriever::Sparse {
5070 column_id: *column_id,
5071 query: query.clone(),
5072 k: *k,
5073 },
5074 snapshot,
5075 None,
5076 allowed,
5077 )?
5078 .into_iter()
5079 .map(|hit| hit.row_id.0)
5080 .collect(),
5081 ),
5082 Condition::MinHashSimilar {
5083 column_id,
5084 query,
5085 k,
5086 } => match self.minhash.get(column_id) {
5087 Some(index) => {
5088 let candidates = index.candidate_row_ids(query);
5089 let eligible =
5090 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
5091 RowIdSet::from_unsorted(
5092 index
5093 .search_filtered(query, *k, |row_id| {
5094 eligible.contains(&row_id)
5095 && allowed.map_or(true, |allowed| allowed.contains(&row_id))
5096 })
5097 .into_iter()
5098 .map(|(row_id, _)| row_id.0)
5099 .collect(),
5100 )
5101 }
5102 None => RowIdSet::empty(),
5103 },
5104 Condition::Range { column_id, lo, hi } => {
5105 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5114 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5115 } else if self.run_refs.len() == 1 {
5116 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5117 r.range_row_id_set_i64(*column_id, *lo, *hi)?
5118 } else {
5119 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
5120 };
5121 set.remove_many(self.overlay_rid_set(snapshot));
5122 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
5123 set
5124 }
5125 Condition::RangeF64 {
5126 column_id,
5127 lo,
5128 lo_inclusive,
5129 hi,
5130 hi_inclusive,
5131 } => {
5132 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5135 RowIdSet::from_unsorted(
5136 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5137 .into_iter()
5138 .collect(),
5139 )
5140 } else if self.run_refs.len() == 1 {
5141 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5142 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
5143 } else {
5144 return self.range_scan_f64(
5145 *column_id,
5146 *lo,
5147 *lo_inclusive,
5148 *hi,
5149 *hi_inclusive,
5150 snapshot,
5151 );
5152 };
5153 set.remove_many(self.overlay_rid_set(snapshot));
5154 self.range_scan_overlay_f64(
5155 &mut set,
5156 *column_id,
5157 *lo,
5158 *lo_inclusive,
5159 *hi,
5160 *hi_inclusive,
5161 snapshot,
5162 );
5163 set
5164 }
5165 Condition::IsNull { column_id } => {
5166 let mut set = if self.run_refs.len() == 1 {
5167 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5168 r.null_row_id_set(*column_id, true)?
5169 } else {
5170 return self.null_scan(*column_id, true, snapshot);
5171 };
5172 set.remove_many(self.overlay_rid_set(snapshot));
5173 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
5174 set
5175 }
5176 Condition::IsNotNull { column_id } => {
5177 let mut set = if self.run_refs.len() == 1 {
5178 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5179 r.null_row_id_set(*column_id, false)?
5180 } else {
5181 return self.null_scan(*column_id, false, snapshot);
5182 };
5183 set.remove_many(self.overlay_rid_set(snapshot));
5184 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
5185 set
5186 }
5187 })
5188 }
5189
5190 fn range_scan_i64(
5198 &self,
5199 column_id: u16,
5200 lo: i64,
5201 hi: i64,
5202 snapshot: Snapshot,
5203 ) -> Result<RowIdSet> {
5204 let mut row_ids = Vec::new();
5205 let overlay_rids = self.overlay_rid_set(snapshot);
5206 for rr in &self.run_refs {
5207 let mut reader = self.open_reader(rr.run_id)?;
5208 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
5209 for rid in matched {
5210 if !overlay_rids.contains(&rid) {
5211 row_ids.push(rid);
5212 }
5213 }
5214 }
5215 let mut s = RowIdSet::from_unsorted(row_ids);
5216 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
5217 Ok(s)
5218 }
5219
5220 fn range_scan_f64(
5223 &self,
5224 column_id: u16,
5225 lo: f64,
5226 lo_inclusive: bool,
5227 hi: f64,
5228 hi_inclusive: bool,
5229 snapshot: Snapshot,
5230 ) -> Result<RowIdSet> {
5231 let mut row_ids = Vec::new();
5232 let overlay_rids = self.overlay_rid_set(snapshot);
5233 for rr in &self.run_refs {
5234 let mut reader = self.open_reader(rr.run_id)?;
5235 let matched = reader.range_row_ids_visible_f64(
5236 column_id,
5237 lo,
5238 lo_inclusive,
5239 hi,
5240 hi_inclusive,
5241 snapshot.epoch,
5242 )?;
5243 for rid in matched {
5244 if !overlay_rids.contains(&rid) {
5245 row_ids.push(rid);
5246 }
5247 }
5248 }
5249 let mut s = RowIdSet::from_unsorted(row_ids);
5250 self.range_scan_overlay_f64(
5251 &mut s,
5252 column_id,
5253 lo,
5254 lo_inclusive,
5255 hi,
5256 hi_inclusive,
5257 snapshot,
5258 );
5259 Ok(s)
5260 }
5261
5262 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
5264 let mut s = HashSet::new();
5265 for row in self.memtable.visible_versions(snapshot.epoch) {
5266 s.insert(row.row_id.0);
5267 }
5268 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5269 s.insert(row.row_id.0);
5270 }
5271 s
5272 }
5273
5274 fn range_scan_overlay_i64(
5275 &self,
5276 s: &mut RowIdSet,
5277 column_id: u16,
5278 lo: i64,
5279 hi: i64,
5280 snapshot: Snapshot,
5281 ) {
5282 let mut newest: HashMap<u64, &Row> = HashMap::new();
5287 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5288 let memtable = self.memtable.visible_versions(snapshot.epoch);
5289 for r in &mutable {
5290 newest.entry(r.row_id.0).or_insert(r);
5291 }
5292 for r in &memtable {
5293 newest.insert(r.row_id.0, r);
5294 }
5295 for row in newest.values() {
5296 if !row.deleted {
5297 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
5298 if *v >= lo && *v <= hi {
5299 s.insert(row.row_id.0);
5300 }
5301 }
5302 }
5303 }
5304 }
5305
5306 #[allow(clippy::too_many_arguments)]
5307 fn range_scan_overlay_f64(
5308 &self,
5309 s: &mut RowIdSet,
5310 column_id: u16,
5311 lo: f64,
5312 lo_inclusive: bool,
5313 hi: f64,
5314 hi_inclusive: bool,
5315 snapshot: Snapshot,
5316 ) {
5317 let mut newest: HashMap<u64, &Row> = HashMap::new();
5320 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5321 let memtable = self.memtable.visible_versions(snapshot.epoch);
5322 for r in &mutable {
5323 newest.entry(r.row_id.0).or_insert(r);
5324 }
5325 for r in &memtable {
5326 newest.insert(r.row_id.0, r);
5327 }
5328 for row in newest.values() {
5329 if !row.deleted {
5330 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
5331 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
5332 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
5333 if ok_lo && ok_hi {
5334 s.insert(row.row_id.0);
5335 }
5336 }
5337 }
5338 }
5339 }
5340
5341 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
5344 let mut row_ids = Vec::new();
5345 let overlay_rids = self.overlay_rid_set(snapshot);
5346 for rr in &self.run_refs {
5347 let mut reader = self.open_reader(rr.run_id)?;
5348 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
5349 for rid in matched {
5350 if !overlay_rids.contains(&rid) {
5351 row_ids.push(rid);
5352 }
5353 }
5354 }
5355 let mut s = RowIdSet::from_unsorted(row_ids);
5356 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
5357 Ok(s)
5358 }
5359
5360 fn null_scan_overlay(
5364 &self,
5365 s: &mut RowIdSet,
5366 column_id: u16,
5367 want_nulls: bool,
5368 snapshot: Snapshot,
5369 ) {
5370 let mut newest: HashMap<u64, &Row> = HashMap::new();
5371 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5372 let memtable = self.memtable.visible_versions(snapshot.epoch);
5373 for r in &mutable {
5374 newest.entry(r.row_id.0).or_insert(r);
5375 }
5376 for r in &memtable {
5377 newest.insert(r.row_id.0, r);
5378 }
5379 for row in newest.values() {
5380 if row.deleted {
5381 continue;
5382 }
5383 let is_null = !row.columns.contains_key(&column_id)
5384 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
5385 if is_null == want_nulls {
5386 s.insert(row.row_id.0);
5387 }
5388 }
5389 }
5390
5391 pub fn snapshot(&self) -> Snapshot {
5392 Snapshot::at(self.epoch.visible())
5393 }
5394
5395 pub fn data_generation(&self) -> u64 {
5397 self.data_generation
5398 }
5399
5400 pub fn pin_snapshot(&mut self) -> Snapshot {
5403 let e = self.epoch.visible();
5404 *self.pinned.entry(e).or_insert(0) += 1;
5405 Snapshot::at(e)
5406 }
5407
5408 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
5410 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
5411 *count -= 1;
5412 if *count == 0 {
5413 self.pinned.remove(&snap.epoch);
5414 }
5415 }
5416 }
5417
5418 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
5428 let local = self.pinned.keys().next().copied();
5429 let global = self.snapshots.min_pinned();
5430 let history = self.snapshots.history_floor(self.current_epoch());
5431 [local, global, history].into_iter().flatten().min()
5432 }
5433
5434 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
5438 self.ensure_writable()?;
5439 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
5440 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
5441 }
5442
5443 pub fn clear_ttl(&mut self) -> Result<()> {
5444 self.ensure_writable()?;
5445 self.apply_ttl_policy_at(None, self.current_epoch())
5446 }
5447
5448 pub fn ttl(&self) -> Option<TtlPolicy> {
5449 self.ttl
5450 }
5451
5452 pub(crate) fn prepare_ttl_policy(
5453 &self,
5454 column_name: &str,
5455 duration_nanos: u64,
5456 ) -> Result<TtlPolicy> {
5457 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
5458 return Err(MongrelError::InvalidArgument(
5459 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
5460 ));
5461 }
5462 let column = self
5463 .schema
5464 .columns
5465 .iter()
5466 .find(|column| column.name == column_name)
5467 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
5468 if column.ty != TypeId::TimestampNanos {
5469 return Err(MongrelError::Schema(format!(
5470 "TTL column {column_name} must be TimestampNanos, is {:?}",
5471 column.ty
5472 )));
5473 }
5474 Ok(TtlPolicy {
5475 column_id: column.id,
5476 duration_nanos,
5477 })
5478 }
5479
5480 pub(crate) fn apply_ttl_policy_at(
5481 &mut self,
5482 policy: Option<TtlPolicy>,
5483 epoch: Epoch,
5484 ) -> Result<()> {
5485 if let Some(policy) = policy {
5486 let column = self
5487 .schema
5488 .columns
5489 .iter()
5490 .find(|column| column.id == policy.column_id)
5491 .ok_or_else(|| {
5492 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
5493 })?;
5494 if column.ty != TypeId::TimestampNanos
5495 || policy.duration_nanos == 0
5496 || policy.duration_nanos > i64::MAX as u64
5497 {
5498 return Err(MongrelError::Schema("invalid TTL policy".into()));
5499 }
5500 }
5501 self.ttl = policy;
5502 self.agg_cache.clear();
5503 self.clear_result_cache();
5504 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5505 self.persist_manifest(epoch)
5506 }
5507
5508 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
5509 let Some(policy) = self.ttl else {
5510 return false;
5511 };
5512 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
5513 return false;
5514 };
5515 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
5516 }
5517
5518 pub fn current_epoch(&self) -> Epoch {
5519 self.epoch.visible()
5520 }
5521
5522 pub fn memtable_len(&self) -> usize {
5523 self.memtable.len()
5524 }
5525
5526 pub fn count(&self) -> u64 {
5529 if self.ttl.is_none()
5530 && self.pending_put_cols.is_empty()
5531 && self.pending_delete_rids.is_empty()
5532 && self.pending_rows.is_empty()
5533 && self.pending_dels.is_empty()
5534 && self.pending_truncate.is_none()
5535 {
5536 self.live_count
5537 } else {
5538 self.visible_rows(self.snapshot())
5539 .map(|rows| rows.len() as u64)
5540 .unwrap_or(self.live_count)
5541 }
5542 }
5543
5544 pub fn count_conditions(
5548 &mut self,
5549 conditions: &[crate::query::Condition],
5550 snapshot: Snapshot,
5551 ) -> Result<Option<u64>> {
5552 use crate::query::Condition;
5553 if self.ttl.is_some() {
5554 if conditions.is_empty() {
5555 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
5556 }
5557 let mut sets = Vec::with_capacity(conditions.len());
5558 for condition in conditions {
5559 sets.push(self.resolve_condition(condition, snapshot)?);
5560 }
5561 let survivors = RowIdSet::intersect_many(sets);
5562 let rows = self.visible_rows(snapshot)?;
5563 return Ok(Some(
5564 rows.into_iter()
5565 .filter(|row| survivors.contains(row.row_id.0))
5566 .count() as u64,
5567 ));
5568 }
5569 if conditions.is_empty() {
5570 return Ok(Some(self.count()));
5571 }
5572 let served = |c: &Condition| {
5573 matches!(
5574 c,
5575 Condition::Pk(_)
5576 | Condition::BitmapEq { .. }
5577 | Condition::BitmapIn { .. }
5578 | Condition::BytesPrefix { .. }
5579 | Condition::FmContains { .. }
5580 | Condition::FmContainsAll { .. }
5581 | Condition::Ann { .. }
5582 | Condition::Range { .. }
5583 | Condition::RangeF64 { .. }
5584 | Condition::SparseMatch { .. }
5585 | Condition::MinHashSimilar { .. }
5586 | Condition::IsNull { .. }
5587 | Condition::IsNotNull { .. }
5588 )
5589 };
5590 if !conditions.iter().all(served) {
5591 return Ok(None);
5592 }
5593 self.ensure_indexes_complete()?;
5594 if !self.pending_put_cols.is_empty()
5595 || !self.pending_delete_rids.is_empty()
5596 || !self.pending_rows.is_empty()
5597 || !self.pending_dels.is_empty()
5598 || self.pending_truncate.is_some()
5599 {
5600 let mut sets = Vec::with_capacity(conditions.len());
5601 for condition in conditions {
5602 sets.push(self.resolve_condition(condition, snapshot)?);
5603 }
5604 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5605 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
5606 }
5607 let mut sets = Vec::with_capacity(conditions.len());
5608 for condition in conditions {
5609 sets.push(self.resolve_condition(condition, snapshot)?);
5610 }
5611 let mut rids = RowIdSet::intersect_many(sets);
5612 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
5622 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
5623 }
5624 let count = rids.len() as u64;
5625 crate::trace::QueryTrace::record(|t| {
5626 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
5627 t.survivor_count = Some(count as usize);
5628 t.conditions_pushed = conditions.len();
5629 });
5630 Ok(Some(count))
5631 }
5632
5633 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
5638 let mut out = Vec::new();
5639 for row in self.memtable.visible_versions(snapshot.epoch) {
5640 if row.deleted {
5641 out.push(row.row_id.0);
5642 }
5643 }
5644 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5645 if row.deleted {
5646 out.push(row.row_id.0);
5647 }
5648 }
5649 out
5650 }
5651
5652 pub fn bulk_load_columns(
5661 &mut self,
5662 user_columns: Vec<(u16, columnar::NativeColumn)>,
5663 ) -> Result<Epoch> {
5664 self.bulk_load_columns_with(user_columns, 3, false, true)
5665 }
5666
5667 pub fn bulk_load_fast(
5674 &mut self,
5675 user_columns: Vec<(u16, columnar::NativeColumn)>,
5676 ) -> Result<Epoch> {
5677 self.bulk_load_columns_with(user_columns, -1, true, false)
5678 }
5679
5680 fn bulk_load_columns_with(
5681 &mut self,
5682 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
5683 zstd_level: i32,
5684 force_plain: bool,
5685 lz4: bool,
5686 ) -> Result<Epoch> {
5687 let epoch = self.commit()?;
5688 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
5689 if n == 0 {
5690 return Ok(epoch);
5691 }
5692 let live_before = self.live_count;
5693 self.spill_mutable_run(epoch)?;
5695 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
5696 && self.indexes_complete
5697 && self.run_refs.is_empty()
5698 && self.memtable.is_empty()
5699 && self.mutable_run.is_empty();
5700 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
5703 self.validate_columns_not_null(&user_columns, n)?;
5704 let winner_idx = self
5705 .bulk_pk_winner_indices(&user_columns, n)
5706 .filter(|idx| idx.len() != n);
5707 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
5708 match winner_idx.as_deref() {
5709 Some(idx) => {
5710 let compacted = user_columns
5711 .iter()
5712 .map(|(id, c)| (*id, c.gather(idx)))
5713 .collect();
5714 (compacted, idx.len())
5715 }
5716 None => (user_columns, n),
5717 };
5718 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
5719 let first = self.allocator.alloc_range(write_n as u64).0;
5720 for rid in first..first + write_n as u64 {
5721 self.reservoir.offer(rid);
5722 }
5723 let run_id = self.next_run_id;
5724 self.next_run_id += 1;
5725 let path = self.run_path(run_id);
5726 let mut writer =
5727 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
5728 if force_plain {
5729 writer = writer.with_plain();
5730 } else if lz4 {
5731 writer = writer.with_lz4();
5734 } else {
5735 writer = writer.with_zstd_level(zstd_level);
5736 }
5737 if let Some(kek) = &self.kek {
5738 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5739 }
5740 let header = writer.write_native(&path, &write_columns, write_n, first)?;
5741 self.run_refs.push(RunRef {
5742 run_id: run_id as u128,
5743 level: 0,
5744 epoch_created: epoch.0,
5745 row_count: header.row_count,
5746 });
5747 self.live_count = self.live_count.saturating_add(write_n as u64);
5748 if eager_index_build {
5749 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
5750 self.index_columns_bulk(&write_columns, &row_ids);
5751 self.indexes_complete = true;
5752 self.build_learned_ranges()?;
5753 } else {
5754 self.indexes_complete = false;
5758 }
5759 self.mark_flushed(epoch)?;
5760 self.persist_manifest(epoch)?;
5761 if eager_index_build {
5762 self.checkpoint_indexes(epoch);
5763 }
5764 self.clear_result_cache();
5765 self.data_generation = self.data_generation.wrapping_add(1);
5766 Ok(epoch)
5767 }
5768
5769 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
5787 let n = row_ids.len();
5788 if n == 0 {
5789 return;
5790 }
5791 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
5792 columns.iter().map(|(id, c)| (*id, c)).collect();
5793 let ty_of: std::collections::HashMap<u16, TypeId> = self
5794 .schema
5795 .columns
5796 .iter()
5797 .map(|c| (c.id, c.ty.clone()))
5798 .collect();
5799 let pk_id = self.schema.primary_key().map(|c| c.id);
5800
5801 for (i, &rid) in row_ids.iter().enumerate() {
5802 let row_id = RowId(rid);
5803 if let Some(pid) = pk_id {
5804 if let Some(col) = by_id.get(&pid) {
5805 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
5806 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
5807 self.insert_hot_pk(key, row_id);
5808 }
5809 }
5810 }
5811 for idef in &self.schema.indexes {
5812 let Some(col) = by_id.get(&idef.column_id) else {
5813 continue;
5814 };
5815 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
5816 match idef.kind {
5817 IndexKind::Bitmap => {
5818 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5819 if let Some(key) =
5820 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
5821 {
5822 b.insert(key, row_id);
5823 }
5824 }
5825 }
5826 IndexKind::FmIndex => {
5827 if let Some(f) = self.fm.get_mut(&idef.column_id) {
5828 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5829 f.insert(bytes.to_vec(), row_id);
5830 }
5831 }
5832 }
5833 IndexKind::Sparse => {
5834 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
5835 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5836 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
5837 s.insert(&terms, row_id);
5838 }
5839 }
5840 }
5841 }
5842 IndexKind::MinHash => {
5843 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
5844 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5845 let tokens = crate::index::token_hashes_from_bytes(bytes);
5846 mh.insert(&tokens, row_id);
5847 }
5848 }
5849 }
5850 _ => {}
5851 }
5852 }
5853 }
5854 }
5855
5856 pub fn visible_columns_native(
5861 &self,
5862 snapshot: Snapshot,
5863 projection: Option<&[u16]>,
5864 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
5865 let wanted: Vec<u16> = match projection {
5866 Some(p) => p.to_vec(),
5867 None => self.schema.columns.iter().map(|c| c.id).collect(),
5868 };
5869 if self.ttl.is_none()
5870 && self.memtable.is_empty()
5871 && self.mutable_run.is_empty()
5872 && self.run_refs.len() == 1
5873 {
5874 let rr = self.run_refs[0].clone();
5875 let mut reader = self.open_reader(rr.run_id)?;
5876 let idxs = reader.visible_indices_native(snapshot.epoch)?;
5877 let all_visible = idxs.len() == reader.row_count();
5878 if reader.has_mmap() {
5884 use rayon::prelude::*;
5885 let valid: Vec<u16> = wanted
5888 .iter()
5889 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
5890 .copied()
5891 .collect();
5892 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
5894 .par_iter()
5895 .filter_map(|cid| {
5896 reader
5897 .column_native_shared(*cid)
5898 .ok()
5899 .map(|col| (*cid, col))
5900 })
5901 .collect();
5902 let cols = decoded
5903 .into_iter()
5904 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
5905 .collect();
5906 return Ok(cols);
5907 }
5908 let mut cols = Vec::with_capacity(wanted.len());
5909 for cid in &wanted {
5910 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
5911 Some(c) => c,
5912 None => continue,
5913 };
5914 let col = reader.column_native(cdef.id)?;
5915 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
5916 }
5917 return Ok(cols);
5918 }
5919 let vcols = self.visible_columns(snapshot)?;
5920 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
5921 let out: Vec<(u16, columnar::NativeColumn)> = vcols
5922 .into_iter()
5923 .filter(|(id, _)| want_set.contains(id))
5924 .map(|(id, vals)| {
5925 let ty = self
5926 .schema
5927 .columns
5928 .iter()
5929 .find(|c| c.id == id)
5930 .map(|c| c.ty.clone())
5931 .unwrap_or(TypeId::Bytes);
5932 (id, columnar::values_to_native(ty, &vals))
5933 })
5934 .collect();
5935 Ok(out)
5936 }
5937
5938 pub fn run_count(&self) -> usize {
5939 self.run_refs.len()
5940 }
5941
5942 pub fn memtable_is_empty(&self) -> bool {
5944 self.memtable.is_empty()
5945 }
5946
5947 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
5951 self.page_cache.stats()
5952 }
5953
5954 pub fn reset_page_cache_stats(&self) {
5956 self.page_cache.reset_stats();
5957 }
5958
5959 pub fn run_ids(&self) -> Vec<u128> {
5962 self.run_refs.iter().map(|r| r.run_id).collect()
5963 }
5964
5965 pub fn single_run_is_clean(&self) -> bool {
5969 if self.ttl.is_some() || self.run_refs.len() != 1 {
5970 return false;
5971 }
5972 self.open_reader(self.run_refs[0].run_id)
5973 .map(|r| r.is_clean())
5974 .unwrap_or(false)
5975 }
5976
5977 fn resolve_footprint(
5984 &self,
5985 conditions: &[crate::query::Condition],
5986 snapshot: Snapshot,
5987 ) -> roaring::RoaringBitmap {
5988 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
5989 return roaring::RoaringBitmap::new();
5990 }
5991 if self.run_refs.is_empty() {
5992 return roaring::RoaringBitmap::new();
5993 }
5994 if self.run_refs.len() == 1 {
5996 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
5997 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
5998 return rids.to_roaring_lossy();
5999 }
6000 }
6001 }
6002 roaring::RoaringBitmap::new()
6003 }
6004
6005 pub fn query_columns_native_cached(
6016 &mut self,
6017 conditions: &[crate::query::Condition],
6018 projection: Option<&[u16]>,
6019 snapshot: Snapshot,
6020 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
6021 if self.ttl.is_some() {
6024 return self.query_columns_native(conditions, projection, snapshot);
6025 }
6026 if conditions.is_empty() {
6027 return self.query_columns_native(conditions, projection, snapshot);
6028 }
6029 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
6033 if let Some(hit) = self.result_cache.lock().get_columns(key) {
6034 crate::trace::QueryTrace::record(|t| {
6035 t.result_cache_hit = true;
6036 t.scan_mode = crate::trace::ScanMode::NativePushdown;
6037 });
6038 return Ok(Some((*hit).clone()));
6039 }
6040 let res = self.query_columns_native(conditions, projection, snapshot)?;
6041 if let Some(cols) = &res {
6042 let footprint = self.resolve_footprint(conditions, snapshot);
6043 let condition_cols = crate::query::condition_columns(conditions);
6044 self.result_cache.lock().insert(
6045 key,
6046 CachedEntry {
6047 data: CachedData::Columns(Arc::new(cols.clone())),
6048 footprint,
6049 condition_cols,
6050 },
6051 );
6052 }
6053 Ok(res)
6054 }
6055
6056 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
6061 if self.ttl.is_some() {
6062 return self.query(q);
6063 }
6064 if q.conditions.is_empty() {
6065 return self.query(q);
6066 }
6067 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
6068 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15);
6069 if let Some(hit) = self.result_cache.lock().get_rows(key) {
6070 crate::trace::QueryTrace::record(|t| {
6071 t.result_cache_hit = true;
6072 t.scan_mode = crate::trace::ScanMode::Materialized;
6073 });
6074 return Ok((*hit).clone());
6075 }
6076 let rows = self.query(q)?;
6077 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
6078 let condition_cols = crate::query::condition_columns(&q.conditions);
6079 self.result_cache.lock().insert(
6080 key,
6081 CachedEntry {
6082 data: CachedData::Rows(Arc::new(rows.clone())),
6083 footprint,
6084 condition_cols,
6085 },
6086 );
6087 Ok(rows)
6088 }
6089
6090 #[allow(clippy::type_complexity)]
6105 pub fn query_columns_native_traced(
6106 &mut self,
6107 conditions: &[crate::query::Condition],
6108 projection: Option<&[u16]>,
6109 snapshot: Snapshot,
6110 ) -> Result<(
6111 Option<Vec<(u16, columnar::NativeColumn)>>,
6112 crate::trace::QueryTrace,
6113 )> {
6114 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6115 self.query_columns_native(conditions, projection, snapshot)
6116 });
6117 Ok((result?, trace))
6118 }
6119
6120 #[allow(clippy::type_complexity)]
6123 pub fn query_columns_native_cached_traced(
6124 &mut self,
6125 conditions: &[crate::query::Condition],
6126 projection: Option<&[u16]>,
6127 snapshot: Snapshot,
6128 ) -> Result<(
6129 Option<Vec<(u16, columnar::NativeColumn)>>,
6130 crate::trace::QueryTrace,
6131 )> {
6132 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6133 self.query_columns_native_cached(conditions, projection, snapshot)
6134 });
6135 Ok((result?, trace))
6136 }
6137
6138 pub fn native_page_cursor_traced(
6140 &self,
6141 snapshot: Snapshot,
6142 projection: Vec<(u16, TypeId)>,
6143 conditions: &[crate::query::Condition],
6144 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
6145 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6146 self.native_page_cursor(snapshot, projection, conditions)
6147 });
6148 Ok((result?, trace))
6149 }
6150
6151 pub fn native_multi_run_cursor_traced(
6153 &self,
6154 snapshot: Snapshot,
6155 projection: Vec<(u16, TypeId)>,
6156 conditions: &[crate::query::Condition],
6157 ) -> Result<(
6158 Option<crate::cursor::MultiRunCursor>,
6159 crate::trace::QueryTrace,
6160 )> {
6161 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6162 self.native_multi_run_cursor(snapshot, projection, conditions)
6163 });
6164 Ok((result?, trace))
6165 }
6166
6167 pub fn count_conditions_traced(
6169 &mut self,
6170 conditions: &[crate::query::Condition],
6171 snapshot: Snapshot,
6172 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
6173 let (result, trace) =
6174 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
6175 Ok((result?, trace))
6176 }
6177
6178 pub fn query_traced(
6180 &mut self,
6181 q: &crate::query::Query,
6182 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
6183 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
6184 Ok((result?, trace))
6185 }
6186
6187 pub fn query_columns_native(
6192 &mut self,
6193 conditions: &[crate::query::Condition],
6194 projection: Option<&[u16]>,
6195 snapshot: Snapshot,
6196 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
6197 use crate::query::Condition;
6198 if self.ttl.is_some() {
6201 return Ok(None);
6202 }
6203 if conditions.is_empty() {
6204 return Ok(None);
6205 }
6206 self.ensure_indexes_complete()?;
6207
6208 let served = |c: &Condition| {
6213 matches!(
6214 c,
6215 Condition::Pk(_)
6216 | Condition::BitmapEq { .. }
6217 | Condition::BitmapIn { .. }
6218 | Condition::BytesPrefix { .. }
6219 | Condition::FmContains { .. }
6220 | Condition::FmContainsAll { .. }
6221 | Condition::Ann { .. }
6222 | Condition::Range { .. }
6223 | Condition::RangeF64 { .. }
6224 | Condition::SparseMatch { .. }
6225 | Condition::MinHashSimilar { .. }
6226 | Condition::IsNull { .. }
6227 | Condition::IsNotNull { .. }
6228 )
6229 };
6230 if !conditions.iter().all(served) {
6231 return Ok(None);
6232 }
6233 let fast_path =
6234 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
6235 crate::trace::QueryTrace::record(|t| {
6236 t.run_count = self.run_refs.len();
6237 t.memtable_rows = self.memtable.len();
6238 t.mutable_run_rows = self.mutable_run.len();
6239 t.conditions_pushed = conditions.len();
6240 t.learned_range_used = conditions.iter().any(|c| match c {
6241 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
6242 self.learned_range.contains_key(column_id)
6243 }
6244 _ => false,
6245 });
6246 });
6247 let col_ids: Vec<u16> = projection
6249 .map(|p| p.to_vec())
6250 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
6251 let proj_pairs: Vec<(u16, TypeId)> = col_ids
6252 .iter()
6253 .map(|&cid| {
6254 let ty = self
6255 .schema
6256 .columns
6257 .iter()
6258 .find(|c| c.id == cid)
6259 .map(|c| c.ty.clone())
6260 .unwrap_or(TypeId::Bytes);
6261 (cid, ty)
6262 })
6263 .collect();
6264
6265 if fast_path {
6271 let needs_column = conditions.iter().any(|c| match c {
6274 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
6275 Condition::RangeF64 { column_id, .. } => {
6276 !self.learned_range.contains_key(column_id)
6277 }
6278 _ => false,
6279 });
6280 let mut reader_opt: Option<RunReader> = if needs_column {
6281 Some(self.open_reader(self.run_refs[0].run_id)?)
6282 } else {
6283 None
6284 };
6285 let mut sets: Vec<RowIdSet> = Vec::new();
6286 for c in conditions {
6287 let s = match c {
6288 Condition::Range { column_id, lo, hi }
6289 if !self.learned_range.contains_key(column_id) =>
6290 {
6291 if reader_opt.is_none() {
6292 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
6293 }
6294 reader_opt
6295 .as_mut()
6296 .expect("reader opened for range")
6297 .range_row_id_set_i64(*column_id, *lo, *hi)?
6298 }
6299 Condition::RangeF64 {
6300 column_id,
6301 lo,
6302 lo_inclusive,
6303 hi,
6304 hi_inclusive,
6305 } if !self.learned_range.contains_key(column_id) => {
6306 if reader_opt.is_none() {
6307 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
6308 }
6309 reader_opt
6310 .as_mut()
6311 .expect("reader opened for range")
6312 .range_row_id_set_f64(
6313 *column_id,
6314 *lo,
6315 *lo_inclusive,
6316 *hi,
6317 *hi_inclusive,
6318 )?
6319 }
6320 _ => self.resolve_condition(c, snapshot)?,
6321 };
6322 sets.push(s);
6323 }
6324 let candidates = RowIdSet::intersect_many(sets);
6325 crate::trace::QueryTrace::record(|t| {
6326 t.survivor_count = Some(candidates.len());
6327 });
6328 if candidates.is_empty() {
6329 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
6330 .iter()
6331 .map(|&id| {
6332 (
6333 id,
6334 columnar::null_native(
6335 proj_pairs
6336 .iter()
6337 .find(|(c, _)| c == &id)
6338 .map(|(_, t)| t.clone())
6339 .unwrap_or(TypeId::Bytes),
6340 0,
6341 ),
6342 )
6343 })
6344 .collect();
6345 return Ok(Some(cols));
6346 }
6347 let mut reader = match reader_opt.take() {
6348 Some(r) => r,
6349 None => self.open_reader(self.run_refs[0].run_id)?,
6350 };
6351 let candidate_ids = candidates.into_sorted_vec();
6352 let (positions, fast_rid) = if let Some(positions) =
6353 reader.positions_for_row_ids_fast(&candidate_ids)
6354 {
6355 (positions, true)
6356 } else {
6357 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
6358 match col {
6359 columnar::NativeColumn::Int64 { data, .. } => {
6360 let mut p: Vec<usize> = candidate_ids
6361 .iter()
6362 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
6363 .collect();
6364 p.sort_unstable();
6365 (p, false)
6366 }
6367 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
6368 }
6369 };
6370 crate::trace::QueryTrace::record(|t| {
6371 t.scan_mode = crate::trace::ScanMode::NativePushdown;
6372 t.fast_row_id_map = fast_rid;
6373 });
6374 let mut cols = Vec::with_capacity(col_ids.len());
6375 for cid in &col_ids {
6376 let col = reader.column_native(*cid)?;
6377 cols.push((*cid, col.gather(&positions)));
6378 }
6379 return Ok(Some(cols));
6380 }
6381
6382 if !self.run_refs.is_empty() {
6395 use crate::cursor::{drain_cursor_to_columns, Cursor};
6396 let remaining: usize;
6397 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
6398 let c = self
6399 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
6400 .expect("single-run cursor should build when run_refs.len() == 1");
6401 remaining = c.remaining_rows();
6402 Box::new(c)
6403 } else {
6404 let c = self
6405 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
6406 .expect("multi-run cursor should build when run_refs.len() >= 1");
6407 remaining = c.remaining_rows();
6408 Box::new(c)
6409 };
6410 crate::trace::QueryTrace::record(|t| {
6411 if t.survivor_count.is_none() {
6412 t.survivor_count = Some(remaining);
6413 }
6414 });
6415 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
6416 return Ok(Some(cols));
6417 }
6418
6419 crate::trace::QueryTrace::record(|t| {
6424 t.scan_mode = crate::trace::ScanMode::Materialized;
6425 t.row_materialized = true;
6426 });
6427 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6428 for c in conditions {
6429 sets.push(self.resolve_condition(c, snapshot)?);
6430 }
6431 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6432 let rows = self.rows_for_rids(&rids, snapshot)?;
6433 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
6434 for (cid, ty) in &proj_pairs {
6435 let vals: Vec<Value> = rows
6436 .iter()
6437 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6438 .collect();
6439 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
6440 }
6441 Ok(Some(cols))
6442 }
6443
6444 pub fn native_page_cursor(
6459 &self,
6460 snapshot: Snapshot,
6461 projection: Vec<(u16, TypeId)>,
6462 conditions: &[crate::query::Condition],
6463 ) -> Result<Option<NativePageCursor>> {
6464 use crate::cursor::build_page_plans;
6465 if self.ttl.is_some() {
6466 return Ok(None);
6467 }
6468 if !conditions.is_empty() && !self.indexes_complete {
6471 return Ok(None);
6472 }
6473 if self.run_refs.len() != 1 {
6474 return Ok(None);
6475 }
6476 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6477 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6478
6479 let overlay_rids: HashSet<u64> = {
6482 let mut s = HashSet::new();
6483 for row in self.memtable.visible_versions(snapshot.epoch) {
6484 s.insert(row.row_id.0);
6485 }
6486 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6487 s.insert(row.row_id.0);
6488 }
6489 s
6490 };
6491
6492 let survivors = if conditions.is_empty() {
6496 None
6497 } else {
6498 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
6499 };
6500
6501 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
6508 survivors.clone()
6509 } else if let Some(s) = &survivors {
6510 let mut run_set = s.clone();
6511 run_set.remove_many(overlay_rids.iter().copied());
6512 Some(run_set)
6513 } else {
6514 Some(RowIdSet::from_unsorted(
6515 rids.iter()
6516 .map(|&r| r as u64)
6517 .filter(|r| !overlay_rids.contains(r))
6518 .collect(),
6519 ))
6520 };
6521
6522 let overlay_rows = if overlay_rids.is_empty() {
6523 Vec::new()
6524 } else {
6525 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6526 self.overlay_visible_rows(snapshot, bound)
6527 };
6528
6529 let plans = if positions.is_empty() {
6531 Vec::new()
6532 } else {
6533 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
6534 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
6535 };
6536
6537 let overlay = if overlay_rows.is_empty() {
6539 None
6540 } else {
6541 let filtered =
6542 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6543 if filtered.is_empty() {
6544 None
6545 } else {
6546 Some(self.materialize_overlay(&filtered, &projection))
6547 }
6548 };
6549
6550 let overlay_row_count = overlay
6551 .as_ref()
6552 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6553 .unwrap_or(0);
6554 crate::trace::QueryTrace::record(|t| {
6555 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
6556 t.run_count = self.run_refs.len();
6557 t.memtable_rows = self.memtable.len();
6558 t.mutable_run_rows = self.mutable_run.len();
6559 t.overlay_rows = overlay_row_count;
6560 t.conditions_pushed = conditions.len();
6561 t.pages_decoded = plans
6562 .iter()
6563 .map(|p| p.positions.len())
6564 .sum::<usize>()
6565 .min(1);
6566 });
6567
6568 Ok(Some(NativePageCursor::new_with_overlay(
6569 reader, projection, plans, overlay,
6570 )))
6571 }
6572 #[allow(clippy::type_complexity)]
6582 pub fn native_multi_run_cursor(
6583 &self,
6584 snapshot: Snapshot,
6585 projection: Vec<(u16, TypeId)>,
6586 conditions: &[crate::query::Condition],
6587 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
6588 use crate::cursor::{MultiRunCursor, RunStream};
6589 use crate::sorted_run::SYS_ROW_ID;
6590 use std::collections::{BinaryHeap, HashMap, HashSet};
6591 if self.ttl.is_some() {
6592 return Ok(None);
6593 }
6594 if !conditions.is_empty() && !self.indexes_complete {
6597 return Ok(None);
6598 }
6599 if self.run_refs.is_empty() {
6600 return Ok(None);
6601 }
6602
6603 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
6605 Vec::with_capacity(self.run_refs.len());
6606 for rr in &self.run_refs {
6607 let mut reader = self.open_reader(rr.run_id)?;
6608 let (rids, eps, del) = reader.system_columns_native()?;
6609 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
6610 run_meta.push((reader, rids, eps, del, page_rows));
6611 }
6612
6613 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
6617 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
6618 for i in 0..rids.len() {
6619 let rid = rids[i] as u64;
6620 let e = eps[i] as u64;
6621 if e > snapshot.epoch.0 {
6622 continue;
6623 }
6624 let is_del = del[i] != 0;
6625 best.entry(rid)
6626 .and_modify(|cur| {
6627 if e > cur.0 {
6628 *cur = (e, run_idx, i, is_del);
6629 }
6630 })
6631 .or_insert((e, run_idx, i, is_del));
6632 }
6633 }
6634
6635 let overlay_rids: HashSet<u64> = {
6637 let mut s = HashSet::new();
6638 for row in self.memtable.visible_versions(snapshot.epoch) {
6639 s.insert(row.row_id.0);
6640 }
6641 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6642 s.insert(row.row_id.0);
6643 }
6644 s
6645 };
6646
6647 let survivors: Option<RowIdSet> = if conditions.is_empty() {
6649 None
6650 } else {
6651 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6652 for c in conditions {
6653 sets.push(self.resolve_condition(c, snapshot)?);
6654 }
6655 Some(RowIdSet::intersect_many(sets))
6656 };
6657
6658 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
6662 for (rid, (_, run_idx, pos, deleted)) in &best {
6663 if *deleted {
6664 continue;
6665 }
6666 if overlay_rids.contains(rid) {
6667 continue;
6668 }
6669 if let Some(s) = &survivors {
6670 if !s.contains(*rid) {
6671 continue;
6672 }
6673 }
6674 per_run[*run_idx].push((*rid, *pos));
6675 }
6676 for v in per_run.iter_mut() {
6677 v.sort_unstable_by_key(|&(rid, _)| rid);
6678 }
6679
6680 let mut streams = Vec::with_capacity(run_meta.len());
6682 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
6683 let mut total = 0usize;
6684 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
6685 let mut starts = Vec::with_capacity(page_rows.len());
6686 let mut acc = 0usize;
6687 for &r in &page_rows {
6688 starts.push(acc);
6689 acc += r;
6690 }
6691 let mut survivors_vec: Vec<(u64, usize, usize)> =
6692 Vec::with_capacity(per_run[run_idx].len());
6693 for &(rid, pos) in &per_run[run_idx] {
6694 let page_seq = match starts.partition_point(|&s| s <= pos) {
6695 0 => continue,
6696 p => p - 1,
6697 };
6698 let within = pos - starts[page_seq];
6699 survivors_vec.push((rid, page_seq, within));
6700 }
6701 total += survivors_vec.len();
6702 if let Some(&(rid, _, _)) = survivors_vec.first() {
6703 heap.push(std::cmp::Reverse((rid, run_idx)));
6704 }
6705 streams.push(RunStream::new(reader, survivors_vec, page_rows));
6706 }
6707
6708 let overlay_rows = if overlay_rids.is_empty() {
6710 Vec::new()
6711 } else {
6712 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6713 self.overlay_visible_rows(snapshot, bound)
6714 };
6715 let overlay = if overlay_rows.is_empty() {
6716 None
6717 } else {
6718 let filtered =
6719 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6720 if filtered.is_empty() {
6721 None
6722 } else {
6723 Some(self.materialize_overlay(&filtered, &projection))
6724 }
6725 };
6726
6727 let overlay_row_count = overlay
6728 .as_ref()
6729 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6730 .unwrap_or(0);
6731 crate::trace::QueryTrace::record(|t| {
6732 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
6733 t.run_count = self.run_refs.len();
6734 t.memtable_rows = self.memtable.len();
6735 t.mutable_run_rows = self.mutable_run.len();
6736 t.overlay_rows = overlay_row_count;
6737 t.conditions_pushed = conditions.len();
6738 t.survivor_count = Some(total);
6739 });
6740
6741 Ok(Some(MultiRunCursor::new(
6742 streams, projection, heap, total, overlay,
6743 )))
6744 }
6745
6746 fn overlay_materialization_bound<'a>(
6758 conditions: &[crate::query::Condition],
6759 survivors: &'a Option<RowIdSet>,
6760 ) -> Option<&'a RowIdSet> {
6761 use crate::query::Condition;
6762 let has_range = conditions
6763 .iter()
6764 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6765 if has_range {
6766 None
6767 } else {
6768 survivors.as_ref()
6769 }
6770 }
6771
6772 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
6784 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
6785 let mut fold = |row: Row| {
6786 if let Some(b) = bound {
6787 if !b.contains(row.row_id.0) {
6788 return;
6789 }
6790 }
6791 best.entry(row.row_id.0)
6792 .and_modify(|(be, br)| {
6793 if row.committed_epoch > *be {
6794 *be = row.committed_epoch;
6795 *br = row.clone();
6796 }
6797 })
6798 .or_insert_with(|| (row.committed_epoch, row));
6799 };
6800 for row in self.memtable.visible_versions(snapshot.epoch) {
6801 fold(row);
6802 }
6803 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6804 fold(row);
6805 }
6806 let mut out: Vec<Row> = best
6807 .into_values()
6808 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
6809 .collect();
6810 out.sort_by_key(|r| r.row_id);
6811 out
6812 }
6813
6814 fn filter_overlay_rows(
6822 &self,
6823 rows: Vec<Row>,
6824 conditions: &[crate::query::Condition],
6825 survivors: Option<&RowIdSet>,
6826 snapshot: Snapshot,
6827 ) -> Result<Vec<Row>> {
6828 if conditions.is_empty() {
6829 return Ok(rows);
6830 }
6831 use crate::query::Condition;
6832 let all_index_served = !conditions
6836 .iter()
6837 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6838 if all_index_served {
6839 return Ok(rows
6840 .into_iter()
6841 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
6842 .collect());
6843 }
6844 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6847 for c in conditions {
6848 let s = match c {
6849 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
6850 _ => self.resolve_condition(c, snapshot)?,
6851 };
6852 per_cond_sets.push(s);
6853 }
6854 Ok(rows
6855 .into_iter()
6856 .filter(|row| {
6857 conditions.iter().enumerate().all(|(i, c)| match c {
6858 Condition::Range { column_id, lo, hi } => {
6859 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
6860 }
6861 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
6862 match row.columns.get(column_id) {
6863 Some(Value::Float64(v)) => {
6864 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
6865 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
6866 lo_ok && hi_ok
6867 }
6868 _ => false,
6869 }
6870 }
6871 _ => per_cond_sets[i].contains(row.row_id.0),
6872 })
6873 })
6874 .collect())
6875 }
6876
6877 fn materialize_overlay(
6880 &self,
6881 rows: &[Row],
6882 projection: &[(u16, TypeId)],
6883 ) -> Vec<columnar::NativeColumn> {
6884 if projection.is_empty() {
6885 return vec![columnar::null_native(TypeId::Int64, rows.len())];
6886 }
6887 let mut cols = Vec::with_capacity(projection.len());
6888 for (cid, ty) in projection {
6889 let vals: Vec<Value> = rows
6890 .iter()
6891 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6892 .collect();
6893 cols.push(columnar::values_to_native(ty.clone(), &vals));
6894 }
6895 cols
6896 }
6897
6898 fn resolve_survivor_rids(
6903 &self,
6904 conditions: &[crate::query::Condition],
6905 reader: &mut RunReader,
6906 snapshot: Snapshot,
6907 ) -> Result<RowIdSet> {
6908 use crate::query::Condition;
6909 let mut sets: Vec<RowIdSet> = Vec::new();
6910 for c in conditions {
6911 self.validate_condition(c)?;
6912 let s: RowIdSet = match c {
6913 Condition::Pk(key) => {
6914 let lookup = self
6915 .schema
6916 .primary_key()
6917 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
6918 .unwrap_or_else(|| key.clone());
6919 self.hot
6920 .get(&lookup)
6921 .map(|r| RowIdSet::one(r.0))
6922 .unwrap_or_else(RowIdSet::empty)
6923 }
6924 Condition::BitmapEq { column_id, value } => {
6925 let lookup = self.index_lookup_key_bytes(*column_id, value);
6926 self.bitmap
6927 .get(column_id)
6928 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
6929 .unwrap_or_else(RowIdSet::empty)
6930 }
6931 Condition::BitmapIn { column_id, values } => {
6932 let bm = self.bitmap.get(column_id);
6933 let mut acc = roaring::RoaringBitmap::new();
6934 if let Some(b) = bm {
6935 for v in values {
6936 let lookup = self.index_lookup_key_bytes(*column_id, v);
6937 acc |= b.get(&lookup);
6938 }
6939 }
6940 RowIdSet::from_roaring(acc)
6941 }
6942 Condition::BytesPrefix { column_id, prefix } => {
6943 if let Some(b) = self.bitmap.get(column_id) {
6944 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
6945 let mut acc = roaring::RoaringBitmap::new();
6946 for key in b.keys() {
6947 if key.starts_with(&lookup_prefix) {
6948 acc |= b.get(key);
6949 }
6950 }
6951 RowIdSet::from_roaring(acc)
6952 } else {
6953 RowIdSet::empty()
6954 }
6955 }
6956 Condition::FmContains { column_id, pattern } => self
6957 .fm
6958 .get(column_id)
6959 .map(|f| {
6960 RowIdSet::from_unsorted(
6961 f.locate(pattern).into_iter().map(|r| r.0).collect(),
6962 )
6963 })
6964 .unwrap_or_else(RowIdSet::empty),
6965 Condition::FmContainsAll {
6966 column_id,
6967 patterns,
6968 } => {
6969 if let Some(f) = self.fm.get(column_id) {
6970 let sets: Vec<RowIdSet> = patterns
6971 .iter()
6972 .map(|pat| {
6973 RowIdSet::from_unsorted(
6974 f.locate(pat).into_iter().map(|r| r.0).collect(),
6975 )
6976 })
6977 .collect();
6978 RowIdSet::intersect_many(sets)
6979 } else {
6980 RowIdSet::empty()
6981 }
6982 }
6983 Condition::Ann {
6984 column_id,
6985 query,
6986 k,
6987 } => RowIdSet::from_unsorted(
6988 self.retrieve_filtered(
6989 &crate::query::Retriever::Ann {
6990 column_id: *column_id,
6991 query: query.clone(),
6992 k: *k,
6993 },
6994 snapshot,
6995 None,
6996 None,
6997 )?
6998 .into_iter()
6999 .map(|hit| hit.row_id.0)
7000 .collect(),
7001 ),
7002 Condition::SparseMatch {
7003 column_id,
7004 query,
7005 k,
7006 } => RowIdSet::from_unsorted(
7007 self.retrieve_filtered(
7008 &crate::query::Retriever::Sparse {
7009 column_id: *column_id,
7010 query: query.clone(),
7011 k: *k,
7012 },
7013 snapshot,
7014 None,
7015 None,
7016 )?
7017 .into_iter()
7018 .map(|hit| hit.row_id.0)
7019 .collect(),
7020 ),
7021 Condition::MinHashSimilar {
7022 column_id,
7023 query,
7024 k,
7025 } => match self.minhash.get(column_id) {
7026 Some(index) => {
7027 let candidates = index.candidate_row_ids(query);
7028 let eligible =
7029 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
7030 RowIdSet::from_unsorted(
7031 index
7032 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
7033 .into_iter()
7034 .map(|(row_id, _)| row_id.0)
7035 .collect(),
7036 )
7037 }
7038 None => RowIdSet::empty(),
7039 },
7040 Condition::Range { column_id, lo, hi } => {
7041 if let Some(li) = self.learned_range.get(column_id) {
7042 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7043 } else {
7044 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
7045 }
7046 }
7047 Condition::RangeF64 {
7048 column_id,
7049 lo,
7050 lo_inclusive,
7051 hi,
7052 hi_inclusive,
7053 } => {
7054 if let Some(li) = self.learned_range.get(column_id) {
7055 RowIdSet::from_unsorted(
7056 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7057 .into_iter()
7058 .collect(),
7059 )
7060 } else {
7061 reader.range_row_id_set_f64(
7062 *column_id,
7063 *lo,
7064 *lo_inclusive,
7065 *hi,
7066 *hi_inclusive,
7067 )?
7068 }
7069 }
7070 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
7071 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
7072 };
7073 sets.push(s);
7074 }
7075 Ok(RowIdSet::intersect_many(sets))
7076 }
7077
7078 pub fn scan_cursor(
7099 &self,
7100 snapshot: Snapshot,
7101 projection: Vec<(u16, TypeId)>,
7102 conditions: &[crate::query::Condition],
7103 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
7104 if self.ttl.is_some() {
7105 return Ok(None);
7106 }
7107 if !conditions.is_empty() && !self.indexes_complete {
7113 return Ok(None);
7114 }
7115 if self.run_refs.len() == 1 {
7116 Ok(self
7117 .native_page_cursor(snapshot, projection, conditions)?
7118 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7119 } else {
7120 Ok(self
7121 .native_multi_run_cursor(snapshot, projection, conditions)?
7122 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7123 }
7124 }
7125
7126 pub fn aggregate_native(
7140 &self,
7141 snapshot: Snapshot,
7142 column: Option<u16>,
7143 conditions: &[crate::query::Condition],
7144 agg: NativeAgg,
7145 ) -> Result<Option<NativeAggResult>> {
7146 if self.ttl.is_some() {
7147 return Ok(None);
7148 }
7149 if self.run_refs.len() == 1 && conditions.is_empty() {
7151 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
7152 return Ok(Some(res));
7153 }
7154 }
7155 if matches!(agg, NativeAgg::Count) && column.is_none() {
7157 return Ok(self
7158 .scan_cursor(snapshot, Vec::new(), conditions)?
7159 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
7160 }
7161 let cid = match column {
7164 Some(c) => c,
7165 None => return Ok(None),
7166 };
7167 let ty = self.column_type(cid);
7168 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
7169 else {
7170 return Ok(None);
7171 };
7172 match ty {
7173 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
7174 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
7175 Ok(Some(pack_int(agg, count, sum, mn, mx)))
7176 }
7177 TypeId::Float64 => {
7178 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
7179 Ok(Some(pack_float(agg, count, sum, mn, mx)))
7180 }
7181 _ => Ok(None),
7182 }
7183 }
7184
7185 fn aggregate_from_stats(
7193 &self,
7194 snapshot: Snapshot,
7195 column: Option<u16>,
7196 agg: NativeAgg,
7197 ) -> Result<Option<NativeAggResult>> {
7198 let cid = match (agg, column) {
7199 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
7200 _ => return Ok(None), };
7202 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
7203 return Ok(None);
7204 };
7205 let Some(cs) = stats.get(&cid) else {
7206 return Ok(None);
7207 };
7208 match agg {
7209 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
7211 self.live_count.saturating_sub(cs.null_count),
7212 ))),
7213 NativeAgg::Min | NativeAgg::Max => {
7214 let bound = if agg == NativeAgg::Min {
7215 &cs.min
7216 } else {
7217 &cs.max
7218 };
7219 match bound {
7220 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
7221 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
7222 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
7227 None => Ok(None),
7228 }
7229 }
7230 _ => Ok(None),
7231 }
7232 }
7233
7234 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
7243 if self.ttl.is_some() {
7244 return Ok(None);
7245 }
7246 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
7247 return Ok(None);
7248 }
7249 self.ensure_indexes_complete()?;
7252 let reader = self.open_reader(self.run_refs[0].run_id)?;
7253 if self.live_count != reader.row_count() as u64 {
7254 return Ok(None);
7255 }
7256 let Some(bm) = self.bitmap.get(&column_id) else {
7257 return Ok(None); };
7259 let mut distinct = bm.value_count() as u64;
7260 if !bm.get(&Value::Null.encode_key()).is_empty() {
7263 distinct = distinct.saturating_sub(1);
7264 }
7265 Ok(Some(distinct))
7266 }
7267
7268 pub fn aggregate_incremental(
7280 &mut self,
7281 cache_key: u64,
7282 conditions: &[crate::query::Condition],
7283 column: Option<u16>,
7284 agg: NativeAgg,
7285 ) -> Result<IncrementalAggResult> {
7286 let snap = self.snapshot();
7287 let cur_wm = self.allocator.current().0;
7288 let cur_epoch = snap.epoch.0;
7289 let incremental_ok = self.ttl.is_none()
7296 && !self.had_deletes
7297 && self.memtable.is_empty()
7298 && self.mutable_run.is_empty();
7299
7300 if incremental_ok {
7303 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
7304 if cached.epoch == cur_epoch {
7305 return Ok(IncrementalAggResult {
7306 state: cached.state,
7307 incremental: true,
7308 delta_rows: 0,
7309 });
7310 }
7311 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
7312 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
7313 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
7314 let index_sets = self.resolve_index_conditions(conditions, snap)?;
7315 let delta_state = agg_state_from_rows(
7316 &delta_rows,
7317 conditions,
7318 &index_sets,
7319 column,
7320 agg,
7321 &self.schema,
7322 )?;
7323 let merged = cached.state.merge(delta_state);
7324 let delta_n = delta_rids.len() as u64;
7325 self.agg_cache.insert(
7326 cache_key,
7327 CachedAgg {
7328 state: merged.clone(),
7329 watermark: cur_wm,
7330 epoch: cur_epoch,
7331 },
7332 );
7333 return Ok(IncrementalAggResult {
7334 state: merged,
7335 incremental: true,
7336 delta_rows: delta_n,
7337 });
7338 }
7339 }
7340 }
7341
7342 let cursor_ok =
7347 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
7348 let state = if cursor_ok && agg != NativeAgg::Avg {
7349 match self.aggregate_native(snap, column, conditions, agg)? {
7350 Some(result) => {
7351 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
7352 }
7353 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
7354 }
7355 } else {
7356 self.agg_state_full_scan(conditions, column, agg, snap)?
7357 };
7358 if incremental_ok {
7360 self.agg_cache.insert(
7361 cache_key,
7362 CachedAgg {
7363 state: state.clone(),
7364 watermark: cur_wm,
7365 epoch: cur_epoch,
7366 },
7367 );
7368 }
7369 Ok(IncrementalAggResult {
7370 state,
7371 incremental: false,
7372 delta_rows: 0,
7373 })
7374 }
7375
7376 fn agg_state_full_scan(
7379 &self,
7380 conditions: &[crate::query::Condition],
7381 column: Option<u16>,
7382 agg: NativeAgg,
7383 snap: Snapshot,
7384 ) -> Result<AggState> {
7385 let rows = self.visible_rows(snap)?;
7386 let index_sets = self.resolve_index_conditions(conditions, snap)?;
7387 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
7388 }
7389
7390 fn resolve_index_conditions(
7393 &self,
7394 conditions: &[crate::query::Condition],
7395 snapshot: Snapshot,
7396 ) -> Result<Vec<RowIdSet>> {
7397 use crate::query::Condition;
7398 let mut sets = Vec::new();
7399 for c in conditions {
7400 if matches!(
7401 c,
7402 Condition::Ann { .. }
7403 | Condition::SparseMatch { .. }
7404 | Condition::MinHashSimilar { .. }
7405 ) {
7406 sets.push(self.resolve_condition(c, snapshot)?);
7407 }
7408 }
7409 Ok(sets)
7410 }
7411
7412 fn column_type(&self, cid: u16) -> TypeId {
7413 self.schema
7414 .columns
7415 .iter()
7416 .find(|c| c.id == cid)
7417 .map(|c| c.ty.clone())
7418 .unwrap_or(TypeId::Bytes)
7419 }
7420
7421 pub fn approx_aggregate(
7430 &mut self,
7431 conditions: &[crate::query::Condition],
7432 column: Option<u16>,
7433 agg: ApproxAgg,
7434 z: f64,
7435 ) -> Result<Option<ApproxResult>> {
7436 use crate::query::Condition;
7437 self.ensure_reservoir_complete()?;
7438 let snapshot = self.snapshot();
7439 let n_pop = self.count();
7440 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
7441 if sample_rids.is_empty() {
7442 return Ok(None);
7443 }
7444 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
7446 let s = live_sample.len();
7447 if s == 0 {
7448 return Ok(None);
7449 }
7450
7451 let mut index_sets: Vec<RowIdSet> = Vec::new();
7454 for c in conditions {
7455 if matches!(
7456 c,
7457 Condition::Ann { .. }
7458 | Condition::SparseMatch { .. }
7459 | Condition::MinHashSimilar { .. }
7460 ) {
7461 index_sets.push(self.resolve_condition(c, snapshot)?);
7462 }
7463 }
7464
7465 let cid = match (agg, column) {
7467 (ApproxAgg::Count, _) => None,
7468 (_, Some(c)) => Some(c),
7469 _ => return Ok(None),
7470 };
7471 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
7472 for r in &live_sample {
7473 if !conditions
7475 .iter()
7476 .all(|c| condition_matches_row(c, r, &self.schema))
7477 {
7478 continue;
7479 }
7480 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
7482 continue;
7483 }
7484 if let Some(cid) = cid {
7485 if let Some(v) = as_f64(r.columns.get(&cid)) {
7486 passing_vals.push(v);
7487 } } else {
7489 passing_vals.push(0.0); }
7491 }
7492 let m = passing_vals.len();
7493
7494 let (point, half) = match agg {
7495 ApproxAgg::Count => {
7496 let p = m as f64 / s as f64;
7498 let point = n_pop as f64 * p;
7499 let var = if s > 1 {
7500 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
7501 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7502 } else {
7503 0.0
7504 };
7505 (point, z * var.sqrt())
7506 }
7507 ApproxAgg::Sum => {
7508 let y: Vec<f64> = live_sample
7510 .iter()
7511 .map(|r| {
7512 let passes_row = conditions
7513 .iter()
7514 .all(|c| condition_matches_row(c, r, &self.schema))
7515 && index_sets.iter().all(|set| set.contains(r.row_id.0));
7516 if passes_row {
7517 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
7518 } else {
7519 0.0
7520 }
7521 })
7522 .collect();
7523 let mean_y = y.iter().sum::<f64>() / s as f64;
7524 let point = n_pop as f64 * mean_y;
7525 let var = if s > 1 {
7526 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
7527 let var_y = ss / (s - 1) as f64;
7528 n_pop as f64 * n_pop as f64 * var_y / s as f64
7529 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7530 } else {
7531 0.0
7532 };
7533 (point, z * var.sqrt())
7534 }
7535 ApproxAgg::Avg => {
7536 if m == 0 {
7537 return Ok(Some(ApproxResult {
7538 point: 0.0,
7539 ci_low: 0.0,
7540 ci_high: 0.0,
7541 n_population: n_pop,
7542 n_sample_live: s,
7543 n_passing: 0,
7544 }));
7545 }
7546 let mean = passing_vals.iter().sum::<f64>() / m as f64;
7547 let half = if m > 1 {
7548 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
7549 let sd = (ss / (m - 1) as f64).sqrt();
7550 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
7551 z * sd / (m as f64).sqrt() * fpc.sqrt()
7552 } else {
7553 0.0
7554 };
7555 (mean, half)
7556 }
7557 };
7558
7559 Ok(Some(ApproxResult {
7560 point,
7561 ci_low: point - half,
7562 ci_high: point + half,
7563 n_population: n_pop,
7564 n_sample_live: s,
7565 n_passing: m,
7566 }))
7567 }
7568
7569 pub fn exact_column_stats(
7577 &self,
7578 _snapshot: Snapshot,
7579 projection: &[u16],
7580 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
7581 if self.ttl.is_some()
7582 || !(self.memtable.is_empty()
7583 && self.mutable_run.is_empty()
7584 && self.run_refs.len() == 1)
7585 {
7586 return Ok(None);
7587 }
7588 let reader = self.open_reader(self.run_refs[0].run_id)?;
7589 if self.live_count != reader.row_count() as u64 {
7590 return Ok(None);
7591 }
7592 let mut out = HashMap::new();
7593 for &cid in projection {
7594 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
7595 Some(c) => c,
7596 None => continue,
7597 };
7598 let Some(stats) = reader.column_page_stats(cid) else {
7600 out.insert(
7601 cid,
7602 ColumnStat {
7603 min: None,
7604 max: None,
7605 null_count: self.live_count,
7606 },
7607 );
7608 continue;
7609 };
7610 let stat = match cdef.ty {
7611 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
7612 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
7613 min: mn.map(Value::Int64),
7614 max: mx.map(Value::Int64),
7615 null_count: n,
7616 })
7617 }
7618 TypeId::Float64 => {
7619 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
7620 min: mn.map(Value::Float64),
7621 max: mx.map(Value::Float64),
7622 null_count: n,
7623 })
7624 }
7625 _ => None,
7626 };
7627 if let Some(s) = stat {
7628 out.insert(cid, s);
7629 }
7630 }
7631 Ok(Some(out))
7632 }
7633
7634 pub fn dir(&self) -> &Path {
7635 &self.dir
7636 }
7637
7638 pub fn schema(&self) -> &Schema {
7639 &self.schema
7640 }
7641
7642 pub(crate) fn set_catalog_name(&mut self, name: String) {
7643 self.name = name;
7644 }
7645
7646 pub(crate) fn prepare_alter_column(
7647 &mut self,
7648 column_name: &str,
7649 change: &AlterColumn,
7650 ) -> Result<ColumnDef> {
7651 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
7652 return Err(MongrelError::InvalidArgument(
7653 "ALTER COLUMN requires committing staged writes first".into(),
7654 ));
7655 }
7656 let old = self
7657 .schema
7658 .columns
7659 .iter()
7660 .find(|c| c.name == column_name)
7661 .cloned()
7662 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7663 let mut next = old.clone();
7664
7665 if let Some(name) = &change.name {
7666 let trimmed = name.trim();
7667 if trimmed.is_empty() {
7668 return Err(MongrelError::InvalidArgument(
7669 "ALTER COLUMN name must not be empty".into(),
7670 ));
7671 }
7672 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
7673 return Err(MongrelError::Schema(format!(
7674 "column {trimmed} already exists"
7675 )));
7676 }
7677 next.name = trimmed.to_string();
7678 }
7679
7680 if let Some(ty) = &change.ty {
7681 next.ty = ty.clone();
7682 }
7683 if let Some(flags) = change.flags {
7684 validate_alter_column_flags(old.flags, flags)?;
7685 next.flags = flags;
7686 }
7687
7688 if let Some(default_change) = &change.default_value {
7689 next.default_value = default_change.clone();
7690 }
7691
7692 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
7693 if old.flags.contains(ColumnFlags::NULLABLE)
7694 && !next.flags.contains(ColumnFlags::NULLABLE)
7695 && self.column_has_nulls(old.id)?
7696 {
7697 return Err(MongrelError::InvalidArgument(format!(
7698 "column '{}' contains NULL values",
7699 old.name
7700 )));
7701 }
7702 Ok(next)
7703 }
7704
7705 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
7706 let idx = self
7707 .schema
7708 .columns
7709 .iter()
7710 .position(|c| c.id == column.id)
7711 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
7712 if self.schema.columns[idx] == column {
7713 return Ok(());
7714 }
7715 self.schema.columns[idx] = column;
7716 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7717 self.schema.validate_auto_increment()?;
7718 self.schema.validate_defaults()?;
7719 self.auto_inc = resolve_auto_inc(&self.schema);
7720 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
7721 write_schema(&self.dir, &self.schema)?;
7722 self.clear_result_cache();
7723 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7724 self.persist_manifest(self.current_epoch())?;
7725 Ok(())
7726 }
7727
7728 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
7729 self.ensure_writable()?;
7730 let column = self.prepare_alter_column(column_name, &change)?;
7731 self.apply_altered_column(column.clone())?;
7732 Ok(column)
7733 }
7734
7735 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
7736 if self.live_count == 0 {
7737 return Ok(false);
7738 }
7739 let snap = self.snapshot();
7740 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
7741 Ok(columns
7742 .first()
7743 .map(|(_, col)| col.null_count(col.len()) != 0)
7744 .unwrap_or(true))
7745 }
7746
7747 fn has_stored_versions(&self) -> bool {
7748 !self.memtable.is_empty()
7749 || !self.mutable_run.is_empty()
7750 || self.run_refs.iter().any(|r| r.row_count > 0)
7751 || !self.retiring.is_empty()
7752 }
7753
7754 pub fn add_column(
7759 &mut self,
7760 name: &str,
7761 ty: TypeId,
7762 flags: ColumnFlags,
7763 default_value: Option<crate::schema::DefaultExpr>,
7764 ) -> Result<u16> {
7765 self.add_column_with_id(name, ty, flags, default_value, None)
7766 }
7767
7768 pub fn add_column_with_id(
7769 &mut self,
7770 name: &str,
7771 ty: TypeId,
7772 flags: ColumnFlags,
7773 default_value: Option<crate::schema::DefaultExpr>,
7774 requested_id: Option<u16>,
7775 ) -> Result<u16> {
7776 self.ensure_writable()?;
7777 if self.schema.columns.iter().any(|c| c.name == name) {
7778 return Err(MongrelError::Schema(format!(
7779 "column {name} already exists"
7780 )));
7781 }
7782 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
7783 if self.schema.columns.iter().any(|c| c.id == id) {
7784 return Err(MongrelError::Schema(format!(
7785 "column id {id} already exists"
7786 )));
7787 }
7788 id
7789 } else {
7790 self.schema
7791 .columns
7792 .iter()
7793 .map(|c| c.id)
7794 .max()
7795 .unwrap_or(0)
7796 .checked_add(1)
7797 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
7798 };
7799 self.schema.columns.push(ColumnDef {
7800 id,
7801 name: name.to_string(),
7802 ty,
7803 flags,
7804 default_value,
7805 });
7806 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7807 self.schema.validate_auto_increment()?;
7808 self.schema.validate_defaults()?;
7809 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
7810 self.auto_inc = resolve_auto_inc(&self.schema);
7811 }
7812 write_schema(&self.dir, &self.schema)?;
7813 self.clear_result_cache();
7814 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7816 self.persist_manifest(self.current_epoch())?;
7817 Ok(id)
7818 }
7819
7820 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
7829 self.ensure_writable()?;
7830 let cid = self
7831 .schema
7832 .columns
7833 .iter()
7834 .find(|c| c.name == column_name)
7835 .map(|c| c.id)
7836 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7837 let ty = self
7838 .schema
7839 .columns
7840 .iter()
7841 .find(|c| c.id == cid)
7842 .map(|c| c.ty.clone())
7843 .unwrap_or(TypeId::Int64);
7844 if !matches!(
7845 ty,
7846 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
7847 ) {
7848 return Err(MongrelError::Schema(format!(
7849 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
7850 )));
7851 }
7852 if self
7853 .schema
7854 .indexes
7855 .iter()
7856 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
7857 {
7858 return Ok(()); }
7860 self.schema.indexes.push(IndexDef {
7861 name: format!("{}_learned_range", column_name),
7862 column_id: cid,
7863 kind: IndexKind::LearnedRange,
7864 predicate: None,
7865 options: Default::default(),
7866 });
7867 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7868 write_schema(&self.dir, &self.schema)?;
7869 self.build_learned_ranges()?;
7870 Ok(())
7871 }
7872
7873 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
7876 self.sync_byte_threshold = threshold;
7877 if let WalSink::Private(w) = &mut self.wal {
7878 w.set_sync_byte_threshold(threshold);
7879 }
7880 }
7881
7882 pub fn page_cache_flush(&self) {
7886 self.page_cache.flush_to_disk();
7887 }
7888
7889 pub fn page_cache_len(&self) -> usize {
7891 self.page_cache.len()
7892 }
7893
7894 pub fn decoded_cache_len(&self) -> usize {
7897 self.decoded_cache.len()
7898 }
7899
7900 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
7903 self.memtable.drain_sorted()
7904 }
7905
7906 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
7907 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
7908 }
7909
7910 pub(crate) fn table_dir(&self) -> &Path {
7911 &self.dir
7912 }
7913
7914 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
7915 &self.schema
7916 }
7917
7918 pub(crate) fn alloc_run_id(&mut self) -> u64 {
7919 let id = self.next_run_id;
7920 self.next_run_id += 1;
7921 id
7922 }
7923
7924 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
7925 self.run_refs.push(run_ref);
7926 }
7927
7928 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
7938 self.retiring.push(crate::manifest::RetiredRun {
7939 run_id,
7940 retire_epoch,
7941 });
7942 }
7943
7944 pub(crate) fn reap_retiring(
7948 &mut self,
7949 min_active: Epoch,
7950 backup_pinned: &std::collections::HashSet<u128>,
7951 ) -> Result<usize> {
7952 if self.retiring.is_empty() {
7953 return Ok(0);
7954 }
7955 let mut reaped = 0;
7956 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
7957 for r in std::mem::take(&mut self.retiring) {
7963 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
7964 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
7965 reaped += 1;
7966 } else {
7967 kept.push(r);
7968 }
7969 }
7970 self.retiring = kept;
7971 if reaped > 0 {
7972 self.persist_manifest(self.current_epoch())?;
7973 }
7974 Ok(reaped)
7975 }
7976
7977 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
7978 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
7979 return false;
7980 }
7981 self.live_count = self.live_count.saturating_add(run_ref.row_count);
7982 self.run_refs.push(run_ref);
7983 self.indexes_complete = false;
7984 true
7985 }
7986
7987 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
7988 self.kek.as_ref()
7989 }
7990
7991 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
7992 let mut reader = RunReader::open_with_cache(
7993 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
7994 self.schema.clone(),
7995 self.kek.clone(),
7996 Some(self.page_cache.clone()),
7997 Some(self.decoded_cache.clone()),
7998 self.table_id,
7999 Some(&self.verified_runs),
8000 )?;
8001 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
8005 reader.set_uniform_epoch(Epoch(rr.epoch_created));
8006 }
8007 Ok(reader)
8008 }
8009
8010 pub(crate) fn run_refs(&self) -> &[RunRef] {
8011 &self.run_refs
8012 }
8013
8014 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
8015 self.retiring.iter().map(|run| run.run_id)
8016 }
8017
8018 pub(crate) fn runs_dir(&self) -> PathBuf {
8019 self.dir.join(RUNS_DIR)
8020 }
8021
8022 pub(crate) fn wal_dir(&self) -> PathBuf {
8023 self.dir.join(WAL_DIR)
8024 }
8025
8026 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
8027 self.run_refs = refs;
8028 }
8029
8030 pub(crate) fn next_run_id(&self) -> u64 {
8031 self.next_run_id
8032 }
8033
8034 pub(crate) fn compaction_zstd_level(&self) -> i32 {
8035 self.compaction_zstd_level
8036 }
8037
8038 pub(crate) fn bump_next_run_id(&mut self) {
8039 self.next_run_id += 1;
8040 }
8041
8042 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
8043 self.kek.clone()
8044 }
8045
8046 #[cfg(feature = "encryption")]
8050 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8051 self.kek.as_ref().map(|k| k.derive_idx_key())
8052 }
8053
8054 #[cfg(not(feature = "encryption"))]
8055 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8056 None
8057 }
8058
8059 #[cfg(feature = "encryption")]
8063 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8064 self.kek.as_ref().map(|k| *k.derive_meta_key())
8065 }
8066
8067 #[cfg(not(feature = "encryption"))]
8068 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8069 None
8070 }
8071
8072 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
8075 self.column_keys
8076 .iter()
8077 .map(|(&id, &(_, scheme))| (id, scheme))
8078 .collect()
8079 }
8080
8081 #[cfg(feature = "encryption")]
8086 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
8087 self.tokenize_value_enc(column_id, v)
8088 }
8089
8090 #[cfg(feature = "encryption")]
8091 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
8092 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8093 let (key, scheme) = self.column_keys.get(&column_id)?;
8094 let token: Vec<u8> = match (*scheme, v) {
8095 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
8096 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
8097 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
8098 _ => hmac_token(key, &v.encode_key()).to_vec(),
8099 };
8100 Some(Value::Bytes(token))
8101 }
8102
8103 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
8105 self.index_lookup_key_bytes(column_id, &v.encode_key())
8106 }
8107
8108 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
8111 #[cfg(feature = "encryption")]
8112 {
8113 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
8114 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
8115 if *scheme == SCHEME_HMAC_EQ {
8116 return hmac_token(key, encoded).to_vec();
8117 }
8118 }
8119 }
8120 let _ = column_id;
8121 encoded.to_vec()
8122 }
8123}
8124
8125fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
8126 let columnar::NativeColumn::Int64 { data, validity } = col else {
8127 return false;
8128 };
8129 if data.len() < n || !columnar::all_non_null(validity, n) {
8130 return false;
8131 }
8132 data.iter()
8133 .take(n)
8134 .zip(data.iter().skip(1))
8135 .all(|(a, b)| a < b)
8136}
8137
8138#[derive(Debug, Clone)]
8142pub struct ColumnStat {
8143 pub min: Option<Value>,
8144 pub max: Option<Value>,
8145 pub null_count: u64,
8146}
8147
8148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8150pub enum NativeAgg {
8151 Count,
8152 Sum,
8153 Min,
8154 Max,
8155 Avg,
8156}
8157
8158#[derive(Debug, Clone, PartialEq)]
8160pub enum NativeAggResult {
8161 Count(u64),
8162 Int(i64),
8163 Float(f64),
8164 Null,
8166}
8167
8168#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8170pub enum ApproxAgg {
8171 Count,
8172 Sum,
8173 Avg,
8174}
8175
8176#[derive(Debug, Clone)]
8180pub struct ApproxResult {
8181 pub point: f64,
8183 pub ci_low: f64,
8185 pub ci_high: f64,
8187 pub n_population: u64,
8189 pub n_sample_live: usize,
8191 pub n_passing: usize,
8193}
8194
8195#[derive(Debug, Clone, PartialEq)]
8200pub enum AggState {
8201 Count(u64),
8203 SumI {
8205 sum: i128,
8206 count: u64,
8207 },
8208 SumF {
8210 sum: f64,
8211 count: u64,
8212 },
8213 AvgI {
8215 sum: i128,
8216 count: u64,
8217 },
8218 AvgF {
8220 sum: f64,
8221 count: u64,
8222 },
8223 MinI(i64),
8225 MaxI(i64),
8226 MinF(f64),
8228 MaxF(f64),
8229 Empty,
8231}
8232
8233impl AggState {
8234 pub fn merge(self, other: AggState) -> AggState {
8236 use AggState::*;
8237 match (self, other) {
8238 (Empty, x) | (x, Empty) => x,
8239 (Count(a), Count(b)) => Count(a + b),
8240 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
8241 sum: sa + sb,
8242 count: ca + cb,
8243 },
8244 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
8245 sum: sa + sb,
8246 count: ca + cb,
8247 },
8248 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
8249 sum: sa + sb,
8250 count: ca + cb,
8251 },
8252 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
8253 sum: sa + sb,
8254 count: ca + cb,
8255 },
8256 (MinI(a), MinI(b)) => MinI(a.min(b)),
8257 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
8258 (MinF(a), MinF(b)) => MinF(a.min(b)),
8259 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
8260 _ => Empty, }
8262 }
8263
8264 pub fn point(&self) -> Option<f64> {
8266 match self {
8267 AggState::Count(n) => Some(*n as f64),
8268 AggState::SumI { sum, .. } => Some(*sum as f64),
8269 AggState::SumF { sum, .. } => Some(*sum),
8270 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
8271 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
8272 AggState::MinI(n) => Some(*n as f64),
8273 AggState::MaxI(n) => Some(*n as f64),
8274 AggState::MinF(n) => Some(*n),
8275 AggState::MaxF(n) => Some(*n),
8276 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
8277 }
8278 }
8279
8280 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
8284 let is_float = matches!(ty, Some(TypeId::Float64));
8285 match (agg, result) {
8286 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
8287 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
8288 sum: x as i128,
8289 count: 1, },
8291 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
8292 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
8293 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
8294 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
8295 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
8296 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
8297 (NativeAgg::Count, _) => AggState::Empty,
8298 (_, NativeAggResult::Null) => AggState::Empty,
8299 _ => {
8300 let _ = is_float;
8301 AggState::Empty
8302 }
8303 }
8304 }
8305}
8306
8307#[derive(Debug, Clone)]
8310pub struct CachedAgg {
8311 pub state: AggState,
8312 pub watermark: u64,
8313 pub epoch: u64,
8314}
8315
8316#[derive(Debug, Clone)]
8318pub struct IncrementalAggResult {
8319 pub state: AggState,
8321 pub incremental: bool,
8324 pub delta_rows: u64,
8326}
8327
8328fn agg_state_from_rows(
8332 rows: &[Row],
8333 conditions: &[crate::query::Condition],
8334 index_sets: &[RowIdSet],
8335 column: Option<u16>,
8336 agg: NativeAgg,
8337 schema: &Schema,
8338) -> Result<AggState> {
8339 let mut count: u64 = 0;
8340 let mut sum_i: i128 = 0;
8341 let mut sum_f: f64 = 0.0;
8342 let mut mn_i: i64 = i64::MAX;
8343 let mut mx_i: i64 = i64::MIN;
8344 let mut mn_f: f64 = f64::INFINITY;
8345 let mut mx_f: f64 = f64::NEG_INFINITY;
8346 let mut saw_int = false;
8347 let mut saw_float = false;
8348 for r in rows {
8349 if !conditions
8350 .iter()
8351 .all(|c| condition_matches_row(c, r, schema))
8352 {
8353 continue;
8354 }
8355 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
8356 continue;
8357 }
8358 match agg {
8359 NativeAgg::Count => match column {
8360 None => count += 1,
8362 Some(cid) => match r.columns.get(&cid) {
8365 None | Some(Value::Null) => {}
8366 Some(_) => count += 1,
8367 },
8368 },
8369 _ => match column.and_then(|cid| r.columns.get(&cid)) {
8370 Some(Value::Int64(n)) => {
8371 count += 1;
8372 sum_i += *n as i128;
8373 mn_i = mn_i.min(*n);
8374 mx_i = mx_i.max(*n);
8375 saw_int = true;
8376 }
8377 Some(Value::Float64(f)) => {
8378 count += 1;
8379 sum_f += f;
8380 mn_f = mn_f.min(*f);
8381 mx_f = mx_f.max(*f);
8382 saw_float = true;
8383 }
8384 _ => {}
8385 },
8386 }
8387 }
8388 Ok(match agg {
8389 NativeAgg::Count => {
8390 if count == 0 {
8391 AggState::Empty
8392 } else {
8393 AggState::Count(count)
8394 }
8395 }
8396 NativeAgg::Sum => {
8397 if count == 0 {
8398 AggState::Empty
8399 } else if saw_int {
8400 AggState::SumI { sum: sum_i, count }
8401 } else {
8402 AggState::SumF { sum: sum_f, count }
8403 }
8404 }
8405 NativeAgg::Avg => {
8406 if count == 0 {
8407 AggState::Empty
8408 } else if saw_int {
8409 AggState::AvgI { sum: sum_i, count }
8410 } else {
8411 AggState::AvgF { sum: sum_f, count }
8412 }
8413 }
8414 NativeAgg::Min => {
8415 if !saw_int && !saw_float {
8416 AggState::Empty
8417 } else if saw_int {
8418 AggState::MinI(mn_i)
8419 } else {
8420 AggState::MinF(mn_f)
8421 }
8422 }
8423 NativeAgg::Max => {
8424 if !saw_int && !saw_float {
8425 AggState::Empty
8426 } else if saw_int {
8427 AggState::MaxI(mx_i)
8428 } else {
8429 AggState::MaxF(mx_f)
8430 }
8431 }
8432 })
8433}
8434
8435fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
8439 use crate::query::Condition;
8440 match c {
8441 Condition::Pk(key) => match schema.primary_key() {
8442 Some(pk) => row
8443 .columns
8444 .get(&pk.id)
8445 .map(|v| v.encode_key() == *key)
8446 .unwrap_or(false),
8447 None => false,
8448 },
8449 Condition::BitmapEq { column_id, value } => row
8450 .columns
8451 .get(column_id)
8452 .map(|v| v.encode_key() == *value)
8453 .unwrap_or(false),
8454 Condition::BitmapIn { column_id, values } => {
8455 let key = row.columns.get(column_id).map(|v| v.encode_key());
8456 match key {
8457 Some(k) => values.contains(&k),
8458 None => false,
8459 }
8460 }
8461 Condition::BytesPrefix { column_id, prefix } => row
8462 .columns
8463 .get(column_id)
8464 .map(|v| v.encode_key().starts_with(prefix))
8465 .unwrap_or(false),
8466 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
8467 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
8468 _ => false,
8469 },
8470 Condition::RangeF64 {
8471 column_id,
8472 lo,
8473 lo_inclusive,
8474 hi,
8475 hi_inclusive,
8476 } => match row.columns.get(column_id) {
8477 Some(Value::Float64(n)) => {
8478 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
8479 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
8480 lo_ok && hi_ok
8481 }
8482 _ => false,
8483 },
8484 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
8485 Some(Value::Bytes(b)) => {
8486 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
8487 }
8488 _ => false,
8489 },
8490 Condition::FmContainsAll {
8491 column_id,
8492 patterns,
8493 } => match row.columns.get(column_id) {
8494 Some(Value::Bytes(b)) => patterns
8495 .iter()
8496 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
8497 _ => false,
8498 },
8499 Condition::Ann { .. }
8500 | Condition::SparseMatch { .. }
8501 | Condition::MinHashSimilar { .. } => true,
8502 Condition::IsNull { column_id } => {
8503 matches!(row.columns.get(column_id), Some(Value::Null) | None)
8504 }
8505 Condition::IsNotNull { column_id } => {
8506 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
8507 }
8508 }
8509}
8510
8511fn as_f64(v: Option<&Value>) -> Option<f64> {
8513 match v {
8514 Some(Value::Int64(n)) => Some(*n as f64),
8515 Some(Value::Float64(f)) => Some(*f),
8516 _ => None,
8517 }
8518}
8519
8520fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
8524 let mut count: u64 = 0;
8525 let mut sum: i128 = 0;
8526 let mut mn: i64 = i64::MAX;
8527 let mut mx: i64 = i64::MIN;
8528 while let Some(cols) = cursor.next_batch()? {
8529 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
8530 if crate::columnar::all_non_null(validity, data.len()) {
8531 count += data.len() as u64;
8533 sum += data.iter().map(|&v| v as i128).sum::<i128>();
8534 mn = mn.min(*data.iter().min().unwrap_or(&mn));
8535 mx = mx.max(*data.iter().max().unwrap_or(&mx));
8536 } else {
8537 for (i, &v) in data.iter().enumerate() {
8538 if crate::columnar::validity_bit(validity, i) {
8539 count += 1;
8540 sum += v as i128;
8541 mn = mn.min(v);
8542 mx = mx.max(v);
8543 }
8544 }
8545 }
8546 }
8547 }
8548 Ok((count, sum, mn, mx))
8549}
8550
8551fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
8553 let mut count: u64 = 0;
8554 let mut sum: f64 = 0.0;
8555 let mut mn: f64 = f64::INFINITY;
8556 let mut mx: f64 = f64::NEG_INFINITY;
8557 while let Some(cols) = cursor.next_batch()? {
8558 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
8559 if crate::columnar::all_non_null(validity, data.len()) {
8560 count += data.len() as u64;
8561 sum += data.iter().sum::<f64>();
8562 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
8563 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
8564 } else {
8565 for (i, &v) in data.iter().enumerate() {
8566 if crate::columnar::validity_bit(validity, i) {
8567 count += 1;
8568 sum += v;
8569 mn = mn.min(v);
8570 mx = mx.max(v);
8571 }
8572 }
8573 }
8574 }
8575 }
8576 Ok((count, sum, mn, mx))
8577}
8578
8579fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
8580 if count == 0 && !matches!(agg, NativeAgg::Count) {
8581 return NativeAggResult::Null;
8582 }
8583 match agg {
8584 NativeAgg::Count => NativeAggResult::Count(count),
8585 NativeAgg::Sum => match sum.try_into() {
8588 Ok(v) => NativeAggResult::Int(v),
8589 Err(_) => NativeAggResult::Null,
8590 },
8591 NativeAgg::Min => NativeAggResult::Int(mn),
8592 NativeAgg::Max => NativeAggResult::Int(mx),
8593 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
8594 }
8595}
8596
8597fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
8598 if count == 0 && !matches!(agg, NativeAgg::Count) {
8599 return NativeAggResult::Null;
8600 }
8601 match agg {
8602 NativeAgg::Count => NativeAggResult::Count(count),
8603 NativeAgg::Sum => NativeAggResult::Float(sum),
8604 NativeAgg::Min => NativeAggResult::Float(mn),
8605 NativeAgg::Max => NativeAggResult::Float(mx),
8606 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
8607 }
8608}
8609
8610fn agg_int(
8613 stats: &[crate::page::PageStat],
8614 decode: fn(Option<&[u8]>) -> Option<i64>,
8615) -> Option<(Option<i64>, Option<i64>, u64)> {
8616 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
8617 let mut any = false;
8618 for s in stats {
8619 if let Some(v) = decode(s.min.as_deref()) {
8620 mn = mn.min(v);
8621 any = true;
8622 }
8623 if let Some(v) = decode(s.max.as_deref()) {
8624 mx = mx.max(v);
8625 any = true;
8626 }
8627 nulls += s.null_count;
8628 }
8629 any.then_some((Some(mn), Some(mx), nulls))
8630}
8631
8632fn agg_float(
8634 stats: &[crate::page::PageStat],
8635 decode: fn(Option<&[u8]>) -> Option<f64>,
8636) -> Option<(Option<f64>, Option<f64>, u64)> {
8637 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
8638 let mut any = false;
8639 for s in stats {
8640 if let Some(v) = decode(s.min.as_deref()) {
8641 mn = mn.min(v);
8642 any = true;
8643 }
8644 if let Some(v) = decode(s.max.as_deref()) {
8645 mx = mx.max(v);
8646 any = true;
8647 }
8648 nulls += s.null_count;
8649 }
8650 any.then_some((Some(mn), Some(mx), nulls))
8651}
8652
8653type SecondaryIndexes = (
8655 HashMap<u16, BitmapIndex>,
8656 HashMap<u16, AnnIndex>,
8657 HashMap<u16, FmIndex>,
8658 HashMap<u16, SparseIndex>,
8659 HashMap<u16, MinHashIndex>,
8660);
8661
8662fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
8663 let mut bitmap = HashMap::new();
8664 let mut ann = HashMap::new();
8665 let mut fm = HashMap::new();
8666 let mut sparse = HashMap::new();
8667 let mut minhash = HashMap::new();
8668 for idef in &schema.indexes {
8669 match idef.kind {
8670 IndexKind::Bitmap => {
8671 bitmap.insert(idef.column_id, BitmapIndex::new());
8672 }
8673 IndexKind::Ann => {
8674 let dim = schema
8675 .columns
8676 .iter()
8677 .find(|c| c.id == idef.column_id)
8678 .and_then(|c| match c.ty {
8679 TypeId::Embedding { dim } => Some(dim as usize),
8680 _ => None,
8681 })
8682 .unwrap_or(0);
8683 let options = idef.options.ann.clone().unwrap_or_default();
8684 ann.insert(
8685 idef.column_id,
8686 AnnIndex::with_options(
8687 dim,
8688 options.m,
8689 options.ef_construction,
8690 options.ef_search,
8691 ),
8692 );
8693 }
8694 IndexKind::FmIndex => {
8695 fm.insert(idef.column_id, FmIndex::new());
8696 }
8697 IndexKind::Sparse => {
8698 sparse.insert(idef.column_id, SparseIndex::new());
8699 }
8700 IndexKind::MinHash => {
8701 let options = idef.options.minhash.clone().unwrap_or_default();
8702 minhash.insert(
8703 idef.column_id,
8704 MinHashIndex::with_options(options.permutations, options.bands),
8705 );
8706 }
8707 _ => {}
8708 }
8709 }
8710 (bitmap, ann, fm, sparse, minhash)
8711}
8712
8713const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
8714 | ColumnFlags::AUTO_INCREMENT
8715 | ColumnFlags::ENCRYPTED
8716 | ColumnFlags::ENCRYPTED_INDEXABLE
8717 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
8718
8719fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
8720 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
8721 return Err(MongrelError::Schema(
8722 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
8723 ));
8724 }
8725 Ok(())
8726}
8727
8728fn validate_alter_column_type(
8729 schema: &Schema,
8730 old: &ColumnDef,
8731 next: &ColumnDef,
8732 has_stored_versions: bool,
8733) -> Result<()> {
8734 if old.ty == next.ty {
8735 return Ok(());
8736 }
8737 if schema.indexes.iter().any(|i| i.column_id == old.id) {
8738 return Err(MongrelError::Schema(format!(
8739 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
8740 old.name
8741 )));
8742 }
8743 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
8744 return Ok(());
8745 }
8746 Err(MongrelError::Schema(format!(
8747 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
8748 old.ty, next.ty
8749 )))
8750}
8751
8752fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
8753 matches!(
8754 (old, new),
8755 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
8756 )
8757}
8758
8759fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
8765 let mut prev: Option<i64> = None;
8766 for r in rows {
8767 match r.columns.get(&pk_id) {
8768 Some(Value::Int64(v)) => {
8769 if prev.is_some_and(|p| p >= *v) {
8770 return false;
8771 }
8772 prev = Some(*v);
8773 }
8774 _ => return false,
8775 }
8776 }
8777 true
8778}
8779
8780#[allow(clippy::too_many_arguments)]
8781fn index_into(
8782 schema: &Schema,
8783 row: &Row,
8784 hot: &mut HotIndex,
8785 bitmap: &mut HashMap<u16, BitmapIndex>,
8786 ann: &mut HashMap<u16, AnnIndex>,
8787 fm: &mut HashMap<u16, FmIndex>,
8788 sparse: &mut HashMap<u16, SparseIndex>,
8789 minhash: &mut HashMap<u16, MinHashIndex>,
8790) {
8791 for idef in &schema.indexes {
8792 let Some(val) = row.columns.get(&idef.column_id) else {
8793 continue;
8794 };
8795 match idef.kind {
8796 IndexKind::Bitmap => {
8797 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8798 b.insert(val.encode_key(), row.row_id);
8799 }
8800 }
8801 IndexKind::Ann => {
8802 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8803 a.insert_validated(v, row.row_id);
8804 }
8805 }
8806 IndexKind::FmIndex => {
8807 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8808 f.insert(b.clone(), row.row_id);
8809 }
8810 }
8811 IndexKind::Sparse => {
8812 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8813 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8816 s.insert(&terms, row.row_id);
8817 }
8818 }
8819 }
8820 IndexKind::MinHash => {
8821 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8822 let tokens = crate::index::token_hashes_from_bytes(b);
8825 mh.insert(&tokens, row.row_id);
8826 }
8827 }
8828 _ => {}
8829 }
8830 }
8831 if let Some(pk_col) = schema.primary_key() {
8832 if let Some(pk_val) = row.columns.get(&pk_col.id) {
8833 hot.insert(pk_val.encode_key(), row.row_id);
8834 }
8835 }
8836}
8837
8838#[allow(clippy::too_many_arguments)]
8841fn index_into_single(
8842 idef: &IndexDef,
8843 _schema: &Schema,
8844 row: &Row,
8845 _hot: &mut HotIndex,
8846 bitmap: &mut HashMap<u16, BitmapIndex>,
8847 ann: &mut HashMap<u16, AnnIndex>,
8848 fm: &mut HashMap<u16, FmIndex>,
8849 sparse: &mut HashMap<u16, SparseIndex>,
8850 minhash: &mut HashMap<u16, MinHashIndex>,
8851) {
8852 let Some(val) = row.columns.get(&idef.column_id) else {
8853 return;
8854 };
8855 match idef.kind {
8856 IndexKind::Bitmap => {
8857 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8858 b.insert(val.encode_key(), row.row_id);
8859 }
8860 }
8861 IndexKind::Ann => {
8862 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8863 a.insert_validated(v, row.row_id);
8864 }
8865 }
8866 IndexKind::FmIndex => {
8867 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8868 f.insert(b.clone(), row.row_id);
8869 }
8870 }
8871 IndexKind::Sparse => {
8872 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8873 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8874 s.insert(&terms, row.row_id);
8875 }
8876 }
8877 }
8878 IndexKind::MinHash => {
8879 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8880 let tokens = crate::index::token_hashes_from_bytes(b);
8881 mh.insert(&tokens, row.row_id);
8882 }
8883 }
8884 _ => {}
8885 }
8886}
8887
8888fn eval_partial_predicate(
8894 pred: &str,
8895 columns_map: &HashMap<u16, &Value>,
8896 name_to_id: &HashMap<&str, u16>,
8897) -> bool {
8898 let lower = pred.trim().to_ascii_lowercase();
8899 if let Some(rest) = lower.strip_suffix(" is not null") {
8901 let col_name = rest.trim();
8902 if let Some(col_id) = name_to_id.get(col_name) {
8903 return columns_map
8904 .get(col_id)
8905 .is_some_and(|v| !matches!(v, Value::Null));
8906 }
8907 }
8908 if let Some(rest) = lower.strip_suffix(" is null") {
8910 let col_name = rest.trim();
8911 if let Some(col_id) = name_to_id.get(col_name) {
8912 return columns_map
8913 .get(col_id)
8914 .map_or(true, |v| matches!(v, Value::Null));
8915 }
8916 }
8917 true
8920}
8921
8922#[allow(dead_code)]
8928fn bulk_index_key(
8929 column_keys: &HashMap<u16, ([u8; 32], u8)>,
8930 column_id: u16,
8931 ty: TypeId,
8932 col: &columnar::NativeColumn,
8933 i: usize,
8934) -> Option<Vec<u8>> {
8935 let encoded = columnar::encode_key_native(ty, col, i)?;
8936 #[cfg(feature = "encryption")]
8937 {
8938 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8939 if let Some((key, scheme)) = column_keys.get(&column_id) {
8940 return Some(match (*scheme, col) {
8941 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
8942 (_, columnar::NativeColumn::Int64 { data, .. }) => {
8943 ope_token_i64(key, data[i]).to_vec()
8944 }
8945 (_, columnar::NativeColumn::Float64 { data, .. }) => {
8946 ope_token_f64(key, data[i]).to_vec()
8947 }
8948 _ => hmac_token(key, &encoded).to_vec(),
8949 });
8950 }
8951 }
8952 #[cfg(not(feature = "encryption"))]
8953 {
8954 let _ = (column_id, column_keys, col);
8955 }
8956 Some(encoded)
8957}
8958
8959pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
8960 let json = serde_json::to_string_pretty(schema)
8961 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
8962 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
8963 Ok(())
8964}
8965
8966fn read_schema(dir: &Path) -> Result<Schema> {
8967 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
8968 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
8969}
8970
8971fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
8972 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
8973}
8974
8975fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
8976 let n = list_wal_numbers(wal_dir)?;
8977 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
8978}
8979
8980fn next_wal_number(wal_dir: &Path) -> Result<u32> {
8981 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
8982}
8983
8984fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
8985 let _ = std::fs::create_dir_all(wal_dir);
8986 let mut max_n = None;
8987 for entry in std::fs::read_dir(wal_dir)? {
8988 let entry = entry?;
8989 let fname = entry.file_name();
8990 let Some(s) = fname.to_str() else {
8991 continue;
8992 };
8993 let Some(stripped) = s.strip_prefix("seg-") else {
8994 continue;
8995 };
8996 let Some(stripped) = stripped.strip_suffix(".wal") else {
8997 continue;
8998 };
8999 if let Ok(n) = stripped.parse::<u32>() {
9000 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
9001 }
9002 }
9003 Ok(max_n)
9004}