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
139#[derive(Clone)]
141pub struct Table {
142 dir: PathBuf,
143 table_id: u64,
144 name: String,
148 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
153 read_only: bool,
156 wal: WalSink,
157 memtable: Memtable,
158 mutable_run: MutableRun,
163 mutable_run_spill_bytes: u64,
165 compaction_zstd_level: i32,
168 allocator: RowIdAllocator,
169 epoch: Arc<EpochAuthority>,
170 persisted_epoch: u64,
173 data_generation: u64,
176 schema: Schema,
177 hot: HotIndex,
178 kek: Option<Arc<Kek>>,
181 column_keys: HashMap<u16, ([u8; 32], u8)>,
185 run_refs: Vec<RunRef>,
186 retiring: Vec<crate::manifest::RetiredRun>,
189 next_run_id: u64,
190 sync_byte_threshold: u64,
191 current_txn_id: u64,
196 bitmap: HashMap<u16, BitmapIndex>,
197 ann: HashMap<u16, AnnIndex>,
198 fm: HashMap<u16, FmIndex>,
199 sparse: HashMap<u16, SparseIndex>,
200 minhash: HashMap<u16, MinHashIndex>,
201 learned_range: HashMap<u16, ColumnLearnedRange>,
204 pk_by_row: HashMap<RowId, Vec<u8>>,
206 pinned: BTreeMap<Epoch, usize>,
209 pub(crate) live_count: u64,
212 reservoir: crate::reservoir::Reservoir,
215 reservoir_complete: bool,
223 had_deletes: bool,
227 agg_cache: HashMap<u64, CachedAgg>,
231 global_idx_epoch: u64,
235 indexes_complete: bool,
240 index_build_policy: IndexBuildPolicy,
242 pk_by_row_complete: bool,
249 flushed_epoch: u64,
252 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
255 snapshots: Arc<crate::retention::SnapshotRegistry>,
258 commit_lock: Arc<parking_lot::Mutex<()>>,
260 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
263 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
273 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
282 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
284 pending_delete_rids: roaring::RoaringBitmap,
287 pending_put_cols: std::collections::HashSet<u16>,
290 pending_rows: Vec<Row>,
296 pending_rows_auto_inc: Vec<bool>,
297 pending_dels: Vec<RowId>,
300 pending_truncate: Option<Epoch>,
304 auto_inc: Option<AutoIncState>,
307 ttl: Option<TtlPolicy>,
310}
311
312const _: () = {
319 const fn assert_sync<T: ?Sized + Sync>() {}
320 assert_sync::<Table>();
321};
322
323enum CachedData {
329 Rows(Arc<Vec<Row>>),
330 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
331}
332
333impl CachedData {
334 fn approx_bytes(&self) -> u64 {
335 match self {
336 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
337 CachedData::Columns(c) => c
338 .iter()
339 .map(|(_, c)| c.approx_bytes())
340 .sum::<u64>()
341 .saturating_add(c.len() as u64 * 16),
342 }
343 }
344}
345
346struct CachedEntry {
350 data: CachedData,
351 footprint: roaring::RoaringBitmap,
352 condition_cols: Vec<u16>,
353}
354
355struct ResultCache {
366 entries: std::collections::HashMap<u64, CachedEntry>,
367 order: std::collections::VecDeque<u64>,
368 bytes: u64,
369 max_bytes: u64,
370 dir: Option<std::path::PathBuf>,
371 #[allow(dead_code)]
372 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
373}
374
375#[derive(serde::Serialize, serde::Deserialize)]
377struct SerializedEntry {
378 condition_cols: Vec<u16>,
379 footprint_bits: Vec<u32>,
380 data: SerializedData,
381}
382
383#[derive(serde::Serialize, serde::Deserialize)]
384enum SerializedData {
385 Rows(Vec<Row>),
386 Columns(Vec<(u16, columnar::NativeColumn)>),
387}
388
389impl SerializedEntry {
390 fn from_entry(entry: &CachedEntry) -> Self {
391 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
392 let data = match &entry.data {
393 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
394 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
395 };
396 Self {
397 condition_cols: entry.condition_cols.clone(),
398 footprint_bits,
399 data,
400 }
401 }
402
403 fn into_entry(self) -> Option<CachedEntry> {
404 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
405 let data = match self.data {
406 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
407 SerializedData::Columns(c) => {
408 if !c.iter().all(|(_, col)| col.validate()) {
411 return None;
412 }
413 CachedData::Columns(Arc::new(c))
414 }
415 };
416 Some(CachedEntry {
417 data,
418 footprint,
419 condition_cols: self.condition_cols,
420 })
421 }
422}
423
424impl ResultCache {
425 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
426
427 fn new() -> Self {
428 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
429 }
430
431 fn with_max_bytes(max_bytes: u64) -> Self {
432 Self {
433 entries: std::collections::HashMap::new(),
434 order: std::collections::VecDeque::new(),
435 bytes: 0,
436 max_bytes,
437 dir: None,
438 cache_dek: None,
439 }
440 }
441
442 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
443 let _ = std::fs::create_dir_all(&dir);
444 self.dir = Some(dir);
445 self
446 }
447
448 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
449 self.cache_dek = dek;
450 self
451 }
452
453 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
454 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
455 }
456
457 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
461 let Some(path) = self.disk_path(key) else {
462 return;
463 };
464 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
465 Ok(s) => s,
466 Err(_) => return,
467 };
468 let on_disk = if let Some(dek) = &self.cache_dek {
470 match self.encrypt_cache(&serialized, dek) {
471 Some(b) => b,
472 None => return,
473 }
474 } else {
475 serialized
476 };
477 let tmp = path.with_extension("tmp");
478 use std::io::Write;
479 let write = || -> std::io::Result<()> {
480 let mut f = std::fs::File::create(&tmp)?;
481 f.write_all(&on_disk)?;
482 f.flush()?;
483 Ok(())
484 };
485 if write().is_err() {
486 let _ = std::fs::remove_file(&tmp);
487 return;
488 }
489 let _ = std::fs::rename(&tmp, &path);
490 }
491
492 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
494 let path = self.disk_path(key)?;
495 let bytes = std::fs::read(&path).ok()?;
496 let plaintext = if let Some(dek) = &self.cache_dek {
497 self.decrypt_cache(&bytes, dek)?
498 } else {
499 bytes
500 };
501 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
502 serialized.into_entry()
503 }
504
505 fn remove_from_disk(&self, key: u64) {
507 if let Some(path) = self.disk_path(key) {
508 let _ = std::fs::remove_file(&path);
509 }
510 }
511
512 #[cfg(feature = "encryption")]
514 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
515 use crate::encryption::Cipher;
516 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
517 let mut nonce = [0u8; 12];
518 crate::encryption::fill_random(&mut nonce);
519 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
520 let mut out = Vec::with_capacity(12 + ct.len());
521 out.extend_from_slice(&nonce);
522 out.extend_from_slice(&ct);
523 Some(out)
524 }
525
526 #[cfg(not(feature = "encryption"))]
527 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
528 None
529 }
530
531 #[cfg(feature = "encryption")]
533 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
534 use crate::encryption::Cipher;
535 if bytes.len() < 28 {
536 return None;
537 }
538 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
539 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
540 let ct = &bytes[12..];
541 cipher.decrypt_page(&nonce, ct).ok()
542 }
543
544 #[cfg(not(feature = "encryption"))]
545 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
546 None
547 }
548
549 fn load_persistent(&mut self) {
552 let Some(dir) = self.dir.as_ref().cloned() else {
553 return;
554 };
555 let entries = match std::fs::read_dir(&dir) {
556 Ok(e) => e,
557 Err(_) => return,
558 };
559 for entry in entries.flatten() {
560 let path = entry.path();
561 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
563 let _ = std::fs::remove_file(&path);
564 continue;
565 }
566 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
567 continue;
568 }
569 let stem = match path.file_stem().and_then(|s| s.to_str()) {
570 Some(s) => s,
571 None => continue,
572 };
573 let key = match u64::from_str_radix(stem, 16) {
574 Ok(k) => k,
575 Err(_) => continue,
576 };
577 let bytes = match std::fs::read(&path) {
578 Ok(b) => b,
579 Err(_) => continue,
580 };
581 let plaintext = if let Some(dek) = &self.cache_dek {
583 match self.decrypt_cache(&bytes, dek) {
584 Some(p) => p,
585 None => {
586 let _ = std::fs::remove_file(&path);
587 continue;
588 }
589 }
590 } else {
591 bytes
592 };
593 match bincode::deserialize::<SerializedEntry>(&plaintext) {
594 Ok(serialized) => {
595 if let Some(entry) = serialized.into_entry() {
596 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
597 self.entries.insert(key, entry);
598 self.order.push_back(key);
599 } else {
600 let _ = std::fs::remove_file(&path);
601 }
602 }
603 Err(_) => {
604 let _ = std::fs::remove_file(&path);
605 }
606 }
607 }
608 self.evict();
609 }
610
611 fn set_max_bytes(&mut self, max_bytes: u64) {
612 self.max_bytes = max_bytes;
613 self.evict();
614 }
615
616 fn touch(&mut self, key: u64) {
618 self.order.retain(|k| *k != key);
619 self.order.push_back(key);
620 }
621
622 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
623 let res = self.entries.get(&key).and_then(|e| match &e.data {
624 CachedData::Rows(r) => Some(r.clone()),
625 CachedData::Columns(_) => None,
626 });
627 if res.is_some() {
628 self.touch(key);
629 return res;
630 }
631 if let Some(entry) = self.load_from_disk(key) {
633 let res = match &entry.data {
634 CachedData::Rows(r) => Some(r.clone()),
635 CachedData::Columns(_) => None,
636 };
637 if res.is_some() {
638 let approx = entry.data.approx_bytes();
639 self.bytes = self.bytes.saturating_add(approx);
640 self.entries.insert(key, entry);
641 self.order.push_back(key);
642 self.evict();
643 return res;
644 }
645 }
646 None
647 }
648
649 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
650 let res = self.entries.get(&key).and_then(|e| match &e.data {
651 CachedData::Columns(c) => Some(c.clone()),
652 CachedData::Rows(_) => None,
653 });
654 if res.is_some() {
655 self.touch(key);
656 return res;
657 }
658 if let Some(entry) = self.load_from_disk(key) {
660 let res = match &entry.data {
661 CachedData::Columns(c) => Some(c.clone()),
662 CachedData::Rows(_) => None,
663 };
664 if res.is_some() {
665 let approx = entry.data.approx_bytes();
666 self.bytes = self.bytes.saturating_add(approx);
667 self.entries.insert(key, entry);
668 self.order.push_back(key);
669 self.evict();
670 return res;
671 }
672 }
673 None
674 }
675
676 fn insert(&mut self, key: u64, entry: CachedEntry) {
677 let approx = entry.data.approx_bytes();
678 if self.entries.remove(&key).is_some() {
679 self.order.retain(|k| *k != key);
680 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
681 }
682 self.store_to_disk(key, &entry);
684 self.bytes = self.bytes.saturating_add(approx);
685 self.entries.insert(key, entry);
686 self.order.push_back(key);
687 self.evict();
688 }
689
690 fn invalidate(
699 &mut self,
700 delete_rids: &roaring::RoaringBitmap,
701 put_cols: &std::collections::HashSet<u16>,
702 ) {
703 if self.entries.is_empty() {
704 return;
705 }
706 let has_deletes = !delete_rids.is_empty();
707 let to_remove: std::collections::HashSet<u64> = self
708 .entries
709 .iter()
710 .filter(|(_, e)| {
711 let delete_hit = if e.footprint.is_empty() {
712 has_deletes
713 } else {
714 e.footprint.intersection_len(delete_rids) > 0
715 };
716 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
717 delete_hit || col_hit
718 })
719 .map(|(&k, _)| k)
720 .collect();
721 for key in &to_remove {
722 if let Some(e) = self.entries.remove(key) {
723 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
724 }
725 self.remove_from_disk(*key);
726 }
727 if !to_remove.is_empty() {
728 self.order.retain(|k| !to_remove.contains(k));
729 }
730 }
731
732 fn clear(&mut self) {
733 if let Some(dir) = &self.dir {
735 if let Ok(entries) = std::fs::read_dir(dir) {
736 for entry in entries.flatten() {
737 let path = entry.path();
738 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
739 let _ = std::fs::remove_file(&path);
740 }
741 }
742 }
743 }
744 self.entries.clear();
745 self.order.clear();
746 self.bytes = 0;
747 }
748
749 fn evict(&mut self) {
750 while self.bytes > self.max_bytes {
751 let Some(k) = self.order.pop_front() else {
752 break;
753 };
754 if let Some(e) = self.entries.remove(&k) {
755 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
756 self.remove_from_disk(k);
760 }
761 }
762 }
763}
764
765type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
772
773fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
774 let _ = kek;
775 #[cfg(feature = "encryption")]
776 {
777 if let Some(k) = kek {
778 return (
779 Some(k.derive_table_wal_key(_table_id)),
780 Some(k.derive_cache_key()),
781 );
782 }
783 }
784 (None, None)
785}
786
787#[cfg(feature = "encryption")]
789fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
790 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
791}
792
793#[cfg(not(feature = "encryption"))]
794fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
795 Box::new(crate::encryption::PlaintextCipher)
796}
797
798fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
799 let Some(kek) = kek else {
800 return HashMap::new();
801 };
802 #[cfg(feature = "encryption")]
803 {
804 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
805 schema
806 .columns
807 .iter()
808 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
809 .map(|c| {
810 let scheme = if schema
811 .indexes
812 .iter()
813 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
814 {
815 SCHEME_OPE_RANGE
816 } else {
817 SCHEME_HMAC_EQ
818 };
819 let key: [u8; 32] = *kek.derive_column_key(c.id);
820 (c.id, (key, scheme))
821 })
822 .collect()
823 }
824 #[cfg(not(feature = "encryption"))]
825 {
826 let _ = (kek, schema);
827 HashMap::new()
828 }
829}
830
831pub(crate) struct SharedCtx {
836 pub epoch: Arc<EpochAuthority>,
837 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
838 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
839 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
840 pub kek: Option<Arc<Kek>>,
841 pub commit_lock: Arc<parking_lot::Mutex<()>>,
847 pub shared: Option<SharedWalCtx>,
851 pub table_name: Option<String>,
854 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
857 pub read_only: bool,
859}
860
861#[derive(Clone)]
867pub(crate) struct SharedWalCtx {
868 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
869 pub group: Arc<GroupCommit>,
870 pub poisoned: Arc<AtomicBool>,
871 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
872 pub change_wake: tokio::sync::broadcast::Sender<()>,
873}
874
875enum WalSink {
878 Private(Wal),
879 Shared(SharedWalCtx),
880 ReadOnly,
881}
882
883impl Clone for WalSink {
884 fn clone(&self) -> Self {
885 match self {
886 Self::Shared(shared) => Self::Shared(shared.clone()),
887 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
888 }
889 }
890}
891
892impl SharedCtx {
893 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
897 let n_shards = if cache_dir.is_some() {
901 1
902 } else {
903 crate::cache::CACHE_SHARDS
904 };
905 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
906 let page_cache = if let Some(d) = cache_dir {
907 Arc::new(crate::cache::Sharded::new(1, || {
908 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
909 }))
910 } else {
911 Arc::new(crate::cache::Sharded::new(n_shards, || {
912 crate::cache::PageCache::new(per_shard)
913 }))
914 };
915 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
916 let decoded_cache = Arc::new(crate::cache::Sharded::new(
917 crate::cache::CACHE_SHARDS,
918 || crate::cache::DecodedPageCache::new(decoded_per_shard),
919 ));
920 Self {
921 epoch: Arc::new(EpochAuthority::new(0)),
922 page_cache,
923 decoded_cache,
924 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
925 kek,
926 commit_lock: Arc::new(parking_lot::Mutex::new(())),
927 shared: None,
928 table_name: None,
929 auth: None,
930 read_only: false,
931 }
932 }
933}
934
935fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
939 use crate::query::Condition;
940 match c {
941 Condition::Pk(_)
943 | Condition::BitmapEq { .. }
944 | Condition::BitmapIn { .. }
945 | Condition::BytesPrefix { .. }
946 | Condition::IsNull { .. }
947 | Condition::IsNotNull { .. } => 0,
948 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
950 1
951 }
952 Condition::FmContains { .. }
954 | Condition::FmContainsAll { .. }
955 | Condition::Ann { .. }
956 | Condition::SparseMatch { .. } => 2,
957 }
958}
959
960impl Table {
961 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
962 let dir = dir.as_ref().to_path_buf();
963 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
964 Self::create_in(&dir, schema, table_id, ctx)
965 }
966
967 #[cfg(feature = "encryption")]
978 pub fn create_encrypted(
979 dir: impl AsRef<Path>,
980 schema: Schema,
981 table_id: u64,
982 passphrase: &str,
983 ) -> Result<Self> {
984 let dir = dir.as_ref();
985 std::fs::create_dir_all(dir.join(META_DIR))?;
986 let salt = crate::encryption::random_salt();
987 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
988 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
989 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
990 Self::create_in(dir, schema, table_id, ctx)
991 }
992
993 #[cfg(feature = "encryption")]
998 pub fn create_with_key(
999 dir: impl AsRef<Path>,
1000 schema: Schema,
1001 table_id: u64,
1002 key: &[u8],
1003 ) -> Result<Self> {
1004 let dir = dir.as_ref();
1005 std::fs::create_dir_all(dir.join(META_DIR))?;
1006 let salt = crate::encryption::random_salt();
1007 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
1008 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1009 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1010 Self::create_in(dir, schema, table_id, ctx)
1011 }
1012
1013 #[cfg(feature = "encryption")]
1015 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1016 let dir = dir.as_ref();
1017 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1018 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1019 MongrelError::NotFound(format!(
1020 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1021 salt_path
1022 ))
1023 })?;
1024 if salt_bytes.len() != crate::encryption::SALT_LEN {
1025 return Err(MongrelError::InvalidArgument(format!(
1026 "salt file is {} bytes, expected {}",
1027 salt_bytes.len(),
1028 crate::encryption::SALT_LEN
1029 )));
1030 }
1031 let mut salt = [0u8; crate::encryption::SALT_LEN];
1032 salt.copy_from_slice(&salt_bytes);
1033 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1034 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1035 Self::open_in(dir, ctx)
1036 }
1037
1038 pub(crate) fn create_in(
1039 dir: impl AsRef<Path>,
1040 schema: Schema,
1041 table_id: u64,
1042 ctx: SharedCtx,
1043 ) -> Result<Self> {
1044 schema.validate_auto_increment()?;
1045 schema.validate_defaults()?;
1046 schema.validate_ai()?;
1047 for index in &schema.indexes {
1048 index.validate_options()?;
1049 }
1050 let dir = dir.as_ref().to_path_buf();
1051 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1052 write_schema(&dir, &schema)?;
1053 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1054 let (wal, current_txn_id) = match ctx.shared.clone() {
1057 Some(s) => (WalSink::Shared(s), 0),
1058 None => {
1059 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1060 let mut w = if let Some(ref dk) = wal_dek {
1061 Wal::create_with_cipher(
1062 dir.join(WAL_DIR).join("seg-000000.wal"),
1063 Epoch(0),
1064 Some(make_cipher(dk)),
1065 0,
1066 )?
1067 } else {
1068 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1069 };
1070 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1071 (WalSink::Private(w), 1)
1072 }
1073 };
1074 let mut manifest = Manifest::new(table_id, schema.schema_id);
1075 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1080 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1081 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1082 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1083 let auto_inc = resolve_auto_inc(&schema);
1084 let rcache_dir = dir.join(RCACHE_DIR);
1085 Ok(Self {
1086 dir,
1087 table_id,
1088 name: ctx.table_name.unwrap_or_default(),
1089 auth: ctx.auth,
1090 read_only: ctx.read_only,
1091 wal,
1092 memtable: Memtable::new(),
1093 mutable_run: MutableRun::new(),
1094 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1095 compaction_zstd_level: 3,
1096 allocator: RowIdAllocator::new(0),
1097 epoch: ctx.epoch,
1098 persisted_epoch: 0,
1099 data_generation: 0,
1100 schema,
1101 hot: HotIndex::new(),
1102 kek: ctx.kek,
1103 column_keys,
1104 run_refs: Vec::new(),
1105 retiring: Vec::new(),
1106 next_run_id: 1,
1107 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1108 current_txn_id,
1109 bitmap,
1110 ann,
1111 fm,
1112 sparse,
1113 minhash,
1114 learned_range: HashMap::new(),
1115 pk_by_row: HashMap::new(),
1116 pinned: BTreeMap::new(),
1117 live_count: 0,
1118 reservoir: crate::reservoir::Reservoir::default(),
1119 reservoir_complete: true,
1120 had_deletes: false,
1121 agg_cache: HashMap::new(),
1122 global_idx_epoch: 0,
1123 indexes_complete: true,
1124 index_build_policy: IndexBuildPolicy::default(),
1125 pk_by_row_complete: false,
1126 flushed_epoch: 0,
1127 page_cache: ctx.page_cache,
1128 decoded_cache: ctx.decoded_cache,
1129 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1130 snapshots: ctx.snapshots,
1131 commit_lock: ctx.commit_lock,
1132 result_cache: Arc::new(parking_lot::Mutex::new(
1133 ResultCache::new()
1134 .with_dir(rcache_dir)
1135 .with_cache_dek(cache_dek.clone()),
1136 )),
1137 pending_delete_rids: roaring::RoaringBitmap::new(),
1138 pending_put_cols: std::collections::HashSet::new(),
1139 pending_rows: Vec::new(),
1140 pending_rows_auto_inc: Vec::new(),
1141 pending_dels: Vec::new(),
1142 pending_truncate: None,
1143 wal_dek,
1144 auto_inc,
1145 ttl: None,
1146 })
1147 }
1148
1149 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1153 let dir = dir.as_ref();
1154 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1155 Self::open_in(dir, ctx)
1156 }
1157
1158 #[cfg(feature = "encryption")]
1161 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1162 let dir = dir.as_ref();
1163 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1164 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1165 MongrelError::NotFound(format!(
1166 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1167 salt_path
1168 ))
1169 })?;
1170 let salt_len = crate::encryption::SALT_LEN;
1171 if salt_bytes.len() != salt_len {
1172 return Err(MongrelError::InvalidArgument(format!(
1173 "encryption salt is {} bytes, expected {salt_len}",
1174 salt_bytes.len()
1175 )));
1176 }
1177 let mut salt = [0u8; 16];
1178 salt.copy_from_slice(&salt_bytes);
1179 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1180 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1181 let t = Self::open_in(dir, ctx)?;
1182 Ok(t)
1183 }
1184
1185 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1186 let dir = dir.as_ref().to_path_buf();
1187 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1188 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1189 let schema: Schema = read_schema(&dir)?;
1190 schema.validate_ai()?;
1191 for index in &schema.indexes {
1192 index.validate_options()?;
1193 }
1194 let replay_epoch = Epoch(manifest.current_epoch);
1195 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1196 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1200 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1201 None => {
1202 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1203 let replayed = match &active {
1205 Some(path) => {
1206 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1207 crate::wal::replay_with_cipher(path, cipher)?
1208 }
1209 None => Vec::new(),
1210 };
1211 let mut w = match &active {
1212 Some(path) => Wal::create_with_cipher(
1213 path,
1214 replay_epoch,
1215 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1216 0,
1217 )?,
1218 None => Wal::create_with_cipher(
1219 dir.join(WAL_DIR).join("seg-000000.wal"),
1220 replay_epoch,
1221 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1222 0,
1223 )?,
1224 };
1225 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1226 (WalSink::Private(w), replayed, 1)
1227 }
1228 };
1229
1230 let mut memtable = Memtable::new();
1231 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1232 let persisted_epoch = manifest.current_epoch;
1233 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1240 s.next = manifest.auto_inc_next;
1241 s.seeded = manifest.auto_inc_next > 0;
1242 s
1243 });
1244
1245 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1252 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1253 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1254 std::collections::BTreeMap::new();
1255 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1256 let mut saw_delete = false;
1257 for record in replayed {
1258 let txn_id = record.txn_id;
1259 match record.op {
1260 Op::Put { rows, .. } => {
1261 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1262 for row in &rows {
1263 allocator.advance_to(row.row_id);
1264 if let Some(ai) = auto_inc.as_mut() {
1265 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1266 if *n + 1 > ai.next {
1267 ai.next = *n + 1;
1268 }
1269 }
1270 }
1271 }
1272 staged_puts.entry(txn_id).or_default().extend(rows);
1273 }
1274 Op::Delete { row_ids, .. } => {
1275 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1276 }
1277 Op::TxnCommit { epoch, .. } => {
1278 let commit_epoch = Epoch(epoch);
1279 if let Some(puts) = staged_puts.remove(&txn_id) {
1280 if commit_epoch.0 > manifest.flushed_epoch {
1281 for row in &puts {
1282 memtable.upsert(row.clone());
1283 }
1284 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1285 }
1286 }
1287 if let Some(dels) = staged_deletes.remove(&txn_id) {
1288 saw_delete = true;
1289 if commit_epoch.0 > manifest.flushed_epoch {
1290 for rid in dels {
1291 memtable.tombstone(rid, commit_epoch);
1292 replayed_deletes.push((rid, commit_epoch));
1293 }
1294 }
1295 }
1296 }
1297 Op::TxnAbort => {
1298 staged_puts.remove(&txn_id);
1299 staged_deletes.remove(&txn_id);
1300 }
1301 Op::TruncateTable { .. }
1302 | Op::ExternalTableState { .. }
1303 | Op::Flush { .. }
1304 | Op::Ddl(_)
1305 | Op::BeforeImage { .. }
1306 | Op::CommitTimestamp { .. } => {}
1307 }
1308 }
1309
1310 let rcache_dir = dir.join(RCACHE_DIR);
1311 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1312 let mut db = Self {
1313 dir,
1314 table_id: manifest.table_id,
1315 name: ctx.table_name.unwrap_or_default(),
1316 auth: ctx.auth,
1317 read_only: ctx.read_only,
1318 wal,
1319 memtable,
1320 mutable_run: MutableRun::new(),
1321 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1322 compaction_zstd_level: 3,
1323 allocator,
1324 epoch: ctx.epoch,
1325 persisted_epoch,
1326 data_generation: persisted_epoch,
1327 schema,
1328 hot: HotIndex::new(),
1329 kek: ctx.kek,
1330 column_keys,
1331 run_refs: manifest.runs.clone(),
1332 retiring: manifest.retiring.clone(),
1333 next_run_id: manifest
1334 .runs
1335 .iter()
1336 .map(|r| r.run_id as u64 + 1)
1337 .max()
1338 .unwrap_or(1),
1339 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1340 current_txn_id,
1341 bitmap: HashMap::new(),
1342 ann: HashMap::new(),
1343 fm: HashMap::new(),
1344 sparse: HashMap::new(),
1345 minhash: HashMap::new(),
1346 learned_range: HashMap::new(),
1347 pk_by_row: HashMap::new(),
1348 pinned: BTreeMap::new(),
1349 live_count: manifest.live_count,
1350 reservoir: crate::reservoir::Reservoir::default(),
1351 reservoir_complete: false,
1352 had_deletes: saw_delete
1353 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1354 != manifest.live_count,
1355 agg_cache: HashMap::new(),
1356 global_idx_epoch: manifest.global_idx_epoch,
1357 indexes_complete: true,
1358 index_build_policy: IndexBuildPolicy::default(),
1359 pk_by_row_complete: false,
1360 flushed_epoch: manifest.flushed_epoch,
1361 page_cache: ctx.page_cache,
1362 decoded_cache: ctx.decoded_cache,
1363 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1364 snapshots: ctx.snapshots,
1365 commit_lock: ctx.commit_lock,
1366 result_cache: Arc::new(parking_lot::Mutex::new(
1367 ResultCache::new()
1368 .with_dir(rcache_dir)
1369 .with_cache_dek(cache_dek.clone()),
1370 )),
1371 pending_delete_rids: roaring::RoaringBitmap::new(),
1372 pending_put_cols: std::collections::HashSet::new(),
1373 pending_rows: Vec::new(),
1374 pending_rows_auto_inc: Vec::new(),
1375 pending_dels: Vec::new(),
1376 pending_truncate: None,
1377 wal_dek,
1378 auto_inc,
1379 ttl: manifest.ttl,
1380 };
1381
1382 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1385
1386 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1391 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1392 c.epoch_built == manifest.global_idx_epoch
1393 && manifest.global_idx_epoch > 0
1394 && manifest
1395 .runs
1396 .iter()
1397 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1398 });
1399 if let Some(loaded) = checkpoint {
1400 if checkpoint_valid {
1401 db.hot = loaded.hot;
1402 db.bitmap = loaded.bitmap;
1403 db.ann = loaded.ann;
1404 db.fm = loaded.fm;
1405 db.sparse = loaded.sparse;
1406 db.minhash = loaded.minhash;
1407 db.learned_range = loaded.learned_range;
1408 }
1411 }
1412 if !checkpoint_valid {
1413 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1414 db.bitmap = bitmap;
1415 db.ann = ann;
1416 db.fm = fm;
1417 db.sparse = sparse;
1418 db.minhash = minhash;
1419 db.rebuild_indexes_from_runs()?;
1420 db.build_learned_ranges()?;
1421 }
1422
1423 for (epoch, group) in replayed_puts {
1428 let (losers, winner_pks) = db.partition_pk_winners(&group);
1429 for (key, &row_id) in &winner_pks {
1430 if let Some(old_rid) = db.hot.get(key) {
1431 if old_rid != row_id {
1432 db.tombstone_row(old_rid, epoch, false);
1433 }
1434 }
1435 }
1436 for &loser_rid in &losers {
1437 db.tombstone_row(loser_rid, epoch, false);
1438 }
1439 for (key, row_id) in winner_pks {
1440 db.insert_hot_pk(key, row_id);
1441 }
1442 if db.schema.primary_key().is_none() {
1443 for r in &group {
1444 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1445 }
1446 }
1447 for r in &group {
1448 if !losers.contains(&r.row_id) {
1449 db.index_row(r);
1450 }
1451 }
1452 }
1453 for (rid, epoch) in &replayed_deletes {
1457 db.remove_hot_for_row(*rid, *epoch);
1458 }
1459
1460 db.result_cache.lock().load_persistent();
1467 Ok(db)
1468 }
1469
1470 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1476 if self.reservoir_complete {
1477 return Ok(());
1478 }
1479 self.rebuild_reservoir()?;
1480 self.reservoir_complete = true;
1481 Ok(())
1482 }
1483
1484 fn rebuild_reservoir(&mut self) -> Result<()> {
1487 let snap = self.snapshot();
1488 let rows = self.visible_rows(snap)?;
1489 self.reservoir.reset();
1490 for r in rows {
1491 self.reservoir.offer(r.row_id.0);
1492 }
1493 Ok(())
1494 }
1495
1496 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1497 self.hot = HotIndex::new();
1498 self.pk_by_row.clear();
1499 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1500 self.bitmap = bitmap;
1501 self.ann = ann;
1502 self.fm = fm;
1503 self.sparse = sparse;
1504 self.minhash = minhash;
1505 let snapshot = Epoch(u64::MAX);
1506 let ttl_now = unix_nanos_now();
1507 for rr in self.run_refs.clone() {
1508 let mut reader = self.open_reader(rr.run_id)?;
1509 for row in reader.visible_rows(snapshot)? {
1510 if self.row_expired_at(&row, ttl_now) {
1511 continue;
1512 }
1513 let tok_row = self.tokenized_for_indexes(&row);
1514 index_into(
1515 &self.schema,
1516 &tok_row,
1517 &mut self.hot,
1518 &mut self.bitmap,
1519 &mut self.ann,
1520 &mut self.fm,
1521 &mut self.sparse,
1522 &mut self.minhash,
1523 );
1524 }
1525 }
1526 for row in self.mutable_run.visible_versions(snapshot) {
1527 if row.deleted {
1528 self.remove_hot_for_row(row.row_id, snapshot);
1529 } else if !self.row_expired_at(&row, ttl_now) {
1530 self.index_row(&row);
1531 }
1532 }
1533 for row in self.memtable.visible_versions(snapshot) {
1534 if row.deleted {
1535 self.remove_hot_for_row(row.row_id, snapshot);
1536 } else if !self.row_expired_at(&row, ttl_now) {
1537 self.index_row(&row);
1538 }
1539 }
1540 self.refresh_pk_by_row_from_hot();
1541 Ok(())
1542 }
1543
1544 fn refresh_pk_by_row_from_hot(&mut self) {
1545 self.pk_by_row_complete = true;
1546 if self.schema.primary_key().is_none() {
1547 self.pk_by_row.clear();
1548 return;
1549 }
1550 self.pk_by_row = self
1556 .hot
1557 .entries()
1558 .into_iter()
1559 .map(|(key, row_id)| (row_id, key))
1560 .collect();
1561 }
1562
1563 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1564 if self.schema.primary_key().is_some() {
1565 self.pk_by_row.insert(row_id, key.clone());
1566 }
1567 self.hot.insert(key, row_id);
1568 }
1569
1570 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1574 self.learned_range.clear();
1575 if self.run_refs.len() != 1 {
1576 return Ok(());
1577 }
1578 let cols: Vec<(u16, usize)> = self
1579 .schema
1580 .indexes
1581 .iter()
1582 .filter(|i| i.kind == IndexKind::LearnedRange)
1583 .map(|i| {
1584 (
1585 i.column_id,
1586 i.options
1587 .learned_range
1588 .as_ref()
1589 .map(|options| options.epsilon)
1590 .unwrap_or(16),
1591 )
1592 })
1593 .collect();
1594 if cols.is_empty() {
1595 return Ok(());
1596 }
1597 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1598 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1599 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1600 _ => return Ok(()),
1601 };
1602 for (cid, epsilon) in cols {
1603 let ty = self
1604 .schema
1605 .columns
1606 .iter()
1607 .find(|c| c.id == cid)
1608 .map(|c| c.ty.clone())
1609 .unwrap_or(TypeId::Int64);
1610 match ty {
1611 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1612 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1613 let pairs: Vec<(i64, u64)> = data
1614 .iter()
1615 .zip(row_ids.iter())
1616 .map(|(v, r)| (*v, *r))
1617 .collect();
1618 self.learned_range.insert(
1619 cid,
1620 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
1621 );
1622 }
1623 }
1624 TypeId::Float64 => {
1625 if let columnar::NativeColumn::Float64 { data, .. } =
1626 reader.column_native(cid)?
1627 {
1628 let pairs: Vec<(f64, u64)> = data
1629 .iter()
1630 .zip(row_ids.iter())
1631 .map(|(v, r)| (*v, *r))
1632 .collect();
1633 self.learned_range.insert(
1634 cid,
1635 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
1636 );
1637 }
1638 }
1639 _ => {}
1640 }
1641 }
1642 Ok(())
1643 }
1644
1645 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1652 if self.indexes_complete {
1653 crate::trace::QueryTrace::record(|t| {
1654 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1655 });
1656 return Ok(());
1657 }
1658 crate::trace::QueryTrace::record(|t| {
1659 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1660 });
1661 self.rebuild_indexes_from_runs()?;
1662 self.build_learned_ranges()?;
1663 self.indexes_complete = true;
1664 let epoch = self.current_epoch();
1665 self.checkpoint_indexes(epoch);
1666 Ok(())
1667 }
1668
1669 fn pending_epoch(&self) -> Epoch {
1670 Epoch(self.epoch.visible().0 + 1)
1671 }
1672
1673 fn is_shared(&self) -> bool {
1676 matches!(self.wal, WalSink::Shared(_))
1677 }
1678
1679 fn ensure_txn_id(&mut self) -> u64 {
1683 if self.current_txn_id == 0 {
1684 let id = match &self.wal {
1685 WalSink::Shared(s) => {
1686 let mut g = s.txn_ids.lock();
1687 let v = *g;
1688 *g = g.wrapping_add(1);
1689 v
1690 }
1691 WalSink::Private(_) => 1,
1692 WalSink::ReadOnly => 1,
1693 };
1694 self.current_txn_id = id;
1695 }
1696 self.current_txn_id
1697 }
1698
1699 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1702 self.ensure_writable()?;
1703 let txn_id = self.ensure_txn_id();
1704 let table_id = self.table_id;
1705 match &mut self.wal {
1706 WalSink::Private(w) => {
1707 w.append_txn(txn_id, op)?;
1708 }
1709 WalSink::Shared(s) => {
1710 s.wal.lock().append(txn_id, table_id, op)?;
1711 }
1712 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
1713 }
1714 Ok(())
1715 }
1716
1717 fn ensure_writable(&self) -> Result<()> {
1718 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
1719 Err(MongrelError::ReadOnlyReplica)
1720 } else {
1721 Ok(())
1722 }
1723 }
1724
1725 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1736 match &self.auth {
1737 Some(checker) => checker.check(&self.name, perm),
1738 None => Ok(()),
1739 }
1740 }
1741 pub fn require_select(&self) -> Result<()> {
1746 self.require(crate::auth_state::RequiredPermission::Select)
1747 }
1748 fn require_insert(&self) -> Result<()> {
1749 self.require(crate::auth_state::RequiredPermission::Insert)
1750 }
1751 #[allow(dead_code)]
1755 fn require_update(&self) -> Result<()> {
1756 self.require(crate::auth_state::RequiredPermission::Update)
1757 }
1758 fn require_delete(&self) -> Result<()> {
1759 self.require(crate::auth_state::RequiredPermission::Delete)
1760 }
1761
1762 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1765 self.require_insert()?;
1766 Ok(self.put_returning(columns)?.0)
1767 }
1768
1769 pub fn put_returning(
1774 &mut self,
1775 mut columns: Vec<(u16, Value)>,
1776 ) -> Result<(RowId, Option<i64>)> {
1777 self.require_insert()?;
1778 let assigned = self.fill_auto_inc(&mut columns)?;
1779 self.apply_defaults(&mut columns)?;
1780 self.schema.validate_values(&columns)?;
1781 let row_id = if self.schema.clustered {
1786 self.derive_clustered_row_id(&columns)?
1787 } else {
1788 self.allocator.alloc()
1789 };
1790 let epoch = self.pending_epoch();
1791 let mut row = Row::new(row_id, epoch);
1792 for (col_id, val) in columns {
1793 row.columns.insert(col_id, val);
1794 }
1795 self.commit_rows(vec![row], assigned.is_some())?;
1796 Ok((row_id, assigned))
1797 }
1798
1799 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1802 self.require_insert()?;
1803 Ok(self
1804 .put_batch_returning(batch)?
1805 .into_iter()
1806 .map(|(r, _)| r)
1807 .collect())
1808 }
1809
1810 pub fn put_batch_returning(
1813 &mut self,
1814 batch: Vec<Vec<(u16, Value)>>,
1815 ) -> Result<Vec<(RowId, Option<i64>)>> {
1816 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1817 for mut cols in batch {
1818 let assigned = self.fill_auto_inc(&mut cols)?;
1819 self.apply_defaults(&mut cols)?;
1820 filled.push((cols, assigned));
1821 }
1822 for (cols, _) in &filled {
1823 self.schema.validate_values(cols)?;
1824 }
1825 let epoch = self.pending_epoch();
1826 let mut rows = Vec::with_capacity(filled.len());
1827 let mut ids = Vec::with_capacity(filled.len());
1828 for (cols, assigned) in filled {
1829 let row_id = if self.schema.clustered {
1830 self.derive_clustered_row_id(&cols)?
1831 } else {
1832 self.allocator.alloc()
1833 };
1834 let mut row = Row::new(row_id, epoch);
1835 for (c, v) in cols {
1836 row.columns.insert(c, v);
1837 }
1838 ids.push((row_id, assigned));
1839 rows.push(row);
1840 }
1841 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1842 self.commit_rows(rows, all_auto_generated)?;
1843 Ok(ids)
1844 }
1845
1846 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1852 self.ensure_writable()?;
1853 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1854 return Ok(None);
1855 };
1856 let pos = columns.iter().position(|(c, _)| *c == cid);
1857 let assigned = match pos {
1858 Some(i) => match &columns[i].1 {
1859 Value::Null => {
1860 let next = self.alloc_auto_inc_value()?;
1861 columns[i].1 = Value::Int64(next);
1862 Some(next)
1863 }
1864 Value::Int64(n) => {
1865 self.advance_auto_inc_past(*n)?;
1866 None
1867 }
1868 other => {
1869 return Err(MongrelError::InvalidArgument(format!(
1870 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1871 other
1872 )))
1873 }
1874 },
1875 None => {
1876 let next = self.alloc_auto_inc_value()?;
1877 columns.push((cid, Value::Int64(next)));
1878 Some(next)
1879 }
1880 };
1881 Ok(assigned)
1882 }
1883
1884 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1890 for col in &self.schema.columns {
1891 let Some(expr) = &col.default_value else {
1892 continue;
1893 };
1894 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
1896 continue;
1897 }
1898 let pos = columns.iter().position(|(c, _)| *c == col.id);
1899 let needs_default = match pos {
1900 None => true,
1901 Some(i) => matches!(columns[i].1, Value::Null),
1902 };
1903 if !needs_default {
1904 continue;
1905 }
1906 let v = match expr {
1907 crate::schema::DefaultExpr::Static(v) => v.clone(),
1908 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
1909 crate::schema::DefaultExpr::Uuid => {
1910 let mut buf = [0u8; 16];
1911 getrandom::getrandom(&mut buf)
1912 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
1913 Value::Uuid(buf)
1914 }
1915 };
1916 match pos {
1917 None => columns.push((col.id, v)),
1918 Some(i) => columns[i].1 = v,
1919 }
1920 }
1921 Ok(())
1922 }
1923
1924 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1926 self.ensure_auto_inc_seeded()?;
1927 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1929 let v = ai.next;
1930 ai.next = ai.next.saturating_add(1);
1931 Ok(v)
1932 }
1933
1934 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1937 self.ensure_auto_inc_seeded()?;
1938 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1939 let floor = used.saturating_add(1).max(1);
1940 if ai.next < floor {
1941 ai.next = floor;
1942 }
1943 Ok(())
1944 }
1945
1946 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1951 let needs_seed = match self.auto_inc {
1952 Some(ai) => !ai.seeded,
1953 None => return Ok(()),
1954 };
1955 if !needs_seed {
1956 return Ok(());
1957 }
1958 if self.seed_empty_auto_inc() {
1959 return Ok(());
1960 }
1961 let cid = self
1962 .auto_inc
1963 .as_ref()
1964 .expect("auto-inc column present")
1965 .column_id;
1966 let max = self.scan_max_int64(cid)?;
1967 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1968 let floor = max.saturating_add(1).max(1);
1969 if ai.next < floor {
1970 ai.next = floor;
1971 }
1972 ai.seeded = true;
1973 Ok(())
1974 }
1975
1976 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1977 if n == 0 || self.auto_inc.is_none() {
1978 return Ok(None);
1979 }
1980 self.ensure_auto_inc_seeded()?;
1981 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1982 let start = ai.next;
1983 ai.next = ai.next.saturating_add(n as i64);
1984 Ok(Some(start))
1985 }
1986
1987 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1991 let mut max: i64 = 0;
1992 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1993 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1994 if *n > max {
1995 max = *n;
1996 }
1997 }
1998 }
1999 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2000 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2001 if *n > max {
2002 max = *n;
2003 }
2004 }
2005 }
2006 for rr in self.run_refs.clone() {
2007 let reader = self.open_reader(rr.run_id)?;
2008 if let Some(stats) = reader.column_page_stats(column_id) {
2009 for s in stats {
2010 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2011 if n > max {
2012 max = n;
2013 }
2014 }
2015 }
2016 } else if reader.has_column(column_id) {
2017 if let columnar::NativeColumn::Int64 { data, validity } =
2018 reader.column_native_shared(column_id)?
2019 {
2020 for (i, n) in data.iter().enumerate() {
2021 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2022 {
2023 max = *n;
2024 }
2025 }
2026 }
2027 }
2028 }
2029 Ok(max)
2030 }
2031
2032 fn seed_empty_auto_inc(&mut self) -> bool {
2033 let Some(ai) = self.auto_inc.as_mut() else {
2034 return false;
2035 };
2036 if ai.seeded || self.live_count != 0 {
2037 return false;
2038 }
2039 if ai.next < 1 {
2040 ai.next = 1;
2041 }
2042 ai.seeded = true;
2043 true
2044 }
2045
2046 fn advance_auto_inc_from_native_columns(
2047 &mut self,
2048 columns: &[(u16, columnar::NativeColumn)],
2049 n: usize,
2050 live_before: u64,
2051 ) -> Result<()> {
2052 let Some(ai) = self.auto_inc.as_mut() else {
2053 return Ok(());
2054 };
2055 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2056 return Ok(());
2057 };
2058 let columnar::NativeColumn::Int64 { data, validity } = col else {
2059 return Err(MongrelError::InvalidArgument(format!(
2060 "AUTO_INCREMENT column {} must be Int64",
2061 ai.column_id
2062 )));
2063 };
2064 let max = if native_int64_strictly_increasing(col, n) {
2065 data.get(n.saturating_sub(1)).copied()
2066 } else {
2067 data.iter()
2068 .take(n)
2069 .enumerate()
2070 .filter_map(|(i, v)| {
2071 if validity.is_empty() || columnar::validity_bit(validity, i) {
2072 Some(*v)
2073 } else {
2074 None
2075 }
2076 })
2077 .max()
2078 };
2079 if let Some(max) = max {
2080 let floor = max.saturating_add(1).max(1);
2081 if ai.next < floor {
2082 ai.next = floor;
2083 }
2084 if ai.seeded || live_before == 0 {
2085 ai.seeded = true;
2086 }
2087 }
2088 Ok(())
2089 }
2090
2091 fn fill_auto_inc_native_columns(
2092 &mut self,
2093 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2094 n: usize,
2095 ) -> Result<()> {
2096 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2097 return Ok(());
2098 };
2099 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2100 if let Some(start) = self.alloc_auto_inc_range(n)? {
2101 columns.push((
2102 cid,
2103 columnar::NativeColumn::Int64 {
2104 data: (start..start.saturating_add(n as i64)).collect(),
2105 validity: vec![0xFF; n.div_ceil(8)],
2106 },
2107 ));
2108 }
2109 return Ok(());
2110 };
2111
2112 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2113 return Err(MongrelError::InvalidArgument(format!(
2114 "AUTO_INCREMENT column {cid} must be Int64"
2115 )));
2116 };
2117 if data.len() < n {
2118 return Err(MongrelError::InvalidArgument(format!(
2119 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2120 data.len()
2121 )));
2122 }
2123 if columnar::all_non_null(validity, n) {
2124 return Ok(());
2125 }
2126 if validity.iter().all(|b| *b == 0) {
2127 if let Some(start) = self.alloc_auto_inc_range(n)? {
2128 for (i, slot) in data.iter_mut().take(n).enumerate() {
2129 *slot = start.saturating_add(i as i64);
2130 }
2131 *validity = vec![0xFF; n.div_ceil(8)];
2132 }
2133 return Ok(());
2134 }
2135
2136 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2137 for (i, slot) in data.iter_mut().enumerate().take(n) {
2138 if columnar::validity_bit(validity, i) {
2139 self.advance_auto_inc_past(*slot)?;
2140 } else {
2141 *slot = self.alloc_auto_inc_value()?;
2142 }
2143 }
2144 *validity = new_validity;
2145 Ok(())
2146 }
2147
2148 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2162 self.ensure_writable()?;
2163 if self.auto_inc.is_none() {
2164 return Ok(None);
2165 }
2166 Ok(Some(self.alloc_auto_inc_value()?))
2167 }
2168
2169 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2175 let payload = bincode::serialize(&rows)?;
2176 self.wal_append_data(Op::Put {
2177 table_id: self.table_id,
2178 rows: payload,
2179 })?;
2180 if self.is_shared() {
2181 self.pending_rows_auto_inc
2182 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2183 self.pending_rows.extend(rows);
2184 } else {
2185 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2186 }
2187 Ok(())
2188 }
2189
2190 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2195 self.apply_put_rows_inner(rows, true)
2196 }
2197
2198 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2199 if check_existing_pk {
2200 self.ensure_indexes_complete()?;
2201 }
2202 if rows.len() == 1 {
2206 let row = rows.into_iter().next().expect("len checked");
2207 return self.apply_put_row_single(row, check_existing_pk);
2208 }
2209 let pk_id = self.schema.primary_key().map(|c| c.id);
2226 let probe = match pk_id {
2227 Some(pid) => {
2228 check_existing_pk
2229 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2230 }
2231 None => false,
2232 };
2233 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2236 for r in rows {
2237 for &cid in r.columns.keys() {
2238 self.pending_put_cols.insert(cid);
2239 }
2240 match pk_id {
2241 Some(pid) if probe || maintain_pk_by_row => {
2242 if let Some(pk_val) = r.columns.get(&pid) {
2243 let key = self.index_lookup_key(pid, pk_val);
2244 if probe {
2245 if let Some(old_rid) = self.hot.get(&key) {
2246 if old_rid != r.row_id {
2247 self.tombstone_row(old_rid, r.committed_epoch, true);
2248 }
2249 }
2250 }
2251 if maintain_pk_by_row {
2252 self.pk_by_row.insert(r.row_id, key);
2253 }
2254 }
2255 }
2256 Some(_) => {}
2257 None => {
2258 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2259 }
2260 }
2261 self.index_row(&r);
2262 self.reservoir.offer(r.row_id.0);
2263 self.memtable.upsert(r);
2264 self.live_count = self.live_count.saturating_add(1);
2267 }
2268 self.data_generation = self.data_generation.wrapping_add(1);
2269 Ok(())
2270 }
2271
2272 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2276 for &cid in row.columns.keys() {
2277 self.pending_put_cols.insert(cid);
2278 }
2279 let epoch = row.committed_epoch;
2280 if let Some(pk_col) = self.schema.primary_key() {
2281 let pk_id = pk_col.id;
2282 if let Some(pk_val) = row.columns.get(&pk_id) {
2283 let maintain_pk_by_row = self.pk_by_row_complete;
2287 if check_existing_pk || maintain_pk_by_row {
2288 let key = self.index_lookup_key(pk_id, pk_val);
2289 if check_existing_pk {
2290 if let Some(old_rid) = self.hot.get(&key) {
2291 if old_rid != row.row_id {
2292 self.tombstone_row(old_rid, epoch, true);
2293 }
2294 }
2295 }
2296 if maintain_pk_by_row {
2297 self.pk_by_row.insert(row.row_id, key);
2298 }
2299 }
2300 }
2301 } else {
2302 self.hot
2303 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2304 }
2305 self.index_row(&row);
2306 self.reservoir.offer(row.row_id.0);
2307 self.memtable.upsert(row);
2308 self.live_count = self.live_count.saturating_add(1);
2309 self.data_generation = self.data_generation.wrapping_add(1);
2310 Ok(())
2311 }
2312
2313 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2316 self.allocator.alloc()
2317 }
2318
2319 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2325 let pk = self.schema.primary_key().ok_or_else(|| {
2326 MongrelError::Schema("clustered table requires a single-column primary key".into())
2327 })?;
2328 let pk_val = columns
2329 .iter()
2330 .find(|(id, _)| *id == pk.id)
2331 .map(|(_, v)| v)
2332 .ok_or_else(|| {
2333 MongrelError::Schema(format!(
2334 "clustered table missing primary key column {} ({})",
2335 pk.id, pk.name
2336 ))
2337 })?;
2338 let key_bytes = pk_val.encode_key();
2339 let mut hash: u64 = 0xcbf29ce484222325;
2341 for &b in &key_bytes {
2342 hash ^= b as u64;
2343 hash = hash.wrapping_mul(0x100000001b3);
2344 }
2345 Ok(RowId(hash.max(1)))
2348 }
2349
2350 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2358 self.ensure_indexes_complete()?;
2359 let n = rows.len();
2360 for r in rows {
2361 for &cid in r.columns.keys() {
2362 self.pending_put_cols.insert(cid);
2363 }
2364 }
2365 let (losers, winner_pks) = self.partition_pk_winners(rows);
2366 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2367 for (key, &row_id) in &winner_pks {
2369 if let Some(old_rid) = self.hot.get(key) {
2370 if old_rid != row_id {
2371 self.tombstone_row(old_rid, epoch, true);
2372 }
2373 }
2374 }
2375 for &loser_rid in &losers {
2378 self.tombstone_row(loser_rid, epoch, false);
2379 }
2380 for (key, row_id) in winner_pks {
2382 self.insert_hot_pk(key, row_id);
2383 }
2384 if self.schema.primary_key().is_none() {
2385 for r in rows {
2386 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2387 }
2388 }
2389 for r in rows {
2390 self.allocator.advance_to(r.row_id);
2391 if !losers.contains(&r.row_id) {
2392 self.index_row(r);
2393 }
2394 }
2395 for r in rows {
2396 if !losers.contains(&r.row_id) {
2397 self.reservoir.offer(r.row_id.0);
2398 }
2399 }
2400 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2401 self.data_generation = self.data_generation.wrapping_add(1);
2402 Ok(())
2403 }
2404
2405 pub(crate) fn recover_apply(
2410 &mut self,
2411 rows: Vec<Row>,
2412 deletes: Vec<(RowId, Epoch)>,
2413 ) -> Result<()> {
2414 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2418 std::collections::BTreeMap::new();
2419 for row in rows {
2420 self.allocator.advance_to(row.row_id);
2421 if let Some(ai) = self.auto_inc.as_mut() {
2426 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2427 if *n + 1 > ai.next {
2428 ai.next = *n + 1;
2429 }
2430 }
2431 }
2432 by_epoch.entry(row.committed_epoch).or_default().push(row);
2433 }
2434 for (epoch, group) in by_epoch {
2435 let (losers, winner_pks) = self.partition_pk_winners(&group);
2436 for (key, &row_id) in &winner_pks {
2438 if let Some(old_rid) = self.hot.get(key) {
2439 if old_rid != row_id {
2440 self.tombstone_row(old_rid, epoch, false);
2441 }
2442 }
2443 }
2444 for (key, row_id) in winner_pks {
2445 self.insert_hot_pk(key, row_id);
2446 }
2447 if self.schema.primary_key().is_none() {
2448 for r in &group {
2449 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2450 }
2451 }
2452 for r in &group {
2453 if !losers.contains(&r.row_id) {
2454 self.memtable.upsert(r.clone());
2455 self.index_row(r);
2456 }
2457 }
2458 }
2459 for (rid, epoch) in deletes {
2460 self.memtable.tombstone(rid, epoch);
2461 self.remove_hot_for_row(rid, epoch);
2462 }
2463 self.reservoir_complete = false;
2466 Ok(())
2467 }
2468
2469 pub(crate) fn flushed_epoch(&self) -> u64 {
2471 self.flushed_epoch
2472 }
2473
2474 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2475 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2476 }
2477
2478 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2480 self.schema.validate_values(cells)
2481 }
2482
2483 fn validate_columns_not_null(
2487 &self,
2488 columns: &[(u16, columnar::NativeColumn)],
2489 n: usize,
2490 ) -> Result<()> {
2491 let by_id: HashMap<u16, &columnar::NativeColumn> =
2492 columns.iter().map(|(id, c)| (*id, c)).collect();
2493 for col in &self.schema.columns {
2494 if !col.flags.contains(ColumnFlags::NULLABLE) {
2495 match by_id.get(&col.id) {
2496 None => {
2497 return Err(MongrelError::InvalidArgument(format!(
2498 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2499 col.name, col.id
2500 )));
2501 }
2502 Some(c) => {
2503 if c.null_count(n) != 0 {
2504 return Err(MongrelError::InvalidArgument(format!(
2505 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2506 col.name, col.id
2507 )));
2508 }
2509 }
2510 }
2511 }
2512 if let TypeId::Enum { variants } = &col.ty {
2513 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2514 if by_id.contains_key(&col.id) {
2515 return Err(MongrelError::InvalidArgument(format!(
2516 "column '{}' ({}) enum requires a bytes column",
2517 col.name, col.id
2518 )));
2519 }
2520 continue;
2521 };
2522 for index in 0..n {
2523 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2524 continue;
2525 };
2526 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2527 return Err(MongrelError::InvalidArgument(format!(
2528 "column '{}' ({}) enum value {:?} is not one of {:?}",
2529 col.name,
2530 col.id,
2531 String::from_utf8_lossy(value),
2532 variants
2533 )));
2534 }
2535 }
2536 }
2537 }
2538 Ok(())
2539 }
2540
2541 fn bulk_pk_winner_indices(
2546 &self,
2547 columns: &[(u16, columnar::NativeColumn)],
2548 n: usize,
2549 ) -> Option<Vec<usize>> {
2550 let pk_col = self.schema.primary_key()?;
2551 let pk_id = pk_col.id;
2552 let pk_ty = pk_col.ty.clone();
2553 let by_id: HashMap<u16, &columnar::NativeColumn> =
2554 columns.iter().map(|(id, c)| (*id, c)).collect();
2555 let pk_native = by_id.get(&pk_id)?;
2556 if native_int64_strictly_increasing(pk_native, n) {
2557 return None;
2558 }
2559 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2561 let mut null_pk_rows: Vec<usize> = Vec::new();
2562 for i in 0..n {
2563 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2564 Some(key) => {
2565 last.insert(key, i);
2566 }
2567 None => null_pk_rows.push(i),
2568 }
2569 }
2570 let mut winners: HashSet<usize> = last.values().copied().collect();
2571 for i in null_pk_rows {
2572 winners.insert(i);
2573 }
2574 Some((0..n).filter(|i| winners.contains(i)).collect())
2575 }
2576
2577 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2579 self.require_delete()?;
2580 let epoch = self.pending_epoch();
2581 self.wal_append_data(Op::Delete {
2582 table_id: self.table_id,
2583 row_ids: vec![row_id],
2584 })?;
2585 if self.is_shared() {
2586 self.pending_dels.push(row_id);
2587 } else {
2588 self.apply_delete(row_id, epoch);
2589 }
2590 Ok(())
2591 }
2592
2593 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2594 let pre = self.get(row_id, self.snapshot());
2595 self.delete(row_id)?;
2596 Ok(pre.map(|row| {
2597 let mut columns: Vec<_> = row.columns.into_iter().collect();
2598 columns.sort_by_key(|(id, _)| *id);
2599 OwnedRow { columns }
2600 }))
2601 }
2602
2603 pub fn truncate(&mut self) -> Result<()> {
2605 self.require_delete()?;
2606 let epoch = self.pending_epoch();
2607 self.wal_append_data(Op::TruncateTable {
2608 table_id: self.table_id,
2609 })?;
2610 self.pending_rows.clear();
2611 self.pending_rows_auto_inc.clear();
2612 self.pending_dels.clear();
2613 self.pending_truncate = Some(epoch);
2614 Ok(())
2615 }
2616
2617 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2619 for rr in std::mem::take(&mut self.run_refs) {
2620 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2621 }
2622 for r in std::mem::take(&mut self.retiring) {
2623 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2624 }
2625 self.memtable = Memtable::new();
2626 self.mutable_run = MutableRun::new();
2627 self.hot = HotIndex::new();
2628 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2629 self.bitmap = bitmap;
2630 self.ann = ann;
2631 self.fm = fm;
2632 self.sparse = sparse;
2633 self.minhash = minhash;
2634 self.learned_range.clear();
2635 self.pk_by_row.clear();
2636 self.pk_by_row_complete = false;
2637 self.live_count = 0;
2638 self.reservoir = crate::reservoir::Reservoir::default();
2639 self.reservoir_complete = true;
2640 self.had_deletes = true;
2641 self.agg_cache.clear();
2642 self.global_idx_epoch = 0;
2643 self.indexes_complete = true;
2644 self.pending_delete_rids.clear();
2645 self.pending_put_cols.clear();
2646 self.pending_rows.clear();
2647 self.pending_rows_auto_inc.clear();
2648 self.pending_dels.clear();
2649 self.clear_result_cache();
2650 self.invalidate_index_checkpoint();
2651 self.data_generation = self.data_generation.wrapping_add(1);
2652 Ok(())
2653 }
2654
2655 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2658 self.remove_hot_for_row(row_id, epoch);
2659 self.tombstone_row(row_id, epoch, true);
2660 self.data_generation = self.data_generation.wrapping_add(1);
2661 }
2662
2663 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2667 let tombstone = Row {
2668 row_id,
2669 committed_epoch: epoch,
2670 columns: std::collections::HashMap::new(),
2671 deleted: true,
2672 };
2673 self.memtable.upsert(tombstone);
2674 self.pk_by_row.remove(&row_id);
2675 if adjust_live_count {
2676 self.live_count = self.live_count.saturating_sub(1);
2677 }
2678 self.pending_delete_rids.insert(row_id.0 as u32);
2680 self.had_deletes = true;
2683 self.agg_cache.clear();
2684 }
2685
2686 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2690 let Some(pk_col) = self.schema.primary_key() else {
2691 return;
2692 };
2693 if self.pk_by_row_complete {
2696 if let Some(key) = self.pk_by_row.remove(&row_id) {
2697 if self.hot.get(&key) == Some(row_id) {
2698 self.hot.remove(&key);
2699 }
2700 }
2701 return;
2702 }
2703 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2722 if self.indexes_complete {
2723 let pk_val = self
2724 .memtable
2725 .get_version(row_id, lookup_epoch)
2726 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2727 .or_else(|| {
2728 self.mutable_run
2729 .get_version(row_id, lookup_epoch)
2730 .filter(|(_, r)| !r.deleted)
2731 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2732 })
2733 .or_else(|| {
2734 self.run_refs.iter().find_map(|rr| {
2735 let mut reader = self.open_reader(rr.run_id).ok()?;
2736 let (_, deleted, val) = reader
2737 .get_version_column(row_id, lookup_epoch, pk_col.id)
2738 .ok()??;
2739 if deleted {
2740 return None;
2741 }
2742 val
2743 })
2744 });
2745 if let Some(pk_val) = pk_val {
2746 let key = self.index_lookup_key(pk_col.id, &pk_val);
2747 if self.hot.get(&key) == Some(row_id) {
2748 self.hot.remove(&key);
2749 }
2750 return;
2751 }
2752 }
2753 self.refresh_pk_by_row_from_hot();
2758 if let Some(key) = self.pk_by_row.remove(&row_id) {
2759 if self.hot.get(&key) == Some(row_id) {
2760 self.hot.remove(&key);
2761 }
2762 }
2763 }
2764
2765 fn partition_pk_winners(
2770 &self,
2771 rows: &[Row],
2772 ) -> (
2773 std::collections::HashSet<RowId>,
2774 std::collections::HashMap<Vec<u8>, RowId>,
2775 ) {
2776 let mut losers = std::collections::HashSet::new();
2777 let Some(pk_col) = self.schema.primary_key() else {
2778 return (losers, std::collections::HashMap::new());
2779 };
2780 let pk_id = pk_col.id;
2781 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2782 std::collections::HashMap::new();
2783 for r in rows {
2784 let Some(pk_val) = r.columns.get(&pk_id) else {
2785 continue;
2786 };
2787 let key = self.index_lookup_key(pk_id, pk_val);
2788 if let Some(&old_rid) = winners.get(&key) {
2789 losers.insert(old_rid);
2790 }
2791 winners.insert(key, r.row_id);
2792 }
2793 (losers, winners)
2794 }
2795
2796 fn index_row(&mut self, row: &Row) {
2797 if row.deleted {
2798 return;
2799 }
2800 let any_predicate = self
2808 .schema
2809 .indexes
2810 .iter()
2811 .any(|idx| idx.predicate.is_some());
2812 if any_predicate {
2813 let columns_map: HashMap<u16, &Value> =
2814 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2815 let name_to_id: HashMap<&str, u16> = self
2816 .schema
2817 .columns
2818 .iter()
2819 .map(|c| (c.name.as_str(), c.id))
2820 .collect();
2821 for idx in &self.schema.indexes {
2822 if let Some(pred) = &idx.predicate {
2823 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2824 continue; }
2826 }
2827 index_into_single(
2829 idx,
2830 &self.schema,
2831 row,
2832 &mut self.hot,
2833 &mut self.bitmap,
2834 &mut self.ann,
2835 &mut self.fm,
2836 &mut self.sparse,
2837 &mut self.minhash,
2838 );
2839 }
2840 return;
2841 }
2842 if self.column_keys.is_empty() {
2846 index_into(
2847 &self.schema,
2848 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 return;
2857 }
2858 let effective_row = self.tokenized_for_indexes(row);
2859 index_into(
2860 &self.schema,
2861 &effective_row,
2862 &mut self.hot,
2863 &mut self.bitmap,
2864 &mut self.ann,
2865 &mut self.fm,
2866 &mut self.sparse,
2867 &mut self.minhash,
2868 );
2869 }
2870
2871 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2877 if self.column_keys.is_empty() {
2878 return row.clone();
2879 }
2880 #[cfg(feature = "encryption")]
2881 {
2882 use crate::encryption::SCHEME_HMAC_EQ;
2883 let mut tok = row.clone();
2884 for (&cid, &(_, scheme)) in &self.column_keys {
2885 if scheme != SCHEME_HMAC_EQ {
2886 continue;
2887 }
2888 if let Some(v) = tok.columns.get(&cid).cloned() {
2889 if let Some(t) = self.tokenize_value(cid, &v) {
2890 tok.columns.insert(cid, t);
2891 }
2892 }
2893 }
2894 tok
2895 }
2896 #[cfg(not(feature = "encryption"))]
2897 {
2898 row.clone()
2899 }
2900 }
2901
2902 pub fn commit(&mut self) -> Result<Epoch> {
2907 self.ensure_writable()?;
2908 if self.is_shared() {
2909 self.commit_shared()
2910 } else {
2911 self.commit_private()
2912 }
2913 }
2914
2915 fn commit_private(&mut self) -> Result<Epoch> {
2917 let commit_lock = Arc::clone(&self.commit_lock);
2921 let _g = commit_lock.lock();
2922 let new_epoch = self.epoch.bump_assigned();
2923 let txn_id = self.current_txn_id;
2924 match &mut self.wal {
2928 WalSink::Private(w) => {
2929 w.append_txn(
2930 txn_id,
2931 Op::TxnCommit {
2932 epoch: new_epoch.0,
2933 added_runs: Vec::new(),
2934 },
2935 )?;
2936 w.sync()?;
2937 }
2938 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2939 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2940 }
2941 if let Some(epoch) = self.pending_truncate.take() {
2943 self.apply_truncate(epoch)?;
2944 }
2945 self.invalidate_pending_cache();
2946 self.persist_manifest(new_epoch)?;
2947 self.epoch.publish_in_order(new_epoch);
2951 self.current_txn_id += 1;
2952 self.data_generation = self.data_generation.wrapping_add(1);
2953 Ok(new_epoch)
2954 }
2955
2956 fn commit_shared(&mut self) -> Result<Epoch> {
2962 use std::sync::atomic::Ordering;
2963 let s = match &self.wal {
2964 WalSink::Shared(s) => s.clone(),
2965 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2966 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2967 };
2968 if s.poisoned.load(Ordering::Relaxed) {
2969 return Err(MongrelError::Other(
2970 "database poisoned by fsync error".into(),
2971 ));
2972 }
2973 let commit_lock = Arc::clone(&self.commit_lock);
2980 let _g = commit_lock.lock();
2981 let txn_id = self.ensure_txn_id();
2984 let (new_epoch, commit_seq) = {
2985 let mut wal = s.wal.lock();
2986 let new_epoch = self.epoch.bump_assigned();
2987 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2988 (new_epoch, seq)
2989 };
2990 s.group
2991 .await_durable(&s.wal, commit_seq)
2992 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2993
2994 if self.pending_truncate.take().is_some() {
2997 self.apply_truncate(new_epoch)?;
2998 }
2999 let mut rows = std::mem::take(&mut self.pending_rows);
3000 if !rows.is_empty() {
3001 for r in &mut rows {
3002 r.committed_epoch = new_epoch;
3003 }
3004 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3005 let all_auto_generated =
3006 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3007 self.apply_put_rows_inner(rows, !all_auto_generated)?;
3008 } else {
3009 self.pending_rows_auto_inc.clear();
3010 }
3011 let dels = std::mem::take(&mut self.pending_dels);
3012 for rid in dels {
3013 self.apply_delete(rid, new_epoch);
3014 }
3015
3016 self.invalidate_pending_cache();
3017 self.persist_manifest(new_epoch)?;
3018 self.epoch.publish_in_order(new_epoch);
3019 let _ = s.change_wake.send(());
3020 self.current_txn_id = 0;
3022 self.data_generation = self.data_generation.wrapping_add(1);
3023 Ok(new_epoch)
3024 }
3025
3026 pub fn flush(&mut self) -> Result<Epoch> {
3034 self.ensure_indexes_complete()?;
3035 let epoch = self.commit()?;
3036 let rows = self.memtable.drain_sorted();
3037 if !rows.is_empty() {
3038 self.mutable_run.insert_many(rows);
3039 }
3040 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3041 self.spill_mutable_run(epoch)?;
3042 self.mark_flushed(epoch)?;
3046 self.persist_manifest(epoch)?;
3047 self.build_learned_ranges()?;
3048 self.checkpoint_indexes(epoch);
3051 }
3052 Ok(epoch)
3055 }
3056
3057 pub fn force_flush(&mut self) -> Result<Epoch> {
3066 let saved = self.mutable_run_spill_bytes;
3067 self.mutable_run_spill_bytes = 1;
3068 let result = self.flush();
3069 self.mutable_run_spill_bytes = saved;
3070 result
3071 }
3072
3073 pub fn close(&mut self) -> Result<()> {
3080 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3081 self.force_flush()?;
3082 }
3083 Ok(())
3084 }
3085
3086 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3093 let op = Op::Flush {
3094 table_id: self.table_id,
3095 flushed_epoch: epoch.0,
3096 };
3097 match &mut self.wal {
3098 WalSink::Private(w) => {
3099 w.append_system(op)?;
3100 w.sync()?;
3101 }
3102 WalSink::Shared(s) => {
3103 s.wal.lock().append_system(op)?;
3108 }
3109 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3110 }
3111 self.flushed_epoch = epoch.0;
3112 if matches!(self.wal, WalSink::Private(_)) {
3113 self.rotate_wal(epoch)?;
3114 }
3115 Ok(())
3116 }
3117
3118 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3122 let rows = self.mutable_run.drain_sorted();
3123 if rows.is_empty() {
3124 return Ok(());
3125 }
3126 let run_id = self.next_run_id;
3127 self.next_run_id += 1;
3128 let path = self.run_path(run_id);
3129 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3130 if let Some(kek) = &self.kek {
3131 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3132 }
3133 let header = writer.write(&path, &rows)?;
3134 self.run_refs.push(RunRef {
3135 run_id: run_id as u128,
3136 level: 0,
3137 epoch_created: epoch.0,
3138 row_count: header.row_count,
3139 });
3140 Ok(())
3141 }
3142
3143 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3147 self.mutable_run_spill_bytes = bytes.max(1);
3148 }
3149
3150 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3154 self.compaction_zstd_level = level;
3155 }
3156
3157 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3161 self.result_cache.lock().set_max_bytes(max_bytes);
3162 }
3163
3164 pub(crate) fn clear_result_cache(&mut self) {
3168 self.result_cache.lock().clear();
3169 }
3170
3171 pub fn mutable_run_len(&self) -> usize {
3173 self.mutable_run.len()
3174 }
3175
3176 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3179 self.mutable_run.drain_sorted()
3180 }
3181
3182 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3187 let epoch = self.commit()?;
3188 let n = batch.len();
3189 if n == 0 {
3190 return Ok(epoch);
3191 }
3192 for row in &batch {
3193 self.schema.validate_values(row)?;
3194 }
3195 let live_before = self.live_count;
3196 self.spill_mutable_run(epoch)?;
3200 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3201 && self.indexes_complete
3202 && self.run_refs.is_empty()
3203 && self.memtable.is_empty()
3204 && self.mutable_run.is_empty();
3205 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3211 use rayon::prelude::*;
3212 self.schema
3213 .columns
3214 .par_iter()
3215 .map(|cdef| {
3216 (
3217 cdef.id,
3218 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3219 )
3220 })
3221 .collect::<Vec<_>>()
3222 };
3223 drop(batch);
3224 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3229 self.validate_columns_not_null(&user_columns, n)?;
3230 let winner_idx = self
3231 .bulk_pk_winner_indices(&user_columns, n)
3232 .filter(|idx| idx.len() != n);
3233 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3234 match winner_idx.as_deref() {
3235 Some(idx) => {
3236 let compacted = user_columns
3237 .iter()
3238 .map(|(id, c)| (*id, c.gather(idx)))
3239 .collect();
3240 (compacted, idx.len())
3241 }
3242 None => (std::mem::take(&mut user_columns), n),
3243 };
3244 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3245 let first = self.allocator.alloc_range(write_n as u64).0;
3246 for rid in first..first + write_n as u64 {
3247 self.reservoir.offer(rid);
3248 }
3249 let run_id = self.next_run_id;
3250 self.next_run_id += 1;
3251 let path = self.run_path(run_id);
3252 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3253 .clean(true)
3254 .with_lz4()
3255 .with_native_endian();
3256 if let Some(kek) = &self.kek {
3257 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3258 }
3259 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3260 self.run_refs.push(RunRef {
3261 run_id: run_id as u128,
3262 level: 0,
3263 epoch_created: epoch.0,
3264 row_count: header.row_count,
3265 });
3266 self.live_count = self.live_count.saturating_add(write_n as u64);
3267 if eager_index_build {
3268 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3269 self.index_columns_bulk(&write_columns, &row_ids);
3270 self.indexes_complete = true;
3271 self.build_learned_ranges()?;
3272 } else {
3273 self.indexes_complete = false;
3274 }
3275 self.mark_flushed(epoch)?;
3276 self.persist_manifest(epoch)?;
3277 if eager_index_build {
3278 self.checkpoint_indexes(epoch);
3279 }
3280 self.clear_result_cache();
3281 Ok(epoch)
3282 }
3283
3284 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3287 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3288 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3289 let segment_no = segment
3292 .file_stem()
3293 .and_then(|s| s.to_str())
3294 .and_then(|s| s.strip_prefix("seg-"))
3295 .and_then(|s| s.parse::<u64>().ok())
3296 .unwrap_or(0);
3297 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3298 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3299 wal.sync()?;
3300 self.wal = WalSink::Private(wal);
3301 Ok(())
3302 }
3303
3304 pub(crate) fn invalidate_pending_cache(&mut self) {
3309 self.result_cache
3310 .lock()
3311 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3312 self.pending_delete_rids.clear();
3313 self.pending_put_cols.clear();
3314 }
3315
3316 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3317 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3318 m.current_epoch = epoch.0;
3319 m.next_row_id = self.allocator.current().0;
3320 m.runs = self.run_refs.clone();
3321 m.live_count = self.live_count;
3322 m.global_idx_epoch = self.global_idx_epoch;
3323 m.flushed_epoch = self.flushed_epoch;
3324 m.retiring = self.retiring.clone();
3325 m.auto_inc_next = match self.auto_inc {
3329 Some(ai) if ai.seeded => ai.next,
3330 _ => 0,
3331 };
3332 m.ttl = self.ttl;
3333 let meta_dek = self.manifest_meta_dek();
3334 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3335 Ok(())
3336 }
3337
3338 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3344 if !self.indexes_complete {
3347 return;
3348 }
3349 let snap = global_idx::IndexSnapshot {
3350 hot: &self.hot,
3351 bitmap: &self.bitmap,
3352 ann: &self.ann,
3353 fm: &self.fm,
3354 sparse: &self.sparse,
3355 minhash: &self.minhash,
3356 learned_range: &self.learned_range,
3357 };
3358 let idx_dek = self.idx_dek();
3360 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3361 .is_ok()
3362 {
3363 self.global_idx_epoch = epoch.0;
3364 let _ = self.persist_manifest(epoch);
3365 }
3366 }
3367
3368 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3371 self.global_idx_epoch = 0;
3372 global_idx::remove(&self.dir);
3373 let _ = self.persist_manifest(self.epoch.visible());
3374 }
3375
3376 pub(crate) fn mark_indexes_incomplete(&mut self) {
3377 self.indexes_complete = false;
3378 self.invalidate_index_checkpoint();
3379 }
3380
3381 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3384 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3385 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3386 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3387 best = Some((epoch, row));
3388 }
3389 }
3390 for rr in &self.run_refs {
3391 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3392 continue;
3393 };
3394 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3395 continue;
3396 };
3397 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3398 best = Some((epoch, row));
3399 }
3400 }
3401 let now_nanos = unix_nanos_now();
3402 match best {
3403 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3404 Some((_, r)) => Some(r),
3405 None => None,
3406 }
3407 }
3408
3409 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3413 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3414 let mut fold = |row: Row| {
3415 best.entry(row.row_id.0)
3416 .and_modify(|e| {
3417 if row.committed_epoch > e.0 {
3418 *e = (row.committed_epoch, row.clone());
3419 }
3420 })
3421 .or_insert_with(|| (row.committed_epoch, row));
3422 };
3423 for row in self.memtable.visible_versions(snapshot.epoch) {
3424 fold(row);
3425 }
3426 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3427 fold(row);
3428 }
3429 for rr in &self.run_refs {
3430 let mut reader = self.open_reader(rr.run_id)?;
3431 for row in reader.visible_versions(snapshot.epoch)? {
3432 fold(row);
3433 }
3434 }
3435 let now_nanos = unix_nanos_now();
3436 let mut out: Vec<Row> = best
3437 .into_values()
3438 .filter_map(|(_, r)| {
3439 if r.deleted || self.row_expired_at(&r, now_nanos) {
3440 None
3441 } else {
3442 Some(r)
3443 }
3444 })
3445 .collect();
3446 out.sort_by_key(|r| r.row_id);
3447 Ok(out)
3448 }
3449
3450 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3457 if self.ttl.is_none()
3458 && self.memtable.is_empty()
3459 && self.mutable_run.is_empty()
3460 && self.run_refs.len() == 1
3461 {
3462 let rr = self.run_refs[0].clone();
3463 let mut reader = self.open_reader(rr.run_id)?;
3464 let idxs = reader.visible_indices(snapshot.epoch)?;
3465 let mut cols = Vec::with_capacity(self.schema.columns.len());
3466 for cdef in &self.schema.columns {
3467 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3468 }
3469 return Ok(cols);
3470 }
3471 let rows = self.visible_rows(snapshot)?;
3473 let mut cols: Vec<(u16, Vec<Value>)> = self
3474 .schema
3475 .columns
3476 .iter()
3477 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3478 .collect();
3479 for r in &rows {
3480 for (cid, vec) in cols.iter_mut() {
3481 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3482 }
3483 }
3484 Ok(cols)
3485 }
3486
3487 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3489 let row_id = self.hot.get(key)?;
3490 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3491 Some(row_id)
3492 } else {
3493 None
3494 }
3495 }
3496
3497 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3502 self.query_at_with_allowed(q, self.snapshot(), None)
3503 }
3504
3505 pub fn query_at_with_allowed(
3508 &mut self,
3509 q: &crate::query::Query,
3510 snapshot: Snapshot,
3511 allowed: Option<&std::collections::HashSet<RowId>>,
3512 ) -> Result<Vec<Row>> {
3513 self.query_at_with_allowed_after(q, snapshot, allowed, None)
3514 }
3515
3516 #[doc(hidden)]
3517 pub fn query_at_with_allowed_after(
3518 &mut self,
3519 q: &crate::query::Query,
3520 snapshot: Snapshot,
3521 allowed: Option<&std::collections::HashSet<RowId>>,
3522 after_row_id: Option<RowId>,
3523 ) -> Result<Vec<Row>> {
3524 self.require_select()?;
3525 self.ensure_indexes_complete()?;
3526 if q.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 if let Some(limit) = q.limit {
3533 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
3534 return Err(MongrelError::InvalidArgument(format!(
3535 "query limit must be between 1 and {}",
3536 crate::query::MAX_FINAL_LIMIT
3537 )));
3538 }
3539 }
3540 if q.offset > crate::query::MAX_QUERY_OFFSET {
3541 return Err(MongrelError::InvalidArgument(format!(
3542 "query offset exceeds {}",
3543 crate::query::MAX_QUERY_OFFSET
3544 )));
3545 }
3546 self.query_conditions_at(
3547 &q.conditions,
3548 snapshot,
3549 allowed,
3550 q.limit,
3551 q.offset,
3552 after_row_id,
3553 )
3554 }
3555
3556 #[doc(hidden)]
3559 pub fn query_all_at(
3560 &mut self,
3561 conditions: &[crate::query::Condition],
3562 snapshot: Snapshot,
3563 ) -> Result<Vec<Row>> {
3564 self.require_select()?;
3565 self.ensure_indexes_complete()?;
3566 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3567 return Err(MongrelError::InvalidArgument(format!(
3568 "query exceeds {} conditions",
3569 crate::query::MAX_HARD_CONDITIONS
3570 )));
3571 }
3572 self.query_conditions_at(conditions, snapshot, None, None, 0, None)
3573 }
3574
3575 fn query_conditions_at(
3576 &self,
3577 conditions: &[crate::query::Condition],
3578 snapshot: Snapshot,
3579 allowed: Option<&std::collections::HashSet<RowId>>,
3580 limit: Option<usize>,
3581 offset: usize,
3582 after_row_id: Option<RowId>,
3583 ) -> Result<Vec<Row>> {
3584 crate::trace::QueryTrace::record(|t| {
3585 t.run_count = self.run_refs.len();
3586 t.memtable_rows = self.memtable.len();
3587 t.mutable_run_rows = self.mutable_run.len();
3588 });
3589 if conditions.is_empty() {
3593 crate::trace::QueryTrace::record(|t| {
3594 t.scan_mode = crate::trace::ScanMode::Materialized;
3595 t.row_materialized = true;
3596 });
3597 let mut rows = self.visible_rows(snapshot)?;
3598 if let Some(allowed) = allowed {
3599 rows.retain(|row| allowed.contains(&row.row_id));
3600 }
3601 if let Some(after_row_id) = after_row_id {
3602 rows.retain(|row| row.row_id > after_row_id);
3603 }
3604 rows.drain(..offset.min(rows.len()));
3605 if let Some(limit) = limit {
3606 rows.truncate(limit);
3607 }
3608 return Ok(rows);
3609 }
3610 crate::trace::QueryTrace::record(|t| {
3611 t.conditions_pushed = conditions.len();
3612 t.scan_mode = crate::trace::ScanMode::Materialized;
3613 t.row_materialized = true;
3614 });
3615 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
3622 ordered.sort_by_key(|c| condition_cost_rank(c));
3623 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3624 for c in &ordered {
3625 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
3626 let empty = s.is_empty();
3627 sets.push(s);
3628 if empty {
3629 break;
3630 }
3631 }
3632 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3633 if let Some(allowed) = allowed {
3634 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
3635 }
3636 if let Some(after_row_id) = after_row_id {
3637 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
3638 rids.drain(..first);
3639 }
3640 rids.drain(..offset.min(rids.len()));
3641 if let Some(limit) = limit {
3642 rids.truncate(limit);
3643 }
3644 self.rows_for_rids(&rids, snapshot)
3645 }
3646
3647 pub fn retrieve(
3649 &mut self,
3650 retriever: &crate::query::Retriever,
3651 ) -> Result<Vec<crate::query::RetrieverHit>> {
3652 self.retrieve_with_allowed(retriever, None)
3653 }
3654
3655 pub fn retrieve_at(
3656 &mut self,
3657 retriever: &crate::query::Retriever,
3658 snapshot: Snapshot,
3659 allowed: Option<&std::collections::HashSet<RowId>>,
3660 ) -> Result<Vec<crate::query::RetrieverHit>> {
3661 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
3662 }
3663
3664 pub fn retrieve_with_allowed(
3667 &mut self,
3668 retriever: &crate::query::Retriever,
3669 allowed: Option<&std::collections::HashSet<RowId>>,
3670 ) -> Result<Vec<crate::query::RetrieverHit>> {
3671 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
3672 }
3673
3674 pub fn retrieve_at_with_allowed(
3675 &mut self,
3676 retriever: &crate::query::Retriever,
3677 snapshot: Snapshot,
3678 allowed: Option<&std::collections::HashSet<RowId>>,
3679 ) -> Result<Vec<crate::query::RetrieverHit>> {
3680 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
3681 }
3682
3683 pub fn retrieve_at_with_allowed_and_context(
3684 &mut self,
3685 retriever: &crate::query::Retriever,
3686 snapshot: Snapshot,
3687 allowed: Option<&std::collections::HashSet<RowId>>,
3688 context: Option<&crate::query::AiExecutionContext>,
3689 ) -> Result<Vec<crate::query::RetrieverHit>> {
3690 self.require_select()?;
3691 self.ensure_indexes_complete()?;
3692 self.validate_retriever(retriever)?;
3693 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
3694 }
3695
3696 pub fn retrieve_at_with_candidate_authorization_and_context(
3697 &mut self,
3698 retriever: &crate::query::Retriever,
3699 snapshot: Snapshot,
3700 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3701 context: Option<&crate::query::AiExecutionContext>,
3702 ) -> Result<Vec<crate::query::RetrieverHit>> {
3703 self.require_select()?;
3704 self.ensure_indexes_complete()?;
3705 self.retrieve_at_with_candidate_authorization_on_generation(
3706 retriever,
3707 snapshot,
3708 authorization,
3709 context,
3710 )
3711 }
3712
3713 #[doc(hidden)]
3714 pub fn retrieve_at_with_candidate_authorization_on_generation(
3715 &self,
3716 retriever: &crate::query::Retriever,
3717 snapshot: Snapshot,
3718 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3719 context: Option<&crate::query::AiExecutionContext>,
3720 ) -> Result<Vec<crate::query::RetrieverHit>> {
3721 self.require_select()?;
3722 self.validate_retriever(retriever)?;
3723 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
3724 }
3725
3726 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3727 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
3728 let (column_id, k) = match retriever {
3729 Retriever::Ann {
3730 column_id,
3731 query,
3732 k,
3733 } => {
3734 let index = self.ann.get(column_id).ok_or_else(|| {
3735 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3736 })?;
3737 if query.len() != index.dim() {
3738 return Err(MongrelError::InvalidArgument(format!(
3739 "ANN query dimension must be {}, got {}",
3740 index.dim(),
3741 query.len()
3742 )));
3743 }
3744 if query.iter().any(|value| !value.is_finite()) {
3745 return Err(MongrelError::InvalidArgument(
3746 "ANN query values must be finite".into(),
3747 ));
3748 }
3749 (*column_id, *k)
3750 }
3751 Retriever::Sparse {
3752 column_id,
3753 query,
3754 k,
3755 } => {
3756 if !self.sparse.contains_key(column_id) {
3757 return Err(MongrelError::InvalidArgument(format!(
3758 "column {column_id} has no Sparse index"
3759 )));
3760 }
3761 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3762 return Err(MongrelError::InvalidArgument(
3763 "Sparse query must be non-empty with finite weights".into(),
3764 ));
3765 }
3766 if query.len() > MAX_SPARSE_TERMS {
3767 return Err(MongrelError::InvalidArgument(format!(
3768 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
3769 )));
3770 }
3771 (*column_id, *k)
3772 }
3773 Retriever::MinHash {
3774 column_id,
3775 members,
3776 k,
3777 } => {
3778 if !self.minhash.contains_key(column_id) {
3779 return Err(MongrelError::InvalidArgument(format!(
3780 "column {column_id} has no MinHash index"
3781 )));
3782 }
3783 if members.is_empty() {
3784 return Err(MongrelError::InvalidArgument(
3785 "MinHash members must not be empty".into(),
3786 ));
3787 }
3788 if members.len() > MAX_SET_MEMBERS {
3789 return Err(MongrelError::InvalidArgument(format!(
3790 "MinHash query exceeds {MAX_SET_MEMBERS} members"
3791 )));
3792 }
3793 let mut total_bytes = 0usize;
3794 for member in members {
3795 let bytes = member.encoded_len();
3796 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
3797 return Err(MongrelError::InvalidArgument(format!(
3798 "MinHash member exceeds {} bytes",
3799 crate::query::MAX_SET_MEMBER_BYTES
3800 )));
3801 }
3802 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
3803 MongrelError::InvalidArgument("MinHash input size overflow".into())
3804 })?;
3805 }
3806 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
3807 return Err(MongrelError::InvalidArgument(format!(
3808 "MinHash input exceeds {} bytes",
3809 crate::query::MAX_SET_INPUT_BYTES
3810 )));
3811 }
3812 (*column_id, *k)
3813 }
3814 };
3815 if k == 0 {
3816 return Err(MongrelError::InvalidArgument(
3817 "retriever k must be > 0".into(),
3818 ));
3819 }
3820 if k > MAX_RETRIEVER_K {
3821 return Err(MongrelError::InvalidArgument(format!(
3822 "retriever k exceeds {MAX_RETRIEVER_K}"
3823 )));
3824 }
3825 debug_assert!(self
3826 .schema
3827 .columns
3828 .iter()
3829 .any(|column| column.id == column_id));
3830 Ok(())
3831 }
3832
3833 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3834 use crate::query::Condition;
3835 match condition {
3836 Condition::Ann {
3837 column_id,
3838 query,
3839 k,
3840 } => self.validate_retriever(&crate::query::Retriever::Ann {
3841 column_id: *column_id,
3842 query: query.clone(),
3843 k: *k,
3844 }),
3845 Condition::SparseMatch {
3846 column_id,
3847 query,
3848 k,
3849 } => self.validate_retriever(&crate::query::Retriever::Sparse {
3850 column_id: *column_id,
3851 query: query.clone(),
3852 k: *k,
3853 }),
3854 Condition::MinHashSimilar {
3855 column_id,
3856 query,
3857 k,
3858 } => {
3859 if !self.minhash.contains_key(column_id) {
3860 return Err(MongrelError::InvalidArgument(format!(
3861 "column {column_id} has no MinHash index"
3862 )));
3863 }
3864 if query.is_empty() || *k == 0 {
3865 return Err(MongrelError::InvalidArgument(
3866 "MinHash query must be non-empty and k must be > 0".into(),
3867 ));
3868 }
3869 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
3870 {
3871 return Err(MongrelError::InvalidArgument(format!(
3872 "MinHash query must have <= {} members and k <= {}",
3873 crate::query::MAX_SET_MEMBERS,
3874 crate::query::MAX_RETRIEVER_K
3875 )));
3876 }
3877 Ok(())
3878 }
3879 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
3880 Err(MongrelError::InvalidArgument(format!(
3881 "bitmap IN exceeds {} values",
3882 crate::query::MAX_SET_MEMBERS
3883 )))
3884 }
3885 Condition::FmContainsAll { patterns, .. }
3886 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
3887 {
3888 Err(MongrelError::InvalidArgument(format!(
3889 "FM query exceeds {} patterns",
3890 crate::query::MAX_HARD_CONDITIONS
3891 )))
3892 }
3893 _ => Ok(()),
3894 }
3895 }
3896
3897 fn retrieve_filtered(
3898 &self,
3899 retriever: &crate::query::Retriever,
3900 snapshot: Snapshot,
3901 hard_filter: Option<&RowIdSet>,
3902 allowed: Option<&std::collections::HashSet<RowId>>,
3903 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3904 context: Option<&crate::query::AiExecutionContext>,
3905 ) -> Result<Vec<crate::query::RetrieverHit>> {
3906 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
3907 let started = std::time::Instant::now();
3908 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
3909 Retriever::Ann {
3910 column_id,
3911 query,
3912 k,
3913 } => {
3914 let Some(index) = self.ann.get(column_id) else {
3915 return Ok(Vec::new());
3916 };
3917 let cap = index.len();
3918 if cap == 0 {
3919 return Ok(Vec::new());
3920 }
3921 let mut breadth = (*k).max(1).min(cap);
3922 let mut eligibility = std::collections::HashMap::new();
3923 let mut filtered = loop {
3924 let mut seen = std::collections::HashSet::new();
3925 if let Some(context) = context {
3926 context.checkpoint()?;
3927 }
3928 let raw = index.search_with_context(query, breadth, context)?;
3929 let unchecked: Vec<_> = raw
3930 .iter()
3931 .map(|(row_id, _)| *row_id)
3932 .filter(|row_id| !eligibility.contains_key(row_id))
3933 .filter(|row_id| {
3934 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3935 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3936 })
3937 .collect();
3938 let eligible = self.eligible_and_authorized_candidate_ids(
3939 &unchecked,
3940 *column_id,
3941 snapshot,
3942 candidate_authorization,
3943 context,
3944 )?;
3945 for row_id in unchecked {
3946 eligibility.insert(row_id, eligible.contains(&row_id));
3947 }
3948 let filtered: Vec<_> = raw
3949 .into_iter()
3950 .filter(|(row_id, _)| {
3951 seen.insert(*row_id)
3952 && eligibility.get(row_id).copied().unwrap_or(false)
3953 })
3954 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
3955 .collect();
3956 if filtered.len() >= *k || breadth >= cap {
3957 break filtered;
3958 }
3959 breadth = breadth.saturating_mul(2).min(cap);
3960 };
3961 filtered.truncate(*k);
3962 filtered
3963 }
3964 Retriever::Sparse {
3965 column_id,
3966 query,
3967 k,
3968 } => self
3969 .sparse
3970 .get(column_id)
3971 .map(|index| -> Result<Vec<_>> {
3972 let mut breadth = (*k).max(1);
3973 let mut eligibility = std::collections::HashMap::new();
3974 loop {
3975 if let Some(context) = context {
3976 context.checkpoint()?;
3977 }
3978 let raw = index.search_with_context(query, breadth, context)?;
3979 let unchecked: Vec<_> = raw
3980 .iter()
3981 .map(|(row_id, _)| *row_id)
3982 .filter(|row_id| !eligibility.contains_key(row_id))
3983 .filter(|row_id| {
3984 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3985 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3986 })
3987 .collect();
3988 let eligible = self.eligible_and_authorized_candidate_ids(
3989 &unchecked,
3990 *column_id,
3991 snapshot,
3992 candidate_authorization,
3993 context,
3994 )?;
3995 for row_id in unchecked {
3996 eligibility.insert(row_id, eligible.contains(&row_id));
3997 }
3998 let filtered: Vec<_> = raw
3999 .iter()
4000 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4001 .take(*k)
4002 .map(|(row_id, score)| {
4003 (*row_id, RetrieverScore::SparseDotProduct(*score))
4004 })
4005 .collect();
4006 if filtered.len() >= *k || raw.len() < breadth {
4007 break Ok(filtered);
4008 }
4009 let next = breadth.saturating_mul(2);
4010 if next == breadth {
4011 break Ok(filtered);
4012 }
4013 breadth = next;
4014 }
4015 })
4016 .transpose()?
4017 .unwrap_or_default(),
4018 Retriever::MinHash {
4019 column_id,
4020 members,
4021 k,
4022 } => self
4023 .minhash
4024 .get(column_id)
4025 .map(|index| -> Result<Vec<_>> {
4026 let mut hashes = Vec::with_capacity(members.len());
4027 for member in members {
4028 if let Some(context) = context {
4029 context.consume(crate::query::work_units(
4030 member.encoded_len(),
4031 crate::query::PARSE_WORK_QUANTUM,
4032 ))?;
4033 }
4034 hashes.push(member.hash_v1());
4035 }
4036 let mut breadth = (*k).max(1);
4037 let mut eligibility = std::collections::HashMap::new();
4038 loop {
4039 if let Some(context) = context {
4040 context.checkpoint()?;
4041 }
4042 let raw = index.search_with_context(&hashes, breadth, context)?;
4043 let unchecked: Vec<_> = raw
4044 .iter()
4045 .map(|(row_id, _)| *row_id)
4046 .filter(|row_id| !eligibility.contains_key(row_id))
4047 .filter(|row_id| {
4048 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4049 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4050 })
4051 .collect();
4052 let eligible = self.eligible_and_authorized_candidate_ids(
4053 &unchecked,
4054 *column_id,
4055 snapshot,
4056 candidate_authorization,
4057 context,
4058 )?;
4059 for row_id in unchecked {
4060 eligibility.insert(row_id, eligible.contains(&row_id));
4061 }
4062 let filtered: Vec<_> = raw
4063 .iter()
4064 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4065 .take(*k)
4066 .map(|(row_id, score)| {
4067 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
4068 })
4069 .collect();
4070 if filtered.len() >= *k || raw.len() < breadth {
4071 break Ok(filtered);
4072 }
4073 let next = breadth.saturating_mul(2);
4074 if next == breadth {
4075 break Ok(filtered);
4076 }
4077 breadth = next;
4078 }
4079 })
4080 .transpose()?
4081 .unwrap_or_default(),
4082 };
4083 let elapsed = started.elapsed().as_nanos() as u64;
4084 crate::trace::QueryTrace::record(|trace| {
4085 match retriever {
4086 Retriever::Ann { .. } => {
4087 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
4088 }
4089 Retriever::Sparse { .. } => {
4090 trace.sparse_candidate_nanos =
4091 trace.sparse_candidate_nanos.saturating_add(elapsed)
4092 }
4093 Retriever::MinHash { .. } => {
4094 trace.minhash_candidate_nanos =
4095 trace.minhash_candidate_nanos.saturating_add(elapsed)
4096 }
4097 }
4098 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
4099 });
4100 Ok(scored
4101 .into_iter()
4102 .enumerate()
4103 .map(|(rank, (row_id, score))| RetrieverHit {
4104 row_id,
4105 rank: rank + 1,
4106 score,
4107 })
4108 .collect())
4109 }
4110
4111 fn eligible_candidate_ids(
4112 &self,
4113 candidates: &[RowId],
4114 _column_id: u16,
4115 snapshot: Snapshot,
4116 context: Option<&crate::query::AiExecutionContext>,
4117 ) -> Result<std::collections::HashSet<RowId>> {
4118 if !self.had_deletes
4119 && self.ttl.is_none()
4120 && self.pending_put_cols.is_empty()
4121 && snapshot.epoch == self.snapshot().epoch
4122 {
4123 return Ok(candidates.iter().copied().collect());
4124 }
4125 let mut readers: Vec<_> = self
4126 .run_refs
4127 .iter()
4128 .map(|run| self.open_reader(run.run_id))
4129 .collect::<Result<_>>()?;
4130 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4131 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
4132 for &row_id in candidates {
4133 if let Some(context) = context {
4134 context.consume(1)?;
4135 }
4136 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4137 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4138 let overlay = match (mem, mutable) {
4139 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
4140 (Some(value), None) | (None, Some(value)) => Some(value),
4141 (None, None) => None,
4142 };
4143 if let Some((_, row)) = overlay {
4144 if !row.deleted && !self.row_expired_at(&row, now) {
4145 eligible.insert(row_id);
4146 }
4147 continue;
4148 }
4149 let mut best: Option<(Epoch, bool, usize)> = None;
4150 for (index, reader) in readers.iter_mut().enumerate() {
4151 if let Some((epoch, deleted)) =
4152 reader.get_version_visibility(row_id, snapshot.epoch)?
4153 {
4154 if best
4155 .as_ref()
4156 .map(|(best_epoch, ..)| epoch > *best_epoch)
4157 .unwrap_or(true)
4158 {
4159 best = Some((epoch, deleted, index));
4160 }
4161 }
4162 }
4163 let Some((_, false, reader_index)) = best else {
4164 continue;
4165 };
4166 if let Some(ttl) = self.ttl {
4167 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4168 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4169 {
4170 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4171 continue;
4172 }
4173 }
4174 }
4175 eligible.insert(row_id);
4176 }
4177 Ok(eligible)
4178 }
4179
4180 fn eligible_and_authorized_candidate_ids(
4181 &self,
4182 candidates: &[RowId],
4183 column_id: u16,
4184 snapshot: Snapshot,
4185 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4186 context: Option<&crate::query::AiExecutionContext>,
4187 ) -> Result<std::collections::HashSet<RowId>> {
4188 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
4189 let Some(authorization) = authorization else {
4190 return Ok(eligible);
4191 };
4192 let candidates: Vec<_> = eligible.into_iter().collect();
4193 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
4194 }
4195
4196 fn policy_allowed_candidate_ids(
4197 &self,
4198 candidates: &[RowId],
4199 snapshot: Snapshot,
4200 authorization: &crate::security::CandidateAuthorization<'_>,
4201 context: Option<&crate::query::AiExecutionContext>,
4202 ) -> Result<std::collections::HashSet<RowId>> {
4203 let started = std::time::Instant::now();
4204 if candidates.is_empty()
4205 || authorization.principal.is_admin
4206 || !authorization.security.rls_enabled(authorization.table)
4207 {
4208 return Ok(candidates.iter().copied().collect());
4209 }
4210 if let Some(context) = context {
4211 context.checkpoint()?;
4212 }
4213 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
4214 let mut rows: std::collections::HashMap<RowId, Row> = candidates
4215 .iter()
4216 .map(|row_id| {
4217 (
4218 *row_id,
4219 Row {
4220 row_id: *row_id,
4221 committed_epoch: snapshot.epoch,
4222 columns: std::collections::HashMap::new(),
4223 deleted: false,
4224 },
4225 )
4226 })
4227 .collect();
4228 let columns = authorization
4229 .security
4230 .select_policy_columns(authorization.table, authorization.principal);
4231 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4232 let mut decoded = 0usize;
4233 for column_id in &columns {
4234 if let Some(context) = context {
4235 context.checkpoint()?;
4236 }
4237 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4238 &row_ids, *column_id, snapshot, query_now, context,
4239 )? {
4240 if let Some(row) = rows.get_mut(&row_id) {
4241 row.columns.insert(*column_id, value);
4242 decoded = decoded.saturating_add(1);
4243 }
4244 }
4245 }
4246 if let Some(context) = context {
4247 context.consume(candidates.len().saturating_add(decoded))?;
4248 }
4249 let allowed = rows
4250 .into_values()
4251 .filter_map(|row| {
4252 authorization
4253 .security
4254 .row_allowed(
4255 authorization.table,
4256 crate::security::PolicyCommand::Select,
4257 &row,
4258 authorization.principal,
4259 false,
4260 )
4261 .then_some(row.row_id)
4262 })
4263 .collect();
4264 crate::trace::QueryTrace::record(|trace| {
4265 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
4266 trace.rls_policy_columns_decoded =
4267 trace.rls_policy_columns_decoded.saturating_add(decoded);
4268 trace.authorization_nanos = trace
4269 .authorization_nanos
4270 .saturating_add(started.elapsed().as_nanos() as u64);
4271 });
4272 Ok(allowed)
4273 }
4274
4275 pub fn search(
4277 &mut self,
4278 request: &crate::query::SearchRequest,
4279 ) -> Result<Vec<crate::query::SearchHit>> {
4280 self.search_with_allowed(request, None)
4281 }
4282
4283 pub fn search_at(
4284 &mut self,
4285 request: &crate::query::SearchRequest,
4286 snapshot: Snapshot,
4287 authorized: Option<&std::collections::HashSet<RowId>>,
4288 ) -> Result<Vec<crate::query::SearchHit>> {
4289 self.search_at_with_allowed(request, snapshot, authorized)
4290 }
4291
4292 pub fn search_with_allowed(
4293 &mut self,
4294 request: &crate::query::SearchRequest,
4295 authorized: Option<&std::collections::HashSet<RowId>>,
4296 ) -> Result<Vec<crate::query::SearchHit>> {
4297 self.search_at_with_allowed(request, self.snapshot(), authorized)
4298 }
4299
4300 pub fn search_at_with_allowed(
4301 &mut self,
4302 request: &crate::query::SearchRequest,
4303 snapshot: Snapshot,
4304 authorized: Option<&std::collections::HashSet<RowId>>,
4305 ) -> Result<Vec<crate::query::SearchHit>> {
4306 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
4307 }
4308
4309 pub fn search_at_with_allowed_and_context(
4310 &mut self,
4311 request: &crate::query::SearchRequest,
4312 snapshot: Snapshot,
4313 authorized: Option<&std::collections::HashSet<RowId>>,
4314 context: Option<&crate::query::AiExecutionContext>,
4315 ) -> Result<Vec<crate::query::SearchHit>> {
4316 self.ensure_indexes_complete()?;
4317 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
4318 }
4319
4320 pub fn search_at_with_candidate_authorization_and_context(
4321 &mut self,
4322 request: &crate::query::SearchRequest,
4323 snapshot: Snapshot,
4324 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4325 context: Option<&crate::query::AiExecutionContext>,
4326 ) -> Result<Vec<crate::query::SearchHit>> {
4327 self.ensure_indexes_complete()?;
4328 self.search_at_with_filters_and_context(
4329 request,
4330 snapshot,
4331 None,
4332 authorization,
4333 context,
4334 None,
4335 )
4336 }
4337
4338 #[doc(hidden)]
4339 pub fn search_at_with_candidate_authorization_on_generation(
4340 &self,
4341 request: &crate::query::SearchRequest,
4342 snapshot: Snapshot,
4343 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4344 context: Option<&crate::query::AiExecutionContext>,
4345 ) -> Result<Vec<crate::query::SearchHit>> {
4346 self.search_at_with_filters_and_context(
4347 request,
4348 snapshot,
4349 None,
4350 authorization,
4351 context,
4352 None,
4353 )
4354 }
4355
4356 #[doc(hidden)]
4357 pub fn search_at_with_candidate_authorization_on_generation_after(
4358 &self,
4359 request: &crate::query::SearchRequest,
4360 snapshot: Snapshot,
4361 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4362 context: Option<&crate::query::AiExecutionContext>,
4363 after: Option<crate::query::SearchAfter>,
4364 ) -> Result<Vec<crate::query::SearchHit>> {
4365 self.search_at_with_filters_and_context(
4366 request,
4367 snapshot,
4368 None,
4369 authorization,
4370 context,
4371 after,
4372 )
4373 }
4374
4375 fn search_at_with_filters_and_context(
4376 &self,
4377 request: &crate::query::SearchRequest,
4378 snapshot: Snapshot,
4379 authorized: Option<&std::collections::HashSet<RowId>>,
4380 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4381 context: Option<&crate::query::AiExecutionContext>,
4382 after: Option<crate::query::SearchAfter>,
4383 ) -> Result<Vec<crate::query::SearchHit>> {
4384 use crate::query::{
4385 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
4386 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
4387 };
4388 let total_started = std::time::Instant::now();
4389 self.require_select()?;
4390 if request.limit == 0 {
4391 return Err(MongrelError::InvalidArgument(
4392 "search limit must be > 0".into(),
4393 ));
4394 }
4395 if request.limit > MAX_FINAL_LIMIT {
4396 return Err(MongrelError::InvalidArgument(format!(
4397 "search limit exceeds {MAX_FINAL_LIMIT}"
4398 )));
4399 }
4400 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
4401 return Err(MongrelError::InvalidArgument(
4402 "search-after score must be finite".into(),
4403 ));
4404 }
4405 if request.retrievers.is_empty() {
4406 return Err(MongrelError::InvalidArgument(
4407 "search requires at least one retriever".into(),
4408 ));
4409 }
4410 if request.retrievers.len() > MAX_RETRIEVERS {
4411 return Err(MongrelError::InvalidArgument(format!(
4412 "search exceeds {MAX_RETRIEVERS} retrievers"
4413 )));
4414 }
4415 if request.must.len() > MAX_HARD_CONDITIONS {
4416 return Err(MongrelError::InvalidArgument(format!(
4417 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
4418 )));
4419 }
4420 for condition in &request.must {
4421 self.validate_condition(condition)?;
4422 }
4423 if request.must.iter().any(|condition| {
4424 matches!(
4425 condition,
4426 Condition::Ann { .. }
4427 | Condition::SparseMatch { .. }
4428 | Condition::MinHashSimilar { .. }
4429 )
4430 }) {
4431 return Err(MongrelError::InvalidArgument(
4432 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
4433 .into(),
4434 ));
4435 }
4436 let mut names = std::collections::HashSet::new();
4437 for named in &request.retrievers {
4438 if named.name.is_empty()
4439 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
4440 || !names.insert(named.name.as_str())
4441 {
4442 return Err(MongrelError::InvalidArgument(format!(
4443 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
4444 crate::query::MAX_RETRIEVER_NAME_BYTES
4445 )));
4446 }
4447 if !named.weight.is_finite()
4448 || named.weight < 0.0
4449 || named.weight > MAX_RETRIEVER_WEIGHT
4450 {
4451 return Err(MongrelError::InvalidArgument(format!(
4452 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4453 )));
4454 }
4455 self.validate_retriever(&named.retriever)?;
4456 }
4457 let projection = request
4458 .projection
4459 .clone()
4460 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
4461 if projection.len() > MAX_PROJECTION_COLUMNS {
4462 return Err(MongrelError::InvalidArgument(format!(
4463 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
4464 )));
4465 }
4466 for column_id in &projection {
4467 if !self
4468 .schema
4469 .columns
4470 .iter()
4471 .any(|column| column.id == *column_id)
4472 {
4473 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
4474 }
4475 }
4476 if let Some(crate::query::Rerank::ExactVector {
4477 embedding_column,
4478 query,
4479 candidate_limit,
4480 weight,
4481 ..
4482 }) = &request.rerank
4483 {
4484 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
4485 {
4486 return Err(MongrelError::InvalidArgument(format!(
4487 "rerank candidate_limit must be between search limit and {}",
4488 crate::query::MAX_RETRIEVER_K
4489 )));
4490 }
4491 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
4492 return Err(MongrelError::InvalidArgument(format!(
4493 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4494 )));
4495 }
4496 let column = self
4497 .schema
4498 .columns
4499 .iter()
4500 .find(|column| column.id == *embedding_column)
4501 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
4502 let crate::schema::TypeId::Embedding { dim } = column.ty else {
4503 return Err(MongrelError::InvalidArgument(format!(
4504 "rerank column {embedding_column} is not an embedding"
4505 )));
4506 };
4507 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
4508 return Err(MongrelError::InvalidArgument(format!(
4509 "rerank query must contain {dim} finite values"
4510 )));
4511 }
4512 }
4513
4514 let hard_filter_started = std::time::Instant::now();
4515 let hard_filter = if request.must.is_empty() {
4516 None
4517 } else {
4518 let mut sets = Vec::with_capacity(request.must.len());
4519 for condition in &request.must {
4520 if let Some(context) = context {
4521 context.checkpoint()?;
4522 }
4523 sets.push(self.resolve_condition(condition, snapshot)?);
4524 }
4525 Some(RowIdSet::intersect_many(sets))
4526 };
4527 crate::trace::QueryTrace::record(|trace| {
4528 trace.hard_filter_nanos = trace
4529 .hard_filter_nanos
4530 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
4531 });
4532 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
4533 return Ok(Vec::new());
4534 }
4535
4536 let constant = match request.fusion {
4537 Fusion::ReciprocalRank { constant } => constant,
4538 };
4539 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
4540 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
4541 let mut fusion_nanos = 0u64;
4542 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
4543 std::collections::HashMap::new();
4544 for named in retrievers {
4545 if named.weight == 0.0 {
4546 continue;
4547 }
4548 if let Some(context) = context {
4549 context.checkpoint()?;
4550 }
4551 let hits = self.retrieve_filtered(
4552 &named.retriever,
4553 snapshot,
4554 hard_filter.as_ref(),
4555 authorized,
4556 candidate_authorization,
4557 context,
4558 )?;
4559 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
4560 let fusion_started = std::time::Instant::now();
4561 for hit in hits {
4562 if let Some(context) = context {
4563 context.consume(1)?;
4564 }
4565 let contribution = named.weight / (constant as f64 + hit.rank as f64);
4566 if !contribution.is_finite() {
4567 return Err(MongrelError::InvalidArgument(
4568 "retriever contribution must be finite".into(),
4569 ));
4570 }
4571 let max_fused_candidates = context.map_or(
4572 crate::query::MAX_FUSED_CANDIDATES,
4573 crate::query::AiExecutionContext::max_fused_candidates,
4574 );
4575 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
4576 return Err(MongrelError::WorkBudgetExceeded);
4577 }
4578 let entry = fused.entry(hit.row_id).or_default();
4579 entry.0 += contribution;
4580 if !entry.0.is_finite() {
4581 return Err(MongrelError::InvalidArgument(
4582 "fused score must be finite".into(),
4583 ));
4584 }
4585 entry.1.push(ComponentScore {
4586 retriever_name: retriever_name.clone(),
4587 rank: hit.rank,
4588 raw_score: hit.score,
4589 contribution,
4590 });
4591 }
4592 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
4593 }
4594 let union_size = fused.len();
4595 let mut ranked: Vec<_> = fused
4596 .into_iter()
4597 .map(|(row_id, (fused_score, components))| {
4598 (row_id, fused_score, components, None, fused_score)
4599 })
4600 .collect();
4601 let order = |(a_row, _, _, _, a_score): &(
4602 RowId,
4603 f64,
4604 Vec<ComponentScore>,
4605 Option<f32>,
4606 f64,
4607 ),
4608 (b_row, _, _, _, b_score): &(
4609 RowId,
4610 f64,
4611 Vec<ComponentScore>,
4612 Option<f32>,
4613 f64,
4614 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
4615 if let Some(crate::query::Rerank::ExactVector {
4616 embedding_column,
4617 query,
4618 metric,
4619 candidate_limit,
4620 weight,
4621 }) = &request.rerank
4622 {
4623 let fused_order = |(a_row, a_score, ..): &(
4624 RowId,
4625 f64,
4626 Vec<ComponentScore>,
4627 Option<f32>,
4628 f64,
4629 ),
4630 (b_row, b_score, ..): &(
4631 RowId,
4632 f64,
4633 Vec<ComponentScore>,
4634 Option<f32>,
4635 f64,
4636 )| {
4637 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
4638 };
4639 let selection_started = std::time::Instant::now();
4640 if let Some(context) = context {
4641 context.consume(ranked.len())?;
4642 }
4643 if ranked.len() > *candidate_limit {
4644 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
4645 ranked.truncate(*candidate_limit);
4646 }
4647 ranked.sort_by(fused_order);
4648 fusion_nanos =
4649 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4650 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4651 if let Some(context) = context {
4652 context.consume(row_ids.len())?;
4653 }
4654 let query_now =
4655 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4656 let gather_started = std::time::Instant::now();
4657 let vectors = self.values_for_rids_batch_at_with_context(
4658 &row_ids,
4659 *embedding_column,
4660 snapshot,
4661 query_now,
4662 context,
4663 )?;
4664 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
4665 let vector_work =
4666 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
4667 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
4668 if let Some(context) = context {
4669 context.consume(vector_work)?;
4670 }
4671 query
4672 .iter()
4673 .map(|value| f64::from(*value).powi(2))
4674 .sum::<f64>()
4675 .sqrt()
4676 } else {
4677 0.0
4678 };
4679 let score_started = std::time::Instant::now();
4680 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
4681 for (row_id, value) in vectors {
4682 let Value::Embedding(vector) = value else {
4683 continue;
4684 };
4685 let score = match metric {
4686 crate::query::VectorMetric::DotProduct => {
4687 if let Some(context) = context {
4688 context.consume(vector_work)?;
4689 }
4690 query
4691 .iter()
4692 .zip(&vector)
4693 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4694 .sum::<f64>()
4695 }
4696 crate::query::VectorMetric::Cosine => {
4697 if let Some(context) = context {
4698 context.consume(vector_work.saturating_mul(2))?;
4699 }
4700 let dot = query
4701 .iter()
4702 .zip(&vector)
4703 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4704 .sum::<f64>();
4705 let norm = vector
4706 .iter()
4707 .map(|value| f64::from(*value).powi(2))
4708 .sum::<f64>()
4709 .sqrt();
4710 if query_norm == 0.0 || norm == 0.0 {
4711 0.0
4712 } else {
4713 dot / (query_norm * norm)
4714 }
4715 }
4716 crate::query::VectorMetric::Euclidean => {
4717 if let Some(context) = context {
4718 context.consume(vector_work)?;
4719 }
4720 query
4721 .iter()
4722 .zip(&vector)
4723 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
4724 .sum::<f64>()
4725 .sqrt()
4726 }
4727 };
4728 if !score.is_finite() {
4729 return Err(MongrelError::InvalidArgument(
4730 "exact rerank score must be finite".into(),
4731 ));
4732 }
4733 scores.insert(row_id, score as f32);
4734 }
4735 let mut reranked = Vec::with_capacity(ranked.len());
4736 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
4737 let Some(score) = scores.get(&row_id).copied() else {
4738 continue;
4739 };
4740 let ordering_score = match metric {
4741 crate::query::VectorMetric::Euclidean => -f64::from(score),
4742 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
4743 f64::from(score)
4744 }
4745 };
4746 let final_score = fused_score + *weight * ordering_score;
4747 if !final_score.is_finite() {
4748 return Err(MongrelError::InvalidArgument(
4749 "final rerank score must be finite".into(),
4750 ));
4751 }
4752 reranked.push((row_id, fused_score, components, Some(score), final_score));
4753 }
4754 ranked = reranked;
4755 ranked.sort_by(order);
4756 crate::trace::QueryTrace::record(|trace| {
4757 trace.exact_vector_gather_nanos =
4758 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
4759 trace.exact_vector_score_nanos = trace
4760 .exact_vector_score_nanos
4761 .saturating_add(score_started.elapsed().as_nanos() as u64);
4762 });
4763 }
4764 if let Some(after) = after {
4765 ranked.retain(|(row_id, _, _, _, final_score)| {
4766 final_score.total_cmp(&after.final_score).is_lt()
4767 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
4768 });
4769 }
4770 let projection_started = std::time::Instant::now();
4771 let sentinel = projection
4772 .first()
4773 .copied()
4774 .or_else(|| self.schema.columns.first().map(|column| column.id));
4775 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4776 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
4777 let mut projection_rows = 0usize;
4778 let mut projection_cells = 0usize;
4779 while out.len() < request.limit && !ranked.is_empty() {
4780 if let Some(context) = context {
4781 context.checkpoint()?;
4782 context.consume(ranked.len())?;
4783 }
4784 let needed = request.limit - out.len();
4785 let window_size = ranked
4786 .len()
4787 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
4788 let selection_started = std::time::Instant::now();
4789 let mut remainder = if ranked.len() > window_size {
4790 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
4791 ranked.split_off(window_size)
4792 } else {
4793 Vec::new()
4794 };
4795 ranked.sort_by(order);
4796 fusion_nanos =
4797 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4798 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4799 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
4800 if let Some(context) = context {
4801 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
4802 }
4803 projection_rows = projection_rows.saturating_add(row_ids.len());
4804 projection_cells =
4805 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
4806 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
4807 std::collections::HashMap::new();
4808 if let Some(column_id) = sentinel {
4809 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4810 &row_ids, column_id, snapshot, query_now, context,
4811 )? {
4812 cells.entry(row_id).or_default().insert(column_id, value);
4813 }
4814 }
4815 for &column_id in &projection {
4816 if Some(column_id) == sentinel {
4817 continue;
4818 }
4819 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4820 &row_ids, column_id, snapshot, query_now, context,
4821 )? {
4822 cells.entry(row_id).or_default().insert(column_id, value);
4823 }
4824 }
4825 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
4826 ranked.drain(..)
4827 {
4828 let Some(row_cells) = cells.remove(&row_id) else {
4829 continue;
4830 };
4831 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4832 let final_rank = out.len() + 1;
4833 out.push(SearchHit {
4834 row_id,
4835 cells: projection
4836 .iter()
4837 .filter_map(|column_id| {
4838 row_cells
4839 .get(column_id)
4840 .cloned()
4841 .map(|value| (*column_id, value))
4842 })
4843 .collect(),
4844 components,
4845 fused_score,
4846 exact_rerank_score,
4847 final_score,
4848 final_rank,
4849 });
4850 if out.len() == request.limit {
4851 break;
4852 }
4853 }
4854 ranked.append(&mut remainder);
4855 }
4856 crate::trace::QueryTrace::record(|trace| {
4857 trace.union_size = union_size;
4858 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
4859 trace.projection_nanos = trace
4860 .projection_nanos
4861 .saturating_add(projection_started.elapsed().as_nanos() as u64);
4862 trace.total_nanos = trace
4863 .total_nanos
4864 .saturating_add(total_started.elapsed().as_nanos() as u64);
4865 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
4866 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
4867 if let Some(context) = context {
4868 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
4869 }
4870 });
4871 Ok(out)
4872 }
4873
4874 pub fn set_similarity(
4877 &mut self,
4878 request: &crate::query::SetSimilarityRequest,
4879 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4880 self.set_similarity_with_allowed(request, None)
4881 }
4882
4883 pub fn set_similarity_at(
4884 &mut self,
4885 request: &crate::query::SetSimilarityRequest,
4886 snapshot: Snapshot,
4887 allowed: Option<&std::collections::HashSet<RowId>>,
4888 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4889 self.set_similarity_explained_at(request, snapshot, allowed)
4890 .map(|(hits, _)| hits)
4891 }
4892
4893 pub fn ann_rerank(
4895 &mut self,
4896 request: &crate::query::AnnRerankRequest,
4897 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4898 self.ann_rerank_with_allowed(request, None)
4899 }
4900
4901 pub fn ann_rerank_with_allowed(
4902 &mut self,
4903 request: &crate::query::AnnRerankRequest,
4904 allowed: Option<&std::collections::HashSet<RowId>>,
4905 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4906 self.ann_rerank_at(request, self.snapshot(), allowed)
4907 }
4908
4909 pub fn ann_rerank_at(
4910 &mut self,
4911 request: &crate::query::AnnRerankRequest,
4912 snapshot: Snapshot,
4913 allowed: Option<&std::collections::HashSet<RowId>>,
4914 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4915 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
4916 }
4917
4918 pub fn ann_rerank_at_with_context(
4919 &mut self,
4920 request: &crate::query::AnnRerankRequest,
4921 snapshot: Snapshot,
4922 allowed: Option<&std::collections::HashSet<RowId>>,
4923 context: Option<&crate::query::AiExecutionContext>,
4924 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4925 self.ensure_indexes_complete()?;
4926 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
4927 }
4928
4929 pub fn ann_rerank_at_with_candidate_authorization_and_context(
4930 &mut self,
4931 request: &crate::query::AnnRerankRequest,
4932 snapshot: Snapshot,
4933 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4934 context: Option<&crate::query::AiExecutionContext>,
4935 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4936 self.ensure_indexes_complete()?;
4937 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
4938 }
4939
4940 #[doc(hidden)]
4941 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
4942 &self,
4943 request: &crate::query::AnnRerankRequest,
4944 snapshot: Snapshot,
4945 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4946 context: Option<&crate::query::AiExecutionContext>,
4947 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4948 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
4949 }
4950
4951 fn ann_rerank_at_with_filters_and_context(
4952 &self,
4953 request: &crate::query::AnnRerankRequest,
4954 snapshot: Snapshot,
4955 allowed: Option<&std::collections::HashSet<RowId>>,
4956 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4957 context: Option<&crate::query::AiExecutionContext>,
4958 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4959 use crate::query::{
4960 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
4961 };
4962 if request.candidate_k == 0 || request.limit == 0 {
4963 return Err(MongrelError::InvalidArgument(
4964 "candidate_k and limit must be > 0".into(),
4965 ));
4966 }
4967 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
4968 return Err(MongrelError::InvalidArgument(format!(
4969 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
4970 )));
4971 }
4972 let retriever = Retriever::Ann {
4973 column_id: request.column_id,
4974 query: request.query.clone(),
4975 k: request.candidate_k,
4976 };
4977 self.require_select()?;
4978 self.validate_retriever(&retriever)?;
4979 let hits = self.retrieve_filtered(
4980 &retriever,
4981 snapshot,
4982 None,
4983 allowed,
4984 candidate_authorization,
4985 context,
4986 )?;
4987 let distances: std::collections::HashMap<_, _> = hits
4988 .iter()
4989 .filter_map(|hit| match hit.score {
4990 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
4991 _ => None,
4992 })
4993 .collect();
4994 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4995 if let Some(context) = context {
4996 context.consume(row_ids.len())?;
4997 }
4998 let gather_started = std::time::Instant::now();
4999 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5000 let values = self.values_for_rids_batch_at_with_context(
5001 &row_ids,
5002 request.column_id,
5003 snapshot,
5004 query_now,
5005 context,
5006 )?;
5007 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5008 let score_started = std::time::Instant::now();
5009 let vector_work =
5010 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
5011 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
5012 if let Some(context) = context {
5013 context.consume(vector_work)?;
5014 }
5015 request
5016 .query
5017 .iter()
5018 .map(|value| f64::from(*value).powi(2))
5019 .sum::<f64>()
5020 .sqrt()
5021 } else {
5022 0.0
5023 };
5024 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
5025 for (row_id, value) in values {
5026 let Value::Embedding(vector) = value else {
5027 continue;
5028 };
5029 let exact_score = match request.metric {
5030 VectorMetric::DotProduct => {
5031 if let Some(context) = context {
5032 context.consume(vector_work)?;
5033 }
5034 request
5035 .query
5036 .iter()
5037 .zip(&vector)
5038 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5039 .sum::<f64>()
5040 }
5041 VectorMetric::Cosine => {
5042 if let Some(context) = context {
5043 context.consume(vector_work.saturating_mul(2))?;
5044 }
5045 let dot = request
5046 .query
5047 .iter()
5048 .zip(&vector)
5049 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5050 .sum::<f64>();
5051 let norm = vector
5052 .iter()
5053 .map(|value| f64::from(*value).powi(2))
5054 .sum::<f64>()
5055 .sqrt();
5056 if query_norm == 0.0 || norm == 0.0 {
5057 0.0
5058 } else {
5059 dot / (query_norm * norm)
5060 }
5061 }
5062 VectorMetric::Euclidean => {
5063 if let Some(context) = context {
5064 context.consume(vector_work)?;
5065 }
5066 request
5067 .query
5068 .iter()
5069 .zip(&vector)
5070 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
5071 .sum::<f64>()
5072 .sqrt()
5073 }
5074 };
5075 let exact_score = exact_score as f32;
5076 if !exact_score.is_finite() {
5077 return Err(MongrelError::InvalidArgument(
5078 "exact ANN score must be finite".into(),
5079 ));
5080 }
5081 reranked.push(AnnRerankHit {
5082 row_id,
5083 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
5084 exact_score,
5085 });
5086 }
5087 reranked.sort_by(|left, right| {
5088 let score = match request.metric {
5089 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
5090 VectorMetric::Cosine | VectorMetric::DotProduct => {
5091 right.exact_score.total_cmp(&left.exact_score)
5092 }
5093 };
5094 score.then_with(|| left.row_id.cmp(&right.row_id))
5095 });
5096 reranked.truncate(request.limit);
5097 crate::trace::QueryTrace::record(|trace| {
5098 trace.exact_vector_gather_nanos =
5099 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
5100 trace.exact_vector_score_nanos = trace
5101 .exact_vector_score_nanos
5102 .saturating_add(score_started.elapsed().as_nanos() as u64);
5103 });
5104 Ok(reranked)
5105 }
5106
5107 pub fn set_similarity_with_allowed(
5108 &mut self,
5109 request: &crate::query::SetSimilarityRequest,
5110 allowed: Option<&std::collections::HashSet<RowId>>,
5111 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5112 self.set_similarity_explained_at(request, self.snapshot(), allowed)
5113 .map(|(hits, _)| hits)
5114 }
5115
5116 pub fn set_similarity_explained(
5117 &mut self,
5118 request: &crate::query::SetSimilarityRequest,
5119 ) -> Result<(
5120 Vec<crate::query::SetSimilarityHit>,
5121 crate::query::SetSimilarityTrace,
5122 )> {
5123 self.set_similarity_explained_at(request, self.snapshot(), None)
5124 }
5125
5126 fn set_similarity_explained_at(
5127 &mut self,
5128 request: &crate::query::SetSimilarityRequest,
5129 snapshot: Snapshot,
5130 allowed: Option<&std::collections::HashSet<RowId>>,
5131 ) -> Result<(
5132 Vec<crate::query::SetSimilarityHit>,
5133 crate::query::SetSimilarityTrace,
5134 )> {
5135 self.ensure_indexes_complete()?;
5136 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
5137 }
5138
5139 pub fn set_similarity_at_with_context(
5140 &mut self,
5141 request: &crate::query::SetSimilarityRequest,
5142 snapshot: Snapshot,
5143 allowed: Option<&std::collections::HashSet<RowId>>,
5144 context: Option<&crate::query::AiExecutionContext>,
5145 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5146 self.ensure_indexes_complete()?;
5147 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
5148 .map(|(hits, _)| hits)
5149 }
5150
5151 pub fn set_similarity_at_with_candidate_authorization_and_context(
5152 &mut self,
5153 request: &crate::query::SetSimilarityRequest,
5154 snapshot: Snapshot,
5155 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5156 context: Option<&crate::query::AiExecutionContext>,
5157 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5158 self.ensure_indexes_complete()?;
5159 self.set_similarity_explained_at_with_context(
5160 request,
5161 snapshot,
5162 None,
5163 authorization,
5164 context,
5165 )
5166 .map(|(hits, _)| hits)
5167 }
5168
5169 #[doc(hidden)]
5170 pub fn set_similarity_at_with_candidate_authorization_on_generation(
5171 &self,
5172 request: &crate::query::SetSimilarityRequest,
5173 snapshot: Snapshot,
5174 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5175 context: Option<&crate::query::AiExecutionContext>,
5176 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5177 self.set_similarity_explained_at_with_context(
5178 request,
5179 snapshot,
5180 None,
5181 authorization,
5182 context,
5183 )
5184 .map(|(hits, _)| hits)
5185 }
5186
5187 fn set_similarity_explained_at_with_context(
5188 &self,
5189 request: &crate::query::SetSimilarityRequest,
5190 snapshot: Snapshot,
5191 allowed: Option<&std::collections::HashSet<RowId>>,
5192 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5193 context: Option<&crate::query::AiExecutionContext>,
5194 ) -> Result<(
5195 Vec<crate::query::SetSimilarityHit>,
5196 crate::query::SetSimilarityTrace,
5197 )> {
5198 use crate::query::{
5199 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
5200 MAX_SET_MEMBERS,
5201 };
5202 let mut trace = crate::query::SetSimilarityTrace::default();
5203 if request.members.is_empty() {
5204 return Ok((Vec::new(), trace));
5205 }
5206 if request.candidate_k == 0 || request.limit == 0 {
5207 return Err(MongrelError::InvalidArgument(
5208 "candidate_k and limit must be > 0".into(),
5209 ));
5210 }
5211 if request.candidate_k > MAX_RETRIEVER_K
5212 || request.limit > MAX_FINAL_LIMIT
5213 || request.members.len() > MAX_SET_MEMBERS
5214 {
5215 return Err(MongrelError::InvalidArgument(format!(
5216 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
5217 )));
5218 }
5219 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
5220 return Err(MongrelError::InvalidArgument(
5221 "min_jaccard must be finite and between 0 and 1".into(),
5222 ));
5223 }
5224 let started = std::time::Instant::now();
5225 let retriever = Retriever::MinHash {
5226 column_id: request.column_id,
5227 members: request.members.clone(),
5228 k: request.candidate_k,
5229 };
5230 self.require_select()?;
5231 self.validate_retriever(&retriever)?;
5232 let hits = self.retrieve_filtered(
5233 &retriever,
5234 snapshot,
5235 None,
5236 allowed,
5237 candidate_authorization,
5238 context,
5239 )?;
5240 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
5241 trace.candidate_count = hits.len();
5242 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
5243 if let Some(context) = context {
5244 context.consume(row_ids.len())?;
5245 }
5246 let started = std::time::Instant::now();
5247 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5248 let values = self.values_for_rids_batch_at_with_context(
5249 &row_ids,
5250 request.column_id,
5251 snapshot,
5252 query_now,
5253 context,
5254 )?;
5255 trace.gather_us = started.elapsed().as_micros() as u64;
5256 if let Some(context) = context {
5257 context.consume(request.members.len())?;
5258 }
5259 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
5260 let estimates: std::collections::HashMap<_, _> = hits
5261 .into_iter()
5262 .filter_map(|hit| match hit.score {
5263 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
5264 _ => None,
5265 })
5266 .collect();
5267 let started = std::time::Instant::now();
5268 let mut parsed = Vec::with_capacity(values.len());
5269 for (row_id, value) in values {
5270 let Value::Bytes(bytes) = value else {
5271 continue;
5272 };
5273 if let Some(context) = context {
5274 context.consume(crate::query::work_units(
5275 bytes.len(),
5276 crate::query::PARSE_WORK_QUANTUM,
5277 ))?;
5278 }
5279 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
5280 continue;
5281 };
5282 if let Some(context) = context {
5283 context.consume(members.len())?;
5284 }
5285 let stored = members
5286 .into_iter()
5287 .filter_map(|member| match member {
5288 serde_json::Value::String(value) => {
5289 Some(crate::query::SetMember::String(value))
5290 }
5291 serde_json::Value::Number(value) => {
5292 Some(crate::query::SetMember::Number(value))
5293 }
5294 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
5295 _ => None,
5296 })
5297 .collect::<std::collections::HashSet<_>>();
5298 parsed.push((row_id, stored));
5299 }
5300 trace.parse_us = started.elapsed().as_micros() as u64;
5301 trace.verified_count = parsed.len();
5302 let started = std::time::Instant::now();
5303 let mut exact = Vec::new();
5304 for (row_id, stored) in parsed {
5305 if let Some(context) = context {
5306 context.consume(query.len().saturating_add(stored.len()))?;
5307 }
5308 let union = query.union(&stored).count();
5309 let score = if union == 0 {
5310 1.0
5311 } else {
5312 query.intersection(&stored).count() as f32 / union as f32
5313 };
5314 if score >= request.min_jaccard {
5315 exact.push(SetSimilarityHit {
5316 row_id,
5317 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
5318 exact_jaccard: score,
5319 });
5320 }
5321 }
5322 exact.sort_by(|a, b| {
5323 b.exact_jaccard
5324 .total_cmp(&a.exact_jaccard)
5325 .then_with(|| a.row_id.cmp(&b.row_id))
5326 });
5327 exact.truncate(request.limit);
5328 trace.score_us = started.elapsed().as_micros() as u64;
5329 crate::trace::QueryTrace::record(|query_trace| {
5330 query_trace.exact_set_gather_nanos = query_trace
5331 .exact_set_gather_nanos
5332 .saturating_add(trace.gather_us.saturating_mul(1_000));
5333 query_trace.exact_set_parse_nanos = query_trace
5334 .exact_set_parse_nanos
5335 .saturating_add(trace.parse_us.saturating_mul(1_000));
5336 query_trace.exact_set_score_nanos = query_trace
5337 .exact_set_score_nanos
5338 .saturating_add(trace.score_us.saturating_mul(1_000));
5339 });
5340 Ok((exact, trace))
5341 }
5342
5343 fn values_for_rids_batch_at(
5345 &self,
5346 row_ids: &[u64],
5347 column_id: u16,
5348 snapshot: Snapshot,
5349 now: i64,
5350 ) -> Result<Vec<(RowId, Value)>> {
5351 if self.ttl.is_none()
5352 && self.memtable.is_empty()
5353 && self.mutable_run.is_empty()
5354 && self.run_refs.len() == 1
5355 {
5356 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5357 if row_ids.len().saturating_mul(24) < reader.row_count() {
5362 let mut values = Vec::with_capacity(row_ids.len());
5363 for &raw_row_id in row_ids {
5364 let row_id = RowId(raw_row_id);
5365 if let Some((_, false, Some(value))) =
5366 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5367 {
5368 values.push((row_id, value));
5369 }
5370 }
5371 return Ok(values);
5372 }
5373 let (positions, visible_row_ids) =
5374 reader.visible_positions_with_rids(snapshot.epoch)?;
5375 let requested: Vec<(RowId, usize)> = row_ids
5376 .iter()
5377 .filter_map(|raw| {
5378 visible_row_ids
5379 .binary_search(&(*raw as i64))
5380 .ok()
5381 .map(|index| (RowId(*raw), positions[index]))
5382 })
5383 .collect();
5384 let values = reader.gather_column(
5385 column_id,
5386 &requested
5387 .iter()
5388 .map(|(_, position)| *position)
5389 .collect::<Vec<_>>(),
5390 )?;
5391 return Ok(requested
5392 .into_iter()
5393 .zip(values)
5394 .map(|((row_id, _), value)| (row_id, value))
5395 .collect());
5396 }
5397 self.values_for_rids_at(row_ids, column_id, snapshot, now)
5398 }
5399
5400 fn values_for_rids_batch_at_with_context(
5401 &self,
5402 row_ids: &[u64],
5403 column_id: u16,
5404 snapshot: Snapshot,
5405 now: i64,
5406 context: Option<&crate::query::AiExecutionContext>,
5407 ) -> Result<Vec<(RowId, Value)>> {
5408 let Some(context) = context else {
5409 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
5410 };
5411 let mut values = Vec::with_capacity(row_ids.len());
5412 for chunk in row_ids.chunks(256) {
5413 context.checkpoint()?;
5414 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
5415 }
5416 Ok(values)
5417 }
5418
5419 fn values_for_rids_at(
5421 &self,
5422 row_ids: &[u64],
5423 column_id: u16,
5424 snapshot: Snapshot,
5425 now: i64,
5426 ) -> Result<Vec<(RowId, Value)>> {
5427 let mut readers: Vec<_> = self
5428 .run_refs
5429 .iter()
5430 .map(|run| self.open_reader(run.run_id))
5431 .collect::<Result<_>>()?;
5432 let mut out = Vec::with_capacity(row_ids.len());
5433 for &raw_row_id in row_ids {
5434 let row_id = RowId(raw_row_id);
5435 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5436 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5437 let overlay = match (mem, mutable) {
5438 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
5439 (a_epoch, a)
5440 } else {
5441 (b_epoch, b)
5442 }),
5443 (Some(value), None) | (None, Some(value)) => Some(value),
5444 (None, None) => None,
5445 };
5446 if let Some((_, row)) = overlay {
5447 if !row.deleted && !self.row_expired_at(&row, now) {
5448 if let Some(value) = row.columns.get(&column_id) {
5449 out.push((row_id, value.clone()));
5450 }
5451 }
5452 continue;
5453 }
5454
5455 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
5456 for (index, reader) in readers.iter_mut().enumerate() {
5457 if let Some((epoch, deleted, value)) =
5458 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5459 {
5460 if best
5461 .as_ref()
5462 .map(|(best_epoch, ..)| epoch > *best_epoch)
5463 .unwrap_or(true)
5464 {
5465 best = Some((epoch, deleted, value, index));
5466 }
5467 }
5468 }
5469 let Some((_, false, Some(value), reader_index)) = best else {
5470 continue;
5471 };
5472 if let Some(ttl) = self.ttl {
5473 if ttl.column_id != column_id {
5474 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5475 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5476 {
5477 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5478 continue;
5479 }
5480 }
5481 } else if let Value::Int64(timestamp) = value {
5482 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5483 continue;
5484 }
5485 }
5486 }
5487 out.push((row_id, value));
5488 }
5489 Ok(out)
5490 }
5491
5492 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
5497 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now())
5498 }
5499
5500 pub fn rows_for_rids_with_context(
5501 &self,
5502 rids: &[u64],
5503 snapshot: Snapshot,
5504 context: &crate::query::AiExecutionContext,
5505 ) -> Result<Vec<Row>> {
5506 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
5507 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos())
5508 }
5509
5510 fn rows_for_rids_at_time(
5511 &self,
5512 rids: &[u64],
5513 snapshot: Snapshot,
5514 ttl_now: i64,
5515 ) -> Result<Vec<Row>> {
5516 use std::collections::HashMap;
5517 let mut rows = Vec::with_capacity(rids.len());
5518 let tier_size = self.memtable.len() + self.mutable_run.len();
5535 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
5536 if rids.len().saturating_mul(24) < tier_size {
5537 for &rid in rids {
5538 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
5539 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
5540 let newest = match (mem, mrun) {
5541 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
5542 (Some((_, mr)), None) => Some(mr),
5543 (None, Some((_, rr))) => Some(rr),
5544 (None, None) => None,
5545 };
5546 if let Some(row) = newest {
5547 overlay.insert(rid, row);
5548 }
5549 }
5550 } else {
5551 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
5552 overlay
5553 .entry(row.row_id.0)
5554 .and_modify(|e| {
5555 if row.committed_epoch > e.committed_epoch {
5556 *e = row.clone();
5557 }
5558 })
5559 .or_insert(row);
5560 };
5561 for row in self.memtable.visible_versions(snapshot.epoch) {
5562 fold_newest(row, &mut overlay);
5563 }
5564 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5565 fold_newest(row, &mut overlay);
5566 }
5567 }
5568 if self.run_refs.len() == 1 {
5569 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5570 if rids.len().saturating_mul(24) < reader.row_count() {
5578 for &rid in rids {
5579 if let Some(r) = overlay.get(&rid) {
5580 if !r.deleted {
5581 rows.push(r.clone());
5582 }
5583 continue;
5584 }
5585 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
5586 if !row.deleted {
5587 rows.push(row);
5588 }
5589 }
5590 }
5591 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5592 return Ok(rows);
5593 }
5594 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
5603 enum Src {
5606 Overlay,
5607 Run,
5608 }
5609 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
5610 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
5611 for rid in rids {
5612 if overlay.contains_key(rid) {
5613 plan.push(Src::Overlay);
5614 continue;
5615 }
5616 match vis_rids.binary_search(&(*rid as i64)) {
5617 Ok(i) => {
5618 plan.push(Src::Run);
5619 fetch.push(positions[i]);
5620 }
5621 Err(_) => { }
5622 }
5623 }
5624 let fetched = reader.materialize_batch(&fetch)?;
5625 let mut fetched_iter = fetched.into_iter();
5626 for (rid, src) in rids.iter().zip(plan) {
5627 match src {
5628 Src::Overlay => {
5629 if let Some(r) = overlay.get(rid) {
5630 if !r.deleted {
5631 rows.push(r.clone());
5632 }
5633 }
5634 }
5635 Src::Run => {
5636 if let Some(row) = fetched_iter.next() {
5637 if !row.deleted {
5638 rows.push(row);
5639 }
5640 }
5641 }
5642 }
5643 }
5644 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5645 return Ok(rows);
5646 }
5647 let mut readers: Vec<_> = self
5651 .run_refs
5652 .iter()
5653 .map(|rr| self.open_reader(rr.run_id))
5654 .collect::<Result<Vec<_>>>()?;
5655 for rid in rids {
5656 if let Some(r) = overlay.get(rid) {
5657 if !r.deleted {
5658 rows.push(r.clone());
5659 }
5660 continue;
5661 }
5662 let mut best: Option<(Epoch, Row)> = None;
5663 for reader in readers.iter_mut() {
5664 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
5665 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
5666 best = Some((epoch, row));
5667 }
5668 }
5669 }
5670 if let Some((_, r)) = best {
5671 if !r.deleted {
5672 rows.push(r);
5673 }
5674 }
5675 }
5676 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5677 Ok(rows)
5678 }
5679
5680 pub fn indexes_complete(&self) -> bool {
5690 self.indexes_complete
5691 }
5692
5693 pub fn index_build_policy(&self) -> IndexBuildPolicy {
5695 self.index_build_policy
5696 }
5697
5698 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
5702 self.index_build_policy = policy;
5703 }
5704
5705 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
5710 if !self.indexes_complete {
5714 return None;
5715 }
5716 let b = self.bitmap.get(&column_id)?;
5717 let result: Vec<Vec<u8>> = b
5718 .keys()
5719 .into_iter()
5720 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
5721 .cloned()
5722 .collect();
5723 Some(result)
5724 }
5725
5726 pub fn fk_join_row_ids(
5727 &self,
5728 fk_column_id: u16,
5729 pk_values: &[Vec<u8>],
5730 fk_conditions: &[crate::query::Condition],
5731 snapshot: Snapshot,
5732 ) -> Result<Vec<u64>> {
5733 let Some(b) = self.bitmap.get(&fk_column_id) else {
5734 return Ok(Vec::new());
5735 };
5736 let mut join_set = {
5737 let mut acc = roaring::RoaringBitmap::new();
5738 for v in pk_values {
5739 acc |= b.get(v);
5740 }
5741 RowIdSet::from_roaring(acc)
5742 };
5743 if !fk_conditions.is_empty() {
5744 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5745 sets.push(join_set);
5746 for c in fk_conditions {
5747 sets.push(self.resolve_condition(c, snapshot)?);
5748 }
5749 join_set = RowIdSet::intersect_many(sets);
5750 }
5751 Ok(join_set.into_sorted_vec())
5752 }
5753
5754 pub fn fk_join_count(
5760 &self,
5761 fk_column_id: u16,
5762 pk_values: &[Vec<u8>],
5763 fk_conditions: &[crate::query::Condition],
5764 snapshot: Snapshot,
5765 ) -> Result<u64> {
5766 let Some(b) = self.bitmap.get(&fk_column_id) else {
5767 return Ok(0);
5768 };
5769 let mut acc = roaring::RoaringBitmap::new();
5770 for v in pk_values {
5771 acc |= b.get(v);
5772 }
5773 if fk_conditions.is_empty() {
5774 return Ok(acc.len());
5775 }
5776 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5777 sets.push(RowIdSet::from_roaring(acc));
5778 for c in fk_conditions {
5779 sets.push(self.resolve_condition(c, snapshot)?);
5780 }
5781 Ok(RowIdSet::intersect_many(sets).len() as u64)
5782 }
5783
5784 fn resolve_condition(
5789 &self,
5790 c: &crate::query::Condition,
5791 snapshot: Snapshot,
5792 ) -> Result<RowIdSet> {
5793 self.resolve_condition_with_allowed(c, snapshot, None)
5794 }
5795
5796 fn resolve_condition_with_allowed(
5797 &self,
5798 c: &crate::query::Condition,
5799 snapshot: Snapshot,
5800 allowed: Option<&std::collections::HashSet<RowId>>,
5801 ) -> Result<RowIdSet> {
5802 use crate::query::Condition;
5803 self.validate_condition(c)?;
5804 Ok(match c {
5805 Condition::Pk(key) => {
5806 let lookup = self
5807 .schema
5808 .primary_key()
5809 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5810 .unwrap_or_else(|| key.clone());
5811 self.hot
5812 .get(&lookup)
5813 .map(|r| RowIdSet::one(r.0))
5814 .unwrap_or_else(RowIdSet::empty)
5815 }
5816 Condition::BitmapEq { column_id, value } => {
5817 let lookup = self.index_lookup_key_bytes(*column_id, value);
5818 self.bitmap
5819 .get(column_id)
5820 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5821 .unwrap_or_else(RowIdSet::empty)
5822 }
5823 Condition::BitmapIn { column_id, values } => {
5824 let bm = self.bitmap.get(column_id);
5825 let mut acc = roaring::RoaringBitmap::new();
5826 if let Some(b) = bm {
5827 for v in values {
5828 let lookup = self.index_lookup_key_bytes(*column_id, v);
5829 acc |= b.get(&lookup);
5830 }
5831 }
5832 RowIdSet::from_roaring(acc)
5833 }
5834 Condition::BytesPrefix { column_id, prefix } => {
5835 if let Some(b) = self.bitmap.get(column_id) {
5840 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5841 let mut acc = roaring::RoaringBitmap::new();
5842 for key in b.keys() {
5843 if key.starts_with(&lookup_prefix) {
5844 acc |= b.get(key);
5845 }
5846 }
5847 RowIdSet::from_roaring(acc)
5848 } else {
5849 RowIdSet::empty()
5850 }
5851 }
5852 Condition::FmContains { column_id, pattern } => self
5853 .fm
5854 .get(column_id)
5855 .map(|f| {
5856 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
5857 })
5858 .unwrap_or_else(RowIdSet::empty),
5859 Condition::FmContainsAll {
5860 column_id,
5861 patterns,
5862 } => {
5863 if let Some(f) = self.fm.get(column_id) {
5866 let sets: Vec<RowIdSet> = patterns
5867 .iter()
5868 .map(|pat| {
5869 RowIdSet::from_unsorted(
5870 f.locate(pat).into_iter().map(|r| r.0).collect(),
5871 )
5872 })
5873 .collect();
5874 RowIdSet::intersect_many(sets)
5875 } else {
5876 RowIdSet::empty()
5877 }
5878 }
5879 Condition::Ann {
5880 column_id,
5881 query,
5882 k,
5883 } => RowIdSet::from_unsorted(
5884 self.retrieve_filtered(
5885 &crate::query::Retriever::Ann {
5886 column_id: *column_id,
5887 query: query.clone(),
5888 k: *k,
5889 },
5890 snapshot,
5891 None,
5892 allowed,
5893 None,
5894 None,
5895 )?
5896 .into_iter()
5897 .map(|hit| hit.row_id.0)
5898 .collect(),
5899 ),
5900 Condition::SparseMatch {
5901 column_id,
5902 query,
5903 k,
5904 } => RowIdSet::from_unsorted(
5905 self.retrieve_filtered(
5906 &crate::query::Retriever::Sparse {
5907 column_id: *column_id,
5908 query: query.clone(),
5909 k: *k,
5910 },
5911 snapshot,
5912 None,
5913 allowed,
5914 None,
5915 None,
5916 )?
5917 .into_iter()
5918 .map(|hit| hit.row_id.0)
5919 .collect(),
5920 ),
5921 Condition::MinHashSimilar {
5922 column_id,
5923 query,
5924 k,
5925 } => match self.minhash.get(column_id) {
5926 Some(index) => {
5927 let candidates = index.candidate_row_ids(query);
5928 let eligible =
5929 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
5930 RowIdSet::from_unsorted(
5931 index
5932 .search_filtered(query, *k, |row_id| {
5933 eligible.contains(&row_id)
5934 && allowed.map_or(true, |allowed| allowed.contains(&row_id))
5935 })
5936 .into_iter()
5937 .map(|(row_id, _)| row_id.0)
5938 .collect(),
5939 )
5940 }
5941 None => RowIdSet::empty(),
5942 },
5943 Condition::Range { column_id, lo, hi } => {
5944 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5953 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5954 } else if self.run_refs.len() == 1 {
5955 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5956 r.range_row_id_set_i64(*column_id, *lo, *hi)?
5957 } else {
5958 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
5959 };
5960 set.remove_many(self.overlay_rid_set(snapshot));
5961 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
5962 set
5963 }
5964 Condition::RangeF64 {
5965 column_id,
5966 lo,
5967 lo_inclusive,
5968 hi,
5969 hi_inclusive,
5970 } => {
5971 let mut set = if let Some(li) = self.learned_range.get(column_id) {
5974 RowIdSet::from_unsorted(
5975 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5976 .into_iter()
5977 .collect(),
5978 )
5979 } else if self.run_refs.len() == 1 {
5980 let mut r = self.open_reader(self.run_refs[0].run_id)?;
5981 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
5982 } else {
5983 return self.range_scan_f64(
5984 *column_id,
5985 *lo,
5986 *lo_inclusive,
5987 *hi,
5988 *hi_inclusive,
5989 snapshot,
5990 );
5991 };
5992 set.remove_many(self.overlay_rid_set(snapshot));
5993 self.range_scan_overlay_f64(
5994 &mut set,
5995 *column_id,
5996 *lo,
5997 *lo_inclusive,
5998 *hi,
5999 *hi_inclusive,
6000 snapshot,
6001 );
6002 set
6003 }
6004 Condition::IsNull { column_id } => {
6005 let mut set = if self.run_refs.len() == 1 {
6006 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6007 r.null_row_id_set(*column_id, true)?
6008 } else {
6009 return self.null_scan(*column_id, true, snapshot);
6010 };
6011 set.remove_many(self.overlay_rid_set(snapshot));
6012 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
6013 set
6014 }
6015 Condition::IsNotNull { column_id } => {
6016 let mut set = if self.run_refs.len() == 1 {
6017 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6018 r.null_row_id_set(*column_id, false)?
6019 } else {
6020 return self.null_scan(*column_id, false, snapshot);
6021 };
6022 set.remove_many(self.overlay_rid_set(snapshot));
6023 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
6024 set
6025 }
6026 })
6027 }
6028
6029 fn range_scan_i64(
6037 &self,
6038 column_id: u16,
6039 lo: i64,
6040 hi: i64,
6041 snapshot: Snapshot,
6042 ) -> Result<RowIdSet> {
6043 let mut row_ids = Vec::new();
6044 let overlay_rids = self.overlay_rid_set(snapshot);
6045 for rr in &self.run_refs {
6046 let mut reader = self.open_reader(rr.run_id)?;
6047 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
6048 for rid in matched {
6049 if !overlay_rids.contains(&rid) {
6050 row_ids.push(rid);
6051 }
6052 }
6053 }
6054 let mut s = RowIdSet::from_unsorted(row_ids);
6055 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
6056 Ok(s)
6057 }
6058
6059 fn range_scan_f64(
6062 &self,
6063 column_id: u16,
6064 lo: f64,
6065 lo_inclusive: bool,
6066 hi: f64,
6067 hi_inclusive: bool,
6068 snapshot: Snapshot,
6069 ) -> Result<RowIdSet> {
6070 let mut row_ids = Vec::new();
6071 let overlay_rids = self.overlay_rid_set(snapshot);
6072 for rr in &self.run_refs {
6073 let mut reader = self.open_reader(rr.run_id)?;
6074 let matched = reader.range_row_ids_visible_f64(
6075 column_id,
6076 lo,
6077 lo_inclusive,
6078 hi,
6079 hi_inclusive,
6080 snapshot.epoch,
6081 )?;
6082 for rid in matched {
6083 if !overlay_rids.contains(&rid) {
6084 row_ids.push(rid);
6085 }
6086 }
6087 }
6088 let mut s = RowIdSet::from_unsorted(row_ids);
6089 self.range_scan_overlay_f64(
6090 &mut s,
6091 column_id,
6092 lo,
6093 lo_inclusive,
6094 hi,
6095 hi_inclusive,
6096 snapshot,
6097 );
6098 Ok(s)
6099 }
6100
6101 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
6103 let mut s = HashSet::new();
6104 for row in self.memtable.visible_versions(snapshot.epoch) {
6105 s.insert(row.row_id.0);
6106 }
6107 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6108 s.insert(row.row_id.0);
6109 }
6110 s
6111 }
6112
6113 fn range_scan_overlay_i64(
6114 &self,
6115 s: &mut RowIdSet,
6116 column_id: u16,
6117 lo: i64,
6118 hi: i64,
6119 snapshot: Snapshot,
6120 ) {
6121 let mut newest: HashMap<u64, &Row> = HashMap::new();
6126 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6127 let memtable = self.memtable.visible_versions(snapshot.epoch);
6128 for r in &mutable {
6129 newest.entry(r.row_id.0).or_insert(r);
6130 }
6131 for r in &memtable {
6132 newest.insert(r.row_id.0, r);
6133 }
6134 for row in newest.values() {
6135 if !row.deleted {
6136 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
6137 if *v >= lo && *v <= hi {
6138 s.insert(row.row_id.0);
6139 }
6140 }
6141 }
6142 }
6143 }
6144
6145 #[allow(clippy::too_many_arguments)]
6146 fn range_scan_overlay_f64(
6147 &self,
6148 s: &mut RowIdSet,
6149 column_id: u16,
6150 lo: f64,
6151 lo_inclusive: bool,
6152 hi: f64,
6153 hi_inclusive: bool,
6154 snapshot: Snapshot,
6155 ) {
6156 let mut newest: HashMap<u64, &Row> = HashMap::new();
6159 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6160 let memtable = self.memtable.visible_versions(snapshot.epoch);
6161 for r in &mutable {
6162 newest.entry(r.row_id.0).or_insert(r);
6163 }
6164 for r in &memtable {
6165 newest.insert(r.row_id.0, r);
6166 }
6167 for row in newest.values() {
6168 if !row.deleted {
6169 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
6170 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
6171 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
6172 if ok_lo && ok_hi {
6173 s.insert(row.row_id.0);
6174 }
6175 }
6176 }
6177 }
6178 }
6179
6180 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
6183 let mut row_ids = Vec::new();
6184 let overlay_rids = self.overlay_rid_set(snapshot);
6185 for rr in &self.run_refs {
6186 let mut reader = self.open_reader(rr.run_id)?;
6187 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
6188 for rid in matched {
6189 if !overlay_rids.contains(&rid) {
6190 row_ids.push(rid);
6191 }
6192 }
6193 }
6194 let mut s = RowIdSet::from_unsorted(row_ids);
6195 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
6196 Ok(s)
6197 }
6198
6199 fn null_scan_overlay(
6203 &self,
6204 s: &mut RowIdSet,
6205 column_id: u16,
6206 want_nulls: bool,
6207 snapshot: Snapshot,
6208 ) {
6209 let mut newest: HashMap<u64, &Row> = HashMap::new();
6210 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6211 let memtable = self.memtable.visible_versions(snapshot.epoch);
6212 for r in &mutable {
6213 newest.entry(r.row_id.0).or_insert(r);
6214 }
6215 for r in &memtable {
6216 newest.insert(r.row_id.0, r);
6217 }
6218 for row in newest.values() {
6219 if row.deleted {
6220 continue;
6221 }
6222 let is_null = !row.columns.contains_key(&column_id)
6223 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
6224 if is_null == want_nulls {
6225 s.insert(row.row_id.0);
6226 }
6227 }
6228 }
6229
6230 pub fn snapshot(&self) -> Snapshot {
6231 Snapshot::at(self.epoch.visible())
6232 }
6233
6234 pub fn data_generation(&self) -> u64 {
6236 self.data_generation
6237 }
6238
6239 pub(crate) fn table_id(&self) -> u64 {
6240 self.table_id
6241 }
6242
6243 #[doc(hidden)]
6244 pub fn clone_read_generation(&mut self) -> Result<Self> {
6245 self.ensure_indexes_complete()?;
6246 let mut generation = self.clone();
6247 generation.read_only = true;
6248 generation.wal = WalSink::ReadOnly;
6249 generation.pending_rows.clear();
6250 generation.pending_rows_auto_inc.clear();
6251 generation.pending_dels.clear();
6252 generation.pending_truncate = None;
6253 Ok(generation)
6254 }
6255
6256 pub fn pin_snapshot(&mut self) -> Snapshot {
6259 let e = self.epoch.visible();
6260 *self.pinned.entry(e).or_insert(0) += 1;
6261 Snapshot::at(e)
6262 }
6263
6264 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
6266 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
6267 *count -= 1;
6268 if *count == 0 {
6269 self.pinned.remove(&snap.epoch);
6270 }
6271 }
6272 }
6273
6274 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
6284 let local = self.pinned.keys().next().copied();
6285 let global = self.snapshots.min_pinned();
6286 let history = self.snapshots.history_floor(self.current_epoch());
6287 [local, global, history].into_iter().flatten().min()
6288 }
6289
6290 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
6294 self.ensure_writable()?;
6295 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
6296 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
6297 }
6298
6299 pub fn clear_ttl(&mut self) -> Result<()> {
6300 self.ensure_writable()?;
6301 self.apply_ttl_policy_at(None, self.current_epoch())
6302 }
6303
6304 pub fn ttl(&self) -> Option<TtlPolicy> {
6305 self.ttl
6306 }
6307
6308 pub(crate) fn prepare_ttl_policy(
6309 &self,
6310 column_name: &str,
6311 duration_nanos: u64,
6312 ) -> Result<TtlPolicy> {
6313 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
6314 return Err(MongrelError::InvalidArgument(
6315 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
6316 ));
6317 }
6318 let column = self
6319 .schema
6320 .columns
6321 .iter()
6322 .find(|column| column.name == column_name)
6323 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
6324 if column.ty != TypeId::TimestampNanos {
6325 return Err(MongrelError::Schema(format!(
6326 "TTL column {column_name} must be TimestampNanos, is {:?}",
6327 column.ty
6328 )));
6329 }
6330 Ok(TtlPolicy {
6331 column_id: column.id,
6332 duration_nanos,
6333 })
6334 }
6335
6336 pub(crate) fn apply_ttl_policy_at(
6337 &mut self,
6338 policy: Option<TtlPolicy>,
6339 epoch: Epoch,
6340 ) -> Result<()> {
6341 if let Some(policy) = policy {
6342 let column = self
6343 .schema
6344 .columns
6345 .iter()
6346 .find(|column| column.id == policy.column_id)
6347 .ok_or_else(|| {
6348 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
6349 })?;
6350 if column.ty != TypeId::TimestampNanos
6351 || policy.duration_nanos == 0
6352 || policy.duration_nanos > i64::MAX as u64
6353 {
6354 return Err(MongrelError::Schema("invalid TTL policy".into()));
6355 }
6356 }
6357 self.ttl = policy;
6358 self.agg_cache.clear();
6359 self.clear_result_cache();
6360 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6361 self.persist_manifest(epoch)
6362 }
6363
6364 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
6365 let Some(policy) = self.ttl else {
6366 return false;
6367 };
6368 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
6369 return false;
6370 };
6371 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
6372 }
6373
6374 pub fn current_epoch(&self) -> Epoch {
6375 self.epoch.visible()
6376 }
6377
6378 pub fn memtable_len(&self) -> usize {
6379 self.memtable.len()
6380 }
6381
6382 pub fn count(&self) -> u64 {
6385 if self.ttl.is_none()
6386 && self.pending_put_cols.is_empty()
6387 && self.pending_delete_rids.is_empty()
6388 && self.pending_rows.is_empty()
6389 && self.pending_dels.is_empty()
6390 && self.pending_truncate.is_none()
6391 {
6392 self.live_count
6393 } else {
6394 self.visible_rows(self.snapshot())
6395 .map(|rows| rows.len() as u64)
6396 .unwrap_or(self.live_count)
6397 }
6398 }
6399
6400 pub fn count_conditions(
6404 &mut self,
6405 conditions: &[crate::query::Condition],
6406 snapshot: Snapshot,
6407 ) -> Result<Option<u64>> {
6408 use crate::query::Condition;
6409 if self.ttl.is_some() {
6410 if conditions.is_empty() {
6411 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
6412 }
6413 let mut sets = Vec::with_capacity(conditions.len());
6414 for condition in conditions {
6415 sets.push(self.resolve_condition(condition, snapshot)?);
6416 }
6417 let survivors = RowIdSet::intersect_many(sets);
6418 let rows = self.visible_rows(snapshot)?;
6419 return Ok(Some(
6420 rows.into_iter()
6421 .filter(|row| survivors.contains(row.row_id.0))
6422 .count() as u64,
6423 ));
6424 }
6425 if conditions.is_empty() {
6426 return Ok(Some(self.count()));
6427 }
6428 let served = |c: &Condition| {
6429 matches!(
6430 c,
6431 Condition::Pk(_)
6432 | Condition::BitmapEq { .. }
6433 | Condition::BitmapIn { .. }
6434 | Condition::BytesPrefix { .. }
6435 | Condition::FmContains { .. }
6436 | Condition::FmContainsAll { .. }
6437 | Condition::Ann { .. }
6438 | Condition::Range { .. }
6439 | Condition::RangeF64 { .. }
6440 | Condition::SparseMatch { .. }
6441 | Condition::MinHashSimilar { .. }
6442 | Condition::IsNull { .. }
6443 | Condition::IsNotNull { .. }
6444 )
6445 };
6446 if !conditions.iter().all(served) {
6447 return Ok(None);
6448 }
6449 self.ensure_indexes_complete()?;
6450 if !self.pending_put_cols.is_empty()
6451 || !self.pending_delete_rids.is_empty()
6452 || !self.pending_rows.is_empty()
6453 || !self.pending_dels.is_empty()
6454 || self.pending_truncate.is_some()
6455 {
6456 let mut sets = Vec::with_capacity(conditions.len());
6457 for condition in conditions {
6458 sets.push(self.resolve_condition(condition, snapshot)?);
6459 }
6460 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6461 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
6462 }
6463 let mut sets = Vec::with_capacity(conditions.len());
6464 for condition in conditions {
6465 sets.push(self.resolve_condition(condition, snapshot)?);
6466 }
6467 let mut rids = RowIdSet::intersect_many(sets);
6468 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6478 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
6479 }
6480 let count = rids.len() as u64;
6481 crate::trace::QueryTrace::record(|t| {
6482 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
6483 t.survivor_count = Some(count as usize);
6484 t.conditions_pushed = conditions.len();
6485 });
6486 Ok(Some(count))
6487 }
6488
6489 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
6494 let mut out = Vec::new();
6495 for row in self.memtable.visible_versions(snapshot.epoch) {
6496 if row.deleted {
6497 out.push(row.row_id.0);
6498 }
6499 }
6500 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6501 if row.deleted {
6502 out.push(row.row_id.0);
6503 }
6504 }
6505 out
6506 }
6507
6508 pub fn bulk_load_columns(
6517 &mut self,
6518 user_columns: Vec<(u16, columnar::NativeColumn)>,
6519 ) -> Result<Epoch> {
6520 self.bulk_load_columns_with(user_columns, 3, false, true)
6521 }
6522
6523 pub fn bulk_load_fast(
6530 &mut self,
6531 user_columns: Vec<(u16, columnar::NativeColumn)>,
6532 ) -> Result<Epoch> {
6533 self.bulk_load_columns_with(user_columns, -1, true, false)
6534 }
6535
6536 fn bulk_load_columns_with(
6537 &mut self,
6538 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
6539 zstd_level: i32,
6540 force_plain: bool,
6541 lz4: bool,
6542 ) -> Result<Epoch> {
6543 let epoch = self.commit()?;
6544 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
6545 if n == 0 {
6546 return Ok(epoch);
6547 }
6548 let live_before = self.live_count;
6549 self.spill_mutable_run(epoch)?;
6551 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6552 && self.indexes_complete
6553 && self.run_refs.is_empty()
6554 && self.memtable.is_empty()
6555 && self.mutable_run.is_empty();
6556 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6559 self.validate_columns_not_null(&user_columns, n)?;
6560 let winner_idx = self
6561 .bulk_pk_winner_indices(&user_columns, n)
6562 .filter(|idx| idx.len() != n);
6563 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6564 match winner_idx.as_deref() {
6565 Some(idx) => {
6566 let compacted = user_columns
6567 .iter()
6568 .map(|(id, c)| (*id, c.gather(idx)))
6569 .collect();
6570 (compacted, idx.len())
6571 }
6572 None => (user_columns, n),
6573 };
6574 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6575 let first = self.allocator.alloc_range(write_n as u64).0;
6576 for rid in first..first + write_n as u64 {
6577 self.reservoir.offer(rid);
6578 }
6579 let run_id = self.next_run_id;
6580 self.next_run_id += 1;
6581 let path = self.run_path(run_id);
6582 let mut writer =
6583 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
6584 if force_plain {
6585 writer = writer.with_plain();
6586 } else if lz4 {
6587 writer = writer.with_lz4();
6590 } else {
6591 writer = writer.with_zstd_level(zstd_level);
6592 }
6593 if let Some(kek) = &self.kek {
6594 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6595 }
6596 let header = writer.write_native(&path, &write_columns, write_n, first)?;
6597 self.run_refs.push(RunRef {
6598 run_id: run_id as u128,
6599 level: 0,
6600 epoch_created: epoch.0,
6601 row_count: header.row_count,
6602 });
6603 self.live_count = self.live_count.saturating_add(write_n as u64);
6604 if eager_index_build {
6605 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6606 self.index_columns_bulk(&write_columns, &row_ids);
6607 self.indexes_complete = true;
6608 self.build_learned_ranges()?;
6609 } else {
6610 self.indexes_complete = false;
6614 }
6615 self.mark_flushed(epoch)?;
6616 self.persist_manifest(epoch)?;
6617 if eager_index_build {
6618 self.checkpoint_indexes(epoch);
6619 }
6620 self.clear_result_cache();
6621 self.data_generation = self.data_generation.wrapping_add(1);
6622 Ok(epoch)
6623 }
6624
6625 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
6643 let n = row_ids.len();
6644 if n == 0 {
6645 return;
6646 }
6647 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
6648 columns.iter().map(|(id, c)| (*id, c)).collect();
6649 let ty_of: std::collections::HashMap<u16, TypeId> = self
6650 .schema
6651 .columns
6652 .iter()
6653 .map(|c| (c.id, c.ty.clone()))
6654 .collect();
6655 let pk_id = self.schema.primary_key().map(|c| c.id);
6656
6657 for (i, &rid) in row_ids.iter().enumerate() {
6658 let row_id = RowId(rid);
6659 if let Some(pid) = pk_id {
6660 if let Some(col) = by_id.get(&pid) {
6661 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
6662 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
6663 self.insert_hot_pk(key, row_id);
6664 }
6665 }
6666 }
6667 for idef in &self.schema.indexes {
6668 let Some(col) = by_id.get(&idef.column_id) else {
6669 continue;
6670 };
6671 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
6672 match idef.kind {
6673 IndexKind::Bitmap => {
6674 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
6675 if let Some(key) =
6676 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
6677 {
6678 b.insert(key, row_id);
6679 }
6680 }
6681 }
6682 IndexKind::FmIndex => {
6683 if let Some(f) = self.fm.get_mut(&idef.column_id) {
6684 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6685 f.insert(bytes.to_vec(), row_id);
6686 }
6687 }
6688 }
6689 IndexKind::Sparse => {
6690 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
6691 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6692 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
6693 s.insert(&terms, row_id);
6694 }
6695 }
6696 }
6697 }
6698 IndexKind::MinHash => {
6699 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
6700 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6701 let tokens = crate::index::token_hashes_from_bytes(bytes);
6702 mh.insert(&tokens, row_id);
6703 }
6704 }
6705 }
6706 _ => {}
6707 }
6708 }
6709 }
6710 }
6711
6712 pub fn visible_columns_native(
6717 &self,
6718 snapshot: Snapshot,
6719 projection: Option<&[u16]>,
6720 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
6721 let wanted: Vec<u16> = match projection {
6722 Some(p) => p.to_vec(),
6723 None => self.schema.columns.iter().map(|c| c.id).collect(),
6724 };
6725 if self.ttl.is_none()
6726 && self.memtable.is_empty()
6727 && self.mutable_run.is_empty()
6728 && self.run_refs.len() == 1
6729 {
6730 let rr = self.run_refs[0].clone();
6731 let mut reader = self.open_reader(rr.run_id)?;
6732 let idxs = reader.visible_indices_native(snapshot.epoch)?;
6733 let all_visible = idxs.len() == reader.row_count();
6734 if reader.has_mmap() {
6740 use rayon::prelude::*;
6741 let valid: Vec<u16> = wanted
6744 .iter()
6745 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
6746 .copied()
6747 .collect();
6748 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
6750 .par_iter()
6751 .filter_map(|cid| {
6752 reader
6753 .column_native_shared(*cid)
6754 .ok()
6755 .map(|col| (*cid, col))
6756 })
6757 .collect();
6758 let cols = decoded
6759 .into_iter()
6760 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
6761 .collect();
6762 return Ok(cols);
6763 }
6764 let mut cols = Vec::with_capacity(wanted.len());
6765 for cid in &wanted {
6766 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
6767 Some(c) => c,
6768 None => continue,
6769 };
6770 let col = reader.column_native(cdef.id)?;
6771 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
6772 }
6773 return Ok(cols);
6774 }
6775 let vcols = self.visible_columns(snapshot)?;
6776 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
6777 let out: Vec<(u16, columnar::NativeColumn)> = vcols
6778 .into_iter()
6779 .filter(|(id, _)| want_set.contains(id))
6780 .map(|(id, vals)| {
6781 let ty = self
6782 .schema
6783 .columns
6784 .iter()
6785 .find(|c| c.id == id)
6786 .map(|c| c.ty.clone())
6787 .unwrap_or(TypeId::Bytes);
6788 (id, columnar::values_to_native(ty, &vals))
6789 })
6790 .collect();
6791 Ok(out)
6792 }
6793
6794 pub fn run_count(&self) -> usize {
6795 self.run_refs.len()
6796 }
6797
6798 pub fn memtable_is_empty(&self) -> bool {
6800 self.memtable.is_empty()
6801 }
6802
6803 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
6807 self.page_cache.stats()
6808 }
6809
6810 pub fn reset_page_cache_stats(&self) {
6812 self.page_cache.reset_stats();
6813 }
6814
6815 pub fn run_ids(&self) -> Vec<u128> {
6818 self.run_refs.iter().map(|r| r.run_id).collect()
6819 }
6820
6821 pub fn single_run_is_clean(&self) -> bool {
6825 if self.ttl.is_some() || self.run_refs.len() != 1 {
6826 return false;
6827 }
6828 self.open_reader(self.run_refs[0].run_id)
6829 .map(|r| r.is_clean())
6830 .unwrap_or(false)
6831 }
6832
6833 fn resolve_footprint(
6840 &self,
6841 conditions: &[crate::query::Condition],
6842 snapshot: Snapshot,
6843 ) -> roaring::RoaringBitmap {
6844 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6845 return roaring::RoaringBitmap::new();
6846 }
6847 if self.run_refs.is_empty() {
6848 return roaring::RoaringBitmap::new();
6849 }
6850 if self.run_refs.len() == 1 {
6852 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
6853 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
6854 return rids.to_roaring_lossy();
6855 }
6856 }
6857 }
6858 roaring::RoaringBitmap::new()
6859 }
6860
6861 pub fn query_columns_native_cached(
6872 &mut self,
6873 conditions: &[crate::query::Condition],
6874 projection: Option<&[u16]>,
6875 snapshot: Snapshot,
6876 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
6877 if self.ttl.is_some() {
6880 return self.query_columns_native(conditions, projection, snapshot);
6881 }
6882 if conditions.is_empty() {
6883 return self.query_columns_native(conditions, projection, snapshot);
6884 }
6885 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
6889 if let Some(hit) = self.result_cache.lock().get_columns(key) {
6890 crate::trace::QueryTrace::record(|t| {
6891 t.result_cache_hit = true;
6892 t.scan_mode = crate::trace::ScanMode::NativePushdown;
6893 });
6894 return Ok(Some((*hit).clone()));
6895 }
6896 let res = self.query_columns_native(conditions, projection, snapshot)?;
6897 if let Some(cols) = &res {
6898 let footprint = self.resolve_footprint(conditions, snapshot);
6899 let condition_cols = crate::query::condition_columns(conditions);
6900 self.result_cache.lock().insert(
6901 key,
6902 CachedEntry {
6903 data: CachedData::Columns(Arc::new(cols.clone())),
6904 footprint,
6905 condition_cols,
6906 },
6907 );
6908 }
6909 Ok(res)
6910 }
6911
6912 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
6917 if self.ttl.is_some() {
6918 return self.query(q);
6919 }
6920 if q.conditions.is_empty() {
6921 return self.query(q);
6922 }
6923 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
6924 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
6925 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
6926 if let Some(hit) = self.result_cache.lock().get_rows(key) {
6927 crate::trace::QueryTrace::record(|t| {
6928 t.result_cache_hit = true;
6929 t.scan_mode = crate::trace::ScanMode::Materialized;
6930 });
6931 return Ok((*hit).clone());
6932 }
6933 let rows = self.query(q)?;
6934 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
6935 let condition_cols = crate::query::condition_columns(&q.conditions);
6936 self.result_cache.lock().insert(
6937 key,
6938 CachedEntry {
6939 data: CachedData::Rows(Arc::new(rows.clone())),
6940 footprint,
6941 condition_cols,
6942 },
6943 );
6944 Ok(rows)
6945 }
6946
6947 #[allow(clippy::type_complexity)]
6962 pub fn query_columns_native_traced(
6963 &mut self,
6964 conditions: &[crate::query::Condition],
6965 projection: Option<&[u16]>,
6966 snapshot: Snapshot,
6967 ) -> Result<(
6968 Option<Vec<(u16, columnar::NativeColumn)>>,
6969 crate::trace::QueryTrace,
6970 )> {
6971 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6972 self.query_columns_native(conditions, projection, snapshot)
6973 });
6974 Ok((result?, trace))
6975 }
6976
6977 #[allow(clippy::type_complexity)]
6980 pub fn query_columns_native_cached_traced(
6981 &mut self,
6982 conditions: &[crate::query::Condition],
6983 projection: Option<&[u16]>,
6984 snapshot: Snapshot,
6985 ) -> Result<(
6986 Option<Vec<(u16, columnar::NativeColumn)>>,
6987 crate::trace::QueryTrace,
6988 )> {
6989 let (result, trace) = crate::trace::QueryTrace::capture(|| {
6990 self.query_columns_native_cached(conditions, projection, snapshot)
6991 });
6992 Ok((result?, trace))
6993 }
6994
6995 pub fn native_page_cursor_traced(
6997 &self,
6998 snapshot: Snapshot,
6999 projection: Vec<(u16, TypeId)>,
7000 conditions: &[crate::query::Condition],
7001 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
7002 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7003 self.native_page_cursor(snapshot, projection, conditions)
7004 });
7005 Ok((result?, trace))
7006 }
7007
7008 pub fn native_multi_run_cursor_traced(
7010 &self,
7011 snapshot: Snapshot,
7012 projection: Vec<(u16, TypeId)>,
7013 conditions: &[crate::query::Condition],
7014 ) -> Result<(
7015 Option<crate::cursor::MultiRunCursor>,
7016 crate::trace::QueryTrace,
7017 )> {
7018 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7019 self.native_multi_run_cursor(snapshot, projection, conditions)
7020 });
7021 Ok((result?, trace))
7022 }
7023
7024 pub fn count_conditions_traced(
7026 &mut self,
7027 conditions: &[crate::query::Condition],
7028 snapshot: Snapshot,
7029 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
7030 let (result, trace) =
7031 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
7032 Ok((result?, trace))
7033 }
7034
7035 pub fn query_traced(
7037 &mut self,
7038 q: &crate::query::Query,
7039 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
7040 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
7041 Ok((result?, trace))
7042 }
7043
7044 pub fn query_columns_native(
7049 &mut self,
7050 conditions: &[crate::query::Condition],
7051 projection: Option<&[u16]>,
7052 snapshot: Snapshot,
7053 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7054 use crate::query::Condition;
7055 if self.ttl.is_some() {
7058 return Ok(None);
7059 }
7060 if conditions.is_empty() {
7061 return Ok(None);
7062 }
7063 self.ensure_indexes_complete()?;
7064
7065 let served = |c: &Condition| {
7070 matches!(
7071 c,
7072 Condition::Pk(_)
7073 | Condition::BitmapEq { .. }
7074 | Condition::BitmapIn { .. }
7075 | Condition::BytesPrefix { .. }
7076 | Condition::FmContains { .. }
7077 | Condition::FmContainsAll { .. }
7078 | Condition::Ann { .. }
7079 | Condition::Range { .. }
7080 | Condition::RangeF64 { .. }
7081 | Condition::SparseMatch { .. }
7082 | Condition::MinHashSimilar { .. }
7083 | Condition::IsNull { .. }
7084 | Condition::IsNotNull { .. }
7085 )
7086 };
7087 if !conditions.iter().all(served) {
7088 return Ok(None);
7089 }
7090 let fast_path =
7091 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
7092 crate::trace::QueryTrace::record(|t| {
7093 t.run_count = self.run_refs.len();
7094 t.memtable_rows = self.memtable.len();
7095 t.mutable_run_rows = self.mutable_run.len();
7096 t.conditions_pushed = conditions.len();
7097 t.learned_range_used = conditions.iter().any(|c| match c {
7098 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
7099 self.learned_range.contains_key(column_id)
7100 }
7101 _ => false,
7102 });
7103 });
7104 let col_ids: Vec<u16> = projection
7106 .map(|p| p.to_vec())
7107 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
7108 let proj_pairs: Vec<(u16, TypeId)> = col_ids
7109 .iter()
7110 .map(|&cid| {
7111 let ty = self
7112 .schema
7113 .columns
7114 .iter()
7115 .find(|c| c.id == cid)
7116 .map(|c| c.ty.clone())
7117 .unwrap_or(TypeId::Bytes);
7118 (cid, ty)
7119 })
7120 .collect();
7121
7122 if fast_path {
7128 let needs_column = conditions.iter().any(|c| match c {
7131 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
7132 Condition::RangeF64 { column_id, .. } => {
7133 !self.learned_range.contains_key(column_id)
7134 }
7135 _ => false,
7136 });
7137 let mut reader_opt: Option<RunReader> = if needs_column {
7138 Some(self.open_reader(self.run_refs[0].run_id)?)
7139 } else {
7140 None
7141 };
7142 let mut sets: Vec<RowIdSet> = Vec::new();
7143 for c in conditions {
7144 let s = match c {
7145 Condition::Range { column_id, lo, hi }
7146 if !self.learned_range.contains_key(column_id) =>
7147 {
7148 if reader_opt.is_none() {
7149 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7150 }
7151 reader_opt
7152 .as_mut()
7153 .expect("reader opened for range")
7154 .range_row_id_set_i64(*column_id, *lo, *hi)?
7155 }
7156 Condition::RangeF64 {
7157 column_id,
7158 lo,
7159 lo_inclusive,
7160 hi,
7161 hi_inclusive,
7162 } if !self.learned_range.contains_key(column_id) => {
7163 if reader_opt.is_none() {
7164 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7165 }
7166 reader_opt
7167 .as_mut()
7168 .expect("reader opened for range")
7169 .range_row_id_set_f64(
7170 *column_id,
7171 *lo,
7172 *lo_inclusive,
7173 *hi,
7174 *hi_inclusive,
7175 )?
7176 }
7177 _ => self.resolve_condition(c, snapshot)?,
7178 };
7179 sets.push(s);
7180 }
7181 let candidates = RowIdSet::intersect_many(sets);
7182 crate::trace::QueryTrace::record(|t| {
7183 t.survivor_count = Some(candidates.len());
7184 });
7185 if candidates.is_empty() {
7186 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
7187 .iter()
7188 .map(|&id| {
7189 (
7190 id,
7191 columnar::null_native(
7192 proj_pairs
7193 .iter()
7194 .find(|(c, _)| c == &id)
7195 .map(|(_, t)| t.clone())
7196 .unwrap_or(TypeId::Bytes),
7197 0,
7198 ),
7199 )
7200 })
7201 .collect();
7202 return Ok(Some(cols));
7203 }
7204 let mut reader = match reader_opt.take() {
7205 Some(r) => r,
7206 None => self.open_reader(self.run_refs[0].run_id)?,
7207 };
7208 let candidate_ids = candidates.into_sorted_vec();
7209 let (positions, fast_rid) = if let Some(positions) =
7210 reader.positions_for_row_ids_fast(&candidate_ids)
7211 {
7212 (positions, true)
7213 } else {
7214 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
7215 match col {
7216 columnar::NativeColumn::Int64 { data, .. } => {
7217 let mut p: Vec<usize> = candidate_ids
7218 .iter()
7219 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
7220 .collect();
7221 p.sort_unstable();
7222 (p, false)
7223 }
7224 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
7225 }
7226 };
7227 crate::trace::QueryTrace::record(|t| {
7228 t.scan_mode = crate::trace::ScanMode::NativePushdown;
7229 t.fast_row_id_map = fast_rid;
7230 });
7231 let mut cols = Vec::with_capacity(col_ids.len());
7232 for cid in &col_ids {
7233 let col = reader.column_native(*cid)?;
7234 cols.push((*cid, col.gather(&positions)));
7235 }
7236 return Ok(Some(cols));
7237 }
7238
7239 if !self.run_refs.is_empty() {
7252 use crate::cursor::{drain_cursor_to_columns, Cursor};
7253 let remaining: usize;
7254 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
7255 let c = self
7256 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
7257 .expect("single-run cursor should build when run_refs.len() == 1");
7258 remaining = c.remaining_rows();
7259 Box::new(c)
7260 } else {
7261 let c = self
7262 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
7263 .expect("multi-run cursor should build when run_refs.len() >= 1");
7264 remaining = c.remaining_rows();
7265 Box::new(c)
7266 };
7267 crate::trace::QueryTrace::record(|t| {
7268 if t.survivor_count.is_none() {
7269 t.survivor_count = Some(remaining);
7270 }
7271 });
7272 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
7273 return Ok(Some(cols));
7274 }
7275
7276 crate::trace::QueryTrace::record(|t| {
7281 t.scan_mode = crate::trace::ScanMode::Materialized;
7282 t.row_materialized = true;
7283 });
7284 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7285 for c in conditions {
7286 sets.push(self.resolve_condition(c, snapshot)?);
7287 }
7288 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7289 let rows = self.rows_for_rids(&rids, snapshot)?;
7290 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
7291 for (cid, ty) in &proj_pairs {
7292 let vals: Vec<Value> = rows
7293 .iter()
7294 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
7295 .collect();
7296 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
7297 }
7298 Ok(Some(cols))
7299 }
7300
7301 pub fn native_page_cursor(
7316 &self,
7317 snapshot: Snapshot,
7318 projection: Vec<(u16, TypeId)>,
7319 conditions: &[crate::query::Condition],
7320 ) -> Result<Option<NativePageCursor>> {
7321 use crate::cursor::build_page_plans;
7322 if self.ttl.is_some() {
7323 return Ok(None);
7324 }
7325 if !conditions.is_empty() && !self.indexes_complete {
7328 return Ok(None);
7329 }
7330 if self.run_refs.len() != 1 {
7331 return Ok(None);
7332 }
7333 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7334 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7335
7336 let overlay_rids: HashSet<u64> = {
7339 let mut s = HashSet::new();
7340 for row in self.memtable.visible_versions(snapshot.epoch) {
7341 s.insert(row.row_id.0);
7342 }
7343 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7344 s.insert(row.row_id.0);
7345 }
7346 s
7347 };
7348
7349 let survivors = if conditions.is_empty() {
7353 None
7354 } else {
7355 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
7356 };
7357
7358 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
7365 survivors.clone()
7366 } else if let Some(s) = &survivors {
7367 let mut run_set = s.clone();
7368 run_set.remove_many(overlay_rids.iter().copied());
7369 Some(run_set)
7370 } else {
7371 Some(RowIdSet::from_unsorted(
7372 rids.iter()
7373 .map(|&r| r as u64)
7374 .filter(|r| !overlay_rids.contains(r))
7375 .collect(),
7376 ))
7377 };
7378
7379 let overlay_rows = if overlay_rids.is_empty() {
7380 Vec::new()
7381 } else {
7382 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7383 self.overlay_visible_rows(snapshot, bound)
7384 };
7385
7386 let plans = if positions.is_empty() {
7388 Vec::new()
7389 } else {
7390 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
7391 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
7392 };
7393
7394 let overlay = if overlay_rows.is_empty() {
7396 None
7397 } else {
7398 let filtered =
7399 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7400 if filtered.is_empty() {
7401 None
7402 } else {
7403 Some(self.materialize_overlay(&filtered, &projection))
7404 }
7405 };
7406
7407 let overlay_row_count = overlay
7408 .as_ref()
7409 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7410 .unwrap_or(0);
7411 crate::trace::QueryTrace::record(|t| {
7412 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
7413 t.run_count = self.run_refs.len();
7414 t.memtable_rows = self.memtable.len();
7415 t.mutable_run_rows = self.mutable_run.len();
7416 t.overlay_rows = overlay_row_count;
7417 t.conditions_pushed = conditions.len();
7418 t.pages_decoded = plans
7419 .iter()
7420 .map(|p| p.positions.len())
7421 .sum::<usize>()
7422 .min(1);
7423 });
7424
7425 Ok(Some(NativePageCursor::new_with_overlay(
7426 reader, projection, plans, overlay,
7427 )))
7428 }
7429 #[allow(clippy::type_complexity)]
7439 pub fn native_multi_run_cursor(
7440 &self,
7441 snapshot: Snapshot,
7442 projection: Vec<(u16, TypeId)>,
7443 conditions: &[crate::query::Condition],
7444 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
7445 use crate::cursor::{MultiRunCursor, RunStream};
7446 use crate::sorted_run::SYS_ROW_ID;
7447 use std::collections::{BinaryHeap, HashMap, HashSet};
7448 if self.ttl.is_some() {
7449 return Ok(None);
7450 }
7451 if !conditions.is_empty() && !self.indexes_complete {
7454 return Ok(None);
7455 }
7456 if self.run_refs.is_empty() {
7457 return Ok(None);
7458 }
7459
7460 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
7462 Vec::with_capacity(self.run_refs.len());
7463 for rr in &self.run_refs {
7464 let mut reader = self.open_reader(rr.run_id)?;
7465 let (rids, eps, del) = reader.system_columns_native()?;
7466 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
7467 run_meta.push((reader, rids, eps, del, page_rows));
7468 }
7469
7470 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
7474 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
7475 for i in 0..rids.len() {
7476 let rid = rids[i] as u64;
7477 let e = eps[i] as u64;
7478 if e > snapshot.epoch.0 {
7479 continue;
7480 }
7481 let is_del = del[i] != 0;
7482 best.entry(rid)
7483 .and_modify(|cur| {
7484 if e > cur.0 {
7485 *cur = (e, run_idx, i, is_del);
7486 }
7487 })
7488 .or_insert((e, run_idx, i, is_del));
7489 }
7490 }
7491
7492 let overlay_rids: HashSet<u64> = {
7494 let mut s = HashSet::new();
7495 for row in self.memtable.visible_versions(snapshot.epoch) {
7496 s.insert(row.row_id.0);
7497 }
7498 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7499 s.insert(row.row_id.0);
7500 }
7501 s
7502 };
7503
7504 let survivors: Option<RowIdSet> = if conditions.is_empty() {
7506 None
7507 } else {
7508 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7509 for c in conditions {
7510 sets.push(self.resolve_condition(c, snapshot)?);
7511 }
7512 Some(RowIdSet::intersect_many(sets))
7513 };
7514
7515 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
7519 for (rid, (_, run_idx, pos, deleted)) in &best {
7520 if *deleted {
7521 continue;
7522 }
7523 if overlay_rids.contains(rid) {
7524 continue;
7525 }
7526 if let Some(s) = &survivors {
7527 if !s.contains(*rid) {
7528 continue;
7529 }
7530 }
7531 per_run[*run_idx].push((*rid, *pos));
7532 }
7533 for v in per_run.iter_mut() {
7534 v.sort_unstable_by_key(|&(rid, _)| rid);
7535 }
7536
7537 let mut streams = Vec::with_capacity(run_meta.len());
7539 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
7540 let mut total = 0usize;
7541 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
7542 let mut starts = Vec::with_capacity(page_rows.len());
7543 let mut acc = 0usize;
7544 for &r in &page_rows {
7545 starts.push(acc);
7546 acc += r;
7547 }
7548 let mut survivors_vec: Vec<(u64, usize, usize)> =
7549 Vec::with_capacity(per_run[run_idx].len());
7550 for &(rid, pos) in &per_run[run_idx] {
7551 let page_seq = match starts.partition_point(|&s| s <= pos) {
7552 0 => continue,
7553 p => p - 1,
7554 };
7555 let within = pos - starts[page_seq];
7556 survivors_vec.push((rid, page_seq, within));
7557 }
7558 total += survivors_vec.len();
7559 if let Some(&(rid, _, _)) = survivors_vec.first() {
7560 heap.push(std::cmp::Reverse((rid, run_idx)));
7561 }
7562 streams.push(RunStream::new(reader, survivors_vec, page_rows));
7563 }
7564
7565 let overlay_rows = if overlay_rids.is_empty() {
7567 Vec::new()
7568 } else {
7569 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7570 self.overlay_visible_rows(snapshot, bound)
7571 };
7572 let overlay = if overlay_rows.is_empty() {
7573 None
7574 } else {
7575 let filtered =
7576 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7577 if filtered.is_empty() {
7578 None
7579 } else {
7580 Some(self.materialize_overlay(&filtered, &projection))
7581 }
7582 };
7583
7584 let overlay_row_count = overlay
7585 .as_ref()
7586 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7587 .unwrap_or(0);
7588 crate::trace::QueryTrace::record(|t| {
7589 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
7590 t.run_count = self.run_refs.len();
7591 t.memtable_rows = self.memtable.len();
7592 t.mutable_run_rows = self.mutable_run.len();
7593 t.overlay_rows = overlay_row_count;
7594 t.conditions_pushed = conditions.len();
7595 t.survivor_count = Some(total);
7596 });
7597
7598 Ok(Some(MultiRunCursor::new(
7599 streams, projection, heap, total, overlay,
7600 )))
7601 }
7602
7603 fn overlay_materialization_bound<'a>(
7615 conditions: &[crate::query::Condition],
7616 survivors: &'a Option<RowIdSet>,
7617 ) -> Option<&'a RowIdSet> {
7618 use crate::query::Condition;
7619 let has_range = conditions
7620 .iter()
7621 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7622 if has_range {
7623 None
7624 } else {
7625 survivors.as_ref()
7626 }
7627 }
7628
7629 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
7641 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
7642 let mut fold = |row: Row| {
7643 if let Some(b) = bound {
7644 if !b.contains(row.row_id.0) {
7645 return;
7646 }
7647 }
7648 best.entry(row.row_id.0)
7649 .and_modify(|(be, br)| {
7650 if row.committed_epoch > *be {
7651 *be = row.committed_epoch;
7652 *br = row.clone();
7653 }
7654 })
7655 .or_insert_with(|| (row.committed_epoch, row));
7656 };
7657 for row in self.memtable.visible_versions(snapshot.epoch) {
7658 fold(row);
7659 }
7660 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7661 fold(row);
7662 }
7663 let mut out: Vec<Row> = best
7664 .into_values()
7665 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
7666 .collect();
7667 out.sort_by_key(|r| r.row_id);
7668 out
7669 }
7670
7671 fn filter_overlay_rows(
7679 &self,
7680 rows: Vec<Row>,
7681 conditions: &[crate::query::Condition],
7682 survivors: Option<&RowIdSet>,
7683 snapshot: Snapshot,
7684 ) -> Result<Vec<Row>> {
7685 if conditions.is_empty() {
7686 return Ok(rows);
7687 }
7688 use crate::query::Condition;
7689 let all_index_served = !conditions
7693 .iter()
7694 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7695 if all_index_served {
7696 return Ok(rows
7697 .into_iter()
7698 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
7699 .collect());
7700 }
7701 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7704 for c in conditions {
7705 let s = match c {
7706 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
7707 _ => self.resolve_condition(c, snapshot)?,
7708 };
7709 per_cond_sets.push(s);
7710 }
7711 Ok(rows
7712 .into_iter()
7713 .filter(|row| {
7714 conditions.iter().enumerate().all(|(i, c)| match c {
7715 Condition::Range { column_id, lo, hi } => {
7716 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
7717 }
7718 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
7719 match row.columns.get(column_id) {
7720 Some(Value::Float64(v)) => {
7721 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
7722 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
7723 lo_ok && hi_ok
7724 }
7725 _ => false,
7726 }
7727 }
7728 _ => per_cond_sets[i].contains(row.row_id.0),
7729 })
7730 })
7731 .collect())
7732 }
7733
7734 fn materialize_overlay(
7737 &self,
7738 rows: &[Row],
7739 projection: &[(u16, TypeId)],
7740 ) -> Vec<columnar::NativeColumn> {
7741 if projection.is_empty() {
7742 return vec![columnar::null_native(TypeId::Int64, rows.len())];
7743 }
7744 let mut cols = Vec::with_capacity(projection.len());
7745 for (cid, ty) in projection {
7746 let vals: Vec<Value> = rows
7747 .iter()
7748 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
7749 .collect();
7750 cols.push(columnar::values_to_native(ty.clone(), &vals));
7751 }
7752 cols
7753 }
7754
7755 fn resolve_survivor_rids(
7760 &self,
7761 conditions: &[crate::query::Condition],
7762 reader: &mut RunReader,
7763 snapshot: Snapshot,
7764 ) -> Result<RowIdSet> {
7765 use crate::query::Condition;
7766 let mut sets: Vec<RowIdSet> = Vec::new();
7767 for c in conditions {
7768 self.validate_condition(c)?;
7769 let s: RowIdSet = match c {
7770 Condition::Pk(key) => {
7771 let lookup = self
7772 .schema
7773 .primary_key()
7774 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7775 .unwrap_or_else(|| key.clone());
7776 self.hot
7777 .get(&lookup)
7778 .map(|r| RowIdSet::one(r.0))
7779 .unwrap_or_else(RowIdSet::empty)
7780 }
7781 Condition::BitmapEq { column_id, value } => {
7782 let lookup = self.index_lookup_key_bytes(*column_id, value);
7783 self.bitmap
7784 .get(column_id)
7785 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7786 .unwrap_or_else(RowIdSet::empty)
7787 }
7788 Condition::BitmapIn { column_id, values } => {
7789 let bm = self.bitmap.get(column_id);
7790 let mut acc = roaring::RoaringBitmap::new();
7791 if let Some(b) = bm {
7792 for v in values {
7793 let lookup = self.index_lookup_key_bytes(*column_id, v);
7794 acc |= b.get(&lookup);
7795 }
7796 }
7797 RowIdSet::from_roaring(acc)
7798 }
7799 Condition::BytesPrefix { column_id, prefix } => {
7800 if let Some(b) = self.bitmap.get(column_id) {
7801 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7802 let mut acc = roaring::RoaringBitmap::new();
7803 for key in b.keys() {
7804 if key.starts_with(&lookup_prefix) {
7805 acc |= b.get(key);
7806 }
7807 }
7808 RowIdSet::from_roaring(acc)
7809 } else {
7810 RowIdSet::empty()
7811 }
7812 }
7813 Condition::FmContains { column_id, pattern } => self
7814 .fm
7815 .get(column_id)
7816 .map(|f| {
7817 RowIdSet::from_unsorted(
7818 f.locate(pattern).into_iter().map(|r| r.0).collect(),
7819 )
7820 })
7821 .unwrap_or_else(RowIdSet::empty),
7822 Condition::FmContainsAll {
7823 column_id,
7824 patterns,
7825 } => {
7826 if let Some(f) = self.fm.get(column_id) {
7827 let sets: Vec<RowIdSet> = patterns
7828 .iter()
7829 .map(|pat| {
7830 RowIdSet::from_unsorted(
7831 f.locate(pat).into_iter().map(|r| r.0).collect(),
7832 )
7833 })
7834 .collect();
7835 RowIdSet::intersect_many(sets)
7836 } else {
7837 RowIdSet::empty()
7838 }
7839 }
7840 Condition::Ann {
7841 column_id,
7842 query,
7843 k,
7844 } => RowIdSet::from_unsorted(
7845 self.retrieve_filtered(
7846 &crate::query::Retriever::Ann {
7847 column_id: *column_id,
7848 query: query.clone(),
7849 k: *k,
7850 },
7851 snapshot,
7852 None,
7853 None,
7854 None,
7855 None,
7856 )?
7857 .into_iter()
7858 .map(|hit| hit.row_id.0)
7859 .collect(),
7860 ),
7861 Condition::SparseMatch {
7862 column_id,
7863 query,
7864 k,
7865 } => RowIdSet::from_unsorted(
7866 self.retrieve_filtered(
7867 &crate::query::Retriever::Sparse {
7868 column_id: *column_id,
7869 query: query.clone(),
7870 k: *k,
7871 },
7872 snapshot,
7873 None,
7874 None,
7875 None,
7876 None,
7877 )?
7878 .into_iter()
7879 .map(|hit| hit.row_id.0)
7880 .collect(),
7881 ),
7882 Condition::MinHashSimilar {
7883 column_id,
7884 query,
7885 k,
7886 } => match self.minhash.get(column_id) {
7887 Some(index) => {
7888 let candidates = index.candidate_row_ids(query);
7889 let eligible =
7890 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7891 RowIdSet::from_unsorted(
7892 index
7893 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
7894 .into_iter()
7895 .map(|(row_id, _)| row_id.0)
7896 .collect(),
7897 )
7898 }
7899 None => RowIdSet::empty(),
7900 },
7901 Condition::Range { column_id, lo, hi } => {
7902 if let Some(li) = self.learned_range.get(column_id) {
7903 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7904 } else {
7905 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
7906 }
7907 }
7908 Condition::RangeF64 {
7909 column_id,
7910 lo,
7911 lo_inclusive,
7912 hi,
7913 hi_inclusive,
7914 } => {
7915 if let Some(li) = self.learned_range.get(column_id) {
7916 RowIdSet::from_unsorted(
7917 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7918 .into_iter()
7919 .collect(),
7920 )
7921 } else {
7922 reader.range_row_id_set_f64(
7923 *column_id,
7924 *lo,
7925 *lo_inclusive,
7926 *hi,
7927 *hi_inclusive,
7928 )?
7929 }
7930 }
7931 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
7932 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
7933 };
7934 sets.push(s);
7935 }
7936 Ok(RowIdSet::intersect_many(sets))
7937 }
7938
7939 pub fn scan_cursor(
7960 &self,
7961 snapshot: Snapshot,
7962 projection: Vec<(u16, TypeId)>,
7963 conditions: &[crate::query::Condition],
7964 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
7965 if self.ttl.is_some() {
7966 return Ok(None);
7967 }
7968 if !conditions.is_empty() && !self.indexes_complete {
7974 return Ok(None);
7975 }
7976 if self.run_refs.len() == 1 {
7977 Ok(self
7978 .native_page_cursor(snapshot, projection, conditions)?
7979 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7980 } else {
7981 Ok(self
7982 .native_multi_run_cursor(snapshot, projection, conditions)?
7983 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
7984 }
7985 }
7986
7987 pub fn aggregate_native(
8001 &self,
8002 snapshot: Snapshot,
8003 column: Option<u16>,
8004 conditions: &[crate::query::Condition],
8005 agg: NativeAgg,
8006 ) -> Result<Option<NativeAggResult>> {
8007 if self.ttl.is_some() {
8008 return Ok(None);
8009 }
8010 if self.run_refs.len() == 1 && conditions.is_empty() {
8012 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
8013 return Ok(Some(res));
8014 }
8015 }
8016 if matches!(agg, NativeAgg::Count) && column.is_none() {
8018 return Ok(self
8019 .scan_cursor(snapshot, Vec::new(), conditions)?
8020 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
8021 }
8022 let cid = match column {
8025 Some(c) => c,
8026 None => return Ok(None),
8027 };
8028 let ty = self.column_type(cid);
8029 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
8030 else {
8031 return Ok(None);
8032 };
8033 match ty {
8034 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8035 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
8036 Ok(Some(pack_int(agg, count, sum, mn, mx)))
8037 }
8038 TypeId::Float64 => {
8039 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
8040 Ok(Some(pack_float(agg, count, sum, mn, mx)))
8041 }
8042 _ => Ok(None),
8043 }
8044 }
8045
8046 fn aggregate_from_stats(
8054 &self,
8055 snapshot: Snapshot,
8056 column: Option<u16>,
8057 agg: NativeAgg,
8058 ) -> Result<Option<NativeAggResult>> {
8059 let cid = match (agg, column) {
8060 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
8061 _ => return Ok(None), };
8063 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
8064 return Ok(None);
8065 };
8066 let Some(cs) = stats.get(&cid) else {
8067 return Ok(None);
8068 };
8069 match agg {
8070 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
8072 self.live_count.saturating_sub(cs.null_count),
8073 ))),
8074 NativeAgg::Min | NativeAgg::Max => {
8075 let bound = if agg == NativeAgg::Min {
8076 &cs.min
8077 } else {
8078 &cs.max
8079 };
8080 match bound {
8081 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
8082 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
8083 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
8088 None => Ok(None),
8089 }
8090 }
8091 _ => Ok(None),
8092 }
8093 }
8094
8095 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
8104 if self.ttl.is_some() {
8105 return Ok(None);
8106 }
8107 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
8108 return Ok(None);
8109 }
8110 self.ensure_indexes_complete()?;
8113 let reader = self.open_reader(self.run_refs[0].run_id)?;
8114 if self.live_count != reader.row_count() as u64 {
8115 return Ok(None);
8116 }
8117 let Some(bm) = self.bitmap.get(&column_id) else {
8118 return Ok(None); };
8120 let mut distinct = bm.value_count() as u64;
8121 if !bm.get(&Value::Null.encode_key()).is_empty() {
8124 distinct = distinct.saturating_sub(1);
8125 }
8126 Ok(Some(distinct))
8127 }
8128
8129 pub fn aggregate_incremental(
8141 &mut self,
8142 cache_key: u64,
8143 conditions: &[crate::query::Condition],
8144 column: Option<u16>,
8145 agg: NativeAgg,
8146 ) -> Result<IncrementalAggResult> {
8147 let snap = self.snapshot();
8148 let cur_wm = self.allocator.current().0;
8149 let cur_epoch = snap.epoch.0;
8150 let incremental_ok = self.ttl.is_none()
8157 && !self.had_deletes
8158 && self.memtable.is_empty()
8159 && self.mutable_run.is_empty();
8160
8161 if incremental_ok {
8164 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
8165 if cached.epoch == cur_epoch {
8166 return Ok(IncrementalAggResult {
8167 state: cached.state,
8168 incremental: true,
8169 delta_rows: 0,
8170 });
8171 }
8172 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
8173 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
8174 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
8175 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8176 let delta_state = agg_state_from_rows(
8177 &delta_rows,
8178 conditions,
8179 &index_sets,
8180 column,
8181 agg,
8182 &self.schema,
8183 )?;
8184 let merged = cached.state.merge(delta_state);
8185 let delta_n = delta_rids.len() as u64;
8186 self.agg_cache.insert(
8187 cache_key,
8188 CachedAgg {
8189 state: merged.clone(),
8190 watermark: cur_wm,
8191 epoch: cur_epoch,
8192 },
8193 );
8194 return Ok(IncrementalAggResult {
8195 state: merged,
8196 incremental: true,
8197 delta_rows: delta_n,
8198 });
8199 }
8200 }
8201 }
8202
8203 let cursor_ok =
8208 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8209 let state = if cursor_ok && agg != NativeAgg::Avg {
8210 match self.aggregate_native(snap, column, conditions, agg)? {
8211 Some(result) => {
8212 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
8213 }
8214 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
8215 }
8216 } else {
8217 self.agg_state_full_scan(conditions, column, agg, snap)?
8218 };
8219 if incremental_ok {
8221 self.agg_cache.insert(
8222 cache_key,
8223 CachedAgg {
8224 state: state.clone(),
8225 watermark: cur_wm,
8226 epoch: cur_epoch,
8227 },
8228 );
8229 }
8230 Ok(IncrementalAggResult {
8231 state,
8232 incremental: false,
8233 delta_rows: 0,
8234 })
8235 }
8236
8237 fn agg_state_full_scan(
8240 &self,
8241 conditions: &[crate::query::Condition],
8242 column: Option<u16>,
8243 agg: NativeAgg,
8244 snap: Snapshot,
8245 ) -> Result<AggState> {
8246 let rows = self.visible_rows(snap)?;
8247 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8248 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
8249 }
8250
8251 fn resolve_index_conditions(
8254 &self,
8255 conditions: &[crate::query::Condition],
8256 snapshot: Snapshot,
8257 ) -> Result<Vec<RowIdSet>> {
8258 use crate::query::Condition;
8259 let mut sets = Vec::new();
8260 for c in conditions {
8261 if matches!(
8262 c,
8263 Condition::Ann { .. }
8264 | Condition::SparseMatch { .. }
8265 | Condition::MinHashSimilar { .. }
8266 ) {
8267 sets.push(self.resolve_condition(c, snapshot)?);
8268 }
8269 }
8270 Ok(sets)
8271 }
8272
8273 fn column_type(&self, cid: u16) -> TypeId {
8274 self.schema
8275 .columns
8276 .iter()
8277 .find(|c| c.id == cid)
8278 .map(|c| c.ty.clone())
8279 .unwrap_or(TypeId::Bytes)
8280 }
8281
8282 pub fn approx_aggregate(
8291 &mut self,
8292 conditions: &[crate::query::Condition],
8293 column: Option<u16>,
8294 agg: ApproxAgg,
8295 z: f64,
8296 ) -> Result<Option<ApproxResult>> {
8297 use crate::query::Condition;
8298 self.ensure_reservoir_complete()?;
8299 let snapshot = self.snapshot();
8300 let n_pop = self.count();
8301 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
8302 if sample_rids.is_empty() {
8303 return Ok(None);
8304 }
8305 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
8307 let s = live_sample.len();
8308 if s == 0 {
8309 return Ok(None);
8310 }
8311
8312 let mut index_sets: Vec<RowIdSet> = Vec::new();
8315 for c in conditions {
8316 if matches!(
8317 c,
8318 Condition::Ann { .. }
8319 | Condition::SparseMatch { .. }
8320 | Condition::MinHashSimilar { .. }
8321 ) {
8322 index_sets.push(self.resolve_condition(c, snapshot)?);
8323 }
8324 }
8325
8326 let cid = match (agg, column) {
8328 (ApproxAgg::Count, _) => None,
8329 (_, Some(c)) => Some(c),
8330 _ => return Ok(None),
8331 };
8332 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
8333 for r in &live_sample {
8334 if !conditions
8336 .iter()
8337 .all(|c| condition_matches_row(c, r, &self.schema))
8338 {
8339 continue;
8340 }
8341 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
8343 continue;
8344 }
8345 if let Some(cid) = cid {
8346 if let Some(v) = as_f64(r.columns.get(&cid)) {
8347 passing_vals.push(v);
8348 } } else {
8350 passing_vals.push(0.0); }
8352 }
8353 let m = passing_vals.len();
8354
8355 let (point, half) = match agg {
8356 ApproxAgg::Count => {
8357 let p = m as f64 / s as f64;
8359 let point = n_pop as f64 * p;
8360 let var = if s > 1 {
8361 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
8362 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8363 } else {
8364 0.0
8365 };
8366 (point, z * var.sqrt())
8367 }
8368 ApproxAgg::Sum => {
8369 let y: Vec<f64> = live_sample
8371 .iter()
8372 .map(|r| {
8373 let passes_row = conditions
8374 .iter()
8375 .all(|c| condition_matches_row(c, r, &self.schema))
8376 && index_sets.iter().all(|set| set.contains(r.row_id.0));
8377 if passes_row {
8378 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
8379 } else {
8380 0.0
8381 }
8382 })
8383 .collect();
8384 let mean_y = y.iter().sum::<f64>() / s as f64;
8385 let point = n_pop as f64 * mean_y;
8386 let var = if s > 1 {
8387 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
8388 let var_y = ss / (s - 1) as f64;
8389 n_pop as f64 * n_pop as f64 * var_y / s as f64
8390 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8391 } else {
8392 0.0
8393 };
8394 (point, z * var.sqrt())
8395 }
8396 ApproxAgg::Avg => {
8397 if m == 0 {
8398 return Ok(Some(ApproxResult {
8399 point: 0.0,
8400 ci_low: 0.0,
8401 ci_high: 0.0,
8402 n_population: n_pop,
8403 n_sample_live: s,
8404 n_passing: 0,
8405 }));
8406 }
8407 let mean = passing_vals.iter().sum::<f64>() / m as f64;
8408 let half = if m > 1 {
8409 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
8410 let sd = (ss / (m - 1) as f64).sqrt();
8411 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
8412 z * sd / (m as f64).sqrt() * fpc.sqrt()
8413 } else {
8414 0.0
8415 };
8416 (mean, half)
8417 }
8418 };
8419
8420 Ok(Some(ApproxResult {
8421 point,
8422 ci_low: point - half,
8423 ci_high: point + half,
8424 n_population: n_pop,
8425 n_sample_live: s,
8426 n_passing: m,
8427 }))
8428 }
8429
8430 pub fn exact_column_stats(
8438 &self,
8439 _snapshot: Snapshot,
8440 projection: &[u16],
8441 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
8442 if self.ttl.is_some()
8443 || !(self.memtable.is_empty()
8444 && self.mutable_run.is_empty()
8445 && self.run_refs.len() == 1)
8446 {
8447 return Ok(None);
8448 }
8449 let reader = self.open_reader(self.run_refs[0].run_id)?;
8450 if self.live_count != reader.row_count() as u64 {
8451 return Ok(None);
8452 }
8453 let mut out = HashMap::new();
8454 for &cid in projection {
8455 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
8456 Some(c) => c,
8457 None => continue,
8458 };
8459 let Some(stats) = reader.column_page_stats(cid) else {
8461 out.insert(
8462 cid,
8463 ColumnStat {
8464 min: None,
8465 max: None,
8466 null_count: self.live_count,
8467 },
8468 );
8469 continue;
8470 };
8471 let stat = match cdef.ty {
8472 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8473 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
8474 min: mn.map(Value::Int64),
8475 max: mx.map(Value::Int64),
8476 null_count: n,
8477 })
8478 }
8479 TypeId::Float64 => {
8480 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
8481 min: mn.map(Value::Float64),
8482 max: mx.map(Value::Float64),
8483 null_count: n,
8484 })
8485 }
8486 _ => None,
8487 };
8488 if let Some(s) = stat {
8489 out.insert(cid, s);
8490 }
8491 }
8492 Ok(Some(out))
8493 }
8494
8495 pub fn dir(&self) -> &Path {
8496 &self.dir
8497 }
8498
8499 pub fn schema(&self) -> &Schema {
8500 &self.schema
8501 }
8502
8503 pub(crate) fn set_catalog_name(&mut self, name: String) {
8504 self.name = name;
8505 }
8506
8507 pub(crate) fn prepare_alter_column(
8508 &mut self,
8509 column_name: &str,
8510 change: &AlterColumn,
8511 ) -> Result<ColumnDef> {
8512 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
8513 return Err(MongrelError::InvalidArgument(
8514 "ALTER COLUMN requires committing staged writes first".into(),
8515 ));
8516 }
8517 let old = self
8518 .schema
8519 .columns
8520 .iter()
8521 .find(|c| c.name == column_name)
8522 .cloned()
8523 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
8524 let mut next = old.clone();
8525
8526 if let Some(name) = &change.name {
8527 let trimmed = name.trim();
8528 if trimmed.is_empty() {
8529 return Err(MongrelError::InvalidArgument(
8530 "ALTER COLUMN name must not be empty".into(),
8531 ));
8532 }
8533 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
8534 return Err(MongrelError::Schema(format!(
8535 "column {trimmed} already exists"
8536 )));
8537 }
8538 next.name = trimmed.to_string();
8539 }
8540
8541 if let Some(ty) = &change.ty {
8542 next.ty = ty.clone();
8543 }
8544 if let Some(flags) = change.flags {
8545 validate_alter_column_flags(old.flags, flags)?;
8546 next.flags = flags;
8547 }
8548
8549 if let Some(default_change) = &change.default_value {
8550 next.default_value = default_change.clone();
8551 }
8552
8553 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
8554 if old.flags.contains(ColumnFlags::NULLABLE)
8555 && !next.flags.contains(ColumnFlags::NULLABLE)
8556 && self.column_has_nulls(old.id)?
8557 {
8558 return Err(MongrelError::InvalidArgument(format!(
8559 "column '{}' contains NULL values",
8560 old.name
8561 )));
8562 }
8563 Ok(next)
8564 }
8565
8566 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
8567 let idx = self
8568 .schema
8569 .columns
8570 .iter()
8571 .position(|c| c.id == column.id)
8572 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
8573 if self.schema.columns[idx] == column {
8574 return Ok(());
8575 }
8576 self.schema.columns[idx] = column;
8577 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8578 self.schema.validate_auto_increment()?;
8579 self.schema.validate_defaults()?;
8580 self.auto_inc = resolve_auto_inc(&self.schema);
8581 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
8582 write_schema(&self.dir, &self.schema)?;
8583 self.clear_result_cache();
8584 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8585 self.persist_manifest(self.current_epoch())?;
8586 Ok(())
8587 }
8588
8589 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
8590 self.ensure_writable()?;
8591 let column = self.prepare_alter_column(column_name, &change)?;
8592 self.apply_altered_column(column.clone())?;
8593 Ok(column)
8594 }
8595
8596 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
8597 if self.live_count == 0 {
8598 return Ok(false);
8599 }
8600 let snap = self.snapshot();
8601 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
8602 Ok(columns
8603 .first()
8604 .map(|(_, col)| col.null_count(col.len()) != 0)
8605 .unwrap_or(true))
8606 }
8607
8608 fn has_stored_versions(&self) -> bool {
8609 !self.memtable.is_empty()
8610 || !self.mutable_run.is_empty()
8611 || self.run_refs.iter().any(|r| r.row_count > 0)
8612 || !self.retiring.is_empty()
8613 }
8614
8615 pub fn add_column(
8620 &mut self,
8621 name: &str,
8622 ty: TypeId,
8623 flags: ColumnFlags,
8624 default_value: Option<crate::schema::DefaultExpr>,
8625 ) -> Result<u16> {
8626 self.add_column_with_id(name, ty, flags, default_value, None)
8627 }
8628
8629 pub fn add_column_with_id(
8630 &mut self,
8631 name: &str,
8632 ty: TypeId,
8633 flags: ColumnFlags,
8634 default_value: Option<crate::schema::DefaultExpr>,
8635 requested_id: Option<u16>,
8636 ) -> Result<u16> {
8637 self.ensure_writable()?;
8638 if self.schema.columns.iter().any(|c| c.name == name) {
8639 return Err(MongrelError::Schema(format!(
8640 "column {name} already exists"
8641 )));
8642 }
8643 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
8644 if self.schema.columns.iter().any(|c| c.id == id) {
8645 return Err(MongrelError::Schema(format!(
8646 "column id {id} already exists"
8647 )));
8648 }
8649 id
8650 } else {
8651 self.schema
8652 .columns
8653 .iter()
8654 .map(|c| c.id)
8655 .max()
8656 .unwrap_or(0)
8657 .checked_add(1)
8658 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
8659 };
8660 self.schema.columns.push(ColumnDef {
8661 id,
8662 name: name.to_string(),
8663 ty,
8664 flags,
8665 default_value,
8666 });
8667 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8668 self.schema.validate_auto_increment()?;
8669 self.schema.validate_defaults()?;
8670 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
8671 self.auto_inc = resolve_auto_inc(&self.schema);
8672 }
8673 write_schema(&self.dir, &self.schema)?;
8674 self.clear_result_cache();
8675 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8677 self.persist_manifest(self.current_epoch())?;
8678 Ok(id)
8679 }
8680
8681 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
8690 self.ensure_writable()?;
8691 let cid = self
8692 .schema
8693 .columns
8694 .iter()
8695 .find(|c| c.name == column_name)
8696 .map(|c| c.id)
8697 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
8698 let ty = self
8699 .schema
8700 .columns
8701 .iter()
8702 .find(|c| c.id == cid)
8703 .map(|c| c.ty.clone())
8704 .unwrap_or(TypeId::Int64);
8705 if !matches!(
8706 ty,
8707 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
8708 ) {
8709 return Err(MongrelError::Schema(format!(
8710 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
8711 )));
8712 }
8713 if self
8714 .schema
8715 .indexes
8716 .iter()
8717 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
8718 {
8719 return Ok(()); }
8721 self.schema.indexes.push(IndexDef {
8722 name: format!("{}_learned_range", column_name),
8723 column_id: cid,
8724 kind: IndexKind::LearnedRange,
8725 predicate: None,
8726 options: Default::default(),
8727 });
8728 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8729 write_schema(&self.dir, &self.schema)?;
8730 self.build_learned_ranges()?;
8731 Ok(())
8732 }
8733
8734 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
8737 self.sync_byte_threshold = threshold;
8738 if let WalSink::Private(w) = &mut self.wal {
8739 w.set_sync_byte_threshold(threshold);
8740 }
8741 }
8742
8743 pub fn page_cache_flush(&self) {
8747 self.page_cache.flush_to_disk();
8748 }
8749
8750 pub fn page_cache_len(&self) -> usize {
8752 self.page_cache.len()
8753 }
8754
8755 pub fn decoded_cache_len(&self) -> usize {
8758 self.decoded_cache.len()
8759 }
8760
8761 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
8764 self.memtable.drain_sorted()
8765 }
8766
8767 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
8768 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
8769 }
8770
8771 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
8772 self.run_refs.iter().map(|run| run.run_id)
8773 }
8774
8775 pub(crate) fn table_dir(&self) -> &Path {
8776 &self.dir
8777 }
8778
8779 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
8780 &self.schema
8781 }
8782
8783 pub(crate) fn alloc_run_id(&mut self) -> u64 {
8784 let id = self.next_run_id;
8785 self.next_run_id += 1;
8786 id
8787 }
8788
8789 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
8790 self.run_refs.push(run_ref);
8791 }
8792
8793 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
8803 self.retiring.push(crate::manifest::RetiredRun {
8804 run_id,
8805 retire_epoch,
8806 });
8807 }
8808
8809 pub(crate) fn reap_retiring(
8813 &mut self,
8814 min_active: Epoch,
8815 backup_pinned: &std::collections::HashSet<u128>,
8816 ) -> Result<usize> {
8817 if self.retiring.is_empty() {
8818 return Ok(0);
8819 }
8820 let mut reaped = 0;
8821 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
8822 for r in std::mem::take(&mut self.retiring) {
8828 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
8829 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
8830 reaped += 1;
8831 } else {
8832 kept.push(r);
8833 }
8834 }
8835 self.retiring = kept;
8836 if reaped > 0 {
8837 self.persist_manifest(self.current_epoch())?;
8838 }
8839 Ok(reaped)
8840 }
8841
8842 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
8843 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
8844 return false;
8845 }
8846 self.live_count = self.live_count.saturating_add(run_ref.row_count);
8847 self.run_refs.push(run_ref);
8848 self.indexes_complete = false;
8849 true
8850 }
8851
8852 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
8853 self.kek.as_ref()
8854 }
8855
8856 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
8857 let mut reader = RunReader::open_with_cache(
8858 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
8859 self.schema.clone(),
8860 self.kek.clone(),
8861 Some(self.page_cache.clone()),
8862 Some(self.decoded_cache.clone()),
8863 self.table_id,
8864 Some(&self.verified_runs),
8865 )?;
8866 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
8870 reader.set_uniform_epoch(Epoch(rr.epoch_created));
8871 }
8872 Ok(reader)
8873 }
8874
8875 pub(crate) fn run_refs(&self) -> &[RunRef] {
8876 &self.run_refs
8877 }
8878
8879 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
8880 self.retiring.iter().map(|run| run.run_id)
8881 }
8882
8883 pub(crate) fn runs_dir(&self) -> PathBuf {
8884 self.dir.join(RUNS_DIR)
8885 }
8886
8887 pub(crate) fn wal_dir(&self) -> PathBuf {
8888 self.dir.join(WAL_DIR)
8889 }
8890
8891 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
8892 self.run_refs = refs;
8893 }
8894
8895 pub(crate) fn next_run_id(&self) -> u64 {
8896 self.next_run_id
8897 }
8898
8899 pub(crate) fn compaction_zstd_level(&self) -> i32 {
8900 self.compaction_zstd_level
8901 }
8902
8903 pub(crate) fn bump_next_run_id(&mut self) {
8904 self.next_run_id += 1;
8905 }
8906
8907 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
8908 self.kek.clone()
8909 }
8910
8911 #[cfg(feature = "encryption")]
8915 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8916 self.kek.as_ref().map(|k| k.derive_idx_key())
8917 }
8918
8919 #[cfg(not(feature = "encryption"))]
8920 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
8921 None
8922 }
8923
8924 #[cfg(feature = "encryption")]
8928 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8929 self.kek.as_ref().map(|k| *k.derive_meta_key())
8930 }
8931
8932 #[cfg(not(feature = "encryption"))]
8933 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
8934 None
8935 }
8936
8937 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
8940 self.column_keys
8941 .iter()
8942 .map(|(&id, &(_, scheme))| (id, scheme))
8943 .collect()
8944 }
8945
8946 #[cfg(feature = "encryption")]
8951 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
8952 self.tokenize_value_enc(column_id, v)
8953 }
8954
8955 #[cfg(feature = "encryption")]
8956 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
8957 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8958 let (key, scheme) = self.column_keys.get(&column_id)?;
8959 let token: Vec<u8> = match (*scheme, v) {
8960 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
8961 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
8962 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
8963 _ => hmac_token(key, &v.encode_key()).to_vec(),
8964 };
8965 Some(Value::Bytes(token))
8966 }
8967
8968 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
8970 self.index_lookup_key_bytes(column_id, &v.encode_key())
8971 }
8972
8973 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
8976 #[cfg(feature = "encryption")]
8977 {
8978 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
8979 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
8980 if *scheme == SCHEME_HMAC_EQ {
8981 return hmac_token(key, encoded).to_vec();
8982 }
8983 }
8984 }
8985 let _ = column_id;
8986 encoded.to_vec()
8987 }
8988}
8989
8990fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
8991 let columnar::NativeColumn::Int64 { data, validity } = col else {
8992 return false;
8993 };
8994 if data.len() < n || !columnar::all_non_null(validity, n) {
8995 return false;
8996 }
8997 data.iter()
8998 .take(n)
8999 .zip(data.iter().skip(1))
9000 .all(|(a, b)| a < b)
9001}
9002
9003#[derive(Debug, Clone)]
9007pub struct ColumnStat {
9008 pub min: Option<Value>,
9009 pub max: Option<Value>,
9010 pub null_count: u64,
9011}
9012
9013#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9015pub enum NativeAgg {
9016 Count,
9017 Sum,
9018 Min,
9019 Max,
9020 Avg,
9021}
9022
9023#[derive(Debug, Clone, PartialEq)]
9025pub enum NativeAggResult {
9026 Count(u64),
9027 Int(i64),
9028 Float(f64),
9029 Null,
9031}
9032
9033#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9035pub enum ApproxAgg {
9036 Count,
9037 Sum,
9038 Avg,
9039}
9040
9041#[derive(Debug, Clone)]
9045pub struct ApproxResult {
9046 pub point: f64,
9048 pub ci_low: f64,
9050 pub ci_high: f64,
9052 pub n_population: u64,
9054 pub n_sample_live: usize,
9056 pub n_passing: usize,
9058}
9059
9060#[derive(Debug, Clone, PartialEq)]
9065pub enum AggState {
9066 Count(u64),
9068 SumI {
9070 sum: i128,
9071 count: u64,
9072 },
9073 SumF {
9075 sum: f64,
9076 count: u64,
9077 },
9078 AvgI {
9080 sum: i128,
9081 count: u64,
9082 },
9083 AvgF {
9085 sum: f64,
9086 count: u64,
9087 },
9088 MinI(i64),
9090 MaxI(i64),
9091 MinF(f64),
9093 MaxF(f64),
9094 Empty,
9096}
9097
9098impl AggState {
9099 pub fn merge(self, other: AggState) -> AggState {
9101 use AggState::*;
9102 match (self, other) {
9103 (Empty, x) | (x, Empty) => x,
9104 (Count(a), Count(b)) => Count(a + b),
9105 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
9106 sum: sa + sb,
9107 count: ca + cb,
9108 },
9109 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
9110 sum: sa + sb,
9111 count: ca + cb,
9112 },
9113 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
9114 sum: sa + sb,
9115 count: ca + cb,
9116 },
9117 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
9118 sum: sa + sb,
9119 count: ca + cb,
9120 },
9121 (MinI(a), MinI(b)) => MinI(a.min(b)),
9122 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
9123 (MinF(a), MinF(b)) => MinF(a.min(b)),
9124 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
9125 _ => Empty, }
9127 }
9128
9129 pub fn point(&self) -> Option<f64> {
9131 match self {
9132 AggState::Count(n) => Some(*n as f64),
9133 AggState::SumI { sum, .. } => Some(*sum as f64),
9134 AggState::SumF { sum, .. } => Some(*sum),
9135 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
9136 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
9137 AggState::MinI(n) => Some(*n as f64),
9138 AggState::MaxI(n) => Some(*n as f64),
9139 AggState::MinF(n) => Some(*n),
9140 AggState::MaxF(n) => Some(*n),
9141 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
9142 }
9143 }
9144
9145 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
9149 let is_float = matches!(ty, Some(TypeId::Float64));
9150 match (agg, result) {
9151 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
9152 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
9153 sum: x as i128,
9154 count: 1, },
9156 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
9157 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
9158 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
9159 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
9160 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
9161 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
9162 (NativeAgg::Count, _) => AggState::Empty,
9163 (_, NativeAggResult::Null) => AggState::Empty,
9164 _ => {
9165 let _ = is_float;
9166 AggState::Empty
9167 }
9168 }
9169 }
9170}
9171
9172#[derive(Debug, Clone)]
9175pub struct CachedAgg {
9176 pub state: AggState,
9177 pub watermark: u64,
9178 pub epoch: u64,
9179}
9180
9181#[derive(Debug, Clone)]
9183pub struct IncrementalAggResult {
9184 pub state: AggState,
9186 pub incremental: bool,
9189 pub delta_rows: u64,
9191}
9192
9193fn agg_state_from_rows(
9197 rows: &[Row],
9198 conditions: &[crate::query::Condition],
9199 index_sets: &[RowIdSet],
9200 column: Option<u16>,
9201 agg: NativeAgg,
9202 schema: &Schema,
9203) -> Result<AggState> {
9204 let mut count: u64 = 0;
9205 let mut sum_i: i128 = 0;
9206 let mut sum_f: f64 = 0.0;
9207 let mut mn_i: i64 = i64::MAX;
9208 let mut mx_i: i64 = i64::MIN;
9209 let mut mn_f: f64 = f64::INFINITY;
9210 let mut mx_f: f64 = f64::NEG_INFINITY;
9211 let mut saw_int = false;
9212 let mut saw_float = false;
9213 for r in rows {
9214 if !conditions
9215 .iter()
9216 .all(|c| condition_matches_row(c, r, schema))
9217 {
9218 continue;
9219 }
9220 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
9221 continue;
9222 }
9223 match agg {
9224 NativeAgg::Count => match column {
9225 None => count += 1,
9227 Some(cid) => match r.columns.get(&cid) {
9230 None | Some(Value::Null) => {}
9231 Some(_) => count += 1,
9232 },
9233 },
9234 _ => match column.and_then(|cid| r.columns.get(&cid)) {
9235 Some(Value::Int64(n)) => {
9236 count += 1;
9237 sum_i += *n as i128;
9238 mn_i = mn_i.min(*n);
9239 mx_i = mx_i.max(*n);
9240 saw_int = true;
9241 }
9242 Some(Value::Float64(f)) => {
9243 count += 1;
9244 sum_f += f;
9245 mn_f = mn_f.min(*f);
9246 mx_f = mx_f.max(*f);
9247 saw_float = true;
9248 }
9249 _ => {}
9250 },
9251 }
9252 }
9253 Ok(match agg {
9254 NativeAgg::Count => {
9255 if count == 0 {
9256 AggState::Empty
9257 } else {
9258 AggState::Count(count)
9259 }
9260 }
9261 NativeAgg::Sum => {
9262 if count == 0 {
9263 AggState::Empty
9264 } else if saw_int {
9265 AggState::SumI { sum: sum_i, count }
9266 } else {
9267 AggState::SumF { sum: sum_f, count }
9268 }
9269 }
9270 NativeAgg::Avg => {
9271 if count == 0 {
9272 AggState::Empty
9273 } else if saw_int {
9274 AggState::AvgI { sum: sum_i, count }
9275 } else {
9276 AggState::AvgF { sum: sum_f, count }
9277 }
9278 }
9279 NativeAgg::Min => {
9280 if !saw_int && !saw_float {
9281 AggState::Empty
9282 } else if saw_int {
9283 AggState::MinI(mn_i)
9284 } else {
9285 AggState::MinF(mn_f)
9286 }
9287 }
9288 NativeAgg::Max => {
9289 if !saw_int && !saw_float {
9290 AggState::Empty
9291 } else if saw_int {
9292 AggState::MaxI(mx_i)
9293 } else {
9294 AggState::MaxF(mx_f)
9295 }
9296 }
9297 })
9298}
9299
9300fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
9304 use crate::query::Condition;
9305 match c {
9306 Condition::Pk(key) => match schema.primary_key() {
9307 Some(pk) => row
9308 .columns
9309 .get(&pk.id)
9310 .map(|v| v.encode_key() == *key)
9311 .unwrap_or(false),
9312 None => false,
9313 },
9314 Condition::BitmapEq { column_id, value } => row
9315 .columns
9316 .get(column_id)
9317 .map(|v| v.encode_key() == *value)
9318 .unwrap_or(false),
9319 Condition::BitmapIn { column_id, values } => {
9320 let key = row.columns.get(column_id).map(|v| v.encode_key());
9321 match key {
9322 Some(k) => values.contains(&k),
9323 None => false,
9324 }
9325 }
9326 Condition::BytesPrefix { column_id, prefix } => row
9327 .columns
9328 .get(column_id)
9329 .map(|v| v.encode_key().starts_with(prefix))
9330 .unwrap_or(false),
9331 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
9332 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
9333 _ => false,
9334 },
9335 Condition::RangeF64 {
9336 column_id,
9337 lo,
9338 lo_inclusive,
9339 hi,
9340 hi_inclusive,
9341 } => match row.columns.get(column_id) {
9342 Some(Value::Float64(n)) => {
9343 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
9344 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
9345 lo_ok && hi_ok
9346 }
9347 _ => false,
9348 },
9349 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
9350 Some(Value::Bytes(b)) => {
9351 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
9352 }
9353 _ => false,
9354 },
9355 Condition::FmContainsAll {
9356 column_id,
9357 patterns,
9358 } => match row.columns.get(column_id) {
9359 Some(Value::Bytes(b)) => patterns
9360 .iter()
9361 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
9362 _ => false,
9363 },
9364 Condition::Ann { .. }
9365 | Condition::SparseMatch { .. }
9366 | Condition::MinHashSimilar { .. } => true,
9367 Condition::IsNull { column_id } => {
9368 matches!(row.columns.get(column_id), Some(Value::Null) | None)
9369 }
9370 Condition::IsNotNull { column_id } => {
9371 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
9372 }
9373 }
9374}
9375
9376fn as_f64(v: Option<&Value>) -> Option<f64> {
9378 match v {
9379 Some(Value::Int64(n)) => Some(*n as f64),
9380 Some(Value::Float64(f)) => Some(*f),
9381 _ => None,
9382 }
9383}
9384
9385fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
9389 let mut count: u64 = 0;
9390 let mut sum: i128 = 0;
9391 let mut mn: i64 = i64::MAX;
9392 let mut mx: i64 = i64::MIN;
9393 while let Some(cols) = cursor.next_batch()? {
9394 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
9395 if crate::columnar::all_non_null(validity, data.len()) {
9396 count += data.len() as u64;
9398 sum += data.iter().map(|&v| v as i128).sum::<i128>();
9399 mn = mn.min(*data.iter().min().unwrap_or(&mn));
9400 mx = mx.max(*data.iter().max().unwrap_or(&mx));
9401 } else {
9402 for (i, &v) in data.iter().enumerate() {
9403 if crate::columnar::validity_bit(validity, i) {
9404 count += 1;
9405 sum += v as i128;
9406 mn = mn.min(v);
9407 mx = mx.max(v);
9408 }
9409 }
9410 }
9411 }
9412 }
9413 Ok((count, sum, mn, mx))
9414}
9415
9416fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
9418 let mut count: u64 = 0;
9419 let mut sum: f64 = 0.0;
9420 let mut mn: f64 = f64::INFINITY;
9421 let mut mx: f64 = f64::NEG_INFINITY;
9422 while let Some(cols) = cursor.next_batch()? {
9423 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
9424 if crate::columnar::all_non_null(validity, data.len()) {
9425 count += data.len() as u64;
9426 sum += data.iter().sum::<f64>();
9427 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
9428 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
9429 } else {
9430 for (i, &v) in data.iter().enumerate() {
9431 if crate::columnar::validity_bit(validity, i) {
9432 count += 1;
9433 sum += v;
9434 mn = mn.min(v);
9435 mx = mx.max(v);
9436 }
9437 }
9438 }
9439 }
9440 }
9441 Ok((count, sum, mn, mx))
9442}
9443
9444fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
9445 if count == 0 && !matches!(agg, NativeAgg::Count) {
9446 return NativeAggResult::Null;
9447 }
9448 match agg {
9449 NativeAgg::Count => NativeAggResult::Count(count),
9450 NativeAgg::Sum => match sum.try_into() {
9453 Ok(v) => NativeAggResult::Int(v),
9454 Err(_) => NativeAggResult::Null,
9455 },
9456 NativeAgg::Min => NativeAggResult::Int(mn),
9457 NativeAgg::Max => NativeAggResult::Int(mx),
9458 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
9459 }
9460}
9461
9462fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
9463 if count == 0 && !matches!(agg, NativeAgg::Count) {
9464 return NativeAggResult::Null;
9465 }
9466 match agg {
9467 NativeAgg::Count => NativeAggResult::Count(count),
9468 NativeAgg::Sum => NativeAggResult::Float(sum),
9469 NativeAgg::Min => NativeAggResult::Float(mn),
9470 NativeAgg::Max => NativeAggResult::Float(mx),
9471 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
9472 }
9473}
9474
9475fn agg_int(
9478 stats: &[crate::page::PageStat],
9479 decode: fn(Option<&[u8]>) -> Option<i64>,
9480) -> Option<(Option<i64>, Option<i64>, u64)> {
9481 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
9482 let mut any = false;
9483 for s in stats {
9484 if let Some(v) = decode(s.min.as_deref()) {
9485 mn = mn.min(v);
9486 any = true;
9487 }
9488 if let Some(v) = decode(s.max.as_deref()) {
9489 mx = mx.max(v);
9490 any = true;
9491 }
9492 nulls += s.null_count;
9493 }
9494 any.then_some((Some(mn), Some(mx), nulls))
9495}
9496
9497fn agg_float(
9499 stats: &[crate::page::PageStat],
9500 decode: fn(Option<&[u8]>) -> Option<f64>,
9501) -> Option<(Option<f64>, Option<f64>, u64)> {
9502 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
9503 let mut any = false;
9504 for s in stats {
9505 if let Some(v) = decode(s.min.as_deref()) {
9506 mn = mn.min(v);
9507 any = true;
9508 }
9509 if let Some(v) = decode(s.max.as_deref()) {
9510 mx = mx.max(v);
9511 any = true;
9512 }
9513 nulls += s.null_count;
9514 }
9515 any.then_some((Some(mn), Some(mx), nulls))
9516}
9517
9518type SecondaryIndexes = (
9520 HashMap<u16, BitmapIndex>,
9521 HashMap<u16, AnnIndex>,
9522 HashMap<u16, FmIndex>,
9523 HashMap<u16, SparseIndex>,
9524 HashMap<u16, MinHashIndex>,
9525);
9526
9527fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
9528 let mut bitmap = HashMap::new();
9529 let mut ann = HashMap::new();
9530 let mut fm = HashMap::new();
9531 let mut sparse = HashMap::new();
9532 let mut minhash = HashMap::new();
9533 for idef in &schema.indexes {
9534 match idef.kind {
9535 IndexKind::Bitmap => {
9536 bitmap.insert(idef.column_id, BitmapIndex::new());
9537 }
9538 IndexKind::Ann => {
9539 let dim = schema
9540 .columns
9541 .iter()
9542 .find(|c| c.id == idef.column_id)
9543 .and_then(|c| match c.ty {
9544 TypeId::Embedding { dim } => Some(dim as usize),
9545 _ => None,
9546 })
9547 .unwrap_or(0);
9548 let options = idef.options.ann.clone().unwrap_or_default();
9549 ann.insert(
9550 idef.column_id,
9551 AnnIndex::with_options(
9552 dim,
9553 options.m,
9554 options.ef_construction,
9555 options.ef_search,
9556 ),
9557 );
9558 }
9559 IndexKind::FmIndex => {
9560 fm.insert(idef.column_id, FmIndex::new());
9561 }
9562 IndexKind::Sparse => {
9563 sparse.insert(idef.column_id, SparseIndex::new());
9564 }
9565 IndexKind::MinHash => {
9566 let options = idef.options.minhash.clone().unwrap_or_default();
9567 minhash.insert(
9568 idef.column_id,
9569 MinHashIndex::with_options(options.permutations, options.bands),
9570 );
9571 }
9572 _ => {}
9573 }
9574 }
9575 (bitmap, ann, fm, sparse, minhash)
9576}
9577
9578const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
9579 | ColumnFlags::AUTO_INCREMENT
9580 | ColumnFlags::ENCRYPTED
9581 | ColumnFlags::ENCRYPTED_INDEXABLE
9582 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
9583
9584fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
9585 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
9586 return Err(MongrelError::Schema(
9587 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
9588 ));
9589 }
9590 Ok(())
9591}
9592
9593fn validate_alter_column_type(
9594 schema: &Schema,
9595 old: &ColumnDef,
9596 next: &ColumnDef,
9597 has_stored_versions: bool,
9598) -> Result<()> {
9599 if old.ty == next.ty {
9600 return Ok(());
9601 }
9602 if schema.indexes.iter().any(|i| i.column_id == old.id) {
9603 return Err(MongrelError::Schema(format!(
9604 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
9605 old.name
9606 )));
9607 }
9608 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
9609 return Ok(());
9610 }
9611 Err(MongrelError::Schema(format!(
9612 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
9613 old.ty, next.ty
9614 )))
9615}
9616
9617fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
9618 matches!(
9619 (old, new),
9620 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
9621 )
9622}
9623
9624fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
9630 let mut prev: Option<i64> = None;
9631 for r in rows {
9632 match r.columns.get(&pk_id) {
9633 Some(Value::Int64(v)) => {
9634 if prev.is_some_and(|p| p >= *v) {
9635 return false;
9636 }
9637 prev = Some(*v);
9638 }
9639 _ => return false,
9640 }
9641 }
9642 true
9643}
9644
9645#[allow(clippy::too_many_arguments)]
9646fn index_into(
9647 schema: &Schema,
9648 row: &Row,
9649 hot: &mut HotIndex,
9650 bitmap: &mut HashMap<u16, BitmapIndex>,
9651 ann: &mut HashMap<u16, AnnIndex>,
9652 fm: &mut HashMap<u16, FmIndex>,
9653 sparse: &mut HashMap<u16, SparseIndex>,
9654 minhash: &mut HashMap<u16, MinHashIndex>,
9655) {
9656 for idef in &schema.indexes {
9657 let Some(val) = row.columns.get(&idef.column_id) else {
9658 continue;
9659 };
9660 match idef.kind {
9661 IndexKind::Bitmap => {
9662 if let Some(b) = bitmap.get_mut(&idef.column_id) {
9663 b.insert(val.encode_key(), row.row_id);
9664 }
9665 }
9666 IndexKind::Ann => {
9667 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
9668 a.insert_validated(v, row.row_id);
9669 }
9670 }
9671 IndexKind::FmIndex => {
9672 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
9673 f.insert(b.clone(), row.row_id);
9674 }
9675 }
9676 IndexKind::Sparse => {
9677 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
9678 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
9681 s.insert(&terms, row.row_id);
9682 }
9683 }
9684 }
9685 IndexKind::MinHash => {
9686 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
9687 let tokens = crate::index::token_hashes_from_bytes(b);
9690 mh.insert(&tokens, row.row_id);
9691 }
9692 }
9693 _ => {}
9694 }
9695 }
9696 if let Some(pk_col) = schema.primary_key() {
9697 if let Some(pk_val) = row.columns.get(&pk_col.id) {
9698 hot.insert(pk_val.encode_key(), row.row_id);
9699 }
9700 }
9701}
9702
9703#[allow(clippy::too_many_arguments)]
9706fn index_into_single(
9707 idef: &IndexDef,
9708 _schema: &Schema,
9709 row: &Row,
9710 _hot: &mut HotIndex,
9711 bitmap: &mut HashMap<u16, BitmapIndex>,
9712 ann: &mut HashMap<u16, AnnIndex>,
9713 fm: &mut HashMap<u16, FmIndex>,
9714 sparse: &mut HashMap<u16, SparseIndex>,
9715 minhash: &mut HashMap<u16, MinHashIndex>,
9716) {
9717 let Some(val) = row.columns.get(&idef.column_id) else {
9718 return;
9719 };
9720 match idef.kind {
9721 IndexKind::Bitmap => {
9722 if let Some(b) = bitmap.get_mut(&idef.column_id) {
9723 b.insert(val.encode_key(), row.row_id);
9724 }
9725 }
9726 IndexKind::Ann => {
9727 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
9728 a.insert_validated(v, row.row_id);
9729 }
9730 }
9731 IndexKind::FmIndex => {
9732 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
9733 f.insert(b.clone(), row.row_id);
9734 }
9735 }
9736 IndexKind::Sparse => {
9737 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
9738 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
9739 s.insert(&terms, row.row_id);
9740 }
9741 }
9742 }
9743 IndexKind::MinHash => {
9744 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
9745 let tokens = crate::index::token_hashes_from_bytes(b);
9746 mh.insert(&tokens, row.row_id);
9747 }
9748 }
9749 _ => {}
9750 }
9751}
9752
9753fn eval_partial_predicate(
9759 pred: &str,
9760 columns_map: &HashMap<u16, &Value>,
9761 name_to_id: &HashMap<&str, u16>,
9762) -> bool {
9763 let lower = pred.trim().to_ascii_lowercase();
9764 if let Some(rest) = lower.strip_suffix(" is not null") {
9766 let col_name = rest.trim();
9767 if let Some(col_id) = name_to_id.get(col_name) {
9768 return columns_map
9769 .get(col_id)
9770 .is_some_and(|v| !matches!(v, Value::Null));
9771 }
9772 }
9773 if let Some(rest) = lower.strip_suffix(" is null") {
9775 let col_name = rest.trim();
9776 if let Some(col_id) = name_to_id.get(col_name) {
9777 return columns_map
9778 .get(col_id)
9779 .map_or(true, |v| matches!(v, Value::Null));
9780 }
9781 }
9782 true
9785}
9786
9787#[allow(dead_code)]
9793fn bulk_index_key(
9794 column_keys: &HashMap<u16, ([u8; 32], u8)>,
9795 column_id: u16,
9796 ty: TypeId,
9797 col: &columnar::NativeColumn,
9798 i: usize,
9799) -> Option<Vec<u8>> {
9800 let encoded = columnar::encode_key_native(ty, col, i)?;
9801 #[cfg(feature = "encryption")]
9802 {
9803 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
9804 if let Some((key, scheme)) = column_keys.get(&column_id) {
9805 return Some(match (*scheme, col) {
9806 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
9807 (_, columnar::NativeColumn::Int64 { data, .. }) => {
9808 ope_token_i64(key, data[i]).to_vec()
9809 }
9810 (_, columnar::NativeColumn::Float64 { data, .. }) => {
9811 ope_token_f64(key, data[i]).to_vec()
9812 }
9813 _ => hmac_token(key, &encoded).to_vec(),
9814 });
9815 }
9816 }
9817 #[cfg(not(feature = "encryption"))]
9818 {
9819 let _ = (column_id, column_keys, col);
9820 }
9821 Some(encoded)
9822}
9823
9824pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
9825 let json = serde_json::to_string_pretty(schema)
9826 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
9827 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
9828 Ok(())
9829}
9830
9831fn read_schema(dir: &Path) -> Result<Schema> {
9832 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
9833 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
9834}
9835
9836fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
9837 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
9838}
9839
9840fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
9841 let n = list_wal_numbers(wal_dir)?;
9842 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
9843}
9844
9845fn next_wal_number(wal_dir: &Path) -> Result<u32> {
9846 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
9847}
9848
9849fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
9850 let _ = std::fs::create_dir_all(wal_dir);
9851 let mut max_n = None;
9852 for entry in std::fs::read_dir(wal_dir)? {
9853 let entry = entry?;
9854 let fname = entry.file_name();
9855 let Some(s) = fname.to_str() else {
9856 continue;
9857 };
9858 let Some(stripped) = s.strip_prefix("seg-") else {
9859 continue;
9860 };
9861 let Some(stripped) = stripped.strip_suffix(".wal") else {
9862 continue;
9863 };
9864 if let Ok(n) = stripped.parse::<u32>() {
9865 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
9866 }
9867 }
9868 Ok(max_n)
9869}