1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::collections::{BTreeMap, HashMap, HashSet};
30use std::path::{Path, PathBuf};
31use std::sync::atomic::AtomicBool;
32use std::sync::Arc;
33use zeroize::Zeroizing;
34
35pub const WAL_DIR: &str = "_wal";
36pub const RUNS_DIR: &str = "_runs";
37pub const CACHE_DIR: &str = "_cache";
38pub const META_DIR: &str = "_meta";
39pub const RCACHE_DIR: &str = "_rcache";
40pub const KEYS_FILENAME: &str = "keys";
41pub const SCHEMA_FILENAME: &str = "schema.json";
42
43fn iso_now_bytes() -> Vec<u8> {
46 let secs = std::time::SystemTime::now()
47 .duration_since(std::time::UNIX_EPOCH)
48 .map(|d| d.as_secs() as i64)
49 .unwrap_or(0);
50 let days = secs.div_euclid(86_400);
51 let rem = secs.rem_euclid(86_400);
52 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
53 let (year, month, day) = civil_from_days(days);
54 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
55}
56
57pub(crate) fn unix_nanos_now() -> i64 {
58 std::time::SystemTime::now()
59 .duration_since(std::time::UNIX_EPOCH)
60 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
61 .unwrap_or(0)
62}
63
64fn civil_from_days(z: i64) -> (i64, u32, u32) {
65 let z = z + 719_468;
66 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
67 let doe = z - era * 146_097;
68 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
69 let y = yoe + era * 400;
70 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
71 let mp = (5 * doy + 2) / 153;
72 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
73 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
74 (if m <= 2 { y + 1 } else { y }, m, d)
75}
76
77const DEFAULT_SYNC_BYTE_THRESHOLD: u64 = 0; pub(crate) const PAGE_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; pub(crate) const DECODED_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; const DEFAULT_MUTABLE_RUN_SPILL_BYTES: u64 = 8 * 1024 * 1024;
84
85#[derive(Clone, Copy, Debug)]
100struct AutoIncState {
101 column_id: u16,
102 next: i64,
103 seeded: bool,
104}
105
106type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
107
108fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
111 schema.auto_increment_column().map(|c| AutoIncState {
112 column_id: c.id,
113 next: 0,
114 seeded: false,
115 })
116}
117
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
131pub enum IndexBuildPolicy {
132 #[default]
134 Deferred,
135 Eager,
137}
138
139pub struct Table {
141 dir: PathBuf,
142 table_id: u64,
143 name: String,
147 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
152 read_only: bool,
155 wal: WalSink,
156 memtable: Memtable,
157 mutable_run: MutableRun,
162 mutable_run_spill_bytes: u64,
164 compaction_zstd_level: i32,
167 allocator: RowIdAllocator,
168 epoch: Arc<EpochAuthority>,
169 persisted_epoch: u64,
172 schema: Schema,
173 hot: HotIndex,
174 kek: Option<Arc<Kek>>,
177 column_keys: HashMap<u16, ([u8; 32], u8)>,
181 run_refs: Vec<RunRef>,
182 retiring: Vec<crate::manifest::RetiredRun>,
185 next_run_id: u64,
186 sync_byte_threshold: u64,
187 current_txn_id: u64,
192 bitmap: HashMap<u16, BitmapIndex>,
193 ann: HashMap<u16, AnnIndex>,
194 fm: HashMap<u16, FmIndex>,
195 sparse: HashMap<u16, SparseIndex>,
196 minhash: HashMap<u16, MinHashIndex>,
197 learned_range: HashMap<u16, ColumnLearnedRange>,
200 pk_by_row: HashMap<RowId, Vec<u8>>,
202 pinned: BTreeMap<Epoch, usize>,
205 pub(crate) live_count: u64,
208 reservoir: crate::reservoir::Reservoir,
211 reservoir_complete: bool,
219 had_deletes: bool,
223 agg_cache: HashMap<u64, CachedAgg>,
227 global_idx_epoch: u64,
231 indexes_complete: bool,
236 index_build_policy: IndexBuildPolicy,
238 pk_by_row_complete: bool,
245 flushed_epoch: u64,
248 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
251 snapshots: Arc<crate::retention::SnapshotRegistry>,
254 commit_lock: Arc<parking_lot::Mutex<()>>,
256 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
259 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
269 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
278 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
280 pending_delete_rids: roaring::RoaringBitmap,
283 pending_put_cols: std::collections::HashSet<u16>,
286 pending_rows: Vec<Row>,
292 pending_rows_auto_inc: Vec<bool>,
293 pending_dels: Vec<RowId>,
296 pending_truncate: Option<Epoch>,
300 auto_inc: Option<AutoIncState>,
303 ttl: Option<TtlPolicy>,
306}
307
308const _: () = {
315 const fn assert_sync<T: ?Sized + Sync>() {}
316 assert_sync::<Table>();
317};
318
319enum CachedData {
325 Rows(Arc<Vec<Row>>),
326 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
327}
328
329impl CachedData {
330 fn approx_bytes(&self) -> u64 {
331 match self {
332 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
333 CachedData::Columns(c) => c
334 .iter()
335 .map(|(_, c)| c.approx_bytes())
336 .sum::<u64>()
337 .saturating_add(c.len() as u64 * 16),
338 }
339 }
340}
341
342struct CachedEntry {
346 data: CachedData,
347 footprint: roaring::RoaringBitmap,
348 condition_cols: Vec<u16>,
349}
350
351struct ResultCache {
362 entries: std::collections::HashMap<u64, CachedEntry>,
363 order: std::collections::VecDeque<u64>,
364 bytes: u64,
365 max_bytes: u64,
366 dir: Option<std::path::PathBuf>,
367 #[allow(dead_code)]
368 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
369}
370
371#[derive(serde::Serialize, serde::Deserialize)]
373struct SerializedEntry {
374 condition_cols: Vec<u16>,
375 footprint_bits: Vec<u32>,
376 data: SerializedData,
377}
378
379#[derive(serde::Serialize, serde::Deserialize)]
380enum SerializedData {
381 Rows(Vec<Row>),
382 Columns(Vec<(u16, columnar::NativeColumn)>),
383}
384
385impl SerializedEntry {
386 fn from_entry(entry: &CachedEntry) -> Self {
387 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
388 let data = match &entry.data {
389 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
390 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
391 };
392 Self {
393 condition_cols: entry.condition_cols.clone(),
394 footprint_bits,
395 data,
396 }
397 }
398
399 fn into_entry(self) -> Option<CachedEntry> {
400 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
401 let data = match self.data {
402 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
403 SerializedData::Columns(c) => {
404 if !c.iter().all(|(_, col)| col.validate()) {
407 return None;
408 }
409 CachedData::Columns(Arc::new(c))
410 }
411 };
412 Some(CachedEntry {
413 data,
414 footprint,
415 condition_cols: self.condition_cols,
416 })
417 }
418}
419
420impl ResultCache {
421 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
422
423 fn new() -> Self {
424 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
425 }
426
427 fn with_max_bytes(max_bytes: u64) -> Self {
428 Self {
429 entries: std::collections::HashMap::new(),
430 order: std::collections::VecDeque::new(),
431 bytes: 0,
432 max_bytes,
433 dir: None,
434 cache_dek: None,
435 }
436 }
437
438 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
439 let _ = std::fs::create_dir_all(&dir);
440 self.dir = Some(dir);
441 self
442 }
443
444 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
445 self.cache_dek = dek;
446 self
447 }
448
449 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
450 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
451 }
452
453 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
457 let Some(path) = self.disk_path(key) else {
458 return;
459 };
460 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
461 Ok(s) => s,
462 Err(_) => return,
463 };
464 let on_disk = if let Some(dek) = &self.cache_dek {
466 match self.encrypt_cache(&serialized, dek) {
467 Some(b) => b,
468 None => return,
469 }
470 } else {
471 serialized
472 };
473 let tmp = path.with_extension("tmp");
474 use std::io::Write;
475 let write = || -> std::io::Result<()> {
476 let mut f = std::fs::File::create(&tmp)?;
477 f.write_all(&on_disk)?;
478 f.flush()?;
479 Ok(())
480 };
481 if write().is_err() {
482 let _ = std::fs::remove_file(&tmp);
483 return;
484 }
485 let _ = std::fs::rename(&tmp, &path);
486 }
487
488 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
490 let path = self.disk_path(key)?;
491 let bytes = std::fs::read(&path).ok()?;
492 let plaintext = if let Some(dek) = &self.cache_dek {
493 self.decrypt_cache(&bytes, dek)?
494 } else {
495 bytes
496 };
497 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
498 serialized.into_entry()
499 }
500
501 fn remove_from_disk(&self, key: u64) {
503 if let Some(path) = self.disk_path(key) {
504 let _ = std::fs::remove_file(&path);
505 }
506 }
507
508 #[cfg(feature = "encryption")]
510 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
511 use crate::encryption::Cipher;
512 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
513 let mut nonce = [0u8; 12];
514 crate::encryption::fill_random(&mut nonce);
515 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
516 let mut out = Vec::with_capacity(12 + ct.len());
517 out.extend_from_slice(&nonce);
518 out.extend_from_slice(&ct);
519 Some(out)
520 }
521
522 #[cfg(not(feature = "encryption"))]
523 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
524 None
525 }
526
527 #[cfg(feature = "encryption")]
529 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
530 use crate::encryption::Cipher;
531 if bytes.len() < 28 {
532 return None;
533 }
534 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
535 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
536 let ct = &bytes[12..];
537 cipher.decrypt_page(&nonce, ct).ok()
538 }
539
540 #[cfg(not(feature = "encryption"))]
541 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
542 None
543 }
544
545 fn load_persistent(&mut self) {
548 let Some(dir) = self.dir.as_ref().cloned() else {
549 return;
550 };
551 let entries = match std::fs::read_dir(&dir) {
552 Ok(e) => e,
553 Err(_) => return,
554 };
555 for entry in entries.flatten() {
556 let path = entry.path();
557 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
559 let _ = std::fs::remove_file(&path);
560 continue;
561 }
562 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
563 continue;
564 }
565 let stem = match path.file_stem().and_then(|s| s.to_str()) {
566 Some(s) => s,
567 None => continue,
568 };
569 let key = match u64::from_str_radix(stem, 16) {
570 Ok(k) => k,
571 Err(_) => continue,
572 };
573 let bytes = match std::fs::read(&path) {
574 Ok(b) => b,
575 Err(_) => continue,
576 };
577 let plaintext = if let Some(dek) = &self.cache_dek {
579 match self.decrypt_cache(&bytes, dek) {
580 Some(p) => p,
581 None => {
582 let _ = std::fs::remove_file(&path);
583 continue;
584 }
585 }
586 } else {
587 bytes
588 };
589 match bincode::deserialize::<SerializedEntry>(&plaintext) {
590 Ok(serialized) => {
591 if let Some(entry) = serialized.into_entry() {
592 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
593 self.entries.insert(key, entry);
594 self.order.push_back(key);
595 } else {
596 let _ = std::fs::remove_file(&path);
597 }
598 }
599 Err(_) => {
600 let _ = std::fs::remove_file(&path);
601 }
602 }
603 }
604 self.evict();
605 }
606
607 fn set_max_bytes(&mut self, max_bytes: u64) {
608 self.max_bytes = max_bytes;
609 self.evict();
610 }
611
612 fn touch(&mut self, key: u64) {
614 self.order.retain(|k| *k != key);
615 self.order.push_back(key);
616 }
617
618 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
619 let res = self.entries.get(&key).and_then(|e| match &e.data {
620 CachedData::Rows(r) => Some(r.clone()),
621 CachedData::Columns(_) => None,
622 });
623 if res.is_some() {
624 self.touch(key);
625 return res;
626 }
627 if let Some(entry) = self.load_from_disk(key) {
629 let res = match &entry.data {
630 CachedData::Rows(r) => Some(r.clone()),
631 CachedData::Columns(_) => None,
632 };
633 if res.is_some() {
634 let approx = entry.data.approx_bytes();
635 self.bytes = self.bytes.saturating_add(approx);
636 self.entries.insert(key, entry);
637 self.order.push_back(key);
638 self.evict();
639 return res;
640 }
641 }
642 None
643 }
644
645 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
646 let res = self.entries.get(&key).and_then(|e| match &e.data {
647 CachedData::Columns(c) => Some(c.clone()),
648 CachedData::Rows(_) => None,
649 });
650 if res.is_some() {
651 self.touch(key);
652 return res;
653 }
654 if let Some(entry) = self.load_from_disk(key) {
656 let res = match &entry.data {
657 CachedData::Columns(c) => Some(c.clone()),
658 CachedData::Rows(_) => None,
659 };
660 if res.is_some() {
661 let approx = entry.data.approx_bytes();
662 self.bytes = self.bytes.saturating_add(approx);
663 self.entries.insert(key, entry);
664 self.order.push_back(key);
665 self.evict();
666 return res;
667 }
668 }
669 None
670 }
671
672 fn insert(&mut self, key: u64, entry: CachedEntry) {
673 let approx = entry.data.approx_bytes();
674 if self.entries.remove(&key).is_some() {
675 self.order.retain(|k| *k != key);
676 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
677 }
678 self.store_to_disk(key, &entry);
680 self.bytes = self.bytes.saturating_add(approx);
681 self.entries.insert(key, entry);
682 self.order.push_back(key);
683 self.evict();
684 }
685
686 fn invalidate(
695 &mut self,
696 delete_rids: &roaring::RoaringBitmap,
697 put_cols: &std::collections::HashSet<u16>,
698 ) {
699 if self.entries.is_empty() {
700 return;
701 }
702 let has_deletes = !delete_rids.is_empty();
703 let to_remove: std::collections::HashSet<u64> = self
704 .entries
705 .iter()
706 .filter(|(_, e)| {
707 let delete_hit = if e.footprint.is_empty() {
708 has_deletes
709 } else {
710 e.footprint.intersection_len(delete_rids) > 0
711 };
712 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
713 delete_hit || col_hit
714 })
715 .map(|(&k, _)| k)
716 .collect();
717 for key in &to_remove {
718 if let Some(e) = self.entries.remove(key) {
719 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
720 }
721 self.remove_from_disk(*key);
722 }
723 if !to_remove.is_empty() {
724 self.order.retain(|k| !to_remove.contains(k));
725 }
726 }
727
728 fn clear(&mut self) {
729 if let Some(dir) = &self.dir {
731 if let Ok(entries) = std::fs::read_dir(dir) {
732 for entry in entries.flatten() {
733 let path = entry.path();
734 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
735 let _ = std::fs::remove_file(&path);
736 }
737 }
738 }
739 }
740 self.entries.clear();
741 self.order.clear();
742 self.bytes = 0;
743 }
744
745 fn evict(&mut self) {
746 while self.bytes > self.max_bytes {
747 let Some(k) = self.order.pop_front() else {
748 break;
749 };
750 if let Some(e) = self.entries.remove(&k) {
751 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
752 self.remove_from_disk(k);
756 }
757 }
758 }
759}
760
761type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
768
769fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
770 let _ = kek;
771 #[cfg(feature = "encryption")]
772 {
773 if let Some(k) = kek {
774 return (
775 Some(k.derive_table_wal_key(_table_id)),
776 Some(k.derive_cache_key()),
777 );
778 }
779 }
780 (None, None)
781}
782
783#[cfg(feature = "encryption")]
785fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
786 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
787}
788
789#[cfg(not(feature = "encryption"))]
790fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
791 Box::new(crate::encryption::PlaintextCipher)
792}
793
794fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
795 let Some(kek) = kek else {
796 return HashMap::new();
797 };
798 #[cfg(feature = "encryption")]
799 {
800 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
801 schema
802 .columns
803 .iter()
804 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
805 .map(|c| {
806 let scheme = if schema
807 .indexes
808 .iter()
809 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
810 {
811 SCHEME_OPE_RANGE
812 } else {
813 SCHEME_HMAC_EQ
814 };
815 let key: [u8; 32] = *kek.derive_column_key(c.id);
816 (c.id, (key, scheme))
817 })
818 .collect()
819 }
820 #[cfg(not(feature = "encryption"))]
821 {
822 let _ = (kek, schema);
823 HashMap::new()
824 }
825}
826
827pub(crate) struct SharedCtx {
832 pub epoch: Arc<EpochAuthority>,
833 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
834 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
835 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
836 pub kek: Option<Arc<Kek>>,
837 pub commit_lock: Arc<parking_lot::Mutex<()>>,
843 pub shared: Option<SharedWalCtx>,
847 pub table_name: Option<String>,
850 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
853 pub read_only: bool,
855}
856
857#[derive(Clone)]
863pub(crate) struct SharedWalCtx {
864 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
865 pub group: Arc<GroupCommit>,
866 pub poisoned: Arc<AtomicBool>,
867 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
868 pub change_wake: tokio::sync::broadcast::Sender<()>,
869}
870
871enum WalSink {
874 Private(Wal),
875 Shared(SharedWalCtx),
876}
877
878impl SharedCtx {
879 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
883 let n_shards = if cache_dir.is_some() {
887 1
888 } else {
889 crate::cache::CACHE_SHARDS
890 };
891 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
892 let page_cache = if let Some(d) = cache_dir {
893 Arc::new(crate::cache::Sharded::new(1, || {
894 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
895 }))
896 } else {
897 Arc::new(crate::cache::Sharded::new(n_shards, || {
898 crate::cache::PageCache::new(per_shard)
899 }))
900 };
901 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
902 let decoded_cache = Arc::new(crate::cache::Sharded::new(
903 crate::cache::CACHE_SHARDS,
904 || crate::cache::DecodedPageCache::new(decoded_per_shard),
905 ));
906 Self {
907 epoch: Arc::new(EpochAuthority::new(0)),
908 page_cache,
909 decoded_cache,
910 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
911 kek,
912 commit_lock: Arc::new(parking_lot::Mutex::new(())),
913 shared: None,
914 table_name: None,
915 auth: None,
916 read_only: false,
917 }
918 }
919}
920
921fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
925 use crate::query::Condition;
926 match c {
927 Condition::Pk(_)
929 | Condition::BitmapEq { .. }
930 | Condition::BitmapIn { .. }
931 | Condition::BytesPrefix { .. }
932 | Condition::IsNull { .. }
933 | Condition::IsNotNull { .. } => 0,
934 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
936 1
937 }
938 Condition::FmContains { .. }
940 | Condition::FmContainsAll { .. }
941 | Condition::Ann { .. }
942 | Condition::SparseMatch { .. } => 2,
943 }
944}
945
946impl Table {
947 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
948 let dir = dir.as_ref().to_path_buf();
949 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
950 Self::create_in(&dir, schema, table_id, ctx)
951 }
952
953 #[cfg(feature = "encryption")]
964 pub fn create_encrypted(
965 dir: impl AsRef<Path>,
966 schema: Schema,
967 table_id: u64,
968 passphrase: &str,
969 ) -> Result<Self> {
970 let dir = dir.as_ref();
971 std::fs::create_dir_all(dir.join(META_DIR))?;
972 let salt = crate::encryption::random_salt();
973 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
974 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
975 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
976 Self::create_in(dir, schema, table_id, ctx)
977 }
978
979 #[cfg(feature = "encryption")]
984 pub fn create_with_key(
985 dir: impl AsRef<Path>,
986 schema: Schema,
987 table_id: u64,
988 key: &[u8],
989 ) -> Result<Self> {
990 let dir = dir.as_ref();
991 std::fs::create_dir_all(dir.join(META_DIR))?;
992 let salt = crate::encryption::random_salt();
993 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
994 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
995 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
996 Self::create_in(dir, schema, table_id, ctx)
997 }
998
999 #[cfg(feature = "encryption")]
1001 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1002 let dir = dir.as_ref();
1003 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1004 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1005 MongrelError::NotFound(format!(
1006 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1007 salt_path
1008 ))
1009 })?;
1010 if salt_bytes.len() != crate::encryption::SALT_LEN {
1011 return Err(MongrelError::InvalidArgument(format!(
1012 "salt file is {} bytes, expected {}",
1013 salt_bytes.len(),
1014 crate::encryption::SALT_LEN
1015 )));
1016 }
1017 let mut salt = [0u8; crate::encryption::SALT_LEN];
1018 salt.copy_from_slice(&salt_bytes);
1019 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1020 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1021 Self::open_in(dir, ctx)
1022 }
1023
1024 pub(crate) fn create_in(
1025 dir: impl AsRef<Path>,
1026 schema: Schema,
1027 table_id: u64,
1028 ctx: SharedCtx,
1029 ) -> Result<Self> {
1030 schema.validate_auto_increment()?;
1031 schema.validate_defaults()?;
1032 schema.validate_ai()?;
1033 for index in &schema.indexes {
1034 index.validate_options()?;
1035 }
1036 let dir = dir.as_ref().to_path_buf();
1037 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1038 write_schema(&dir, &schema)?;
1039 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1040 let (wal, current_txn_id) = match ctx.shared.clone() {
1043 Some(s) => (WalSink::Shared(s), 0),
1044 None => {
1045 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1046 let mut w = if let Some(ref dk) = wal_dek {
1047 Wal::create_with_cipher(
1048 dir.join(WAL_DIR).join("seg-000000.wal"),
1049 Epoch(0),
1050 Some(make_cipher(dk)),
1051 0,
1052 )?
1053 } else {
1054 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1055 };
1056 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1057 (WalSink::Private(w), 1)
1058 }
1059 };
1060 let mut manifest = Manifest::new(table_id, schema.schema_id);
1061 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1066 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1067 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1068 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1069 let auto_inc = resolve_auto_inc(&schema);
1070 let rcache_dir = dir.join(RCACHE_DIR);
1071 Ok(Self {
1072 dir,
1073 table_id,
1074 name: ctx.table_name.unwrap_or_default(),
1075 auth: ctx.auth,
1076 read_only: ctx.read_only,
1077 wal,
1078 memtable: Memtable::new(),
1079 mutable_run: MutableRun::new(),
1080 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1081 compaction_zstd_level: 3,
1082 allocator: RowIdAllocator::new(0),
1083 epoch: ctx.epoch,
1084 persisted_epoch: 0,
1085 schema,
1086 hot: HotIndex::new(),
1087 kek: ctx.kek,
1088 column_keys,
1089 run_refs: Vec::new(),
1090 retiring: Vec::new(),
1091 next_run_id: 1,
1092 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1093 current_txn_id,
1094 bitmap,
1095 ann,
1096 fm,
1097 sparse,
1098 minhash,
1099 learned_range: HashMap::new(),
1100 pk_by_row: HashMap::new(),
1101 pinned: BTreeMap::new(),
1102 live_count: 0,
1103 reservoir: crate::reservoir::Reservoir::default(),
1104 reservoir_complete: true,
1105 had_deletes: false,
1106 agg_cache: HashMap::new(),
1107 global_idx_epoch: 0,
1108 indexes_complete: true,
1109 index_build_policy: IndexBuildPolicy::default(),
1110 pk_by_row_complete: false,
1111 flushed_epoch: 0,
1112 page_cache: ctx.page_cache,
1113 decoded_cache: ctx.decoded_cache,
1114 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1115 snapshots: ctx.snapshots,
1116 commit_lock: ctx.commit_lock,
1117 result_cache: Arc::new(parking_lot::Mutex::new(
1118 ResultCache::new()
1119 .with_dir(rcache_dir)
1120 .with_cache_dek(cache_dek.clone()),
1121 )),
1122 pending_delete_rids: roaring::RoaringBitmap::new(),
1123 pending_put_cols: std::collections::HashSet::new(),
1124 pending_rows: Vec::new(),
1125 pending_rows_auto_inc: Vec::new(),
1126 pending_dels: Vec::new(),
1127 pending_truncate: None,
1128 wal_dek,
1129 auto_inc,
1130 ttl: None,
1131 })
1132 }
1133
1134 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1138 let dir = dir.as_ref();
1139 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1140 Self::open_in(dir, ctx)
1141 }
1142
1143 #[cfg(feature = "encryption")]
1146 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1147 let dir = dir.as_ref();
1148 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1149 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1150 MongrelError::NotFound(format!(
1151 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1152 salt_path
1153 ))
1154 })?;
1155 let salt_len = crate::encryption::SALT_LEN;
1156 if salt_bytes.len() != salt_len {
1157 return Err(MongrelError::InvalidArgument(format!(
1158 "encryption salt is {} bytes, expected {salt_len}",
1159 salt_bytes.len()
1160 )));
1161 }
1162 let mut salt = [0u8; 16];
1163 salt.copy_from_slice(&salt_bytes);
1164 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1165 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1166 let t = Self::open_in(dir, ctx)?;
1167 Ok(t)
1168 }
1169
1170 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1171 let dir = dir.as_ref().to_path_buf();
1172 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1173 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1174 let schema: Schema = read_schema(&dir)?;
1175 schema.validate_ai()?;
1176 for index in &schema.indexes {
1177 index.validate_options()?;
1178 }
1179 let replay_epoch = Epoch(manifest.current_epoch);
1180 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1181 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1185 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1186 None => {
1187 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1188 let replayed = match &active {
1190 Some(path) => {
1191 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1192 crate::wal::replay_with_cipher(path, cipher)?
1193 }
1194 None => Vec::new(),
1195 };
1196 let mut w = match &active {
1197 Some(path) => Wal::create_with_cipher(
1198 path,
1199 replay_epoch,
1200 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1201 0,
1202 )?,
1203 None => Wal::create_with_cipher(
1204 dir.join(WAL_DIR).join("seg-000000.wal"),
1205 replay_epoch,
1206 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1207 0,
1208 )?,
1209 };
1210 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1211 (WalSink::Private(w), replayed, 1)
1212 }
1213 };
1214
1215 let mut memtable = Memtable::new();
1216 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1217 let persisted_epoch = manifest.current_epoch;
1218 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1225 s.next = manifest.auto_inc_next;
1226 s.seeded = manifest.auto_inc_next > 0;
1227 s
1228 });
1229
1230 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1237 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1238 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1239 std::collections::BTreeMap::new();
1240 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1241 let mut saw_delete = false;
1242 for record in replayed {
1243 let txn_id = record.txn_id;
1244 match record.op {
1245 Op::Put { rows, .. } => {
1246 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1247 for row in &rows {
1248 allocator.advance_to(row.row_id);
1249 if let Some(ai) = auto_inc.as_mut() {
1250 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1251 if *n + 1 > ai.next {
1252 ai.next = *n + 1;
1253 }
1254 }
1255 }
1256 }
1257 staged_puts.entry(txn_id).or_default().extend(rows);
1258 }
1259 Op::Delete { row_ids, .. } => {
1260 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1261 }
1262 Op::TxnCommit { epoch, .. } => {
1263 let commit_epoch = Epoch(epoch);
1264 if let Some(puts) = staged_puts.remove(&txn_id) {
1265 if commit_epoch.0 > manifest.flushed_epoch {
1266 for row in &puts {
1267 memtable.upsert(row.clone());
1268 }
1269 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1270 }
1271 }
1272 if let Some(dels) = staged_deletes.remove(&txn_id) {
1273 saw_delete = true;
1274 if commit_epoch.0 > manifest.flushed_epoch {
1275 for rid in dels {
1276 memtable.tombstone(rid, commit_epoch);
1277 replayed_deletes.push((rid, commit_epoch));
1278 }
1279 }
1280 }
1281 }
1282 Op::TxnAbort => {
1283 staged_puts.remove(&txn_id);
1284 staged_deletes.remove(&txn_id);
1285 }
1286 Op::TruncateTable { .. }
1287 | Op::ExternalTableState { .. }
1288 | Op::Flush { .. }
1289 | Op::Ddl(_)
1290 | Op::BeforeImage { .. }
1291 | Op::CommitTimestamp { .. } => {}
1292 }
1293 }
1294
1295 let rcache_dir = dir.join(RCACHE_DIR);
1296 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1297 let mut db = Self {
1298 dir,
1299 table_id: manifest.table_id,
1300 name: ctx.table_name.unwrap_or_default(),
1301 auth: ctx.auth,
1302 read_only: ctx.read_only,
1303 wal,
1304 memtable,
1305 mutable_run: MutableRun::new(),
1306 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1307 compaction_zstd_level: 3,
1308 allocator,
1309 epoch: ctx.epoch,
1310 persisted_epoch,
1311 schema,
1312 hot: HotIndex::new(),
1313 kek: ctx.kek,
1314 column_keys,
1315 run_refs: manifest.runs.clone(),
1316 retiring: manifest.retiring.clone(),
1317 next_run_id: manifest
1318 .runs
1319 .iter()
1320 .map(|r| r.run_id as u64 + 1)
1321 .max()
1322 .unwrap_or(1),
1323 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1324 current_txn_id,
1325 bitmap: HashMap::new(),
1326 ann: HashMap::new(),
1327 fm: HashMap::new(),
1328 sparse: HashMap::new(),
1329 minhash: HashMap::new(),
1330 learned_range: HashMap::new(),
1331 pk_by_row: HashMap::new(),
1332 pinned: BTreeMap::new(),
1333 live_count: manifest.live_count,
1334 reservoir: crate::reservoir::Reservoir::default(),
1335 reservoir_complete: false,
1336 had_deletes: saw_delete
1337 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1338 != manifest.live_count,
1339 agg_cache: HashMap::new(),
1340 global_idx_epoch: manifest.global_idx_epoch,
1341 indexes_complete: true,
1342 index_build_policy: IndexBuildPolicy::default(),
1343 pk_by_row_complete: false,
1344 flushed_epoch: manifest.flushed_epoch,
1345 page_cache: ctx.page_cache,
1346 decoded_cache: ctx.decoded_cache,
1347 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1348 snapshots: ctx.snapshots,
1349 commit_lock: ctx.commit_lock,
1350 result_cache: Arc::new(parking_lot::Mutex::new(
1351 ResultCache::new()
1352 .with_dir(rcache_dir)
1353 .with_cache_dek(cache_dek.clone()),
1354 )),
1355 pending_delete_rids: roaring::RoaringBitmap::new(),
1356 pending_put_cols: std::collections::HashSet::new(),
1357 pending_rows: Vec::new(),
1358 pending_rows_auto_inc: Vec::new(),
1359 pending_dels: Vec::new(),
1360 pending_truncate: None,
1361 wal_dek,
1362 auto_inc,
1363 ttl: manifest.ttl,
1364 };
1365
1366 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1369
1370 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1375 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1376 c.epoch_built == manifest.global_idx_epoch
1377 && manifest.global_idx_epoch > 0
1378 && manifest
1379 .runs
1380 .iter()
1381 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1382 });
1383 if let Some(loaded) = checkpoint {
1384 if checkpoint_valid {
1385 db.hot = loaded.hot;
1386 db.bitmap = loaded.bitmap;
1387 db.ann = loaded.ann;
1388 db.fm = loaded.fm;
1389 db.sparse = loaded.sparse;
1390 db.minhash = loaded.minhash;
1391 db.learned_range = loaded.learned_range;
1392 }
1395 }
1396 if !checkpoint_valid {
1397 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1398 db.bitmap = bitmap;
1399 db.ann = ann;
1400 db.fm = fm;
1401 db.sparse = sparse;
1402 db.minhash = minhash;
1403 db.rebuild_indexes_from_runs()?;
1404 db.build_learned_ranges()?;
1405 }
1406
1407 for (epoch, group) in replayed_puts {
1412 let (losers, winner_pks) = db.partition_pk_winners(&group);
1413 for (key, &row_id) in &winner_pks {
1414 if let Some(old_rid) = db.hot.get(key) {
1415 if old_rid != row_id {
1416 db.tombstone_row(old_rid, epoch, false);
1417 }
1418 }
1419 }
1420 for &loser_rid in &losers {
1421 db.tombstone_row(loser_rid, epoch, false);
1422 }
1423 for (key, row_id) in winner_pks {
1424 db.insert_hot_pk(key, row_id);
1425 }
1426 if db.schema.primary_key().is_none() {
1427 for r in &group {
1428 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1429 }
1430 }
1431 for r in &group {
1432 if !losers.contains(&r.row_id) {
1433 db.index_row(r);
1434 }
1435 }
1436 }
1437 for (rid, epoch) in &replayed_deletes {
1441 db.remove_hot_for_row(*rid, *epoch);
1442 }
1443
1444 db.result_cache.lock().load_persistent();
1451 Ok(db)
1452 }
1453
1454 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1460 if self.reservoir_complete {
1461 return Ok(());
1462 }
1463 self.rebuild_reservoir()?;
1464 self.reservoir_complete = true;
1465 Ok(())
1466 }
1467
1468 fn rebuild_reservoir(&mut self) -> Result<()> {
1471 let snap = self.snapshot();
1472 let rows = self.visible_rows(snap)?;
1473 self.reservoir.reset();
1474 for r in rows {
1475 self.reservoir.offer(r.row_id.0);
1476 }
1477 Ok(())
1478 }
1479
1480 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1481 self.hot = HotIndex::new();
1482 self.pk_by_row.clear();
1483 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1484 self.bitmap = bitmap;
1485 self.ann = ann;
1486 self.fm = fm;
1487 self.sparse = sparse;
1488 self.minhash = minhash;
1489 let snapshot = Epoch(u64::MAX);
1490 let ttl_now = unix_nanos_now();
1491 for rr in self.run_refs.clone() {
1492 let mut reader = self.open_reader(rr.run_id)?;
1493 for row in reader.visible_rows(snapshot)? {
1494 if self.row_expired_at(&row, ttl_now) {
1495 continue;
1496 }
1497 let tok_row = self.tokenized_for_indexes(&row);
1498 index_into(
1499 &self.schema,
1500 &tok_row,
1501 &mut self.hot,
1502 &mut self.bitmap,
1503 &mut self.ann,
1504 &mut self.fm,
1505 &mut self.sparse,
1506 &mut self.minhash,
1507 );
1508 }
1509 }
1510 for row in self.mutable_run.visible_versions(snapshot) {
1511 if row.deleted {
1512 self.remove_hot_for_row(row.row_id, snapshot);
1513 } else if !self.row_expired_at(&row, ttl_now) {
1514 self.index_row(&row);
1515 }
1516 }
1517 for row in self.memtable.visible_versions(snapshot) {
1518 if row.deleted {
1519 self.remove_hot_for_row(row.row_id, snapshot);
1520 } else if !self.row_expired_at(&row, ttl_now) {
1521 self.index_row(&row);
1522 }
1523 }
1524 self.refresh_pk_by_row_from_hot();
1525 Ok(())
1526 }
1527
1528 fn refresh_pk_by_row_from_hot(&mut self) {
1529 self.pk_by_row_complete = true;
1530 if self.schema.primary_key().is_none() {
1531 self.pk_by_row.clear();
1532 return;
1533 }
1534 self.pk_by_row = self
1540 .hot
1541 .entries()
1542 .into_iter()
1543 .map(|(key, row_id)| (row_id, key))
1544 .collect();
1545 }
1546
1547 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1548 if self.schema.primary_key().is_some() {
1549 self.pk_by_row.insert(row_id, key.clone());
1550 }
1551 self.hot.insert(key, row_id);
1552 }
1553
1554 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1558 self.learned_range.clear();
1559 if self.run_refs.len() != 1 {
1560 return Ok(());
1561 }
1562 let cols: Vec<(u16, usize)> = self
1563 .schema
1564 .indexes
1565 .iter()
1566 .filter(|i| i.kind == IndexKind::LearnedRange)
1567 .map(|i| {
1568 (
1569 i.column_id,
1570 i.options
1571 .learned_range
1572 .as_ref()
1573 .map(|options| options.epsilon)
1574 .unwrap_or(16),
1575 )
1576 })
1577 .collect();
1578 if cols.is_empty() {
1579 return Ok(());
1580 }
1581 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1582 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1583 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1584 _ => return Ok(()),
1585 };
1586 for (cid, epsilon) in cols {
1587 let ty = self
1588 .schema
1589 .columns
1590 .iter()
1591 .find(|c| c.id == cid)
1592 .map(|c| c.ty.clone())
1593 .unwrap_or(TypeId::Int64);
1594 match ty {
1595 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1596 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1597 let pairs: Vec<(i64, u64)> = data
1598 .iter()
1599 .zip(row_ids.iter())
1600 .map(|(v, r)| (*v, *r))
1601 .collect();
1602 self.learned_range.insert(
1603 cid,
1604 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
1605 );
1606 }
1607 }
1608 TypeId::Float64 => {
1609 if let columnar::NativeColumn::Float64 { data, .. } =
1610 reader.column_native(cid)?
1611 {
1612 let pairs: Vec<(f64, u64)> = data
1613 .iter()
1614 .zip(row_ids.iter())
1615 .map(|(v, r)| (*v, *r))
1616 .collect();
1617 self.learned_range.insert(
1618 cid,
1619 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
1620 );
1621 }
1622 }
1623 _ => {}
1624 }
1625 }
1626 Ok(())
1627 }
1628
1629 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1636 if self.indexes_complete {
1637 crate::trace::QueryTrace::record(|t| {
1638 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1639 });
1640 return Ok(());
1641 }
1642 crate::trace::QueryTrace::record(|t| {
1643 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1644 });
1645 self.rebuild_indexes_from_runs()?;
1646 self.build_learned_ranges()?;
1647 self.indexes_complete = true;
1648 let epoch = self.current_epoch();
1649 self.checkpoint_indexes(epoch);
1650 Ok(())
1651 }
1652
1653 fn pending_epoch(&self) -> Epoch {
1654 Epoch(self.epoch.visible().0 + 1)
1655 }
1656
1657 fn is_shared(&self) -> bool {
1660 matches!(self.wal, WalSink::Shared(_))
1661 }
1662
1663 fn ensure_txn_id(&mut self) -> u64 {
1667 if self.current_txn_id == 0 {
1668 let id = match &self.wal {
1669 WalSink::Shared(s) => {
1670 let mut g = s.txn_ids.lock();
1671 let v = *g;
1672 *g = g.wrapping_add(1);
1673 v
1674 }
1675 WalSink::Private(_) => 1,
1676 };
1677 self.current_txn_id = id;
1678 }
1679 self.current_txn_id
1680 }
1681
1682 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1685 self.ensure_writable()?;
1686 let txn_id = self.ensure_txn_id();
1687 let table_id = self.table_id;
1688 match &mut self.wal {
1689 WalSink::Private(w) => {
1690 w.append_txn(txn_id, op)?;
1691 }
1692 WalSink::Shared(s) => {
1693 s.wal.lock().append(txn_id, table_id, op)?;
1694 }
1695 }
1696 Ok(())
1697 }
1698
1699 fn ensure_writable(&self) -> Result<()> {
1700 if self.read_only {
1701 Err(MongrelError::ReadOnlyReplica)
1702 } else {
1703 Ok(())
1704 }
1705 }
1706
1707 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1718 match &self.auth {
1719 Some(checker) => checker.check(&self.name, perm),
1720 None => Ok(()),
1721 }
1722 }
1723 pub fn require_select(&self) -> Result<()> {
1728 self.require(crate::auth_state::RequiredPermission::Select)
1729 }
1730 fn require_insert(&self) -> Result<()> {
1731 self.require(crate::auth_state::RequiredPermission::Insert)
1732 }
1733 #[allow(dead_code)]
1737 fn require_update(&self) -> Result<()> {
1738 self.require(crate::auth_state::RequiredPermission::Update)
1739 }
1740 fn require_delete(&self) -> Result<()> {
1741 self.require(crate::auth_state::RequiredPermission::Delete)
1742 }
1743
1744 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1747 self.require_insert()?;
1748 Ok(self.put_returning(columns)?.0)
1749 }
1750
1751 pub fn put_returning(
1756 &mut self,
1757 mut columns: Vec<(u16, Value)>,
1758 ) -> Result<(RowId, Option<i64>)> {
1759 self.require_insert()?;
1760 let assigned = self.fill_auto_inc(&mut columns)?;
1761 self.apply_defaults(&mut columns)?;
1762 self.schema.validate_values(&columns)?;
1763 let row_id = if self.schema.clustered {
1768 self.derive_clustered_row_id(&columns)?
1769 } else {
1770 self.allocator.alloc()
1771 };
1772 let epoch = self.pending_epoch();
1773 let mut row = Row::new(row_id, epoch);
1774 for (col_id, val) in columns {
1775 row.columns.insert(col_id, val);
1776 }
1777 self.commit_rows(vec![row], assigned.is_some())?;
1778 Ok((row_id, assigned))
1779 }
1780
1781 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1784 self.require_insert()?;
1785 Ok(self
1786 .put_batch_returning(batch)?
1787 .into_iter()
1788 .map(|(r, _)| r)
1789 .collect())
1790 }
1791
1792 pub fn put_batch_returning(
1795 &mut self,
1796 batch: Vec<Vec<(u16, Value)>>,
1797 ) -> Result<Vec<(RowId, Option<i64>)>> {
1798 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1799 for mut cols in batch {
1800 let assigned = self.fill_auto_inc(&mut cols)?;
1801 self.apply_defaults(&mut cols)?;
1802 filled.push((cols, assigned));
1803 }
1804 for (cols, _) in &filled {
1805 self.schema.validate_values(cols)?;
1806 }
1807 let epoch = self.pending_epoch();
1808 let mut rows = Vec::with_capacity(filled.len());
1809 let mut ids = Vec::with_capacity(filled.len());
1810 for (cols, assigned) in filled {
1811 let row_id = if self.schema.clustered {
1812 self.derive_clustered_row_id(&cols)?
1813 } else {
1814 self.allocator.alloc()
1815 };
1816 let mut row = Row::new(row_id, epoch);
1817 for (c, v) in cols {
1818 row.columns.insert(c, v);
1819 }
1820 ids.push((row_id, assigned));
1821 rows.push(row);
1822 }
1823 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1824 self.commit_rows(rows, all_auto_generated)?;
1825 Ok(ids)
1826 }
1827
1828 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1834 self.ensure_writable()?;
1835 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1836 return Ok(None);
1837 };
1838 let pos = columns.iter().position(|(c, _)| *c == cid);
1839 let assigned = match pos {
1840 Some(i) => match &columns[i].1 {
1841 Value::Null => {
1842 let next = self.alloc_auto_inc_value()?;
1843 columns[i].1 = Value::Int64(next);
1844 Some(next)
1845 }
1846 Value::Int64(n) => {
1847 self.advance_auto_inc_past(*n)?;
1848 None
1849 }
1850 other => {
1851 return Err(MongrelError::InvalidArgument(format!(
1852 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1853 other
1854 )))
1855 }
1856 },
1857 None => {
1858 let next = self.alloc_auto_inc_value()?;
1859 columns.push((cid, Value::Int64(next)));
1860 Some(next)
1861 }
1862 };
1863 Ok(assigned)
1864 }
1865
1866 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1872 for col in &self.schema.columns {
1873 let Some(expr) = &col.default_value else {
1874 continue;
1875 };
1876 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
1878 continue;
1879 }
1880 let pos = columns.iter().position(|(c, _)| *c == col.id);
1881 let needs_default = match pos {
1882 None => true,
1883 Some(i) => matches!(columns[i].1, Value::Null),
1884 };
1885 if !needs_default {
1886 continue;
1887 }
1888 let v = match expr {
1889 crate::schema::DefaultExpr::Static(v) => v.clone(),
1890 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
1891 crate::schema::DefaultExpr::Uuid => {
1892 let mut buf = [0u8; 16];
1893 getrandom::getrandom(&mut buf)
1894 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
1895 Value::Uuid(buf)
1896 }
1897 };
1898 match pos {
1899 None => columns.push((col.id, v)),
1900 Some(i) => columns[i].1 = v,
1901 }
1902 }
1903 Ok(())
1904 }
1905
1906 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1908 self.ensure_auto_inc_seeded()?;
1909 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1911 let v = ai.next;
1912 ai.next = ai.next.saturating_add(1);
1913 Ok(v)
1914 }
1915
1916 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1919 self.ensure_auto_inc_seeded()?;
1920 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1921 let floor = used.saturating_add(1).max(1);
1922 if ai.next < floor {
1923 ai.next = floor;
1924 }
1925 Ok(())
1926 }
1927
1928 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1933 let needs_seed = match self.auto_inc {
1934 Some(ai) => !ai.seeded,
1935 None => return Ok(()),
1936 };
1937 if !needs_seed {
1938 return Ok(());
1939 }
1940 if self.seed_empty_auto_inc() {
1941 return Ok(());
1942 }
1943 let cid = self
1944 .auto_inc
1945 .as_ref()
1946 .expect("auto-inc column present")
1947 .column_id;
1948 let max = self.scan_max_int64(cid)?;
1949 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1950 let floor = max.saturating_add(1).max(1);
1951 if ai.next < floor {
1952 ai.next = floor;
1953 }
1954 ai.seeded = true;
1955 Ok(())
1956 }
1957
1958 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1959 if n == 0 || self.auto_inc.is_none() {
1960 return Ok(None);
1961 }
1962 self.ensure_auto_inc_seeded()?;
1963 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1964 let start = ai.next;
1965 ai.next = ai.next.saturating_add(n as i64);
1966 Ok(Some(start))
1967 }
1968
1969 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1973 let mut max: i64 = 0;
1974 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1975 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1976 if *n > max {
1977 max = *n;
1978 }
1979 }
1980 }
1981 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1982 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1983 if *n > max {
1984 max = *n;
1985 }
1986 }
1987 }
1988 for rr in self.run_refs.clone() {
1989 let reader = self.open_reader(rr.run_id)?;
1990 if let Some(stats) = reader.column_page_stats(column_id) {
1991 for s in stats {
1992 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1993 if n > max {
1994 max = n;
1995 }
1996 }
1997 }
1998 } else if reader.has_column(column_id) {
1999 if let columnar::NativeColumn::Int64 { data, validity } =
2000 reader.column_native_shared(column_id)?
2001 {
2002 for (i, n) in data.iter().enumerate() {
2003 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2004 {
2005 max = *n;
2006 }
2007 }
2008 }
2009 }
2010 }
2011 Ok(max)
2012 }
2013
2014 fn seed_empty_auto_inc(&mut self) -> bool {
2015 let Some(ai) = self.auto_inc.as_mut() else {
2016 return false;
2017 };
2018 if ai.seeded || self.live_count != 0 {
2019 return false;
2020 }
2021 if ai.next < 1 {
2022 ai.next = 1;
2023 }
2024 ai.seeded = true;
2025 true
2026 }
2027
2028 fn advance_auto_inc_from_native_columns(
2029 &mut self,
2030 columns: &[(u16, columnar::NativeColumn)],
2031 n: usize,
2032 live_before: u64,
2033 ) -> Result<()> {
2034 let Some(ai) = self.auto_inc.as_mut() else {
2035 return Ok(());
2036 };
2037 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2038 return Ok(());
2039 };
2040 let columnar::NativeColumn::Int64 { data, validity } = col else {
2041 return Err(MongrelError::InvalidArgument(format!(
2042 "AUTO_INCREMENT column {} must be Int64",
2043 ai.column_id
2044 )));
2045 };
2046 let max = if native_int64_strictly_increasing(col, n) {
2047 data.get(n.saturating_sub(1)).copied()
2048 } else {
2049 data.iter()
2050 .take(n)
2051 .enumerate()
2052 .filter_map(|(i, v)| {
2053 if validity.is_empty() || columnar::validity_bit(validity, i) {
2054 Some(*v)
2055 } else {
2056 None
2057 }
2058 })
2059 .max()
2060 };
2061 if let Some(max) = max {
2062 let floor = max.saturating_add(1).max(1);
2063 if ai.next < floor {
2064 ai.next = floor;
2065 }
2066 if ai.seeded || live_before == 0 {
2067 ai.seeded = true;
2068 }
2069 }
2070 Ok(())
2071 }
2072
2073 fn fill_auto_inc_native_columns(
2074 &mut self,
2075 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2076 n: usize,
2077 ) -> Result<()> {
2078 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2079 return Ok(());
2080 };
2081 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2082 if let Some(start) = self.alloc_auto_inc_range(n)? {
2083 columns.push((
2084 cid,
2085 columnar::NativeColumn::Int64 {
2086 data: (start..start.saturating_add(n as i64)).collect(),
2087 validity: vec![0xFF; n.div_ceil(8)],
2088 },
2089 ));
2090 }
2091 return Ok(());
2092 };
2093
2094 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2095 return Err(MongrelError::InvalidArgument(format!(
2096 "AUTO_INCREMENT column {cid} must be Int64"
2097 )));
2098 };
2099 if data.len() < n {
2100 return Err(MongrelError::InvalidArgument(format!(
2101 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2102 data.len()
2103 )));
2104 }
2105 if columnar::all_non_null(validity, n) {
2106 return Ok(());
2107 }
2108 if validity.iter().all(|b| *b == 0) {
2109 if let Some(start) = self.alloc_auto_inc_range(n)? {
2110 for (i, slot) in data.iter_mut().take(n).enumerate() {
2111 *slot = start.saturating_add(i as i64);
2112 }
2113 *validity = vec![0xFF; n.div_ceil(8)];
2114 }
2115 return Ok(());
2116 }
2117
2118 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2119 for (i, slot) in data.iter_mut().enumerate().take(n) {
2120 if columnar::validity_bit(validity, i) {
2121 self.advance_auto_inc_past(*slot)?;
2122 } else {
2123 *slot = self.alloc_auto_inc_value()?;
2124 }
2125 }
2126 *validity = new_validity;
2127 Ok(())
2128 }
2129
2130 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2144 self.ensure_writable()?;
2145 if self.auto_inc.is_none() {
2146 return Ok(None);
2147 }
2148 Ok(Some(self.alloc_auto_inc_value()?))
2149 }
2150
2151 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2157 let payload = bincode::serialize(&rows)?;
2158 self.wal_append_data(Op::Put {
2159 table_id: self.table_id,
2160 rows: payload,
2161 })?;
2162 if self.is_shared() {
2163 self.pending_rows_auto_inc
2164 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2165 self.pending_rows.extend(rows);
2166 } else {
2167 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2168 }
2169 Ok(())
2170 }
2171
2172 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2177 self.apply_put_rows_inner(rows, true)
2178 }
2179
2180 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2181 if check_existing_pk {
2182 self.ensure_indexes_complete()?;
2183 }
2184 if rows.len() == 1 {
2188 let row = rows.into_iter().next().expect("len checked");
2189 return self.apply_put_row_single(row, check_existing_pk);
2190 }
2191 let pk_id = self.schema.primary_key().map(|c| c.id);
2208 let probe = match pk_id {
2209 Some(pid) => {
2210 check_existing_pk
2211 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2212 }
2213 None => false,
2214 };
2215 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2218 for r in rows {
2219 for &cid in r.columns.keys() {
2220 self.pending_put_cols.insert(cid);
2221 }
2222 match pk_id {
2223 Some(pid) if probe || maintain_pk_by_row => {
2224 if let Some(pk_val) = r.columns.get(&pid) {
2225 let key = self.index_lookup_key(pid, pk_val);
2226 if probe {
2227 if let Some(old_rid) = self.hot.get(&key) {
2228 if old_rid != r.row_id {
2229 self.tombstone_row(old_rid, r.committed_epoch, true);
2230 }
2231 }
2232 }
2233 if maintain_pk_by_row {
2234 self.pk_by_row.insert(r.row_id, key);
2235 }
2236 }
2237 }
2238 Some(_) => {}
2239 None => {
2240 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2241 }
2242 }
2243 self.index_row(&r);
2244 self.reservoir.offer(r.row_id.0);
2245 self.memtable.upsert(r);
2246 self.live_count = self.live_count.saturating_add(1);
2249 }
2250 Ok(())
2251 }
2252
2253 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2257 for &cid in row.columns.keys() {
2258 self.pending_put_cols.insert(cid);
2259 }
2260 let epoch = row.committed_epoch;
2261 if let Some(pk_col) = self.schema.primary_key() {
2262 let pk_id = pk_col.id;
2263 if let Some(pk_val) = row.columns.get(&pk_id) {
2264 let maintain_pk_by_row = self.pk_by_row_complete;
2268 if check_existing_pk || maintain_pk_by_row {
2269 let key = self.index_lookup_key(pk_id, pk_val);
2270 if check_existing_pk {
2271 if let Some(old_rid) = self.hot.get(&key) {
2272 if old_rid != row.row_id {
2273 self.tombstone_row(old_rid, epoch, true);
2274 }
2275 }
2276 }
2277 if maintain_pk_by_row {
2278 self.pk_by_row.insert(row.row_id, key);
2279 }
2280 }
2281 }
2282 } else {
2283 self.hot
2284 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2285 }
2286 self.index_row(&row);
2287 self.reservoir.offer(row.row_id.0);
2288 self.memtable.upsert(row);
2289 self.live_count = self.live_count.saturating_add(1);
2290 Ok(())
2291 }
2292
2293 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2296 self.allocator.alloc()
2297 }
2298
2299 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2305 let pk = self.schema.primary_key().ok_or_else(|| {
2306 MongrelError::Schema("clustered table requires a single-column primary key".into())
2307 })?;
2308 let pk_val = columns
2309 .iter()
2310 .find(|(id, _)| *id == pk.id)
2311 .map(|(_, v)| v)
2312 .ok_or_else(|| {
2313 MongrelError::Schema(format!(
2314 "clustered table missing primary key column {} ({})",
2315 pk.id, pk.name
2316 ))
2317 })?;
2318 let key_bytes = pk_val.encode_key();
2319 let mut hash: u64 = 0xcbf29ce484222325;
2321 for &b in &key_bytes {
2322 hash ^= b as u64;
2323 hash = hash.wrapping_mul(0x100000001b3);
2324 }
2325 Ok(RowId(hash.max(1)))
2328 }
2329
2330 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2338 self.ensure_indexes_complete()?;
2339 let n = rows.len();
2340 for r in rows {
2341 for &cid in r.columns.keys() {
2342 self.pending_put_cols.insert(cid);
2343 }
2344 }
2345 let (losers, winner_pks) = self.partition_pk_winners(rows);
2346 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2347 for (key, &row_id) in &winner_pks {
2349 if let Some(old_rid) = self.hot.get(key) {
2350 if old_rid != row_id {
2351 self.tombstone_row(old_rid, epoch, true);
2352 }
2353 }
2354 }
2355 for &loser_rid in &losers {
2358 self.tombstone_row(loser_rid, epoch, false);
2359 }
2360 for (key, row_id) in winner_pks {
2362 self.insert_hot_pk(key, row_id);
2363 }
2364 if self.schema.primary_key().is_none() {
2365 for r in rows {
2366 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2367 }
2368 }
2369 for r in rows {
2370 self.allocator.advance_to(r.row_id);
2371 if !losers.contains(&r.row_id) {
2372 self.index_row(r);
2373 }
2374 }
2375 for r in rows {
2376 if !losers.contains(&r.row_id) {
2377 self.reservoir.offer(r.row_id.0);
2378 }
2379 }
2380 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2381 Ok(())
2382 }
2383
2384 pub(crate) fn recover_apply(
2389 &mut self,
2390 rows: Vec<Row>,
2391 deletes: Vec<(RowId, Epoch)>,
2392 ) -> Result<()> {
2393 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2397 std::collections::BTreeMap::new();
2398 for row in rows {
2399 self.allocator.advance_to(row.row_id);
2400 if let Some(ai) = self.auto_inc.as_mut() {
2405 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2406 if *n + 1 > ai.next {
2407 ai.next = *n + 1;
2408 }
2409 }
2410 }
2411 by_epoch.entry(row.committed_epoch).or_default().push(row);
2412 }
2413 for (epoch, group) in by_epoch {
2414 let (losers, winner_pks) = self.partition_pk_winners(&group);
2415 for (key, &row_id) in &winner_pks {
2417 if let Some(old_rid) = self.hot.get(key) {
2418 if old_rid != row_id {
2419 self.tombstone_row(old_rid, epoch, false);
2420 }
2421 }
2422 }
2423 for (key, row_id) in winner_pks {
2424 self.insert_hot_pk(key, row_id);
2425 }
2426 if self.schema.primary_key().is_none() {
2427 for r in &group {
2428 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2429 }
2430 }
2431 for r in &group {
2432 if !losers.contains(&r.row_id) {
2433 self.memtable.upsert(r.clone());
2434 self.index_row(r);
2435 }
2436 }
2437 }
2438 for (rid, epoch) in deletes {
2439 self.memtable.tombstone(rid, epoch);
2440 self.remove_hot_for_row(rid, epoch);
2441 }
2442 self.reservoir_complete = false;
2445 Ok(())
2446 }
2447
2448 pub(crate) fn flushed_epoch(&self) -> u64 {
2450 self.flushed_epoch
2451 }
2452
2453 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2454 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2455 }
2456
2457 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2459 self.schema.validate_values(cells)
2460 }
2461
2462 fn validate_columns_not_null(
2466 &self,
2467 columns: &[(u16, columnar::NativeColumn)],
2468 n: usize,
2469 ) -> Result<()> {
2470 let by_id: HashMap<u16, &columnar::NativeColumn> =
2471 columns.iter().map(|(id, c)| (*id, c)).collect();
2472 for col in &self.schema.columns {
2473 if !col.flags.contains(ColumnFlags::NULLABLE) {
2474 match by_id.get(&col.id) {
2475 None => {
2476 return Err(MongrelError::InvalidArgument(format!(
2477 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2478 col.name, col.id
2479 )));
2480 }
2481 Some(c) => {
2482 if c.null_count(n) != 0 {
2483 return Err(MongrelError::InvalidArgument(format!(
2484 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2485 col.name, col.id
2486 )));
2487 }
2488 }
2489 }
2490 }
2491 if let TypeId::Enum { variants } = &col.ty {
2492 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2493 if by_id.contains_key(&col.id) {
2494 return Err(MongrelError::InvalidArgument(format!(
2495 "column '{}' ({}) enum requires a bytes column",
2496 col.name, col.id
2497 )));
2498 }
2499 continue;
2500 };
2501 for index in 0..n {
2502 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2503 continue;
2504 };
2505 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2506 return Err(MongrelError::InvalidArgument(format!(
2507 "column '{}' ({}) enum value {:?} is not one of {:?}",
2508 col.name,
2509 col.id,
2510 String::from_utf8_lossy(value),
2511 variants
2512 )));
2513 }
2514 }
2515 }
2516 }
2517 Ok(())
2518 }
2519
2520 fn bulk_pk_winner_indices(
2525 &self,
2526 columns: &[(u16, columnar::NativeColumn)],
2527 n: usize,
2528 ) -> Option<Vec<usize>> {
2529 let pk_col = self.schema.primary_key()?;
2530 let pk_id = pk_col.id;
2531 let pk_ty = pk_col.ty.clone();
2532 let by_id: HashMap<u16, &columnar::NativeColumn> =
2533 columns.iter().map(|(id, c)| (*id, c)).collect();
2534 let pk_native = by_id.get(&pk_id)?;
2535 if native_int64_strictly_increasing(pk_native, n) {
2536 return None;
2537 }
2538 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2540 let mut null_pk_rows: Vec<usize> = Vec::new();
2541 for i in 0..n {
2542 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2543 Some(key) => {
2544 last.insert(key, i);
2545 }
2546 None => null_pk_rows.push(i),
2547 }
2548 }
2549 let mut winners: HashSet<usize> = last.values().copied().collect();
2550 for i in null_pk_rows {
2551 winners.insert(i);
2552 }
2553 Some((0..n).filter(|i| winners.contains(i)).collect())
2554 }
2555
2556 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2558 self.require_delete()?;
2559 let epoch = self.pending_epoch();
2560 self.wal_append_data(Op::Delete {
2561 table_id: self.table_id,
2562 row_ids: vec![row_id],
2563 })?;
2564 if self.is_shared() {
2565 self.pending_dels.push(row_id);
2566 } else {
2567 self.apply_delete(row_id, epoch);
2568 }
2569 Ok(())
2570 }
2571
2572 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2573 let pre = self.get(row_id, self.snapshot());
2574 self.delete(row_id)?;
2575 Ok(pre.map(|row| {
2576 let mut columns: Vec<_> = row.columns.into_iter().collect();
2577 columns.sort_by_key(|(id, _)| *id);
2578 OwnedRow { columns }
2579 }))
2580 }
2581
2582 pub fn truncate(&mut self) -> Result<()> {
2584 self.require_delete()?;
2585 let epoch = self.pending_epoch();
2586 self.wal_append_data(Op::TruncateTable {
2587 table_id: self.table_id,
2588 })?;
2589 self.pending_rows.clear();
2590 self.pending_rows_auto_inc.clear();
2591 self.pending_dels.clear();
2592 self.pending_truncate = Some(epoch);
2593 Ok(())
2594 }
2595
2596 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2598 for rr in std::mem::take(&mut self.run_refs) {
2599 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2600 }
2601 for r in std::mem::take(&mut self.retiring) {
2602 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2603 }
2604 self.memtable = Memtable::new();
2605 self.mutable_run = MutableRun::new();
2606 self.hot = HotIndex::new();
2607 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2608 self.bitmap = bitmap;
2609 self.ann = ann;
2610 self.fm = fm;
2611 self.sparse = sparse;
2612 self.minhash = minhash;
2613 self.learned_range.clear();
2614 self.pk_by_row.clear();
2615 self.pk_by_row_complete = false;
2616 self.live_count = 0;
2617 self.reservoir = crate::reservoir::Reservoir::default();
2618 self.reservoir_complete = true;
2619 self.had_deletes = true;
2620 self.agg_cache.clear();
2621 self.global_idx_epoch = 0;
2622 self.indexes_complete = true;
2623 self.pending_delete_rids.clear();
2624 self.pending_put_cols.clear();
2625 self.pending_rows.clear();
2626 self.pending_rows_auto_inc.clear();
2627 self.pending_dels.clear();
2628 self.clear_result_cache();
2629 self.invalidate_index_checkpoint();
2630 Ok(())
2631 }
2632
2633 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2636 self.remove_hot_for_row(row_id, epoch);
2637 self.tombstone_row(row_id, epoch, true);
2638 }
2639
2640 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2644 let tombstone = Row {
2645 row_id,
2646 committed_epoch: epoch,
2647 columns: std::collections::HashMap::new(),
2648 deleted: true,
2649 };
2650 self.memtable.upsert(tombstone);
2651 self.pk_by_row.remove(&row_id);
2652 if adjust_live_count {
2653 self.live_count = self.live_count.saturating_sub(1);
2654 }
2655 self.pending_delete_rids.insert(row_id.0 as u32);
2657 self.had_deletes = true;
2660 self.agg_cache.clear();
2661 }
2662
2663 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2667 let Some(pk_col) = self.schema.primary_key() else {
2668 return;
2669 };
2670 if self.pk_by_row_complete {
2673 if let Some(key) = self.pk_by_row.remove(&row_id) {
2674 if self.hot.get(&key) == Some(row_id) {
2675 self.hot.remove(&key);
2676 }
2677 }
2678 return;
2679 }
2680 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2699 if self.indexes_complete {
2700 let pk_val = self
2701 .memtable
2702 .get_version(row_id, lookup_epoch)
2703 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2704 .or_else(|| {
2705 self.mutable_run
2706 .get_version(row_id, lookup_epoch)
2707 .filter(|(_, r)| !r.deleted)
2708 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2709 })
2710 .or_else(|| {
2711 self.run_refs.iter().find_map(|rr| {
2712 let mut reader = self.open_reader(rr.run_id).ok()?;
2713 let (_, deleted, val) = reader
2714 .get_version_column(row_id, lookup_epoch, pk_col.id)
2715 .ok()??;
2716 if deleted {
2717 return None;
2718 }
2719 val
2720 })
2721 });
2722 if let Some(pk_val) = pk_val {
2723 let key = self.index_lookup_key(pk_col.id, &pk_val);
2724 if self.hot.get(&key) == Some(row_id) {
2725 self.hot.remove(&key);
2726 }
2727 return;
2728 }
2729 }
2730 self.refresh_pk_by_row_from_hot();
2735 if let Some(key) = self.pk_by_row.remove(&row_id) {
2736 if self.hot.get(&key) == Some(row_id) {
2737 self.hot.remove(&key);
2738 }
2739 }
2740 }
2741
2742 fn partition_pk_winners(
2747 &self,
2748 rows: &[Row],
2749 ) -> (
2750 std::collections::HashSet<RowId>,
2751 std::collections::HashMap<Vec<u8>, RowId>,
2752 ) {
2753 let mut losers = std::collections::HashSet::new();
2754 let Some(pk_col) = self.schema.primary_key() else {
2755 return (losers, std::collections::HashMap::new());
2756 };
2757 let pk_id = pk_col.id;
2758 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2759 std::collections::HashMap::new();
2760 for r in rows {
2761 let Some(pk_val) = r.columns.get(&pk_id) else {
2762 continue;
2763 };
2764 let key = self.index_lookup_key(pk_id, pk_val);
2765 if let Some(&old_rid) = winners.get(&key) {
2766 losers.insert(old_rid);
2767 }
2768 winners.insert(key, r.row_id);
2769 }
2770 (losers, winners)
2771 }
2772
2773 fn index_row(&mut self, row: &Row) {
2774 if row.deleted {
2775 return;
2776 }
2777 let any_predicate = self
2785 .schema
2786 .indexes
2787 .iter()
2788 .any(|idx| idx.predicate.is_some());
2789 if any_predicate {
2790 let columns_map: HashMap<u16, &Value> =
2791 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2792 let name_to_id: HashMap<&str, u16> = self
2793 .schema
2794 .columns
2795 .iter()
2796 .map(|c| (c.name.as_str(), c.id))
2797 .collect();
2798 for idx in &self.schema.indexes {
2799 if let Some(pred) = &idx.predicate {
2800 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2801 continue; }
2803 }
2804 index_into_single(
2806 idx,
2807 &self.schema,
2808 row,
2809 &mut self.hot,
2810 &mut self.bitmap,
2811 &mut self.ann,
2812 &mut self.fm,
2813 &mut self.sparse,
2814 &mut self.minhash,
2815 );
2816 }
2817 return;
2818 }
2819 if self.column_keys.is_empty() {
2823 index_into(
2824 &self.schema,
2825 row,
2826 &mut self.hot,
2827 &mut self.bitmap,
2828 &mut self.ann,
2829 &mut self.fm,
2830 &mut self.sparse,
2831 &mut self.minhash,
2832 );
2833 return;
2834 }
2835 let effective_row = self.tokenized_for_indexes(row);
2836 index_into(
2837 &self.schema,
2838 &effective_row,
2839 &mut self.hot,
2840 &mut self.bitmap,
2841 &mut self.ann,
2842 &mut self.fm,
2843 &mut self.sparse,
2844 &mut self.minhash,
2845 );
2846 }
2847
2848 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2854 if self.column_keys.is_empty() {
2855 return row.clone();
2856 }
2857 #[cfg(feature = "encryption")]
2858 {
2859 use crate::encryption::SCHEME_HMAC_EQ;
2860 let mut tok = row.clone();
2861 for (&cid, &(_, scheme)) in &self.column_keys {
2862 if scheme != SCHEME_HMAC_EQ {
2863 continue;
2864 }
2865 if let Some(v) = tok.columns.get(&cid).cloned() {
2866 if let Some(t) = self.tokenize_value(cid, &v) {
2867 tok.columns.insert(cid, t);
2868 }
2869 }
2870 }
2871 tok
2872 }
2873 #[cfg(not(feature = "encryption"))]
2874 {
2875 row.clone()
2876 }
2877 }
2878
2879 pub fn commit(&mut self) -> Result<Epoch> {
2884 self.ensure_writable()?;
2885 if self.is_shared() {
2886 self.commit_shared()
2887 } else {
2888 self.commit_private()
2889 }
2890 }
2891
2892 fn commit_private(&mut self) -> Result<Epoch> {
2894 let commit_lock = Arc::clone(&self.commit_lock);
2898 let _g = commit_lock.lock();
2899 let new_epoch = self.epoch.bump_assigned();
2900 let txn_id = self.current_txn_id;
2901 match &mut self.wal {
2905 WalSink::Private(w) => {
2906 w.append_txn(
2907 txn_id,
2908 Op::TxnCommit {
2909 epoch: new_epoch.0,
2910 added_runs: Vec::new(),
2911 },
2912 )?;
2913 w.sync()?;
2914 }
2915 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2916 }
2917 if let Some(epoch) = self.pending_truncate.take() {
2919 self.apply_truncate(epoch)?;
2920 }
2921 self.invalidate_pending_cache();
2922 self.persist_manifest(new_epoch)?;
2923 self.epoch.publish_in_order(new_epoch);
2927 self.current_txn_id += 1;
2928 Ok(new_epoch)
2929 }
2930
2931 fn commit_shared(&mut self) -> Result<Epoch> {
2937 use std::sync::atomic::Ordering;
2938 let s = match &self.wal {
2939 WalSink::Shared(s) => s.clone(),
2940 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2941 };
2942 if s.poisoned.load(Ordering::Relaxed) {
2943 return Err(MongrelError::Other(
2944 "database poisoned by fsync error".into(),
2945 ));
2946 }
2947 let commit_lock = Arc::clone(&self.commit_lock);
2954 let _g = commit_lock.lock();
2955 let txn_id = self.ensure_txn_id();
2958 let (new_epoch, commit_seq) = {
2959 let mut wal = s.wal.lock();
2960 let new_epoch = self.epoch.bump_assigned();
2961 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2962 (new_epoch, seq)
2963 };
2964 s.group
2965 .await_durable(&s.wal, commit_seq)
2966 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2967
2968 if self.pending_truncate.take().is_some() {
2971 self.apply_truncate(new_epoch)?;
2972 }
2973 let mut rows = std::mem::take(&mut self.pending_rows);
2974 if !rows.is_empty() {
2975 for r in &mut rows {
2976 r.committed_epoch = new_epoch;
2977 }
2978 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2979 let all_auto_generated =
2980 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2981 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2982 } else {
2983 self.pending_rows_auto_inc.clear();
2984 }
2985 let dels = std::mem::take(&mut self.pending_dels);
2986 for rid in dels {
2987 self.apply_delete(rid, new_epoch);
2988 }
2989
2990 self.invalidate_pending_cache();
2991 self.persist_manifest(new_epoch)?;
2992 self.epoch.publish_in_order(new_epoch);
2993 let _ = s.change_wake.send(());
2994 self.current_txn_id = 0;
2996 Ok(new_epoch)
2997 }
2998
2999 pub fn flush(&mut self) -> Result<Epoch> {
3007 self.ensure_indexes_complete()?;
3008 let epoch = self.commit()?;
3009 let rows = self.memtable.drain_sorted();
3010 if !rows.is_empty() {
3011 self.mutable_run.insert_many(rows);
3012 }
3013 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3014 self.spill_mutable_run(epoch)?;
3015 self.mark_flushed(epoch)?;
3019 self.persist_manifest(epoch)?;
3020 self.build_learned_ranges()?;
3021 self.checkpoint_indexes(epoch);
3024 }
3025 Ok(epoch)
3028 }
3029
3030 pub fn force_flush(&mut self) -> Result<Epoch> {
3039 let saved = self.mutable_run_spill_bytes;
3040 self.mutable_run_spill_bytes = 1;
3041 let result = self.flush();
3042 self.mutable_run_spill_bytes = saved;
3043 result
3044 }
3045
3046 pub fn close(&mut self) -> Result<()> {
3053 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3054 self.force_flush()?;
3055 }
3056 Ok(())
3057 }
3058
3059 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3066 let op = Op::Flush {
3067 table_id: self.table_id,
3068 flushed_epoch: epoch.0,
3069 };
3070 match &mut self.wal {
3071 WalSink::Private(w) => {
3072 w.append_system(op)?;
3073 w.sync()?;
3074 }
3075 WalSink::Shared(s) => {
3076 s.wal.lock().append_system(op)?;
3081 }
3082 }
3083 self.flushed_epoch = epoch.0;
3084 if matches!(self.wal, WalSink::Private(_)) {
3085 self.rotate_wal(epoch)?;
3086 }
3087 Ok(())
3088 }
3089
3090 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3094 let rows = self.mutable_run.drain_sorted();
3095 if rows.is_empty() {
3096 return Ok(());
3097 }
3098 let run_id = self.next_run_id;
3099 self.next_run_id += 1;
3100 let path = self.run_path(run_id);
3101 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3102 if let Some(kek) = &self.kek {
3103 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3104 }
3105 let header = writer.write(&path, &rows)?;
3106 self.run_refs.push(RunRef {
3107 run_id: run_id as u128,
3108 level: 0,
3109 epoch_created: epoch.0,
3110 row_count: header.row_count,
3111 });
3112 Ok(())
3113 }
3114
3115 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3119 self.mutable_run_spill_bytes = bytes.max(1);
3120 }
3121
3122 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3126 self.compaction_zstd_level = level;
3127 }
3128
3129 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3133 self.result_cache.lock().set_max_bytes(max_bytes);
3134 }
3135
3136 pub(crate) fn clear_result_cache(&mut self) {
3140 self.result_cache.lock().clear();
3141 }
3142
3143 pub fn mutable_run_len(&self) -> usize {
3145 self.mutable_run.len()
3146 }
3147
3148 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3151 self.mutable_run.drain_sorted()
3152 }
3153
3154 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3159 let epoch = self.commit()?;
3160 let n = batch.len();
3161 if n == 0 {
3162 return Ok(epoch);
3163 }
3164 for row in &batch {
3165 self.schema.validate_values(row)?;
3166 }
3167 let live_before = self.live_count;
3168 self.spill_mutable_run(epoch)?;
3172 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3173 && self.indexes_complete
3174 && self.run_refs.is_empty()
3175 && self.memtable.is_empty()
3176 && self.mutable_run.is_empty();
3177 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3183 use rayon::prelude::*;
3184 self.schema
3185 .columns
3186 .par_iter()
3187 .map(|cdef| {
3188 (
3189 cdef.id,
3190 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3191 )
3192 })
3193 .collect::<Vec<_>>()
3194 };
3195 drop(batch);
3196 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3201 self.validate_columns_not_null(&user_columns, n)?;
3202 let winner_idx = self
3203 .bulk_pk_winner_indices(&user_columns, n)
3204 .filter(|idx| idx.len() != n);
3205 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3206 match winner_idx.as_deref() {
3207 Some(idx) => {
3208 let compacted = user_columns
3209 .iter()
3210 .map(|(id, c)| (*id, c.gather(idx)))
3211 .collect();
3212 (compacted, idx.len())
3213 }
3214 None => (std::mem::take(&mut user_columns), n),
3215 };
3216 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3217 let first = self.allocator.alloc_range(write_n as u64).0;
3218 for rid in first..first + write_n as u64 {
3219 self.reservoir.offer(rid);
3220 }
3221 let run_id = self.next_run_id;
3222 self.next_run_id += 1;
3223 let path = self.run_path(run_id);
3224 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3225 .clean(true)
3226 .with_lz4()
3227 .with_native_endian();
3228 if let Some(kek) = &self.kek {
3229 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3230 }
3231 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3232 self.run_refs.push(RunRef {
3233 run_id: run_id as u128,
3234 level: 0,
3235 epoch_created: epoch.0,
3236 row_count: header.row_count,
3237 });
3238 self.live_count = self.live_count.saturating_add(write_n as u64);
3239 if eager_index_build {
3240 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3241 self.index_columns_bulk(&write_columns, &row_ids);
3242 self.indexes_complete = true;
3243 self.build_learned_ranges()?;
3244 } else {
3245 self.indexes_complete = false;
3246 }
3247 self.mark_flushed(epoch)?;
3248 self.persist_manifest(epoch)?;
3249 if eager_index_build {
3250 self.checkpoint_indexes(epoch);
3251 }
3252 self.clear_result_cache();
3253 Ok(epoch)
3254 }
3255
3256 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3259 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3260 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3261 let segment_no = segment
3264 .file_stem()
3265 .and_then(|s| s.to_str())
3266 .and_then(|s| s.strip_prefix("seg-"))
3267 .and_then(|s| s.parse::<u64>().ok())
3268 .unwrap_or(0);
3269 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3270 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3271 wal.sync()?;
3272 self.wal = WalSink::Private(wal);
3273 Ok(())
3274 }
3275
3276 pub(crate) fn invalidate_pending_cache(&mut self) {
3281 self.result_cache
3282 .lock()
3283 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3284 self.pending_delete_rids.clear();
3285 self.pending_put_cols.clear();
3286 }
3287
3288 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3289 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3290 m.current_epoch = epoch.0;
3291 m.next_row_id = self.allocator.current().0;
3292 m.runs = self.run_refs.clone();
3293 m.live_count = self.live_count;
3294 m.global_idx_epoch = self.global_idx_epoch;
3295 m.flushed_epoch = self.flushed_epoch;
3296 m.retiring = self.retiring.clone();
3297 m.auto_inc_next = match self.auto_inc {
3301 Some(ai) if ai.seeded => ai.next,
3302 _ => 0,
3303 };
3304 m.ttl = self.ttl;
3305 let meta_dek = self.manifest_meta_dek();
3306 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3307 Ok(())
3308 }
3309
3310 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3316 if !self.indexes_complete {
3319 return;
3320 }
3321 let snap = global_idx::IndexSnapshot {
3322 hot: &self.hot,
3323 bitmap: &self.bitmap,
3324 ann: &self.ann,
3325 fm: &self.fm,
3326 sparse: &self.sparse,
3327 minhash: &self.minhash,
3328 learned_range: &self.learned_range,
3329 };
3330 let idx_dek = self.idx_dek();
3332 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3333 .is_ok()
3334 {
3335 self.global_idx_epoch = epoch.0;
3336 let _ = self.persist_manifest(epoch);
3337 }
3338 }
3339
3340 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3343 self.global_idx_epoch = 0;
3344 global_idx::remove(&self.dir);
3345 let _ = self.persist_manifest(self.epoch.visible());
3346 }
3347
3348 pub(crate) fn mark_indexes_incomplete(&mut self) {
3349 self.indexes_complete = false;
3350 self.invalidate_index_checkpoint();
3351 }
3352
3353 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3356 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3357 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3358 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3359 best = Some((epoch, row));
3360 }
3361 }
3362 for rr in &self.run_refs {
3363 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3364 continue;
3365 };
3366 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3367 continue;
3368 };
3369 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3370 best = Some((epoch, row));
3371 }
3372 }
3373 let now_nanos = unix_nanos_now();
3374 match best {
3375 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3376 Some((_, r)) => Some(r),
3377 None => None,
3378 }
3379 }
3380
3381 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3385 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3386 let mut fold = |row: Row| {
3387 best.entry(row.row_id.0)
3388 .and_modify(|e| {
3389 if row.committed_epoch > e.0 {
3390 *e = (row.committed_epoch, row.clone());
3391 }
3392 })
3393 .or_insert_with(|| (row.committed_epoch, row));
3394 };
3395 for row in self.memtable.visible_versions(snapshot.epoch) {
3396 fold(row);
3397 }
3398 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3399 fold(row);
3400 }
3401 for rr in &self.run_refs {
3402 let mut reader = self.open_reader(rr.run_id)?;
3403 for row in reader.visible_versions(snapshot.epoch)? {
3404 fold(row);
3405 }
3406 }
3407 let now_nanos = unix_nanos_now();
3408 let mut out: Vec<Row> = best
3409 .into_values()
3410 .filter_map(|(_, r)| {
3411 if r.deleted || self.row_expired_at(&r, now_nanos) {
3412 None
3413 } else {
3414 Some(r)
3415 }
3416 })
3417 .collect();
3418 out.sort_by_key(|r| r.row_id);
3419 Ok(out)
3420 }
3421
3422 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3429 if self.ttl.is_none()
3430 && self.memtable.is_empty()
3431 && self.mutable_run.is_empty()
3432 && self.run_refs.len() == 1
3433 {
3434 let rr = self.run_refs[0].clone();
3435 let mut reader = self.open_reader(rr.run_id)?;
3436 let idxs = reader.visible_indices(snapshot.epoch)?;
3437 let mut cols = Vec::with_capacity(self.schema.columns.len());
3438 for cdef in &self.schema.columns {
3439 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3440 }
3441 return Ok(cols);
3442 }
3443 let rows = self.visible_rows(snapshot)?;
3445 let mut cols: Vec<(u16, Vec<Value>)> = self
3446 .schema
3447 .columns
3448 .iter()
3449 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3450 .collect();
3451 for r in &rows {
3452 for (cid, vec) in cols.iter_mut() {
3453 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3454 }
3455 }
3456 Ok(cols)
3457 }
3458
3459 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3461 let row_id = self.hot.get(key)?;
3462 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3463 Some(row_id)
3464 } else {
3465 None
3466 }
3467 }
3468
3469 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3474 self.require_select()?;
3475 self.ensure_indexes_complete()?;
3476 let snapshot = self.snapshot();
3477 crate::trace::QueryTrace::record(|t| {
3478 t.run_count = self.run_refs.len();
3479 t.memtable_rows = self.memtable.len();
3480 t.mutable_run_rows = self.mutable_run.len();
3481 });
3482 if q.conditions.is_empty() {
3486 crate::trace::QueryTrace::record(|t| {
3487 t.scan_mode = crate::trace::ScanMode::Materialized;
3488 t.row_materialized = true;
3489 });
3490 return self.visible_rows(snapshot);
3491 }
3492 crate::trace::QueryTrace::record(|t| {
3493 t.conditions_pushed = q.conditions.len();
3494 t.scan_mode = crate::trace::ScanMode::Materialized;
3495 t.row_materialized = true;
3496 });
3497 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3504 ordered.sort_by_key(|c| condition_cost_rank(c));
3505 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3506 for c in &ordered {
3507 let s = self.resolve_condition(c, snapshot)?;
3508 let empty = s.is_empty();
3509 sets.push(s);
3510 if empty {
3511 break;
3512 }
3513 }
3514 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3515 self.rows_for_rids(&rids, snapshot)
3516 }
3517
3518 pub fn retrieve(
3520 &mut self,
3521 retriever: &crate::query::Retriever,
3522 ) -> Result<Vec<crate::query::RetrieverHit>> {
3523 self.retrieve_with_allowed(retriever, None)
3524 }
3525
3526 pub fn retrieve_with_allowed(
3529 &mut self,
3530 retriever: &crate::query::Retriever,
3531 allowed: Option<&std::collections::HashSet<RowId>>,
3532 ) -> Result<Vec<crate::query::RetrieverHit>> {
3533 self.require_select()?;
3534 self.ensure_indexes_complete()?;
3535 self.validate_retriever(retriever)?;
3536 self.retrieve_filtered(retriever, self.snapshot(), None, allowed)
3537 }
3538
3539 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3540 use crate::query::Retriever;
3541 let (column_id, k) = match retriever {
3542 Retriever::Ann {
3543 column_id,
3544 query,
3545 k,
3546 } => {
3547 let index = self.ann.get(column_id).ok_or_else(|| {
3548 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3549 })?;
3550 if query.len() != index.dim() {
3551 return Err(MongrelError::InvalidArgument(format!(
3552 "ANN query dimension must be {}, got {}",
3553 index.dim(),
3554 query.len()
3555 )));
3556 }
3557 if query.iter().any(|value| !value.is_finite()) {
3558 return Err(MongrelError::InvalidArgument(
3559 "ANN query values must be finite".into(),
3560 ));
3561 }
3562 (*column_id, *k)
3563 }
3564 Retriever::Sparse {
3565 column_id,
3566 query,
3567 k,
3568 } => {
3569 if !self.sparse.contains_key(column_id) {
3570 return Err(MongrelError::InvalidArgument(format!(
3571 "column {column_id} has no Sparse index"
3572 )));
3573 }
3574 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3575 return Err(MongrelError::InvalidArgument(
3576 "Sparse query must be non-empty with finite weights".into(),
3577 ));
3578 }
3579 (*column_id, *k)
3580 }
3581 Retriever::MinHash {
3582 column_id,
3583 members,
3584 k,
3585 } => {
3586 if !self.minhash.contains_key(column_id) {
3587 return Err(MongrelError::InvalidArgument(format!(
3588 "column {column_id} has no MinHash index"
3589 )));
3590 }
3591 if members.is_empty() {
3592 return Err(MongrelError::InvalidArgument(
3593 "MinHash members must not be empty".into(),
3594 ));
3595 }
3596 (*column_id, *k)
3597 }
3598 };
3599 if k == 0 {
3600 return Err(MongrelError::InvalidArgument(
3601 "retriever k must be > 0".into(),
3602 ));
3603 }
3604 debug_assert!(self
3605 .schema
3606 .columns
3607 .iter()
3608 .any(|column| column.id == column_id));
3609 Ok(())
3610 }
3611
3612 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3613 use crate::query::Condition;
3614 match condition {
3615 Condition::Ann {
3616 column_id,
3617 query,
3618 k,
3619 } => self.validate_retriever(&crate::query::Retriever::Ann {
3620 column_id: *column_id,
3621 query: query.clone(),
3622 k: *k,
3623 }),
3624 Condition::SparseMatch {
3625 column_id,
3626 query,
3627 k,
3628 } => self.validate_retriever(&crate::query::Retriever::Sparse {
3629 column_id: *column_id,
3630 query: query.clone(),
3631 k: *k,
3632 }),
3633 Condition::MinHashSimilar {
3634 column_id,
3635 query,
3636 k,
3637 } => {
3638 if !self.minhash.contains_key(column_id) {
3639 return Err(MongrelError::InvalidArgument(format!(
3640 "column {column_id} has no MinHash index"
3641 )));
3642 }
3643 if query.is_empty() || *k == 0 {
3644 return Err(MongrelError::InvalidArgument(
3645 "MinHash query must be non-empty and k must be > 0".into(),
3646 ));
3647 }
3648 Ok(())
3649 }
3650 _ => Ok(()),
3651 }
3652 }
3653
3654 fn retrieve_filtered(
3655 &self,
3656 retriever: &crate::query::Retriever,
3657 snapshot: Snapshot,
3658 hard_filter: Option<&RowIdSet>,
3659 allowed: Option<&std::collections::HashSet<RowId>>,
3660 ) -> Result<Vec<crate::query::RetrieverHit>> {
3661 use crate::query::{Retriever, RetrieverHit, RetrieverScore, SetMember};
3662 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
3663 Retriever::Ann {
3664 column_id,
3665 query,
3666 k,
3667 } => {
3668 let Some(index) = self.ann.get(column_id) else {
3669 return Ok(Vec::new());
3670 };
3671 let cap = index.len();
3672 if cap == 0 {
3673 return Ok(Vec::new());
3674 }
3675 let mut breadth = (*k).max(1).min(cap);
3676 let mut eligibility = std::collections::HashMap::new();
3677 let mut filtered = loop {
3678 let mut seen = std::collections::HashSet::new();
3679 let raw = index.search(query, breadth)?;
3680 let unchecked: Vec<_> = raw
3681 .iter()
3682 .map(|(row_id, _)| *row_id)
3683 .filter(|row_id| !eligibility.contains_key(row_id))
3684 .filter(|row_id| {
3685 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3686 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3687 })
3688 .collect();
3689 let eligible = self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3690 for row_id in unchecked {
3691 eligibility.insert(row_id, eligible.contains(&row_id));
3692 }
3693 let filtered: Vec<_> = raw
3694 .into_iter()
3695 .filter(|(row_id, _)| {
3696 seen.insert(*row_id)
3697 && eligibility.get(row_id).copied().unwrap_or(false)
3698 })
3699 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
3700 .collect();
3701 if filtered.len() >= *k || breadth >= cap {
3702 break filtered;
3703 }
3704 breadth = breadth.saturating_mul(2).min(cap);
3705 };
3706 filtered.truncate(*k);
3707 filtered
3708 }
3709 Retriever::Sparse {
3710 column_id,
3711 query,
3712 k,
3713 } => self
3714 .sparse
3715 .get(column_id)
3716 .map(|index| -> Result<Vec<_>> {
3717 let mut breadth = (*k).max(1);
3718 let mut eligibility = std::collections::HashMap::new();
3719 loop {
3720 let raw = index.search(query, breadth);
3721 let unchecked: Vec<_> = raw
3722 .iter()
3723 .map(|(row_id, _)| *row_id)
3724 .filter(|row_id| !eligibility.contains_key(row_id))
3725 .filter(|row_id| {
3726 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3727 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3728 })
3729 .collect();
3730 let eligible =
3731 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3732 for row_id in unchecked {
3733 eligibility.insert(row_id, eligible.contains(&row_id));
3734 }
3735 let filtered: Vec<_> = raw
3736 .iter()
3737 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3738 .take(*k)
3739 .map(|(row_id, score)| {
3740 (*row_id, RetrieverScore::SparseDotProduct(*score))
3741 })
3742 .collect();
3743 if filtered.len() >= *k || raw.len() < breadth {
3744 break Ok(filtered);
3745 }
3746 let next = breadth.saturating_mul(2);
3747 if next == breadth {
3748 break Ok(filtered);
3749 }
3750 breadth = next;
3751 }
3752 })
3753 .transpose()?
3754 .unwrap_or_default(),
3755 Retriever::MinHash {
3756 column_id,
3757 members,
3758 k,
3759 } => self
3760 .minhash
3761 .get(column_id)
3762 .map(|index| -> Result<Vec<_>> {
3763 let hashes: Vec<_> = members.iter().map(SetMember::hash_v1).collect();
3764 let mut breadth = (*k).max(1);
3765 let mut eligibility = std::collections::HashMap::new();
3766 loop {
3767 let raw = index.search(&hashes, breadth);
3768 let unchecked: Vec<_> = raw
3769 .iter()
3770 .map(|(row_id, _)| *row_id)
3771 .filter(|row_id| !eligibility.contains_key(row_id))
3772 .filter(|row_id| {
3773 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
3774 && allowed.map_or(true, |allowed| allowed.contains(row_id))
3775 })
3776 .collect();
3777 let eligible =
3778 self.eligible_candidate_ids(&unchecked, *column_id, snapshot)?;
3779 for row_id in unchecked {
3780 eligibility.insert(row_id, eligible.contains(&row_id));
3781 }
3782 let filtered: Vec<_> = raw
3783 .iter()
3784 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
3785 .take(*k)
3786 .map(|(row_id, score)| {
3787 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
3788 })
3789 .collect();
3790 if filtered.len() >= *k || raw.len() < breadth {
3791 break Ok(filtered);
3792 }
3793 let next = breadth.saturating_mul(2);
3794 if next == breadth {
3795 break Ok(filtered);
3796 }
3797 breadth = next;
3798 }
3799 })
3800 .transpose()?
3801 .unwrap_or_default(),
3802 };
3803 Ok(scored
3804 .into_iter()
3805 .enumerate()
3806 .map(|(rank, (row_id, score))| RetrieverHit {
3807 row_id,
3808 rank: rank + 1,
3809 score,
3810 })
3811 .collect())
3812 }
3813
3814 fn eligible_candidate_ids(
3815 &self,
3816 candidates: &[RowId],
3817 _column_id: u16,
3818 snapshot: Snapshot,
3819 ) -> Result<std::collections::HashSet<RowId>> {
3820 if !self.had_deletes && self.ttl.is_none() && snapshot.epoch == self.snapshot().epoch {
3821 return Ok(candidates.iter().copied().collect());
3822 }
3823 let mut readers: Vec<_> = self
3824 .run_refs
3825 .iter()
3826 .map(|run| self.open_reader(run.run_id))
3827 .collect::<Result<_>>()?;
3828 let now = unix_nanos_now();
3829 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
3830 for &row_id in candidates {
3831 let mem = self.memtable.get_version(row_id, snapshot.epoch);
3832 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
3833 let overlay = match (mem, mutable) {
3834 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
3835 (Some(value), None) | (None, Some(value)) => Some(value),
3836 (None, None) => None,
3837 };
3838 if let Some((_, row)) = overlay {
3839 if !row.deleted && !self.row_expired_at(&row, now) {
3840 eligible.insert(row_id);
3841 }
3842 continue;
3843 }
3844 let mut best: Option<(Epoch, bool, usize)> = None;
3845 for (index, reader) in readers.iter_mut().enumerate() {
3846 if let Some((epoch, deleted)) =
3847 reader.get_version_visibility(row_id, snapshot.epoch)?
3848 {
3849 if best
3850 .as_ref()
3851 .map(|(best_epoch, ..)| epoch > *best_epoch)
3852 .unwrap_or(true)
3853 {
3854 best = Some((epoch, deleted, index));
3855 }
3856 }
3857 }
3858 let Some((_, false, reader_index)) = best else {
3859 continue;
3860 };
3861 if let Some(ttl) = self.ttl {
3862 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
3863 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
3864 {
3865 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
3866 continue;
3867 }
3868 }
3869 }
3870 eligible.insert(row_id);
3871 }
3872 Ok(eligible)
3873 }
3874
3875 pub fn search(
3877 &mut self,
3878 request: &crate::query::SearchRequest,
3879 ) -> Result<Vec<crate::query::SearchHit>> {
3880 self.search_with_allowed(request, None)
3881 }
3882
3883 pub fn search_with_allowed(
3884 &mut self,
3885 request: &crate::query::SearchRequest,
3886 authorized: Option<&std::collections::HashSet<RowId>>,
3887 ) -> Result<Vec<crate::query::SearchHit>> {
3888 use crate::query::{ComponentScore, Fusion, SearchHit};
3889 self.require_select()?;
3890 self.ensure_indexes_complete()?;
3891 if request.limit == 0 {
3892 return Err(MongrelError::InvalidArgument(
3893 "search limit must be > 0".into(),
3894 ));
3895 }
3896 if request.retrievers.is_empty() {
3897 return Err(MongrelError::InvalidArgument(
3898 "search requires at least one retriever".into(),
3899 ));
3900 }
3901 let mut names = std::collections::HashSet::new();
3902 for named in &request.retrievers {
3903 if named.name.is_empty() || !names.insert(named.name.as_str()) {
3904 return Err(MongrelError::InvalidArgument(
3905 "retriever names must be non-empty and unique".into(),
3906 ));
3907 }
3908 if !named.weight.is_finite() || named.weight < 0.0 {
3909 return Err(MongrelError::InvalidArgument(
3910 "retriever weight must be finite and non-negative".into(),
3911 ));
3912 }
3913 self.validate_retriever(&named.retriever)?;
3914 }
3915 let projection = request
3916 .projection
3917 .clone()
3918 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
3919 for column_id in &projection {
3920 if !self
3921 .schema
3922 .columns
3923 .iter()
3924 .any(|column| column.id == *column_id)
3925 {
3926 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
3927 }
3928 }
3929
3930 let snapshot = self.snapshot();
3931 let hard_filter = if request.must.is_empty() {
3932 None
3933 } else {
3934 let mut sets = Vec::with_capacity(request.must.len());
3935 for condition in &request.must {
3936 sets.push(self.resolve_condition(condition, snapshot)?);
3937 }
3938 Some(RowIdSet::intersect_many(sets))
3939 };
3940 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
3941 return Ok(Vec::new());
3942 }
3943
3944 let constant = match request.fusion {
3945 Fusion::ReciprocalRank { constant } => constant,
3946 };
3947 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
3948 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
3949 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
3950 std::collections::HashMap::new();
3951 for named in retrievers {
3952 let hits = self.retrieve_filtered(
3953 &named.retriever,
3954 snapshot,
3955 hard_filter.as_ref(),
3956 authorized,
3957 )?;
3958 for hit in hits {
3959 let contribution = named.weight / (constant as f64 + hit.rank as f64);
3960 let entry = fused.entry(hit.row_id).or_default();
3961 entry.0 += contribution;
3962 entry.1.push(ComponentScore {
3963 retriever_name: named.name.clone(),
3964 rank: hit.rank,
3965 raw_score: hit.score,
3966 contribution,
3967 });
3968 }
3969 }
3970 let mut ranked: Vec<_> = fused.into_iter().collect();
3971 ranked.sort_by(|(a_row, (a_score, _)), (b_row, (b_score, _))| {
3972 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
3973 });
3974
3975 let row_ids: Vec<_> = ranked.iter().map(|(row_id, _)| row_id.0).collect();
3976 let sentinel = projection
3977 .first()
3978 .copied()
3979 .or_else(|| self.schema.columns.first().map(|column| column.id));
3980 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
3981 std::collections::HashMap::new();
3982 if let Some(column_id) = sentinel {
3983 for (row_id, value) in self.values_for_rids(&row_ids, column_id, snapshot)? {
3984 cells.entry(row_id).or_default().insert(column_id, value);
3985 }
3986 }
3987 for &column_id in &projection {
3988 if Some(column_id) == sentinel {
3989 continue;
3990 }
3991 for (row_id, value) in self.values_for_rids(&row_ids, column_id, snapshot)? {
3992 cells.entry(row_id).or_default().insert(column_id, value);
3993 }
3994 }
3995
3996 let mut out = Vec::with_capacity(request.limit);
3997 for (row_id, (fused_score, mut components)) in ranked {
3998 let Some(row_cells) = cells.remove(&row_id) else {
3999 continue;
4000 };
4001 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4002 out.push(SearchHit {
4003 row_id,
4004 cells: projection
4005 .iter()
4006 .filter_map(|column_id| {
4007 row_cells
4008 .get(column_id)
4009 .cloned()
4010 .map(|value| (*column_id, value))
4011 })
4012 .collect(),
4013 components,
4014 fused_score,
4015 });
4016 if out.len() == request.limit {
4017 break;
4018 }
4019 }
4020 Ok(out)
4021 }
4022
4023 pub fn set_similarity(
4026 &mut self,
4027 request: &crate::query::SetSimilarityRequest,
4028 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4029 self.set_similarity_with_allowed(request, None)
4030 }
4031
4032 pub fn ann_rerank(
4034 &mut self,
4035 request: &crate::query::AnnRerankRequest,
4036 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4037 self.ann_rerank_with_allowed(request, None)
4038 }
4039
4040 pub fn ann_rerank_with_allowed(
4041 &mut self,
4042 request: &crate::query::AnnRerankRequest,
4043 allowed: Option<&std::collections::HashSet<RowId>>,
4044 ) -> Result<Vec<crate::query::AnnRerankHit>> {
4045 use crate::query::{AnnRerankHit, Retriever, RetrieverScore, VectorMetric};
4046 if request.candidate_k == 0 || request.limit == 0 {
4047 return Err(MongrelError::InvalidArgument(
4048 "candidate_k and limit must be > 0".into(),
4049 ));
4050 }
4051 let hits = self.retrieve_with_allowed(
4052 &Retriever::Ann {
4053 column_id: request.column_id,
4054 query: request.query.clone(),
4055 k: request.candidate_k,
4056 },
4057 allowed,
4058 )?;
4059 let distances: std::collections::HashMap<_, _> = hits
4060 .iter()
4061 .filter_map(|hit| match hit.score {
4062 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
4063 _ => None,
4064 })
4065 .collect();
4066 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4067 let values = self.values_for_rids(&row_ids, request.column_id, self.snapshot())?;
4068 let query_norm = request
4069 .query
4070 .iter()
4071 .map(|value| value * value)
4072 .sum::<f32>()
4073 .sqrt();
4074 let mut reranked = values
4075 .into_iter()
4076 .filter_map(|(row_id, value)| {
4077 let Value::Embedding(vector) = value else {
4078 return None;
4079 };
4080 let dot = request
4081 .query
4082 .iter()
4083 .zip(&vector)
4084 .map(|(left, right)| left * right)
4085 .sum::<f32>();
4086 let exact_score = match request.metric {
4087 VectorMetric::DotProduct => dot,
4088 VectorMetric::Cosine => {
4089 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
4090 if query_norm == 0.0 || norm == 0.0 {
4091 0.0
4092 } else {
4093 dot / (query_norm * norm)
4094 }
4095 }
4096 VectorMetric::Euclidean => request
4097 .query
4098 .iter()
4099 .zip(&vector)
4100 .map(|(left, right)| (left - right).powi(2))
4101 .sum::<f32>()
4102 .sqrt(),
4103 };
4104 Some(AnnRerankHit {
4105 row_id,
4106 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
4107 exact_score,
4108 })
4109 })
4110 .collect::<Vec<_>>();
4111 reranked.sort_by(|left, right| {
4112 let score = match request.metric {
4113 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
4114 VectorMetric::Cosine | VectorMetric::DotProduct => {
4115 right.exact_score.total_cmp(&left.exact_score)
4116 }
4117 };
4118 score.then_with(|| left.row_id.cmp(&right.row_id))
4119 });
4120 reranked.truncate(request.limit);
4121 Ok(reranked)
4122 }
4123
4124 pub fn set_similarity_with_allowed(
4125 &mut self,
4126 request: &crate::query::SetSimilarityRequest,
4127 allowed: Option<&std::collections::HashSet<RowId>>,
4128 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4129 self.set_similarity_explained_with_allowed(request, allowed)
4130 .map(|(hits, _)| hits)
4131 }
4132
4133 pub fn set_similarity_explained(
4134 &mut self,
4135 request: &crate::query::SetSimilarityRequest,
4136 ) -> Result<(
4137 Vec<crate::query::SetSimilarityHit>,
4138 crate::query::SetSimilarityTrace,
4139 )> {
4140 self.set_similarity_explained_with_allowed(request, None)
4141 }
4142
4143 fn set_similarity_explained_with_allowed(
4144 &mut self,
4145 request: &crate::query::SetSimilarityRequest,
4146 allowed: Option<&std::collections::HashSet<RowId>>,
4147 ) -> Result<(
4148 Vec<crate::query::SetSimilarityHit>,
4149 crate::query::SetSimilarityTrace,
4150 )> {
4151 use crate::query::{Retriever, RetrieverScore, SetSimilarityHit};
4152 let mut trace = crate::query::SetSimilarityTrace::default();
4153 if request.members.is_empty() {
4154 return Ok((Vec::new(), trace));
4155 }
4156 if request.candidate_k == 0 || request.limit == 0 {
4157 return Err(MongrelError::InvalidArgument(
4158 "candidate_k and limit must be > 0".into(),
4159 ));
4160 }
4161 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
4162 return Err(MongrelError::InvalidArgument(
4163 "min_jaccard must be finite and between 0 and 1".into(),
4164 ));
4165 }
4166 let started = std::time::Instant::now();
4167 let hits = self.retrieve_with_allowed(
4168 &Retriever::MinHash {
4169 column_id: request.column_id,
4170 members: request.members.clone(),
4171 k: request.candidate_k,
4172 },
4173 allowed,
4174 )?;
4175 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
4176 trace.candidate_count = hits.len();
4177 let snapshot = self.snapshot();
4178 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
4179 let started = std::time::Instant::now();
4180 let values = self.values_for_rids(&row_ids, request.column_id, snapshot)?;
4181 trace.gather_us = started.elapsed().as_micros() as u64;
4182 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
4183 let estimates: std::collections::HashMap<_, _> = hits
4184 .into_iter()
4185 .filter_map(|hit| match hit.score {
4186 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
4187 _ => None,
4188 })
4189 .collect();
4190 let started = std::time::Instant::now();
4191 let parsed: Vec<_> = values
4192 .into_iter()
4193 .filter_map(|(row_id, value)| {
4194 let Value::Bytes(bytes) = value else {
4195 return None;
4196 };
4197 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
4198 return None;
4199 };
4200 let stored = members
4201 .into_iter()
4202 .filter_map(|member| match member {
4203 serde_json::Value::String(value) => {
4204 Some(crate::query::SetMember::String(value))
4205 }
4206 serde_json::Value::Number(value) => {
4207 Some(crate::query::SetMember::Number(value))
4208 }
4209 serde_json::Value::Bool(value) => {
4210 Some(crate::query::SetMember::Boolean(value))
4211 }
4212 _ => None,
4213 })
4214 .collect::<std::collections::HashSet<_>>();
4215 Some((row_id, stored))
4216 })
4217 .collect();
4218 trace.parse_us = started.elapsed().as_micros() as u64;
4219 trace.verified_count = parsed.len();
4220 let started = std::time::Instant::now();
4221 let mut exact = Vec::new();
4222 for (row_id, stored) in parsed {
4223 let union = query.union(&stored).count();
4224 let score = if union == 0 {
4225 1.0
4226 } else {
4227 query.intersection(&stored).count() as f32 / union as f32
4228 };
4229 if score >= request.min_jaccard {
4230 exact.push(SetSimilarityHit {
4231 row_id,
4232 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
4233 exact_jaccard: score,
4234 });
4235 }
4236 }
4237 exact.sort_by(|a, b| {
4238 b.exact_jaccard
4239 .total_cmp(&a.exact_jaccard)
4240 .then_with(|| a.row_id.cmp(&b.row_id))
4241 });
4242 exact.truncate(request.limit);
4243 trace.score_us = started.elapsed().as_micros() as u64;
4244 Ok((exact, trace))
4245 }
4246
4247 fn values_for_rids(
4249 &self,
4250 row_ids: &[u64],
4251 column_id: u16,
4252 snapshot: Snapshot,
4253 ) -> Result<Vec<(RowId, Value)>> {
4254 let mut readers: Vec<_> = self
4255 .run_refs
4256 .iter()
4257 .map(|run| self.open_reader(run.run_id))
4258 .collect::<Result<_>>()?;
4259 let now = unix_nanos_now();
4260 let mut out = Vec::with_capacity(row_ids.len());
4261 for &raw_row_id in row_ids {
4262 let row_id = RowId(raw_row_id);
4263 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4264 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4265 let overlay = match (mem, mutable) {
4266 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
4267 (a_epoch, a)
4268 } else {
4269 (b_epoch, b)
4270 }),
4271 (Some(value), None) | (None, Some(value)) => Some(value),
4272 (None, None) => None,
4273 };
4274 if let Some((_, row)) = overlay {
4275 if !row.deleted && !self.row_expired_at(&row, now) {
4276 if let Some(value) = row.columns.get(&column_id) {
4277 out.push((row_id, value.clone()));
4278 }
4279 }
4280 continue;
4281 }
4282
4283 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
4284 for (index, reader) in readers.iter_mut().enumerate() {
4285 if let Some((epoch, deleted, value)) =
4286 reader.get_version_column(row_id, snapshot.epoch, column_id)?
4287 {
4288 if best
4289 .as_ref()
4290 .map(|(best_epoch, ..)| epoch > *best_epoch)
4291 .unwrap_or(true)
4292 {
4293 best = Some((epoch, deleted, value, index));
4294 }
4295 }
4296 }
4297 let Some((_, false, Some(value), reader_index)) = best else {
4298 continue;
4299 };
4300 if let Some(ttl) = self.ttl {
4301 if ttl.column_id != column_id {
4302 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4303 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4304 {
4305 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4306 continue;
4307 }
4308 }
4309 } else if let Value::Int64(timestamp) = value {
4310 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4311 continue;
4312 }
4313 }
4314 }
4315 out.push((row_id, value));
4316 }
4317 Ok(out)
4318 }
4319
4320 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
4325 use std::collections::HashMap;
4326 let mut rows = Vec::with_capacity(rids.len());
4327 let ttl_now = unix_nanos_now();
4328 let tier_size = self.memtable.len() + self.mutable_run.len();
4345 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
4346 if rids.len().saturating_mul(24) < tier_size {
4347 for &rid in rids {
4348 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
4349 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
4350 let newest = match (mem, mrun) {
4351 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
4352 (Some((_, mr)), None) => Some(mr),
4353 (None, Some((_, rr))) => Some(rr),
4354 (None, None) => None,
4355 };
4356 if let Some(row) = newest {
4357 overlay.insert(rid, row);
4358 }
4359 }
4360 } else {
4361 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
4362 overlay
4363 .entry(row.row_id.0)
4364 .and_modify(|e| {
4365 if row.committed_epoch > e.committed_epoch {
4366 *e = row.clone();
4367 }
4368 })
4369 .or_insert(row);
4370 };
4371 for row in self.memtable.visible_versions(snapshot.epoch) {
4372 fold_newest(row, &mut overlay);
4373 }
4374 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4375 fold_newest(row, &mut overlay);
4376 }
4377 }
4378 if self.run_refs.len() == 1 {
4379 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4380 if rids.len().saturating_mul(24) < reader.row_count() {
4388 for &rid in rids {
4389 if let Some(r) = overlay.get(&rid) {
4390 if !r.deleted {
4391 rows.push(r.clone());
4392 }
4393 continue;
4394 }
4395 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
4396 if !row.deleted {
4397 rows.push(row);
4398 }
4399 }
4400 }
4401 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4402 return Ok(rows);
4403 }
4404 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4413 enum Src {
4416 Overlay,
4417 Run,
4418 }
4419 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
4420 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
4421 for rid in rids {
4422 if overlay.contains_key(rid) {
4423 plan.push(Src::Overlay);
4424 continue;
4425 }
4426 match vis_rids.binary_search(&(*rid as i64)) {
4427 Ok(i) => {
4428 plan.push(Src::Run);
4429 fetch.push(positions[i]);
4430 }
4431 Err(_) => { }
4432 }
4433 }
4434 let fetched = reader.materialize_batch(&fetch)?;
4435 let mut fetched_iter = fetched.into_iter();
4436 for (rid, src) in rids.iter().zip(plan) {
4437 match src {
4438 Src::Overlay => {
4439 if let Some(r) = overlay.get(rid) {
4440 if !r.deleted {
4441 rows.push(r.clone());
4442 }
4443 }
4444 }
4445 Src::Run => {
4446 if let Some(row) = fetched_iter.next() {
4447 if !row.deleted {
4448 rows.push(row);
4449 }
4450 }
4451 }
4452 }
4453 }
4454 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4455 return Ok(rows);
4456 }
4457 let mut readers: Vec<_> = self
4461 .run_refs
4462 .iter()
4463 .map(|rr| self.open_reader(rr.run_id))
4464 .collect::<Result<Vec<_>>>()?;
4465 for rid in rids {
4466 if let Some(r) = overlay.get(rid) {
4467 if !r.deleted {
4468 rows.push(r.clone());
4469 }
4470 continue;
4471 }
4472 let mut best: Option<(Epoch, Row)> = None;
4473 for reader in readers.iter_mut() {
4474 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
4475 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4476 best = Some((epoch, row));
4477 }
4478 }
4479 }
4480 if let Some((_, r)) = best {
4481 if !r.deleted {
4482 rows.push(r);
4483 }
4484 }
4485 }
4486 rows.retain(|row| !self.row_expired_at(row, ttl_now));
4487 Ok(rows)
4488 }
4489
4490 pub fn indexes_complete(&self) -> bool {
4500 self.indexes_complete
4501 }
4502
4503 pub fn index_build_policy(&self) -> IndexBuildPolicy {
4505 self.index_build_policy
4506 }
4507
4508 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
4512 self.index_build_policy = policy;
4513 }
4514
4515 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
4520 if !self.indexes_complete {
4524 return None;
4525 }
4526 let b = self.bitmap.get(&column_id)?;
4527 let result: Vec<Vec<u8>> = b
4528 .keys()
4529 .into_iter()
4530 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
4531 .cloned()
4532 .collect();
4533 Some(result)
4534 }
4535
4536 pub fn fk_join_row_ids(
4537 &self,
4538 fk_column_id: u16,
4539 pk_values: &[Vec<u8>],
4540 fk_conditions: &[crate::query::Condition],
4541 snapshot: Snapshot,
4542 ) -> Result<Vec<u64>> {
4543 let Some(b) = self.bitmap.get(&fk_column_id) else {
4544 return Ok(Vec::new());
4545 };
4546 let mut join_set = {
4547 let mut acc = roaring::RoaringBitmap::new();
4548 for v in pk_values {
4549 acc |= b.get(v);
4550 }
4551 RowIdSet::from_roaring(acc)
4552 };
4553 if !fk_conditions.is_empty() {
4554 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4555 sets.push(join_set);
4556 for c in fk_conditions {
4557 sets.push(self.resolve_condition(c, snapshot)?);
4558 }
4559 join_set = RowIdSet::intersect_many(sets);
4560 }
4561 Ok(join_set.into_sorted_vec())
4562 }
4563
4564 pub fn fk_join_count(
4570 &self,
4571 fk_column_id: u16,
4572 pk_values: &[Vec<u8>],
4573 fk_conditions: &[crate::query::Condition],
4574 snapshot: Snapshot,
4575 ) -> Result<u64> {
4576 let Some(b) = self.bitmap.get(&fk_column_id) else {
4577 return Ok(0);
4578 };
4579 let mut acc = roaring::RoaringBitmap::new();
4580 for v in pk_values {
4581 acc |= b.get(v);
4582 }
4583 if fk_conditions.is_empty() {
4584 return Ok(acc.len());
4585 }
4586 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
4587 sets.push(RowIdSet::from_roaring(acc));
4588 for c in fk_conditions {
4589 sets.push(self.resolve_condition(c, snapshot)?);
4590 }
4591 Ok(RowIdSet::intersect_many(sets).len() as u64)
4592 }
4593
4594 fn resolve_condition(
4599 &self,
4600 c: &crate::query::Condition,
4601 snapshot: Snapshot,
4602 ) -> Result<RowIdSet> {
4603 use crate::query::Condition;
4604 self.validate_condition(c)?;
4605 Ok(match c {
4606 Condition::Pk(key) => {
4607 let lookup = self
4608 .schema
4609 .primary_key()
4610 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
4611 .unwrap_or_else(|| key.clone());
4612 self.hot
4613 .get(&lookup)
4614 .map(|r| RowIdSet::one(r.0))
4615 .unwrap_or_else(RowIdSet::empty)
4616 }
4617 Condition::BitmapEq { column_id, value } => {
4618 let lookup = self.index_lookup_key_bytes(*column_id, value);
4619 self.bitmap
4620 .get(column_id)
4621 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
4622 .unwrap_or_else(RowIdSet::empty)
4623 }
4624 Condition::BitmapIn { column_id, values } => {
4625 let bm = self.bitmap.get(column_id);
4626 let mut acc = roaring::RoaringBitmap::new();
4627 if let Some(b) = bm {
4628 for v in values {
4629 let lookup = self.index_lookup_key_bytes(*column_id, v);
4630 acc |= b.get(&lookup);
4631 }
4632 }
4633 RowIdSet::from_roaring(acc)
4634 }
4635 Condition::BytesPrefix { column_id, prefix } => {
4636 if let Some(b) = self.bitmap.get(column_id) {
4641 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
4642 let mut acc = roaring::RoaringBitmap::new();
4643 for key in b.keys() {
4644 if key.starts_with(&lookup_prefix) {
4645 acc |= b.get(key);
4646 }
4647 }
4648 RowIdSet::from_roaring(acc)
4649 } else {
4650 RowIdSet::empty()
4651 }
4652 }
4653 Condition::FmContains { column_id, pattern } => self
4654 .fm
4655 .get(column_id)
4656 .map(|f| {
4657 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
4658 })
4659 .unwrap_or_else(RowIdSet::empty),
4660 Condition::FmContainsAll {
4661 column_id,
4662 patterns,
4663 } => {
4664 if let Some(f) = self.fm.get(column_id) {
4667 let sets: Vec<RowIdSet> = patterns
4668 .iter()
4669 .map(|pat| {
4670 RowIdSet::from_unsorted(
4671 f.locate(pat).into_iter().map(|r| r.0).collect(),
4672 )
4673 })
4674 .collect();
4675 RowIdSet::intersect_many(sets)
4676 } else {
4677 RowIdSet::empty()
4678 }
4679 }
4680 Condition::Ann {
4681 column_id,
4682 query,
4683 k,
4684 } => RowIdSet::from_unsorted(
4685 self.retrieve_filtered(
4686 &crate::query::Retriever::Ann {
4687 column_id: *column_id,
4688 query: query.clone(),
4689 k: *k,
4690 },
4691 snapshot,
4692 None,
4693 None,
4694 )?
4695 .into_iter()
4696 .map(|hit| hit.row_id.0)
4697 .collect(),
4698 ),
4699 Condition::SparseMatch {
4700 column_id,
4701 query,
4702 k,
4703 } => RowIdSet::from_unsorted(
4704 self.retrieve_filtered(
4705 &crate::query::Retriever::Sparse {
4706 column_id: *column_id,
4707 query: query.clone(),
4708 k: *k,
4709 },
4710 snapshot,
4711 None,
4712 None,
4713 )?
4714 .into_iter()
4715 .map(|hit| hit.row_id.0)
4716 .collect(),
4717 ),
4718 Condition::MinHashSimilar {
4719 column_id,
4720 query,
4721 k,
4722 } => match self.minhash.get(column_id) {
4723 Some(index) => {
4724 let candidates = index.candidate_row_ids(query);
4725 let eligible =
4726 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
4727 RowIdSet::from_unsorted(
4728 index
4729 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
4730 .into_iter()
4731 .map(|(row_id, _)| row_id.0)
4732 .collect(),
4733 )
4734 }
4735 None => RowIdSet::empty(),
4736 },
4737 Condition::Range { column_id, lo, hi } => {
4738 let mut set = if let Some(li) = self.learned_range.get(column_id) {
4747 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
4748 } else if self.run_refs.len() == 1 {
4749 let mut r = self.open_reader(self.run_refs[0].run_id)?;
4750 r.range_row_id_set_i64(*column_id, *lo, *hi)?
4751 } else {
4752 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
4753 };
4754 set.remove_many(self.overlay_rid_set(snapshot));
4755 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
4756 set
4757 }
4758 Condition::RangeF64 {
4759 column_id,
4760 lo,
4761 lo_inclusive,
4762 hi,
4763 hi_inclusive,
4764 } => {
4765 let mut set = if let Some(li) = self.learned_range.get(column_id) {
4768 RowIdSet::from_unsorted(
4769 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
4770 .into_iter()
4771 .collect(),
4772 )
4773 } else if self.run_refs.len() == 1 {
4774 let mut r = self.open_reader(self.run_refs[0].run_id)?;
4775 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
4776 } else {
4777 return self.range_scan_f64(
4778 *column_id,
4779 *lo,
4780 *lo_inclusive,
4781 *hi,
4782 *hi_inclusive,
4783 snapshot,
4784 );
4785 };
4786 set.remove_many(self.overlay_rid_set(snapshot));
4787 self.range_scan_overlay_f64(
4788 &mut set,
4789 *column_id,
4790 *lo,
4791 *lo_inclusive,
4792 *hi,
4793 *hi_inclusive,
4794 snapshot,
4795 );
4796 set
4797 }
4798 Condition::IsNull { column_id } => {
4799 let mut set = if self.run_refs.len() == 1 {
4800 let mut r = self.open_reader(self.run_refs[0].run_id)?;
4801 r.null_row_id_set(*column_id, true)?
4802 } else {
4803 return self.null_scan(*column_id, true, snapshot);
4804 };
4805 set.remove_many(self.overlay_rid_set(snapshot));
4806 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
4807 set
4808 }
4809 Condition::IsNotNull { column_id } => {
4810 let mut set = if self.run_refs.len() == 1 {
4811 let mut r = self.open_reader(self.run_refs[0].run_id)?;
4812 r.null_row_id_set(*column_id, false)?
4813 } else {
4814 return self.null_scan(*column_id, false, snapshot);
4815 };
4816 set.remove_many(self.overlay_rid_set(snapshot));
4817 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
4818 set
4819 }
4820 })
4821 }
4822
4823 fn range_scan_i64(
4831 &self,
4832 column_id: u16,
4833 lo: i64,
4834 hi: i64,
4835 snapshot: Snapshot,
4836 ) -> Result<RowIdSet> {
4837 let mut row_ids = Vec::new();
4838 let overlay_rids = self.overlay_rid_set(snapshot);
4839 for rr in &self.run_refs {
4840 let mut reader = self.open_reader(rr.run_id)?;
4841 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
4842 for rid in matched {
4843 if !overlay_rids.contains(&rid) {
4844 row_ids.push(rid);
4845 }
4846 }
4847 }
4848 let mut s = RowIdSet::from_unsorted(row_ids);
4849 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
4850 Ok(s)
4851 }
4852
4853 fn range_scan_f64(
4856 &self,
4857 column_id: u16,
4858 lo: f64,
4859 lo_inclusive: bool,
4860 hi: f64,
4861 hi_inclusive: bool,
4862 snapshot: Snapshot,
4863 ) -> Result<RowIdSet> {
4864 let mut row_ids = Vec::new();
4865 let overlay_rids = self.overlay_rid_set(snapshot);
4866 for rr in &self.run_refs {
4867 let mut reader = self.open_reader(rr.run_id)?;
4868 let matched = reader.range_row_ids_visible_f64(
4869 column_id,
4870 lo,
4871 lo_inclusive,
4872 hi,
4873 hi_inclusive,
4874 snapshot.epoch,
4875 )?;
4876 for rid in matched {
4877 if !overlay_rids.contains(&rid) {
4878 row_ids.push(rid);
4879 }
4880 }
4881 }
4882 let mut s = RowIdSet::from_unsorted(row_ids);
4883 self.range_scan_overlay_f64(
4884 &mut s,
4885 column_id,
4886 lo,
4887 lo_inclusive,
4888 hi,
4889 hi_inclusive,
4890 snapshot,
4891 );
4892 Ok(s)
4893 }
4894
4895 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
4897 let mut s = HashSet::new();
4898 for row in self.memtable.visible_versions(snapshot.epoch) {
4899 s.insert(row.row_id.0);
4900 }
4901 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4902 s.insert(row.row_id.0);
4903 }
4904 s
4905 }
4906
4907 fn range_scan_overlay_i64(
4908 &self,
4909 s: &mut RowIdSet,
4910 column_id: u16,
4911 lo: i64,
4912 hi: i64,
4913 snapshot: Snapshot,
4914 ) {
4915 let mut newest: HashMap<u64, &Row> = HashMap::new();
4920 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4921 let memtable = self.memtable.visible_versions(snapshot.epoch);
4922 for r in &mutable {
4923 newest.entry(r.row_id.0).or_insert(r);
4924 }
4925 for r in &memtable {
4926 newest.insert(r.row_id.0, r);
4927 }
4928 for row in newest.values() {
4929 if !row.deleted {
4930 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
4931 if *v >= lo && *v <= hi {
4932 s.insert(row.row_id.0);
4933 }
4934 }
4935 }
4936 }
4937 }
4938
4939 #[allow(clippy::too_many_arguments)]
4940 fn range_scan_overlay_f64(
4941 &self,
4942 s: &mut RowIdSet,
4943 column_id: u16,
4944 lo: f64,
4945 lo_inclusive: bool,
4946 hi: f64,
4947 hi_inclusive: bool,
4948 snapshot: Snapshot,
4949 ) {
4950 let mut newest: HashMap<u64, &Row> = HashMap::new();
4953 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4954 let memtable = self.memtable.visible_versions(snapshot.epoch);
4955 for r in &mutable {
4956 newest.entry(r.row_id.0).or_insert(r);
4957 }
4958 for r in &memtable {
4959 newest.insert(r.row_id.0, r);
4960 }
4961 for row in newest.values() {
4962 if !row.deleted {
4963 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
4964 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
4965 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
4966 if ok_lo && ok_hi {
4967 s.insert(row.row_id.0);
4968 }
4969 }
4970 }
4971 }
4972 }
4973
4974 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
4977 let mut row_ids = Vec::new();
4978 let overlay_rids = self.overlay_rid_set(snapshot);
4979 for rr in &self.run_refs {
4980 let mut reader = self.open_reader(rr.run_id)?;
4981 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
4982 for rid in matched {
4983 if !overlay_rids.contains(&rid) {
4984 row_ids.push(rid);
4985 }
4986 }
4987 }
4988 let mut s = RowIdSet::from_unsorted(row_ids);
4989 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
4990 Ok(s)
4991 }
4992
4993 fn null_scan_overlay(
4997 &self,
4998 s: &mut RowIdSet,
4999 column_id: u16,
5000 want_nulls: bool,
5001 snapshot: Snapshot,
5002 ) {
5003 let mut newest: HashMap<u64, &Row> = HashMap::new();
5004 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
5005 let memtable = self.memtable.visible_versions(snapshot.epoch);
5006 for r in &mutable {
5007 newest.entry(r.row_id.0).or_insert(r);
5008 }
5009 for r in &memtable {
5010 newest.insert(r.row_id.0, r);
5011 }
5012 for row in newest.values() {
5013 if row.deleted {
5014 continue;
5015 }
5016 let is_null = !row.columns.contains_key(&column_id)
5017 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
5018 if is_null == want_nulls {
5019 s.insert(row.row_id.0);
5020 }
5021 }
5022 }
5023
5024 pub fn snapshot(&self) -> Snapshot {
5025 Snapshot::at(self.epoch.visible())
5026 }
5027
5028 pub fn pin_snapshot(&mut self) -> Snapshot {
5031 let e = self.epoch.visible();
5032 *self.pinned.entry(e).or_insert(0) += 1;
5033 Snapshot::at(e)
5034 }
5035
5036 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
5038 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
5039 *count -= 1;
5040 if *count == 0 {
5041 self.pinned.remove(&snap.epoch);
5042 }
5043 }
5044 }
5045
5046 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
5056 let local = self.pinned.keys().next().copied();
5057 let global = self.snapshots.min_pinned();
5058 let history = self.snapshots.history_floor(self.current_epoch());
5059 [local, global, history].into_iter().flatten().min()
5060 }
5061
5062 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
5066 self.ensure_writable()?;
5067 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
5068 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
5069 }
5070
5071 pub fn clear_ttl(&mut self) -> Result<()> {
5072 self.ensure_writable()?;
5073 self.apply_ttl_policy_at(None, self.current_epoch())
5074 }
5075
5076 pub fn ttl(&self) -> Option<TtlPolicy> {
5077 self.ttl
5078 }
5079
5080 pub(crate) fn prepare_ttl_policy(
5081 &self,
5082 column_name: &str,
5083 duration_nanos: u64,
5084 ) -> Result<TtlPolicy> {
5085 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
5086 return Err(MongrelError::InvalidArgument(
5087 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
5088 ));
5089 }
5090 let column = self
5091 .schema
5092 .columns
5093 .iter()
5094 .find(|column| column.name == column_name)
5095 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
5096 if column.ty != TypeId::TimestampNanos {
5097 return Err(MongrelError::Schema(format!(
5098 "TTL column {column_name} must be TimestampNanos, is {:?}",
5099 column.ty
5100 )));
5101 }
5102 Ok(TtlPolicy {
5103 column_id: column.id,
5104 duration_nanos,
5105 })
5106 }
5107
5108 pub(crate) fn apply_ttl_policy_at(
5109 &mut self,
5110 policy: Option<TtlPolicy>,
5111 epoch: Epoch,
5112 ) -> Result<()> {
5113 if let Some(policy) = policy {
5114 let column = self
5115 .schema
5116 .columns
5117 .iter()
5118 .find(|column| column.id == policy.column_id)
5119 .ok_or_else(|| {
5120 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
5121 })?;
5122 if column.ty != TypeId::TimestampNanos
5123 || policy.duration_nanos == 0
5124 || policy.duration_nanos > i64::MAX as u64
5125 {
5126 return Err(MongrelError::Schema("invalid TTL policy".into()));
5127 }
5128 }
5129 self.ttl = policy;
5130 self.agg_cache.clear();
5131 self.clear_result_cache();
5132 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5133 self.persist_manifest(epoch)
5134 }
5135
5136 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
5137 let Some(policy) = self.ttl else {
5138 return false;
5139 };
5140 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
5141 return false;
5142 };
5143 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
5144 }
5145
5146 pub fn current_epoch(&self) -> Epoch {
5147 self.epoch.visible()
5148 }
5149
5150 pub fn memtable_len(&self) -> usize {
5151 self.memtable.len()
5152 }
5153
5154 pub fn count(&self) -> u64 {
5157 if self.ttl.is_none() {
5158 self.live_count
5159 } else {
5160 self.visible_rows(self.snapshot())
5161 .map(|rows| rows.len() as u64)
5162 .unwrap_or(self.live_count)
5163 }
5164 }
5165
5166 pub fn count_conditions(
5170 &mut self,
5171 conditions: &[crate::query::Condition],
5172 snapshot: Snapshot,
5173 ) -> Result<Option<u64>> {
5174 use crate::query::Condition;
5175 if self.ttl.is_some() {
5176 return self
5177 .query(&crate::query::Query {
5178 conditions: conditions.to_vec(),
5179 })
5180 .map(|rows| Some(rows.len() as u64));
5181 }
5182 if conditions.is_empty() {
5183 return Ok(Some(self.live_count));
5184 }
5185 let served = |c: &Condition| {
5186 matches!(
5187 c,
5188 Condition::Pk(_)
5189 | Condition::BitmapEq { .. }
5190 | Condition::BitmapIn { .. }
5191 | Condition::BytesPrefix { .. }
5192 | Condition::FmContains { .. }
5193 | Condition::FmContainsAll { .. }
5194 | Condition::Ann { .. }
5195 | Condition::Range { .. }
5196 | Condition::RangeF64 { .. }
5197 | Condition::SparseMatch { .. }
5198 | Condition::MinHashSimilar { .. }
5199 | Condition::IsNull { .. }
5200 | Condition::IsNotNull { .. }
5201 )
5202 };
5203 if !conditions.iter().all(served) {
5204 return Ok(None);
5205 }
5206 self.ensure_indexes_complete()?;
5207 let mut sets = Vec::with_capacity(conditions.len());
5208 for condition in conditions {
5209 sets.push(self.resolve_condition(condition, snapshot)?);
5210 }
5211 let mut rids = RowIdSet::intersect_many(sets);
5212 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
5222 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
5223 }
5224 let count = rids.len() as u64;
5225 crate::trace::QueryTrace::record(|t| {
5226 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
5227 t.survivor_count = Some(count as usize);
5228 t.conditions_pushed = conditions.len();
5229 });
5230 Ok(Some(count))
5231 }
5232
5233 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
5238 let mut out = Vec::new();
5239 for row in self.memtable.visible_versions(snapshot.epoch) {
5240 if row.deleted {
5241 out.push(row.row_id.0);
5242 }
5243 }
5244 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5245 if row.deleted {
5246 out.push(row.row_id.0);
5247 }
5248 }
5249 out
5250 }
5251
5252 pub fn bulk_load_columns(
5261 &mut self,
5262 user_columns: Vec<(u16, columnar::NativeColumn)>,
5263 ) -> Result<Epoch> {
5264 self.bulk_load_columns_with(user_columns, 3, false, true)
5265 }
5266
5267 pub fn bulk_load_fast(
5274 &mut self,
5275 user_columns: Vec<(u16, columnar::NativeColumn)>,
5276 ) -> Result<Epoch> {
5277 self.bulk_load_columns_with(user_columns, -1, true, false)
5278 }
5279
5280 fn bulk_load_columns_with(
5281 &mut self,
5282 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
5283 zstd_level: i32,
5284 force_plain: bool,
5285 lz4: bool,
5286 ) -> Result<Epoch> {
5287 let epoch = self.commit()?;
5288 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
5289 if n == 0 {
5290 return Ok(epoch);
5291 }
5292 let live_before = self.live_count;
5293 self.spill_mutable_run(epoch)?;
5295 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
5296 && self.indexes_complete
5297 && self.run_refs.is_empty()
5298 && self.memtable.is_empty()
5299 && self.mutable_run.is_empty();
5300 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
5303 self.validate_columns_not_null(&user_columns, n)?;
5304 let winner_idx = self
5305 .bulk_pk_winner_indices(&user_columns, n)
5306 .filter(|idx| idx.len() != n);
5307 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
5308 match winner_idx.as_deref() {
5309 Some(idx) => {
5310 let compacted = user_columns
5311 .iter()
5312 .map(|(id, c)| (*id, c.gather(idx)))
5313 .collect();
5314 (compacted, idx.len())
5315 }
5316 None => (user_columns, n),
5317 };
5318 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
5319 let first = self.allocator.alloc_range(write_n as u64).0;
5320 for rid in first..first + write_n as u64 {
5321 self.reservoir.offer(rid);
5322 }
5323 let run_id = self.next_run_id;
5324 self.next_run_id += 1;
5325 let path = self.run_path(run_id);
5326 let mut writer =
5327 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
5328 if force_plain {
5329 writer = writer.with_plain();
5330 } else if lz4 {
5331 writer = writer.with_lz4();
5334 } else {
5335 writer = writer.with_zstd_level(zstd_level);
5336 }
5337 if let Some(kek) = &self.kek {
5338 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5339 }
5340 let header = writer.write_native(&path, &write_columns, write_n, first)?;
5341 self.run_refs.push(RunRef {
5342 run_id: run_id as u128,
5343 level: 0,
5344 epoch_created: epoch.0,
5345 row_count: header.row_count,
5346 });
5347 self.live_count = self.live_count.saturating_add(write_n as u64);
5348 if eager_index_build {
5349 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
5350 self.index_columns_bulk(&write_columns, &row_ids);
5351 self.indexes_complete = true;
5352 self.build_learned_ranges()?;
5353 } else {
5354 self.indexes_complete = false;
5358 }
5359 self.mark_flushed(epoch)?;
5360 self.persist_manifest(epoch)?;
5361 if eager_index_build {
5362 self.checkpoint_indexes(epoch);
5363 }
5364 self.clear_result_cache();
5365 Ok(epoch)
5366 }
5367
5368 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
5386 let n = row_ids.len();
5387 if n == 0 {
5388 return;
5389 }
5390 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
5391 columns.iter().map(|(id, c)| (*id, c)).collect();
5392 let ty_of: std::collections::HashMap<u16, TypeId> = self
5393 .schema
5394 .columns
5395 .iter()
5396 .map(|c| (c.id, c.ty.clone()))
5397 .collect();
5398 let pk_id = self.schema.primary_key().map(|c| c.id);
5399
5400 for (i, &rid) in row_ids.iter().enumerate() {
5401 let row_id = RowId(rid);
5402 if let Some(pid) = pk_id {
5403 if let Some(col) = by_id.get(&pid) {
5404 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
5405 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
5406 self.insert_hot_pk(key, row_id);
5407 }
5408 }
5409 }
5410 for idef in &self.schema.indexes {
5411 let Some(col) = by_id.get(&idef.column_id) else {
5412 continue;
5413 };
5414 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
5415 match idef.kind {
5416 IndexKind::Bitmap => {
5417 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5418 if let Some(key) =
5419 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
5420 {
5421 b.insert(key, row_id);
5422 }
5423 }
5424 }
5425 IndexKind::FmIndex => {
5426 if let Some(f) = self.fm.get_mut(&idef.column_id) {
5427 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5428 f.insert(bytes.to_vec(), row_id);
5429 }
5430 }
5431 }
5432 IndexKind::Sparse => {
5433 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
5434 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5435 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
5436 s.insert(&terms, row_id);
5437 }
5438 }
5439 }
5440 }
5441 IndexKind::MinHash => {
5442 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
5443 if let Some(bytes) = columnar::native_bytes_at(col, i) {
5444 let tokens = crate::index::token_hashes_from_bytes(bytes);
5445 mh.insert(&tokens, row_id);
5446 }
5447 }
5448 }
5449 _ => {}
5450 }
5451 }
5452 }
5453 }
5454
5455 pub fn visible_columns_native(
5460 &self,
5461 snapshot: Snapshot,
5462 projection: Option<&[u16]>,
5463 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
5464 let wanted: Vec<u16> = match projection {
5465 Some(p) => p.to_vec(),
5466 None => self.schema.columns.iter().map(|c| c.id).collect(),
5467 };
5468 if self.ttl.is_none()
5469 && self.memtable.is_empty()
5470 && self.mutable_run.is_empty()
5471 && self.run_refs.len() == 1
5472 {
5473 let rr = self.run_refs[0].clone();
5474 let mut reader = self.open_reader(rr.run_id)?;
5475 let idxs = reader.visible_indices_native(snapshot.epoch)?;
5476 let all_visible = idxs.len() == reader.row_count();
5477 if reader.has_mmap() {
5483 use rayon::prelude::*;
5484 let valid: Vec<u16> = wanted
5487 .iter()
5488 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
5489 .copied()
5490 .collect();
5491 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
5493 .par_iter()
5494 .filter_map(|cid| {
5495 reader
5496 .column_native_shared(*cid)
5497 .ok()
5498 .map(|col| (*cid, col))
5499 })
5500 .collect();
5501 let cols = decoded
5502 .into_iter()
5503 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
5504 .collect();
5505 return Ok(cols);
5506 }
5507 let mut cols = Vec::with_capacity(wanted.len());
5508 for cid in &wanted {
5509 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
5510 Some(c) => c,
5511 None => continue,
5512 };
5513 let col = reader.column_native(cdef.id)?;
5514 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
5515 }
5516 return Ok(cols);
5517 }
5518 let vcols = self.visible_columns(snapshot)?;
5519 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
5520 let out: Vec<(u16, columnar::NativeColumn)> = vcols
5521 .into_iter()
5522 .filter(|(id, _)| want_set.contains(id))
5523 .map(|(id, vals)| {
5524 let ty = self
5525 .schema
5526 .columns
5527 .iter()
5528 .find(|c| c.id == id)
5529 .map(|c| c.ty.clone())
5530 .unwrap_or(TypeId::Bytes);
5531 (id, columnar::values_to_native(ty, &vals))
5532 })
5533 .collect();
5534 Ok(out)
5535 }
5536
5537 pub fn run_count(&self) -> usize {
5538 self.run_refs.len()
5539 }
5540
5541 pub fn memtable_is_empty(&self) -> bool {
5543 self.memtable.is_empty()
5544 }
5545
5546 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
5550 self.page_cache.stats()
5551 }
5552
5553 pub fn reset_page_cache_stats(&self) {
5555 self.page_cache.reset_stats();
5556 }
5557
5558 pub fn run_ids(&self) -> Vec<u128> {
5561 self.run_refs.iter().map(|r| r.run_id).collect()
5562 }
5563
5564 pub fn single_run_is_clean(&self) -> bool {
5568 if self.ttl.is_some() || self.run_refs.len() != 1 {
5569 return false;
5570 }
5571 self.open_reader(self.run_refs[0].run_id)
5572 .map(|r| r.is_clean())
5573 .unwrap_or(false)
5574 }
5575
5576 fn resolve_footprint(
5583 &self,
5584 conditions: &[crate::query::Condition],
5585 snapshot: Snapshot,
5586 ) -> roaring::RoaringBitmap {
5587 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
5588 return roaring::RoaringBitmap::new();
5589 }
5590 if self.run_refs.is_empty() {
5591 return roaring::RoaringBitmap::new();
5592 }
5593 if self.run_refs.len() == 1 {
5595 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
5596 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
5597 return rids.to_roaring_lossy();
5598 }
5599 }
5600 }
5601 roaring::RoaringBitmap::new()
5602 }
5603
5604 pub fn query_columns_native_cached(
5615 &mut self,
5616 conditions: &[crate::query::Condition],
5617 projection: Option<&[u16]>,
5618 snapshot: Snapshot,
5619 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
5620 if self.ttl.is_some() {
5623 return self.query_columns_native(conditions, projection, snapshot);
5624 }
5625 if conditions.is_empty() {
5626 return self.query_columns_native(conditions, projection, snapshot);
5627 }
5628 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
5632 if let Some(hit) = self.result_cache.lock().get_columns(key) {
5633 crate::trace::QueryTrace::record(|t| {
5634 t.result_cache_hit = true;
5635 t.scan_mode = crate::trace::ScanMode::NativePushdown;
5636 });
5637 return Ok(Some((*hit).clone()));
5638 }
5639 let res = self.query_columns_native(conditions, projection, snapshot)?;
5640 if let Some(cols) = &res {
5641 let footprint = self.resolve_footprint(conditions, snapshot);
5642 let condition_cols = crate::query::condition_columns(conditions);
5643 self.result_cache.lock().insert(
5644 key,
5645 CachedEntry {
5646 data: CachedData::Columns(Arc::new(cols.clone())),
5647 footprint,
5648 condition_cols,
5649 },
5650 );
5651 }
5652 Ok(res)
5653 }
5654
5655 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
5660 if self.ttl.is_some() {
5661 return self.query(q);
5662 }
5663 if q.conditions.is_empty() {
5664 return self.query(q);
5665 }
5666 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
5667 if let Some(hit) = self.result_cache.lock().get_rows(key) {
5668 crate::trace::QueryTrace::record(|t| {
5669 t.result_cache_hit = true;
5670 t.scan_mode = crate::trace::ScanMode::Materialized;
5671 });
5672 return Ok((*hit).clone());
5673 }
5674 let rows = self.query(q)?;
5675 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
5676 let condition_cols = crate::query::condition_columns(&q.conditions);
5677 self.result_cache.lock().insert(
5678 key,
5679 CachedEntry {
5680 data: CachedData::Rows(Arc::new(rows.clone())),
5681 footprint,
5682 condition_cols,
5683 },
5684 );
5685 Ok(rows)
5686 }
5687
5688 #[allow(clippy::type_complexity)]
5703 pub fn query_columns_native_traced(
5704 &mut self,
5705 conditions: &[crate::query::Condition],
5706 projection: Option<&[u16]>,
5707 snapshot: Snapshot,
5708 ) -> Result<(
5709 Option<Vec<(u16, columnar::NativeColumn)>>,
5710 crate::trace::QueryTrace,
5711 )> {
5712 let (result, trace) = crate::trace::QueryTrace::capture(|| {
5713 self.query_columns_native(conditions, projection, snapshot)
5714 });
5715 Ok((result?, trace))
5716 }
5717
5718 #[allow(clippy::type_complexity)]
5721 pub fn query_columns_native_cached_traced(
5722 &mut self,
5723 conditions: &[crate::query::Condition],
5724 projection: Option<&[u16]>,
5725 snapshot: Snapshot,
5726 ) -> Result<(
5727 Option<Vec<(u16, columnar::NativeColumn)>>,
5728 crate::trace::QueryTrace,
5729 )> {
5730 let (result, trace) = crate::trace::QueryTrace::capture(|| {
5731 self.query_columns_native_cached(conditions, projection, snapshot)
5732 });
5733 Ok((result?, trace))
5734 }
5735
5736 pub fn native_page_cursor_traced(
5738 &self,
5739 snapshot: Snapshot,
5740 projection: Vec<(u16, TypeId)>,
5741 conditions: &[crate::query::Condition],
5742 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
5743 let (result, trace) = crate::trace::QueryTrace::capture(|| {
5744 self.native_page_cursor(snapshot, projection, conditions)
5745 });
5746 Ok((result?, trace))
5747 }
5748
5749 pub fn native_multi_run_cursor_traced(
5751 &self,
5752 snapshot: Snapshot,
5753 projection: Vec<(u16, TypeId)>,
5754 conditions: &[crate::query::Condition],
5755 ) -> Result<(
5756 Option<crate::cursor::MultiRunCursor>,
5757 crate::trace::QueryTrace,
5758 )> {
5759 let (result, trace) = crate::trace::QueryTrace::capture(|| {
5760 self.native_multi_run_cursor(snapshot, projection, conditions)
5761 });
5762 Ok((result?, trace))
5763 }
5764
5765 pub fn count_conditions_traced(
5767 &mut self,
5768 conditions: &[crate::query::Condition],
5769 snapshot: Snapshot,
5770 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
5771 let (result, trace) =
5772 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
5773 Ok((result?, trace))
5774 }
5775
5776 pub fn query_traced(
5778 &mut self,
5779 q: &crate::query::Query,
5780 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
5781 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
5782 Ok((result?, trace))
5783 }
5784
5785 pub fn query_columns_native(
5790 &mut self,
5791 conditions: &[crate::query::Condition],
5792 projection: Option<&[u16]>,
5793 snapshot: Snapshot,
5794 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
5795 use crate::query::Condition;
5796 if self.ttl.is_some() {
5799 return Ok(None);
5800 }
5801 if conditions.is_empty() {
5802 return Ok(None);
5803 }
5804 self.ensure_indexes_complete()?;
5805
5806 let served = |c: &Condition| {
5811 matches!(
5812 c,
5813 Condition::Pk(_)
5814 | Condition::BitmapEq { .. }
5815 | Condition::BitmapIn { .. }
5816 | Condition::BytesPrefix { .. }
5817 | Condition::FmContains { .. }
5818 | Condition::FmContainsAll { .. }
5819 | Condition::Ann { .. }
5820 | Condition::Range { .. }
5821 | Condition::RangeF64 { .. }
5822 | Condition::SparseMatch { .. }
5823 | Condition::MinHashSimilar { .. }
5824 | Condition::IsNull { .. }
5825 | Condition::IsNotNull { .. }
5826 )
5827 };
5828 if !conditions.iter().all(served) {
5829 return Ok(None);
5830 }
5831 let fast_path =
5832 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5833 crate::trace::QueryTrace::record(|t| {
5834 t.run_count = self.run_refs.len();
5835 t.memtable_rows = self.memtable.len();
5836 t.mutable_run_rows = self.mutable_run.len();
5837 t.conditions_pushed = conditions.len();
5838 t.learned_range_used = conditions.iter().any(|c| match c {
5839 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
5840 self.learned_range.contains_key(column_id)
5841 }
5842 _ => false,
5843 });
5844 });
5845 let col_ids: Vec<u16> = projection
5847 .map(|p| p.to_vec())
5848 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
5849 let proj_pairs: Vec<(u16, TypeId)> = col_ids
5850 .iter()
5851 .map(|&cid| {
5852 let ty = self
5853 .schema
5854 .columns
5855 .iter()
5856 .find(|c| c.id == cid)
5857 .map(|c| c.ty.clone())
5858 .unwrap_or(TypeId::Bytes);
5859 (cid, ty)
5860 })
5861 .collect();
5862
5863 if fast_path {
5869 let needs_column = conditions.iter().any(|c| match c {
5872 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
5873 Condition::RangeF64 { column_id, .. } => {
5874 !self.learned_range.contains_key(column_id)
5875 }
5876 _ => false,
5877 });
5878 let mut reader_opt: Option<RunReader> = if needs_column {
5879 Some(self.open_reader(self.run_refs[0].run_id)?)
5880 } else {
5881 None
5882 };
5883 let mut sets: Vec<RowIdSet> = Vec::new();
5884 for c in conditions {
5885 let s = match c {
5886 Condition::Range { column_id, lo, hi }
5887 if !self.learned_range.contains_key(column_id) =>
5888 {
5889 if reader_opt.is_none() {
5890 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
5891 }
5892 reader_opt
5893 .as_mut()
5894 .expect("reader opened for range")
5895 .range_row_id_set_i64(*column_id, *lo, *hi)?
5896 }
5897 Condition::RangeF64 {
5898 column_id,
5899 lo,
5900 lo_inclusive,
5901 hi,
5902 hi_inclusive,
5903 } if !self.learned_range.contains_key(column_id) => {
5904 if reader_opt.is_none() {
5905 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
5906 }
5907 reader_opt
5908 .as_mut()
5909 .expect("reader opened for range")
5910 .range_row_id_set_f64(
5911 *column_id,
5912 *lo,
5913 *lo_inclusive,
5914 *hi,
5915 *hi_inclusive,
5916 )?
5917 }
5918 _ => self.resolve_condition(c, snapshot)?,
5919 };
5920 sets.push(s);
5921 }
5922 let candidates = RowIdSet::intersect_many(sets);
5923 crate::trace::QueryTrace::record(|t| {
5924 t.survivor_count = Some(candidates.len());
5925 });
5926 if candidates.is_empty() {
5927 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
5928 .iter()
5929 .map(|&id| {
5930 (
5931 id,
5932 columnar::null_native(
5933 proj_pairs
5934 .iter()
5935 .find(|(c, _)| c == &id)
5936 .map(|(_, t)| t.clone())
5937 .unwrap_or(TypeId::Bytes),
5938 0,
5939 ),
5940 )
5941 })
5942 .collect();
5943 return Ok(Some(cols));
5944 }
5945 let mut reader = match reader_opt.take() {
5946 Some(r) => r,
5947 None => self.open_reader(self.run_refs[0].run_id)?,
5948 };
5949 let candidate_ids = candidates.into_sorted_vec();
5950 let (positions, fast_rid) = if let Some(positions) =
5951 reader.positions_for_row_ids_fast(&candidate_ids)
5952 {
5953 (positions, true)
5954 } else {
5955 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
5956 match col {
5957 columnar::NativeColumn::Int64 { data, .. } => {
5958 let mut p: Vec<usize> = candidate_ids
5959 .iter()
5960 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
5961 .collect();
5962 p.sort_unstable();
5963 (p, false)
5964 }
5965 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
5966 }
5967 };
5968 crate::trace::QueryTrace::record(|t| {
5969 t.scan_mode = crate::trace::ScanMode::NativePushdown;
5970 t.fast_row_id_map = fast_rid;
5971 });
5972 let mut cols = Vec::with_capacity(col_ids.len());
5973 for cid in &col_ids {
5974 let col = reader.column_native(*cid)?;
5975 cols.push((*cid, col.gather(&positions)));
5976 }
5977 return Ok(Some(cols));
5978 }
5979
5980 if !self.run_refs.is_empty() {
5993 use crate::cursor::{drain_cursor_to_columns, Cursor};
5994 let remaining: usize;
5995 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
5996 let c = self
5997 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
5998 .expect("single-run cursor should build when run_refs.len() == 1");
5999 remaining = c.remaining_rows();
6000 Box::new(c)
6001 } else {
6002 let c = self
6003 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
6004 .expect("multi-run cursor should build when run_refs.len() >= 1");
6005 remaining = c.remaining_rows();
6006 Box::new(c)
6007 };
6008 crate::trace::QueryTrace::record(|t| {
6009 if t.survivor_count.is_none() {
6010 t.survivor_count = Some(remaining);
6011 }
6012 });
6013 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
6014 return Ok(Some(cols));
6015 }
6016
6017 crate::trace::QueryTrace::record(|t| {
6022 t.scan_mode = crate::trace::ScanMode::Materialized;
6023 t.row_materialized = true;
6024 });
6025 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6026 for c in conditions {
6027 sets.push(self.resolve_condition(c, snapshot)?);
6028 }
6029 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6030 let rows = self.rows_for_rids(&rids, snapshot)?;
6031 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
6032 for (cid, ty) in &proj_pairs {
6033 let vals: Vec<Value> = rows
6034 .iter()
6035 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6036 .collect();
6037 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
6038 }
6039 Ok(Some(cols))
6040 }
6041
6042 pub fn native_page_cursor(
6057 &self,
6058 snapshot: Snapshot,
6059 projection: Vec<(u16, TypeId)>,
6060 conditions: &[crate::query::Condition],
6061 ) -> Result<Option<NativePageCursor>> {
6062 use crate::cursor::build_page_plans;
6063 if self.ttl.is_some() {
6064 return Ok(None);
6065 }
6066 if !conditions.is_empty() && !self.indexes_complete {
6069 return Ok(None);
6070 }
6071 if self.run_refs.len() != 1 {
6072 return Ok(None);
6073 }
6074 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6075 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6076
6077 let overlay_rids: HashSet<u64> = {
6080 let mut s = HashSet::new();
6081 for row in self.memtable.visible_versions(snapshot.epoch) {
6082 s.insert(row.row_id.0);
6083 }
6084 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6085 s.insert(row.row_id.0);
6086 }
6087 s
6088 };
6089
6090 let survivors = if conditions.is_empty() {
6094 None
6095 } else {
6096 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
6097 };
6098
6099 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
6106 survivors.clone()
6107 } else if let Some(s) = &survivors {
6108 let mut run_set = s.clone();
6109 run_set.remove_many(overlay_rids.iter().copied());
6110 Some(run_set)
6111 } else {
6112 Some(RowIdSet::from_unsorted(
6113 rids.iter()
6114 .map(|&r| r as u64)
6115 .filter(|r| !overlay_rids.contains(r))
6116 .collect(),
6117 ))
6118 };
6119
6120 let overlay_rows = if overlay_rids.is_empty() {
6121 Vec::new()
6122 } else {
6123 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6124 self.overlay_visible_rows(snapshot, bound)
6125 };
6126
6127 let plans = if positions.is_empty() {
6129 Vec::new()
6130 } else {
6131 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
6132 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
6133 };
6134
6135 let overlay = if overlay_rows.is_empty() {
6137 None
6138 } else {
6139 let filtered =
6140 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6141 if filtered.is_empty() {
6142 None
6143 } else {
6144 Some(self.materialize_overlay(&filtered, &projection))
6145 }
6146 };
6147
6148 let overlay_row_count = overlay
6149 .as_ref()
6150 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6151 .unwrap_or(0);
6152 crate::trace::QueryTrace::record(|t| {
6153 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
6154 t.run_count = self.run_refs.len();
6155 t.memtable_rows = self.memtable.len();
6156 t.mutable_run_rows = self.mutable_run.len();
6157 t.overlay_rows = overlay_row_count;
6158 t.conditions_pushed = conditions.len();
6159 t.pages_decoded = plans
6160 .iter()
6161 .map(|p| p.positions.len())
6162 .sum::<usize>()
6163 .min(1);
6164 });
6165
6166 Ok(Some(NativePageCursor::new_with_overlay(
6167 reader, projection, plans, overlay,
6168 )))
6169 }
6170 #[allow(clippy::type_complexity)]
6180 pub fn native_multi_run_cursor(
6181 &self,
6182 snapshot: Snapshot,
6183 projection: Vec<(u16, TypeId)>,
6184 conditions: &[crate::query::Condition],
6185 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
6186 use crate::cursor::{MultiRunCursor, RunStream};
6187 use crate::sorted_run::SYS_ROW_ID;
6188 use std::collections::{BinaryHeap, HashMap, HashSet};
6189 if self.ttl.is_some() {
6190 return Ok(None);
6191 }
6192 if !conditions.is_empty() && !self.indexes_complete {
6195 return Ok(None);
6196 }
6197 if self.run_refs.is_empty() {
6198 return Ok(None);
6199 }
6200
6201 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
6203 Vec::with_capacity(self.run_refs.len());
6204 for rr in &self.run_refs {
6205 let mut reader = self.open_reader(rr.run_id)?;
6206 let (rids, eps, del) = reader.system_columns_native()?;
6207 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
6208 run_meta.push((reader, rids, eps, del, page_rows));
6209 }
6210
6211 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
6215 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
6216 for i in 0..rids.len() {
6217 let rid = rids[i] as u64;
6218 let e = eps[i] as u64;
6219 if e > snapshot.epoch.0 {
6220 continue;
6221 }
6222 let is_del = del[i] != 0;
6223 best.entry(rid)
6224 .and_modify(|cur| {
6225 if e > cur.0 {
6226 *cur = (e, run_idx, i, is_del);
6227 }
6228 })
6229 .or_insert((e, run_idx, i, is_del));
6230 }
6231 }
6232
6233 let overlay_rids: HashSet<u64> = {
6235 let mut s = HashSet::new();
6236 for row in self.memtable.visible_versions(snapshot.epoch) {
6237 s.insert(row.row_id.0);
6238 }
6239 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6240 s.insert(row.row_id.0);
6241 }
6242 s
6243 };
6244
6245 let survivors: Option<RowIdSet> = if conditions.is_empty() {
6247 None
6248 } else {
6249 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6250 for c in conditions {
6251 sets.push(self.resolve_condition(c, snapshot)?);
6252 }
6253 Some(RowIdSet::intersect_many(sets))
6254 };
6255
6256 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
6260 for (rid, (_, run_idx, pos, deleted)) in &best {
6261 if *deleted {
6262 continue;
6263 }
6264 if overlay_rids.contains(rid) {
6265 continue;
6266 }
6267 if let Some(s) = &survivors {
6268 if !s.contains(*rid) {
6269 continue;
6270 }
6271 }
6272 per_run[*run_idx].push((*rid, *pos));
6273 }
6274 for v in per_run.iter_mut() {
6275 v.sort_unstable_by_key(|&(rid, _)| rid);
6276 }
6277
6278 let mut streams = Vec::with_capacity(run_meta.len());
6280 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
6281 let mut total = 0usize;
6282 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
6283 let mut starts = Vec::with_capacity(page_rows.len());
6284 let mut acc = 0usize;
6285 for &r in &page_rows {
6286 starts.push(acc);
6287 acc += r;
6288 }
6289 let mut survivors_vec: Vec<(u64, usize, usize)> =
6290 Vec::with_capacity(per_run[run_idx].len());
6291 for &(rid, pos) in &per_run[run_idx] {
6292 let page_seq = match starts.partition_point(|&s| s <= pos) {
6293 0 => continue,
6294 p => p - 1,
6295 };
6296 let within = pos - starts[page_seq];
6297 survivors_vec.push((rid, page_seq, within));
6298 }
6299 total += survivors_vec.len();
6300 if let Some(&(rid, _, _)) = survivors_vec.first() {
6301 heap.push(std::cmp::Reverse((rid, run_idx)));
6302 }
6303 streams.push(RunStream::new(reader, survivors_vec, page_rows));
6304 }
6305
6306 let overlay_rows = if overlay_rids.is_empty() {
6308 Vec::new()
6309 } else {
6310 let bound = Self::overlay_materialization_bound(conditions, &survivors);
6311 self.overlay_visible_rows(snapshot, bound)
6312 };
6313 let overlay = if overlay_rows.is_empty() {
6314 None
6315 } else {
6316 let filtered =
6317 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
6318 if filtered.is_empty() {
6319 None
6320 } else {
6321 Some(self.materialize_overlay(&filtered, &projection))
6322 }
6323 };
6324
6325 let overlay_row_count = overlay
6326 .as_ref()
6327 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
6328 .unwrap_or(0);
6329 crate::trace::QueryTrace::record(|t| {
6330 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
6331 t.run_count = self.run_refs.len();
6332 t.memtable_rows = self.memtable.len();
6333 t.mutable_run_rows = self.mutable_run.len();
6334 t.overlay_rows = overlay_row_count;
6335 t.conditions_pushed = conditions.len();
6336 t.survivor_count = Some(total);
6337 });
6338
6339 Ok(Some(MultiRunCursor::new(
6340 streams, projection, heap, total, overlay,
6341 )))
6342 }
6343
6344 fn overlay_materialization_bound<'a>(
6356 conditions: &[crate::query::Condition],
6357 survivors: &'a Option<RowIdSet>,
6358 ) -> Option<&'a RowIdSet> {
6359 use crate::query::Condition;
6360 let has_range = conditions
6361 .iter()
6362 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6363 if has_range {
6364 None
6365 } else {
6366 survivors.as_ref()
6367 }
6368 }
6369
6370 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
6382 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
6383 let mut fold = |row: Row| {
6384 if let Some(b) = bound {
6385 if !b.contains(row.row_id.0) {
6386 return;
6387 }
6388 }
6389 best.entry(row.row_id.0)
6390 .and_modify(|(be, br)| {
6391 if row.committed_epoch > *be {
6392 *be = row.committed_epoch;
6393 *br = row.clone();
6394 }
6395 })
6396 .or_insert_with(|| (row.committed_epoch, row));
6397 };
6398 for row in self.memtable.visible_versions(snapshot.epoch) {
6399 fold(row);
6400 }
6401 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6402 fold(row);
6403 }
6404 let mut out: Vec<Row> = best
6405 .into_values()
6406 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
6407 .collect();
6408 out.sort_by_key(|r| r.row_id);
6409 out
6410 }
6411
6412 fn filter_overlay_rows(
6420 &self,
6421 rows: Vec<Row>,
6422 conditions: &[crate::query::Condition],
6423 survivors: Option<&RowIdSet>,
6424 snapshot: Snapshot,
6425 ) -> Result<Vec<Row>> {
6426 if conditions.is_empty() {
6427 return Ok(rows);
6428 }
6429 use crate::query::Condition;
6430 let all_index_served = !conditions
6434 .iter()
6435 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
6436 if all_index_served {
6437 return Ok(rows
6438 .into_iter()
6439 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
6440 .collect());
6441 }
6442 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
6445 for c in conditions {
6446 let s = match c {
6447 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
6448 _ => self.resolve_condition(c, snapshot)?,
6449 };
6450 per_cond_sets.push(s);
6451 }
6452 Ok(rows
6453 .into_iter()
6454 .filter(|row| {
6455 conditions.iter().enumerate().all(|(i, c)| match c {
6456 Condition::Range { column_id, lo, hi } => {
6457 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
6458 }
6459 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
6460 match row.columns.get(column_id) {
6461 Some(Value::Float64(v)) => {
6462 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
6463 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
6464 lo_ok && hi_ok
6465 }
6466 _ => false,
6467 }
6468 }
6469 _ => per_cond_sets[i].contains(row.row_id.0),
6470 })
6471 })
6472 .collect())
6473 }
6474
6475 fn materialize_overlay(
6478 &self,
6479 rows: &[Row],
6480 projection: &[(u16, TypeId)],
6481 ) -> Vec<columnar::NativeColumn> {
6482 if projection.is_empty() {
6483 return vec![columnar::null_native(TypeId::Int64, rows.len())];
6484 }
6485 let mut cols = Vec::with_capacity(projection.len());
6486 for (cid, ty) in projection {
6487 let vals: Vec<Value> = rows
6488 .iter()
6489 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
6490 .collect();
6491 cols.push(columnar::values_to_native(ty.clone(), &vals));
6492 }
6493 cols
6494 }
6495
6496 fn resolve_survivor_rids(
6501 &self,
6502 conditions: &[crate::query::Condition],
6503 reader: &mut RunReader,
6504 snapshot: Snapshot,
6505 ) -> Result<RowIdSet> {
6506 use crate::query::Condition;
6507 let mut sets: Vec<RowIdSet> = Vec::new();
6508 for c in conditions {
6509 self.validate_condition(c)?;
6510 let s: RowIdSet = match c {
6511 Condition::Pk(key) => {
6512 let lookup = self
6513 .schema
6514 .primary_key()
6515 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
6516 .unwrap_or_else(|| key.clone());
6517 self.hot
6518 .get(&lookup)
6519 .map(|r| RowIdSet::one(r.0))
6520 .unwrap_or_else(RowIdSet::empty)
6521 }
6522 Condition::BitmapEq { column_id, value } => {
6523 let lookup = self.index_lookup_key_bytes(*column_id, value);
6524 self.bitmap
6525 .get(column_id)
6526 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
6527 .unwrap_or_else(RowIdSet::empty)
6528 }
6529 Condition::BitmapIn { column_id, values } => {
6530 let bm = self.bitmap.get(column_id);
6531 let mut acc = roaring::RoaringBitmap::new();
6532 if let Some(b) = bm {
6533 for v in values {
6534 let lookup = self.index_lookup_key_bytes(*column_id, v);
6535 acc |= b.get(&lookup);
6536 }
6537 }
6538 RowIdSet::from_roaring(acc)
6539 }
6540 Condition::BytesPrefix { column_id, prefix } => {
6541 if let Some(b) = self.bitmap.get(column_id) {
6542 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
6543 let mut acc = roaring::RoaringBitmap::new();
6544 for key in b.keys() {
6545 if key.starts_with(&lookup_prefix) {
6546 acc |= b.get(key);
6547 }
6548 }
6549 RowIdSet::from_roaring(acc)
6550 } else {
6551 RowIdSet::empty()
6552 }
6553 }
6554 Condition::FmContains { column_id, pattern } => self
6555 .fm
6556 .get(column_id)
6557 .map(|f| {
6558 RowIdSet::from_unsorted(
6559 f.locate(pattern).into_iter().map(|r| r.0).collect(),
6560 )
6561 })
6562 .unwrap_or_else(RowIdSet::empty),
6563 Condition::FmContainsAll {
6564 column_id,
6565 patterns,
6566 } => {
6567 if let Some(f) = self.fm.get(column_id) {
6568 let sets: Vec<RowIdSet> = patterns
6569 .iter()
6570 .map(|pat| {
6571 RowIdSet::from_unsorted(
6572 f.locate(pat).into_iter().map(|r| r.0).collect(),
6573 )
6574 })
6575 .collect();
6576 RowIdSet::intersect_many(sets)
6577 } else {
6578 RowIdSet::empty()
6579 }
6580 }
6581 Condition::Ann {
6582 column_id,
6583 query,
6584 k,
6585 } => RowIdSet::from_unsorted(
6586 self.retrieve_filtered(
6587 &crate::query::Retriever::Ann {
6588 column_id: *column_id,
6589 query: query.clone(),
6590 k: *k,
6591 },
6592 snapshot,
6593 None,
6594 None,
6595 )?
6596 .into_iter()
6597 .map(|hit| hit.row_id.0)
6598 .collect(),
6599 ),
6600 Condition::SparseMatch {
6601 column_id,
6602 query,
6603 k,
6604 } => RowIdSet::from_unsorted(
6605 self.retrieve_filtered(
6606 &crate::query::Retriever::Sparse {
6607 column_id: *column_id,
6608 query: query.clone(),
6609 k: *k,
6610 },
6611 snapshot,
6612 None,
6613 None,
6614 )?
6615 .into_iter()
6616 .map(|hit| hit.row_id.0)
6617 .collect(),
6618 ),
6619 Condition::MinHashSimilar {
6620 column_id,
6621 query,
6622 k,
6623 } => match self.minhash.get(column_id) {
6624 Some(index) => {
6625 let candidates = index.candidate_row_ids(query);
6626 let eligible =
6627 self.eligible_candidate_ids(&candidates, *column_id, snapshot)?;
6628 RowIdSet::from_unsorted(
6629 index
6630 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
6631 .into_iter()
6632 .map(|(row_id, _)| row_id.0)
6633 .collect(),
6634 )
6635 }
6636 None => RowIdSet::empty(),
6637 },
6638 Condition::Range { column_id, lo, hi } => {
6639 if let Some(li) = self.learned_range.get(column_id) {
6640 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
6641 } else {
6642 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
6643 }
6644 }
6645 Condition::RangeF64 {
6646 column_id,
6647 lo,
6648 lo_inclusive,
6649 hi,
6650 hi_inclusive,
6651 } => {
6652 if let Some(li) = self.learned_range.get(column_id) {
6653 RowIdSet::from_unsorted(
6654 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
6655 .into_iter()
6656 .collect(),
6657 )
6658 } else {
6659 reader.range_row_id_set_f64(
6660 *column_id,
6661 *lo,
6662 *lo_inclusive,
6663 *hi,
6664 *hi_inclusive,
6665 )?
6666 }
6667 }
6668 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
6669 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
6670 };
6671 sets.push(s);
6672 }
6673 Ok(RowIdSet::intersect_many(sets))
6674 }
6675
6676 pub fn scan_cursor(
6697 &self,
6698 snapshot: Snapshot,
6699 projection: Vec<(u16, TypeId)>,
6700 conditions: &[crate::query::Condition],
6701 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
6702 if self.ttl.is_some() {
6703 return Ok(None);
6704 }
6705 if !conditions.is_empty() && !self.indexes_complete {
6711 return Ok(None);
6712 }
6713 if self.run_refs.len() == 1 {
6714 Ok(self
6715 .native_page_cursor(snapshot, projection, conditions)?
6716 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
6717 } else {
6718 Ok(self
6719 .native_multi_run_cursor(snapshot, projection, conditions)?
6720 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
6721 }
6722 }
6723
6724 pub fn aggregate_native(
6738 &self,
6739 snapshot: Snapshot,
6740 column: Option<u16>,
6741 conditions: &[crate::query::Condition],
6742 agg: NativeAgg,
6743 ) -> Result<Option<NativeAggResult>> {
6744 if self.ttl.is_some() {
6745 return Ok(None);
6746 }
6747 if self.run_refs.len() == 1 && conditions.is_empty() {
6749 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
6750 return Ok(Some(res));
6751 }
6752 }
6753 if matches!(agg, NativeAgg::Count) && column.is_none() {
6755 return Ok(self
6756 .scan_cursor(snapshot, Vec::new(), conditions)?
6757 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
6758 }
6759 let cid = match column {
6762 Some(c) => c,
6763 None => return Ok(None),
6764 };
6765 let ty = self.column_type(cid);
6766 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
6767 else {
6768 return Ok(None);
6769 };
6770 match ty {
6771 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
6772 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
6773 Ok(Some(pack_int(agg, count, sum, mn, mx)))
6774 }
6775 TypeId::Float64 => {
6776 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
6777 Ok(Some(pack_float(agg, count, sum, mn, mx)))
6778 }
6779 _ => Ok(None),
6780 }
6781 }
6782
6783 fn aggregate_from_stats(
6791 &self,
6792 snapshot: Snapshot,
6793 column: Option<u16>,
6794 agg: NativeAgg,
6795 ) -> Result<Option<NativeAggResult>> {
6796 let cid = match (agg, column) {
6797 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
6798 _ => return Ok(None), };
6800 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
6801 return Ok(None);
6802 };
6803 let Some(cs) = stats.get(&cid) else {
6804 return Ok(None);
6805 };
6806 match agg {
6807 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
6809 self.live_count.saturating_sub(cs.null_count),
6810 ))),
6811 NativeAgg::Min | NativeAgg::Max => {
6812 let bound = if agg == NativeAgg::Min {
6813 &cs.min
6814 } else {
6815 &cs.max
6816 };
6817 match bound {
6818 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
6819 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
6820 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
6825 None => Ok(None),
6826 }
6827 }
6828 _ => Ok(None),
6829 }
6830 }
6831
6832 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
6841 if self.ttl.is_some() {
6842 return Ok(None);
6843 }
6844 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
6845 return Ok(None);
6846 }
6847 self.ensure_indexes_complete()?;
6850 let reader = self.open_reader(self.run_refs[0].run_id)?;
6851 if self.live_count != reader.row_count() as u64 {
6852 return Ok(None);
6853 }
6854 let Some(bm) = self.bitmap.get(&column_id) else {
6855 return Ok(None); };
6857 let mut distinct = bm.value_count() as u64;
6858 if !bm.get(&Value::Null.encode_key()).is_empty() {
6861 distinct = distinct.saturating_sub(1);
6862 }
6863 Ok(Some(distinct))
6864 }
6865
6866 pub fn aggregate_incremental(
6878 &mut self,
6879 cache_key: u64,
6880 conditions: &[crate::query::Condition],
6881 column: Option<u16>,
6882 agg: NativeAgg,
6883 ) -> Result<IncrementalAggResult> {
6884 let snap = self.snapshot();
6885 let cur_wm = self.allocator.current().0;
6886 let cur_epoch = snap.epoch.0;
6887 let incremental_ok = self.ttl.is_none()
6894 && !self.had_deletes
6895 && self.memtable.is_empty()
6896 && self.mutable_run.is_empty();
6897
6898 if incremental_ok {
6901 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
6902 if cached.epoch == cur_epoch {
6903 return Ok(IncrementalAggResult {
6904 state: cached.state,
6905 incremental: true,
6906 delta_rows: 0,
6907 });
6908 }
6909 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
6910 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
6911 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
6912 let index_sets = self.resolve_index_conditions(conditions, snap)?;
6913 let delta_state = agg_state_from_rows(
6914 &delta_rows,
6915 conditions,
6916 &index_sets,
6917 column,
6918 agg,
6919 &self.schema,
6920 )?;
6921 let merged = cached.state.merge(delta_state);
6922 let delta_n = delta_rids.len() as u64;
6923 self.agg_cache.insert(
6924 cache_key,
6925 CachedAgg {
6926 state: merged.clone(),
6927 watermark: cur_wm,
6928 epoch: cur_epoch,
6929 },
6930 );
6931 return Ok(IncrementalAggResult {
6932 state: merged,
6933 incremental: true,
6934 delta_rows: delta_n,
6935 });
6936 }
6937 }
6938 }
6939
6940 let cursor_ok =
6945 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
6946 let state = if cursor_ok && agg != NativeAgg::Avg {
6947 match self.aggregate_native(snap, column, conditions, agg)? {
6948 Some(result) => {
6949 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
6950 }
6951 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
6952 }
6953 } else {
6954 self.agg_state_full_scan(conditions, column, agg, snap)?
6955 };
6956 if incremental_ok {
6958 self.agg_cache.insert(
6959 cache_key,
6960 CachedAgg {
6961 state: state.clone(),
6962 watermark: cur_wm,
6963 epoch: cur_epoch,
6964 },
6965 );
6966 }
6967 Ok(IncrementalAggResult {
6968 state,
6969 incremental: false,
6970 delta_rows: 0,
6971 })
6972 }
6973
6974 fn agg_state_full_scan(
6977 &self,
6978 conditions: &[crate::query::Condition],
6979 column: Option<u16>,
6980 agg: NativeAgg,
6981 snap: Snapshot,
6982 ) -> Result<AggState> {
6983 let rows = self.visible_rows(snap)?;
6984 let index_sets = self.resolve_index_conditions(conditions, snap)?;
6985 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
6986 }
6987
6988 fn resolve_index_conditions(
6991 &self,
6992 conditions: &[crate::query::Condition],
6993 snapshot: Snapshot,
6994 ) -> Result<Vec<RowIdSet>> {
6995 use crate::query::Condition;
6996 let mut sets = Vec::new();
6997 for c in conditions {
6998 if matches!(
6999 c,
7000 Condition::Ann { .. }
7001 | Condition::SparseMatch { .. }
7002 | Condition::MinHashSimilar { .. }
7003 ) {
7004 sets.push(self.resolve_condition(c, snapshot)?);
7005 }
7006 }
7007 Ok(sets)
7008 }
7009
7010 fn column_type(&self, cid: u16) -> TypeId {
7011 self.schema
7012 .columns
7013 .iter()
7014 .find(|c| c.id == cid)
7015 .map(|c| c.ty.clone())
7016 .unwrap_or(TypeId::Bytes)
7017 }
7018
7019 pub fn approx_aggregate(
7028 &mut self,
7029 conditions: &[crate::query::Condition],
7030 column: Option<u16>,
7031 agg: ApproxAgg,
7032 z: f64,
7033 ) -> Result<Option<ApproxResult>> {
7034 use crate::query::Condition;
7035 self.ensure_reservoir_complete()?;
7036 let snapshot = self.snapshot();
7037 let n_pop = self.count();
7038 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
7039 if sample_rids.is_empty() {
7040 return Ok(None);
7041 }
7042 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
7044 let s = live_sample.len();
7045 if s == 0 {
7046 return Ok(None);
7047 }
7048
7049 let mut index_sets: Vec<RowIdSet> = Vec::new();
7052 for c in conditions {
7053 if matches!(
7054 c,
7055 Condition::Ann { .. }
7056 | Condition::SparseMatch { .. }
7057 | Condition::MinHashSimilar { .. }
7058 ) {
7059 index_sets.push(self.resolve_condition(c, snapshot)?);
7060 }
7061 }
7062
7063 let cid = match (agg, column) {
7065 (ApproxAgg::Count, _) => None,
7066 (_, Some(c)) => Some(c),
7067 _ => return Ok(None),
7068 };
7069 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
7070 for r in &live_sample {
7071 if !conditions
7073 .iter()
7074 .all(|c| condition_matches_row(c, r, &self.schema))
7075 {
7076 continue;
7077 }
7078 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
7080 continue;
7081 }
7082 if let Some(cid) = cid {
7083 if let Some(v) = as_f64(r.columns.get(&cid)) {
7084 passing_vals.push(v);
7085 } } else {
7087 passing_vals.push(0.0); }
7089 }
7090 let m = passing_vals.len();
7091
7092 let (point, half) = match agg {
7093 ApproxAgg::Count => {
7094 let p = m as f64 / s as f64;
7096 let point = n_pop as f64 * p;
7097 let var = if s > 1 {
7098 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
7099 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7100 } else {
7101 0.0
7102 };
7103 (point, z * var.sqrt())
7104 }
7105 ApproxAgg::Sum => {
7106 let y: Vec<f64> = live_sample
7108 .iter()
7109 .map(|r| {
7110 let passes_row = conditions
7111 .iter()
7112 .all(|c| condition_matches_row(c, r, &self.schema))
7113 && index_sets.iter().all(|set| set.contains(r.row_id.0));
7114 if passes_row {
7115 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
7116 } else {
7117 0.0
7118 }
7119 })
7120 .collect();
7121 let mean_y = y.iter().sum::<f64>() / s as f64;
7122 let point = n_pop as f64 * mean_y;
7123 let var = if s > 1 {
7124 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
7125 let var_y = ss / (s - 1) as f64;
7126 n_pop as f64 * n_pop as f64 * var_y / s as f64
7127 * (1.0 - s as f64 / n_pop as f64).max(0.0)
7128 } else {
7129 0.0
7130 };
7131 (point, z * var.sqrt())
7132 }
7133 ApproxAgg::Avg => {
7134 if m == 0 {
7135 return Ok(Some(ApproxResult {
7136 point: 0.0,
7137 ci_low: 0.0,
7138 ci_high: 0.0,
7139 n_population: n_pop,
7140 n_sample_live: s,
7141 n_passing: 0,
7142 }));
7143 }
7144 let mean = passing_vals.iter().sum::<f64>() / m as f64;
7145 let half = if m > 1 {
7146 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
7147 let sd = (ss / (m - 1) as f64).sqrt();
7148 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
7149 z * sd / (m as f64).sqrt() * fpc.sqrt()
7150 } else {
7151 0.0
7152 };
7153 (mean, half)
7154 }
7155 };
7156
7157 Ok(Some(ApproxResult {
7158 point,
7159 ci_low: point - half,
7160 ci_high: point + half,
7161 n_population: n_pop,
7162 n_sample_live: s,
7163 n_passing: m,
7164 }))
7165 }
7166
7167 pub fn exact_column_stats(
7175 &self,
7176 _snapshot: Snapshot,
7177 projection: &[u16],
7178 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
7179 if self.ttl.is_some()
7180 || !(self.memtable.is_empty()
7181 && self.mutable_run.is_empty()
7182 && self.run_refs.len() == 1)
7183 {
7184 return Ok(None);
7185 }
7186 let reader = self.open_reader(self.run_refs[0].run_id)?;
7187 if self.live_count != reader.row_count() as u64 {
7188 return Ok(None);
7189 }
7190 let mut out = HashMap::new();
7191 for &cid in projection {
7192 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
7193 Some(c) => c,
7194 None => continue,
7195 };
7196 let Some(stats) = reader.column_page_stats(cid) else {
7198 out.insert(
7199 cid,
7200 ColumnStat {
7201 min: None,
7202 max: None,
7203 null_count: self.live_count,
7204 },
7205 );
7206 continue;
7207 };
7208 let stat = match cdef.ty {
7209 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
7210 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
7211 min: mn.map(Value::Int64),
7212 max: mx.map(Value::Int64),
7213 null_count: n,
7214 })
7215 }
7216 TypeId::Float64 => {
7217 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
7218 min: mn.map(Value::Float64),
7219 max: mx.map(Value::Float64),
7220 null_count: n,
7221 })
7222 }
7223 _ => None,
7224 };
7225 if let Some(s) = stat {
7226 out.insert(cid, s);
7227 }
7228 }
7229 Ok(Some(out))
7230 }
7231
7232 pub fn dir(&self) -> &Path {
7233 &self.dir
7234 }
7235
7236 pub fn schema(&self) -> &Schema {
7237 &self.schema
7238 }
7239
7240 pub(crate) fn set_catalog_name(&mut self, name: String) {
7241 self.name = name;
7242 }
7243
7244 pub(crate) fn prepare_alter_column(
7245 &mut self,
7246 column_name: &str,
7247 change: &AlterColumn,
7248 ) -> Result<ColumnDef> {
7249 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
7250 return Err(MongrelError::InvalidArgument(
7251 "ALTER COLUMN requires committing staged writes first".into(),
7252 ));
7253 }
7254 let old = self
7255 .schema
7256 .columns
7257 .iter()
7258 .find(|c| c.name == column_name)
7259 .cloned()
7260 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7261 let mut next = old.clone();
7262
7263 if let Some(name) = &change.name {
7264 let trimmed = name.trim();
7265 if trimmed.is_empty() {
7266 return Err(MongrelError::InvalidArgument(
7267 "ALTER COLUMN name must not be empty".into(),
7268 ));
7269 }
7270 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
7271 return Err(MongrelError::Schema(format!(
7272 "column {trimmed} already exists"
7273 )));
7274 }
7275 next.name = trimmed.to_string();
7276 }
7277
7278 if let Some(ty) = &change.ty {
7279 next.ty = ty.clone();
7280 }
7281 if let Some(flags) = change.flags {
7282 validate_alter_column_flags(old.flags, flags)?;
7283 next.flags = flags;
7284 }
7285
7286 if let Some(default_change) = &change.default_value {
7287 next.default_value = default_change.clone();
7288 }
7289
7290 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
7291 if old.flags.contains(ColumnFlags::NULLABLE)
7292 && !next.flags.contains(ColumnFlags::NULLABLE)
7293 && self.column_has_nulls(old.id)?
7294 {
7295 return Err(MongrelError::InvalidArgument(format!(
7296 "column '{}' contains NULL values",
7297 old.name
7298 )));
7299 }
7300 Ok(next)
7301 }
7302
7303 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
7304 let idx = self
7305 .schema
7306 .columns
7307 .iter()
7308 .position(|c| c.id == column.id)
7309 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
7310 if self.schema.columns[idx] == column {
7311 return Ok(());
7312 }
7313 self.schema.columns[idx] = column;
7314 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7315 self.schema.validate_auto_increment()?;
7316 self.schema.validate_defaults()?;
7317 self.auto_inc = resolve_auto_inc(&self.schema);
7318 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
7319 write_schema(&self.dir, &self.schema)?;
7320 self.clear_result_cache();
7321 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7322 self.persist_manifest(self.current_epoch())?;
7323 Ok(())
7324 }
7325
7326 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
7327 self.ensure_writable()?;
7328 let column = self.prepare_alter_column(column_name, &change)?;
7329 self.apply_altered_column(column.clone())?;
7330 Ok(column)
7331 }
7332
7333 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
7334 if self.live_count == 0 {
7335 return Ok(false);
7336 }
7337 let snap = self.snapshot();
7338 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
7339 Ok(columns
7340 .first()
7341 .map(|(_, col)| col.null_count(col.len()) != 0)
7342 .unwrap_or(true))
7343 }
7344
7345 fn has_stored_versions(&self) -> bool {
7346 !self.memtable.is_empty()
7347 || !self.mutable_run.is_empty()
7348 || self.run_refs.iter().any(|r| r.row_count > 0)
7349 || !self.retiring.is_empty()
7350 }
7351
7352 pub fn add_column(
7357 &mut self,
7358 name: &str,
7359 ty: TypeId,
7360 flags: ColumnFlags,
7361 default_value: Option<crate::schema::DefaultExpr>,
7362 ) -> Result<u16> {
7363 self.add_column_with_id(name, ty, flags, default_value, None)
7364 }
7365
7366 pub fn add_column_with_id(
7367 &mut self,
7368 name: &str,
7369 ty: TypeId,
7370 flags: ColumnFlags,
7371 default_value: Option<crate::schema::DefaultExpr>,
7372 requested_id: Option<u16>,
7373 ) -> Result<u16> {
7374 self.ensure_writable()?;
7375 if self.schema.columns.iter().any(|c| c.name == name) {
7376 return Err(MongrelError::Schema(format!(
7377 "column {name} already exists"
7378 )));
7379 }
7380 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
7381 if self.schema.columns.iter().any(|c| c.id == id) {
7382 return Err(MongrelError::Schema(format!(
7383 "column id {id} already exists"
7384 )));
7385 }
7386 id
7387 } else {
7388 self.schema
7389 .columns
7390 .iter()
7391 .map(|c| c.id)
7392 .max()
7393 .unwrap_or(0)
7394 .checked_add(1)
7395 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
7396 };
7397 self.schema.columns.push(ColumnDef {
7398 id,
7399 name: name.to_string(),
7400 ty,
7401 flags,
7402 default_value,
7403 });
7404 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7405 self.schema.validate_auto_increment()?;
7406 self.schema.validate_defaults()?;
7407 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
7408 self.auto_inc = resolve_auto_inc(&self.schema);
7409 }
7410 write_schema(&self.dir, &self.schema)?;
7411 self.clear_result_cache();
7412 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7414 self.persist_manifest(self.current_epoch())?;
7415 Ok(id)
7416 }
7417
7418 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
7427 self.ensure_writable()?;
7428 let cid = self
7429 .schema
7430 .columns
7431 .iter()
7432 .find(|c| c.name == column_name)
7433 .map(|c| c.id)
7434 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
7435 let ty = self
7436 .schema
7437 .columns
7438 .iter()
7439 .find(|c| c.id == cid)
7440 .map(|c| c.ty.clone())
7441 .unwrap_or(TypeId::Int64);
7442 if !matches!(
7443 ty,
7444 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
7445 ) {
7446 return Err(MongrelError::Schema(format!(
7447 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
7448 )));
7449 }
7450 if self
7451 .schema
7452 .indexes
7453 .iter()
7454 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
7455 {
7456 return Ok(()); }
7458 self.schema.indexes.push(IndexDef {
7459 name: format!("{}_learned_range", column_name),
7460 column_id: cid,
7461 kind: IndexKind::LearnedRange,
7462 predicate: None,
7463 options: Default::default(),
7464 });
7465 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
7466 write_schema(&self.dir, &self.schema)?;
7467 self.build_learned_ranges()?;
7468 Ok(())
7469 }
7470
7471 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
7474 self.sync_byte_threshold = threshold;
7475 if let WalSink::Private(w) = &mut self.wal {
7476 w.set_sync_byte_threshold(threshold);
7477 }
7478 }
7479
7480 pub fn page_cache_flush(&self) {
7484 self.page_cache.flush_to_disk();
7485 }
7486
7487 pub fn page_cache_len(&self) -> usize {
7489 self.page_cache.len()
7490 }
7491
7492 pub fn decoded_cache_len(&self) -> usize {
7495 self.decoded_cache.len()
7496 }
7497
7498 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
7501 self.memtable.drain_sorted()
7502 }
7503
7504 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
7505 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
7506 }
7507
7508 pub(crate) fn table_dir(&self) -> &Path {
7509 &self.dir
7510 }
7511
7512 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
7513 &self.schema
7514 }
7515
7516 pub(crate) fn alloc_run_id(&mut self) -> u64 {
7517 let id = self.next_run_id;
7518 self.next_run_id += 1;
7519 id
7520 }
7521
7522 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
7523 self.run_refs.push(run_ref);
7524 }
7525
7526 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
7536 self.retiring.push(crate::manifest::RetiredRun {
7537 run_id,
7538 retire_epoch,
7539 });
7540 }
7541
7542 pub(crate) fn reap_retiring(
7546 &mut self,
7547 min_active: Epoch,
7548 backup_pinned: &std::collections::HashSet<u128>,
7549 ) -> Result<usize> {
7550 if self.retiring.is_empty() {
7551 return Ok(0);
7552 }
7553 let mut reaped = 0;
7554 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
7555 for r in std::mem::take(&mut self.retiring) {
7561 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
7562 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
7563 reaped += 1;
7564 } else {
7565 kept.push(r);
7566 }
7567 }
7568 self.retiring = kept;
7569 if reaped > 0 {
7570 self.persist_manifest(self.current_epoch())?;
7571 }
7572 Ok(reaped)
7573 }
7574
7575 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
7576 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
7577 return false;
7578 }
7579 self.live_count = self.live_count.saturating_add(run_ref.row_count);
7580 self.run_refs.push(run_ref);
7581 self.indexes_complete = false;
7582 true
7583 }
7584
7585 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
7586 self.kek.as_ref()
7587 }
7588
7589 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
7590 let mut reader = RunReader::open_with_cache(
7591 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
7592 self.schema.clone(),
7593 self.kek.clone(),
7594 Some(self.page_cache.clone()),
7595 Some(self.decoded_cache.clone()),
7596 self.table_id,
7597 Some(&self.verified_runs),
7598 )?;
7599 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
7603 reader.set_uniform_epoch(Epoch(rr.epoch_created));
7604 }
7605 Ok(reader)
7606 }
7607
7608 pub(crate) fn run_refs(&self) -> &[RunRef] {
7609 &self.run_refs
7610 }
7611
7612 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
7613 self.retiring.iter().map(|run| run.run_id)
7614 }
7615
7616 pub(crate) fn runs_dir(&self) -> PathBuf {
7617 self.dir.join(RUNS_DIR)
7618 }
7619
7620 pub(crate) fn wal_dir(&self) -> PathBuf {
7621 self.dir.join(WAL_DIR)
7622 }
7623
7624 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
7625 self.run_refs = refs;
7626 }
7627
7628 pub(crate) fn next_run_id(&self) -> u64 {
7629 self.next_run_id
7630 }
7631
7632 pub(crate) fn compaction_zstd_level(&self) -> i32 {
7633 self.compaction_zstd_level
7634 }
7635
7636 pub(crate) fn bump_next_run_id(&mut self) {
7637 self.next_run_id += 1;
7638 }
7639
7640 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
7641 self.kek.clone()
7642 }
7643
7644 #[cfg(feature = "encryption")]
7648 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
7649 self.kek.as_ref().map(|k| k.derive_idx_key())
7650 }
7651
7652 #[cfg(not(feature = "encryption"))]
7653 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
7654 None
7655 }
7656
7657 #[cfg(feature = "encryption")]
7661 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
7662 self.kek.as_ref().map(|k| *k.derive_meta_key())
7663 }
7664
7665 #[cfg(not(feature = "encryption"))]
7666 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
7667 None
7668 }
7669
7670 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
7673 self.column_keys
7674 .iter()
7675 .map(|(&id, &(_, scheme))| (id, scheme))
7676 .collect()
7677 }
7678
7679 #[cfg(feature = "encryption")]
7684 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
7685 self.tokenize_value_enc(column_id, v)
7686 }
7687
7688 #[cfg(feature = "encryption")]
7689 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
7690 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7691 let (key, scheme) = self.column_keys.get(&column_id)?;
7692 let token: Vec<u8> = match (*scheme, v) {
7693 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
7694 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
7695 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
7696 _ => hmac_token(key, &v.encode_key()).to_vec(),
7697 };
7698 Some(Value::Bytes(token))
7699 }
7700
7701 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
7703 self.index_lookup_key_bytes(column_id, &v.encode_key())
7704 }
7705
7706 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
7709 #[cfg(feature = "encryption")]
7710 {
7711 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
7712 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
7713 if *scheme == SCHEME_HMAC_EQ {
7714 return hmac_token(key, encoded).to_vec();
7715 }
7716 }
7717 }
7718 let _ = column_id;
7719 encoded.to_vec()
7720 }
7721}
7722
7723fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
7724 let columnar::NativeColumn::Int64 { data, validity } = col else {
7725 return false;
7726 };
7727 if data.len() < n || !columnar::all_non_null(validity, n) {
7728 return false;
7729 }
7730 data.iter()
7731 .take(n)
7732 .zip(data.iter().skip(1))
7733 .all(|(a, b)| a < b)
7734}
7735
7736#[derive(Debug, Clone)]
7740pub struct ColumnStat {
7741 pub min: Option<Value>,
7742 pub max: Option<Value>,
7743 pub null_count: u64,
7744}
7745
7746#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7748pub enum NativeAgg {
7749 Count,
7750 Sum,
7751 Min,
7752 Max,
7753 Avg,
7754}
7755
7756#[derive(Debug, Clone, PartialEq)]
7758pub enum NativeAggResult {
7759 Count(u64),
7760 Int(i64),
7761 Float(f64),
7762 Null,
7764}
7765
7766#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7768pub enum ApproxAgg {
7769 Count,
7770 Sum,
7771 Avg,
7772}
7773
7774#[derive(Debug, Clone)]
7778pub struct ApproxResult {
7779 pub point: f64,
7781 pub ci_low: f64,
7783 pub ci_high: f64,
7785 pub n_population: u64,
7787 pub n_sample_live: usize,
7789 pub n_passing: usize,
7791}
7792
7793#[derive(Debug, Clone, PartialEq)]
7798pub enum AggState {
7799 Count(u64),
7801 SumI {
7803 sum: i128,
7804 count: u64,
7805 },
7806 SumF {
7808 sum: f64,
7809 count: u64,
7810 },
7811 AvgI {
7813 sum: i128,
7814 count: u64,
7815 },
7816 AvgF {
7818 sum: f64,
7819 count: u64,
7820 },
7821 MinI(i64),
7823 MaxI(i64),
7824 MinF(f64),
7826 MaxF(f64),
7827 Empty,
7829}
7830
7831impl AggState {
7832 pub fn merge(self, other: AggState) -> AggState {
7834 use AggState::*;
7835 match (self, other) {
7836 (Empty, x) | (x, Empty) => x,
7837 (Count(a), Count(b)) => Count(a + b),
7838 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
7839 sum: sa + sb,
7840 count: ca + cb,
7841 },
7842 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
7843 sum: sa + sb,
7844 count: ca + cb,
7845 },
7846 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
7847 sum: sa + sb,
7848 count: ca + cb,
7849 },
7850 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
7851 sum: sa + sb,
7852 count: ca + cb,
7853 },
7854 (MinI(a), MinI(b)) => MinI(a.min(b)),
7855 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
7856 (MinF(a), MinF(b)) => MinF(a.min(b)),
7857 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
7858 _ => Empty, }
7860 }
7861
7862 pub fn point(&self) -> Option<f64> {
7864 match self {
7865 AggState::Count(n) => Some(*n as f64),
7866 AggState::SumI { sum, .. } => Some(*sum as f64),
7867 AggState::SumF { sum, .. } => Some(*sum),
7868 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
7869 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
7870 AggState::MinI(n) => Some(*n as f64),
7871 AggState::MaxI(n) => Some(*n as f64),
7872 AggState::MinF(n) => Some(*n),
7873 AggState::MaxF(n) => Some(*n),
7874 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
7875 }
7876 }
7877
7878 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
7882 let is_float = matches!(ty, Some(TypeId::Float64));
7883 match (agg, result) {
7884 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
7885 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
7886 sum: x as i128,
7887 count: 1, },
7889 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
7890 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
7891 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
7892 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
7893 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
7894 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
7895 (NativeAgg::Count, _) => AggState::Empty,
7896 (_, NativeAggResult::Null) => AggState::Empty,
7897 _ => {
7898 let _ = is_float;
7899 AggState::Empty
7900 }
7901 }
7902 }
7903}
7904
7905#[derive(Debug, Clone)]
7908pub struct CachedAgg {
7909 pub state: AggState,
7910 pub watermark: u64,
7911 pub epoch: u64,
7912}
7913
7914#[derive(Debug, Clone)]
7916pub struct IncrementalAggResult {
7917 pub state: AggState,
7919 pub incremental: bool,
7922 pub delta_rows: u64,
7924}
7925
7926fn agg_state_from_rows(
7930 rows: &[Row],
7931 conditions: &[crate::query::Condition],
7932 index_sets: &[RowIdSet],
7933 column: Option<u16>,
7934 agg: NativeAgg,
7935 schema: &Schema,
7936) -> Result<AggState> {
7937 let mut count: u64 = 0;
7938 let mut sum_i: i128 = 0;
7939 let mut sum_f: f64 = 0.0;
7940 let mut mn_i: i64 = i64::MAX;
7941 let mut mx_i: i64 = i64::MIN;
7942 let mut mn_f: f64 = f64::INFINITY;
7943 let mut mx_f: f64 = f64::NEG_INFINITY;
7944 let mut saw_int = false;
7945 let mut saw_float = false;
7946 for r in rows {
7947 if !conditions
7948 .iter()
7949 .all(|c| condition_matches_row(c, r, schema))
7950 {
7951 continue;
7952 }
7953 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
7954 continue;
7955 }
7956 match agg {
7957 NativeAgg::Count => match column {
7958 None => count += 1,
7960 Some(cid) => match r.columns.get(&cid) {
7963 None | Some(Value::Null) => {}
7964 Some(_) => count += 1,
7965 },
7966 },
7967 _ => match column.and_then(|cid| r.columns.get(&cid)) {
7968 Some(Value::Int64(n)) => {
7969 count += 1;
7970 sum_i += *n as i128;
7971 mn_i = mn_i.min(*n);
7972 mx_i = mx_i.max(*n);
7973 saw_int = true;
7974 }
7975 Some(Value::Float64(f)) => {
7976 count += 1;
7977 sum_f += f;
7978 mn_f = mn_f.min(*f);
7979 mx_f = mx_f.max(*f);
7980 saw_float = true;
7981 }
7982 _ => {}
7983 },
7984 }
7985 }
7986 Ok(match agg {
7987 NativeAgg::Count => {
7988 if count == 0 {
7989 AggState::Empty
7990 } else {
7991 AggState::Count(count)
7992 }
7993 }
7994 NativeAgg::Sum => {
7995 if count == 0 {
7996 AggState::Empty
7997 } else if saw_int {
7998 AggState::SumI { sum: sum_i, count }
7999 } else {
8000 AggState::SumF { sum: sum_f, count }
8001 }
8002 }
8003 NativeAgg::Avg => {
8004 if count == 0 {
8005 AggState::Empty
8006 } else if saw_int {
8007 AggState::AvgI { sum: sum_i, count }
8008 } else {
8009 AggState::AvgF { sum: sum_f, count }
8010 }
8011 }
8012 NativeAgg::Min => {
8013 if !saw_int && !saw_float {
8014 AggState::Empty
8015 } else if saw_int {
8016 AggState::MinI(mn_i)
8017 } else {
8018 AggState::MinF(mn_f)
8019 }
8020 }
8021 NativeAgg::Max => {
8022 if !saw_int && !saw_float {
8023 AggState::Empty
8024 } else if saw_int {
8025 AggState::MaxI(mx_i)
8026 } else {
8027 AggState::MaxF(mx_f)
8028 }
8029 }
8030 })
8031}
8032
8033fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
8037 use crate::query::Condition;
8038 match c {
8039 Condition::Pk(key) => match schema.primary_key() {
8040 Some(pk) => row
8041 .columns
8042 .get(&pk.id)
8043 .map(|v| v.encode_key() == *key)
8044 .unwrap_or(false),
8045 None => false,
8046 },
8047 Condition::BitmapEq { column_id, value } => row
8048 .columns
8049 .get(column_id)
8050 .map(|v| v.encode_key() == *value)
8051 .unwrap_or(false),
8052 Condition::BitmapIn { column_id, values } => {
8053 let key = row.columns.get(column_id).map(|v| v.encode_key());
8054 match key {
8055 Some(k) => values.contains(&k),
8056 None => false,
8057 }
8058 }
8059 Condition::BytesPrefix { column_id, prefix } => row
8060 .columns
8061 .get(column_id)
8062 .map(|v| v.encode_key().starts_with(prefix))
8063 .unwrap_or(false),
8064 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
8065 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
8066 _ => false,
8067 },
8068 Condition::RangeF64 {
8069 column_id,
8070 lo,
8071 lo_inclusive,
8072 hi,
8073 hi_inclusive,
8074 } => match row.columns.get(column_id) {
8075 Some(Value::Float64(n)) => {
8076 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
8077 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
8078 lo_ok && hi_ok
8079 }
8080 _ => false,
8081 },
8082 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
8083 Some(Value::Bytes(b)) => {
8084 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
8085 }
8086 _ => false,
8087 },
8088 Condition::FmContainsAll {
8089 column_id,
8090 patterns,
8091 } => match row.columns.get(column_id) {
8092 Some(Value::Bytes(b)) => patterns
8093 .iter()
8094 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
8095 _ => false,
8096 },
8097 Condition::Ann { .. }
8098 | Condition::SparseMatch { .. }
8099 | Condition::MinHashSimilar { .. } => true,
8100 Condition::IsNull { column_id } => {
8101 matches!(row.columns.get(column_id), Some(Value::Null) | None)
8102 }
8103 Condition::IsNotNull { column_id } => {
8104 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
8105 }
8106 }
8107}
8108
8109fn as_f64(v: Option<&Value>) -> Option<f64> {
8111 match v {
8112 Some(Value::Int64(n)) => Some(*n as f64),
8113 Some(Value::Float64(f)) => Some(*f),
8114 _ => None,
8115 }
8116}
8117
8118fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
8122 let mut count: u64 = 0;
8123 let mut sum: i128 = 0;
8124 let mut mn: i64 = i64::MAX;
8125 let mut mx: i64 = i64::MIN;
8126 while let Some(cols) = cursor.next_batch()? {
8127 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
8128 if crate::columnar::all_non_null(validity, data.len()) {
8129 count += data.len() as u64;
8131 sum += data.iter().map(|&v| v as i128).sum::<i128>();
8132 mn = mn.min(*data.iter().min().unwrap_or(&mn));
8133 mx = mx.max(*data.iter().max().unwrap_or(&mx));
8134 } else {
8135 for (i, &v) in data.iter().enumerate() {
8136 if crate::columnar::validity_bit(validity, i) {
8137 count += 1;
8138 sum += v as i128;
8139 mn = mn.min(v);
8140 mx = mx.max(v);
8141 }
8142 }
8143 }
8144 }
8145 }
8146 Ok((count, sum, mn, mx))
8147}
8148
8149fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
8151 let mut count: u64 = 0;
8152 let mut sum: f64 = 0.0;
8153 let mut mn: f64 = f64::INFINITY;
8154 let mut mx: f64 = f64::NEG_INFINITY;
8155 while let Some(cols) = cursor.next_batch()? {
8156 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
8157 if crate::columnar::all_non_null(validity, data.len()) {
8158 count += data.len() as u64;
8159 sum += data.iter().sum::<f64>();
8160 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
8161 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
8162 } else {
8163 for (i, &v) in data.iter().enumerate() {
8164 if crate::columnar::validity_bit(validity, i) {
8165 count += 1;
8166 sum += v;
8167 mn = mn.min(v);
8168 mx = mx.max(v);
8169 }
8170 }
8171 }
8172 }
8173 }
8174 Ok((count, sum, mn, mx))
8175}
8176
8177fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
8178 if count == 0 && !matches!(agg, NativeAgg::Count) {
8179 return NativeAggResult::Null;
8180 }
8181 match agg {
8182 NativeAgg::Count => NativeAggResult::Count(count),
8183 NativeAgg::Sum => match sum.try_into() {
8186 Ok(v) => NativeAggResult::Int(v),
8187 Err(_) => NativeAggResult::Null,
8188 },
8189 NativeAgg::Min => NativeAggResult::Int(mn),
8190 NativeAgg::Max => NativeAggResult::Int(mx),
8191 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
8192 }
8193}
8194
8195fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
8196 if count == 0 && !matches!(agg, NativeAgg::Count) {
8197 return NativeAggResult::Null;
8198 }
8199 match agg {
8200 NativeAgg::Count => NativeAggResult::Count(count),
8201 NativeAgg::Sum => NativeAggResult::Float(sum),
8202 NativeAgg::Min => NativeAggResult::Float(mn),
8203 NativeAgg::Max => NativeAggResult::Float(mx),
8204 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
8205 }
8206}
8207
8208fn agg_int(
8211 stats: &[crate::page::PageStat],
8212 decode: fn(Option<&[u8]>) -> Option<i64>,
8213) -> Option<(Option<i64>, Option<i64>, u64)> {
8214 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
8215 let mut any = false;
8216 for s in stats {
8217 if let Some(v) = decode(s.min.as_deref()) {
8218 mn = mn.min(v);
8219 any = true;
8220 }
8221 if let Some(v) = decode(s.max.as_deref()) {
8222 mx = mx.max(v);
8223 any = true;
8224 }
8225 nulls += s.null_count;
8226 }
8227 any.then_some((Some(mn), Some(mx), nulls))
8228}
8229
8230fn agg_float(
8232 stats: &[crate::page::PageStat],
8233 decode: fn(Option<&[u8]>) -> Option<f64>,
8234) -> Option<(Option<f64>, Option<f64>, u64)> {
8235 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
8236 let mut any = false;
8237 for s in stats {
8238 if let Some(v) = decode(s.min.as_deref()) {
8239 mn = mn.min(v);
8240 any = true;
8241 }
8242 if let Some(v) = decode(s.max.as_deref()) {
8243 mx = mx.max(v);
8244 any = true;
8245 }
8246 nulls += s.null_count;
8247 }
8248 any.then_some((Some(mn), Some(mx), nulls))
8249}
8250
8251type SecondaryIndexes = (
8253 HashMap<u16, BitmapIndex>,
8254 HashMap<u16, AnnIndex>,
8255 HashMap<u16, FmIndex>,
8256 HashMap<u16, SparseIndex>,
8257 HashMap<u16, MinHashIndex>,
8258);
8259
8260fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
8261 let mut bitmap = HashMap::new();
8262 let mut ann = HashMap::new();
8263 let mut fm = HashMap::new();
8264 let mut sparse = HashMap::new();
8265 let mut minhash = HashMap::new();
8266 for idef in &schema.indexes {
8267 match idef.kind {
8268 IndexKind::Bitmap => {
8269 bitmap.insert(idef.column_id, BitmapIndex::new());
8270 }
8271 IndexKind::Ann => {
8272 let dim = schema
8273 .columns
8274 .iter()
8275 .find(|c| c.id == idef.column_id)
8276 .and_then(|c| match c.ty {
8277 TypeId::Embedding { dim } => Some(dim as usize),
8278 _ => None,
8279 })
8280 .unwrap_or(0);
8281 let options = idef.options.ann.clone().unwrap_or_default();
8282 ann.insert(
8283 idef.column_id,
8284 AnnIndex::with_options(
8285 dim,
8286 options.m,
8287 options.ef_construction,
8288 options.ef_search,
8289 ),
8290 );
8291 }
8292 IndexKind::FmIndex => {
8293 fm.insert(idef.column_id, FmIndex::new());
8294 }
8295 IndexKind::Sparse => {
8296 sparse.insert(idef.column_id, SparseIndex::new());
8297 }
8298 IndexKind::MinHash => {
8299 let options = idef.options.minhash.clone().unwrap_or_default();
8300 minhash.insert(
8301 idef.column_id,
8302 MinHashIndex::with_options(options.permutations, options.bands),
8303 );
8304 }
8305 _ => {}
8306 }
8307 }
8308 (bitmap, ann, fm, sparse, minhash)
8309}
8310
8311const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
8312 | ColumnFlags::AUTO_INCREMENT
8313 | ColumnFlags::ENCRYPTED
8314 | ColumnFlags::ENCRYPTED_INDEXABLE
8315 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
8316
8317fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
8318 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
8319 return Err(MongrelError::Schema(
8320 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
8321 ));
8322 }
8323 Ok(())
8324}
8325
8326fn validate_alter_column_type(
8327 schema: &Schema,
8328 old: &ColumnDef,
8329 next: &ColumnDef,
8330 has_stored_versions: bool,
8331) -> Result<()> {
8332 if old.ty == next.ty {
8333 return Ok(());
8334 }
8335 if schema.indexes.iter().any(|i| i.column_id == old.id) {
8336 return Err(MongrelError::Schema(format!(
8337 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
8338 old.name
8339 )));
8340 }
8341 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
8342 return Ok(());
8343 }
8344 Err(MongrelError::Schema(format!(
8345 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
8346 old.ty, next.ty
8347 )))
8348}
8349
8350fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
8351 matches!(
8352 (old, new),
8353 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
8354 )
8355}
8356
8357fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
8363 let mut prev: Option<i64> = None;
8364 for r in rows {
8365 match r.columns.get(&pk_id) {
8366 Some(Value::Int64(v)) => {
8367 if prev.is_some_and(|p| p >= *v) {
8368 return false;
8369 }
8370 prev = Some(*v);
8371 }
8372 _ => return false,
8373 }
8374 }
8375 true
8376}
8377
8378#[allow(clippy::too_many_arguments)]
8379fn index_into(
8380 schema: &Schema,
8381 row: &Row,
8382 hot: &mut HotIndex,
8383 bitmap: &mut HashMap<u16, BitmapIndex>,
8384 ann: &mut HashMap<u16, AnnIndex>,
8385 fm: &mut HashMap<u16, FmIndex>,
8386 sparse: &mut HashMap<u16, SparseIndex>,
8387 minhash: &mut HashMap<u16, MinHashIndex>,
8388) {
8389 for idef in &schema.indexes {
8390 let Some(val) = row.columns.get(&idef.column_id) else {
8391 continue;
8392 };
8393 match idef.kind {
8394 IndexKind::Bitmap => {
8395 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8396 b.insert(val.encode_key(), row.row_id);
8397 }
8398 }
8399 IndexKind::Ann => {
8400 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8401 a.insert_validated(v, row.row_id);
8402 }
8403 }
8404 IndexKind::FmIndex => {
8405 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8406 f.insert(b.clone(), row.row_id);
8407 }
8408 }
8409 IndexKind::Sparse => {
8410 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8411 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8414 s.insert(&terms, row.row_id);
8415 }
8416 }
8417 }
8418 IndexKind::MinHash => {
8419 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8420 let tokens = crate::index::token_hashes_from_bytes(b);
8423 mh.insert(&tokens, row.row_id);
8424 }
8425 }
8426 _ => {}
8427 }
8428 }
8429 if let Some(pk_col) = schema.primary_key() {
8430 if let Some(pk_val) = row.columns.get(&pk_col.id) {
8431 hot.insert(pk_val.encode_key(), row.row_id);
8432 }
8433 }
8434}
8435
8436#[allow(clippy::too_many_arguments)]
8439fn index_into_single(
8440 idef: &IndexDef,
8441 _schema: &Schema,
8442 row: &Row,
8443 _hot: &mut HotIndex,
8444 bitmap: &mut HashMap<u16, BitmapIndex>,
8445 ann: &mut HashMap<u16, AnnIndex>,
8446 fm: &mut HashMap<u16, FmIndex>,
8447 sparse: &mut HashMap<u16, SparseIndex>,
8448 minhash: &mut HashMap<u16, MinHashIndex>,
8449) {
8450 let Some(val) = row.columns.get(&idef.column_id) else {
8451 return;
8452 };
8453 match idef.kind {
8454 IndexKind::Bitmap => {
8455 if let Some(b) = bitmap.get_mut(&idef.column_id) {
8456 b.insert(val.encode_key(), row.row_id);
8457 }
8458 }
8459 IndexKind::Ann => {
8460 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
8461 a.insert_validated(v, row.row_id);
8462 }
8463 }
8464 IndexKind::FmIndex => {
8465 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
8466 f.insert(b.clone(), row.row_id);
8467 }
8468 }
8469 IndexKind::Sparse => {
8470 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
8471 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
8472 s.insert(&terms, row.row_id);
8473 }
8474 }
8475 }
8476 IndexKind::MinHash => {
8477 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
8478 let tokens = crate::index::token_hashes_from_bytes(b);
8479 mh.insert(&tokens, row.row_id);
8480 }
8481 }
8482 _ => {}
8483 }
8484}
8485
8486fn eval_partial_predicate(
8492 pred: &str,
8493 columns_map: &HashMap<u16, &Value>,
8494 name_to_id: &HashMap<&str, u16>,
8495) -> bool {
8496 let lower = pred.trim().to_ascii_lowercase();
8497 if let Some(rest) = lower.strip_suffix(" is not null") {
8499 let col_name = rest.trim();
8500 if let Some(col_id) = name_to_id.get(col_name) {
8501 return columns_map
8502 .get(col_id)
8503 .is_some_and(|v| !matches!(v, Value::Null));
8504 }
8505 }
8506 if let Some(rest) = lower.strip_suffix(" is null") {
8508 let col_name = rest.trim();
8509 if let Some(col_id) = name_to_id.get(col_name) {
8510 return columns_map
8511 .get(col_id)
8512 .map_or(true, |v| matches!(v, Value::Null));
8513 }
8514 }
8515 true
8518}
8519
8520#[allow(dead_code)]
8526fn bulk_index_key(
8527 column_keys: &HashMap<u16, ([u8; 32], u8)>,
8528 column_id: u16,
8529 ty: TypeId,
8530 col: &columnar::NativeColumn,
8531 i: usize,
8532) -> Option<Vec<u8>> {
8533 let encoded = columnar::encode_key_native(ty, col, i)?;
8534 #[cfg(feature = "encryption")]
8535 {
8536 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
8537 if let Some((key, scheme)) = column_keys.get(&column_id) {
8538 return Some(match (*scheme, col) {
8539 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
8540 (_, columnar::NativeColumn::Int64 { data, .. }) => {
8541 ope_token_i64(key, data[i]).to_vec()
8542 }
8543 (_, columnar::NativeColumn::Float64 { data, .. }) => {
8544 ope_token_f64(key, data[i]).to_vec()
8545 }
8546 _ => hmac_token(key, &encoded).to_vec(),
8547 });
8548 }
8549 }
8550 #[cfg(not(feature = "encryption"))]
8551 {
8552 let _ = (column_id, column_keys, col);
8553 }
8554 Some(encoded)
8555}
8556
8557pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
8558 let json = serde_json::to_string_pretty(schema)
8559 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
8560 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
8561 Ok(())
8562}
8563
8564fn read_schema(dir: &Path) -> Result<Schema> {
8565 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
8566 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
8567}
8568
8569fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
8570 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
8571}
8572
8573fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
8574 let n = list_wal_numbers(wal_dir)?;
8575 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
8576}
8577
8578fn next_wal_number(wal_dir: &Path) -> Result<u32> {
8579 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
8580}
8581
8582fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
8583 let _ = std::fs::create_dir_all(wal_dir);
8584 let mut max_n = None;
8585 for entry in std::fs::read_dir(wal_dir)? {
8586 let entry = entry?;
8587 let fname = entry.file_name();
8588 let Some(s) = fname.to_str() else {
8589 continue;
8590 };
8591 let Some(stripped) = s.strip_prefix("seg-") else {
8592 continue;
8593 };
8594 let Some(stripped) = stripped.strip_suffix(".wal") else {
8595 continue;
8596 };
8597 if let Ok(n) = stripped.parse::<u32>() {
8598 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
8599 }
8600 }
8601 Ok(max_n)
8602}