1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::collections::{BTreeMap, HashMap, HashSet};
30use std::path::{Path, PathBuf};
31use std::sync::atomic::AtomicBool;
32use std::sync::Arc;
33use zeroize::Zeroizing;
34
35pub const WAL_DIR: &str = "_wal";
36pub const RUNS_DIR: &str = "_runs";
37pub const CACHE_DIR: &str = "_cache";
38pub const META_DIR: &str = "_meta";
39pub const RCACHE_DIR: &str = "_rcache";
40pub const KEYS_FILENAME: &str = "keys";
41pub const SCHEMA_FILENAME: &str = "schema.json";
42
43fn iso_now_bytes() -> Vec<u8> {
46 let secs = std::time::SystemTime::now()
47 .duration_since(std::time::UNIX_EPOCH)
48 .map(|d| d.as_secs() as i64)
49 .unwrap_or(0);
50 let days = secs.div_euclid(86_400);
51 let rem = secs.rem_euclid(86_400);
52 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
53 let (year, month, day) = civil_from_days(days);
54 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
55}
56
57pub(crate) fn unix_nanos_now() -> i64 {
58 std::time::SystemTime::now()
59 .duration_since(std::time::UNIX_EPOCH)
60 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
61 .unwrap_or(0)
62}
63
64fn civil_from_days(z: i64) -> (i64, u32, u32) {
65 let z = z + 719_468;
66 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
67 let doe = z - era * 146_097;
68 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
69 let y = yoe + era * 400;
70 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
71 let mp = (5 * doy + 2) / 153;
72 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
73 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
74 (if m <= 2 { y + 1 } else { y }, m, d)
75}
76
77const DEFAULT_SYNC_BYTE_THRESHOLD: u64 = 0; pub(crate) const PAGE_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; pub(crate) const DECODED_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; const DEFAULT_MUTABLE_RUN_SPILL_BYTES: u64 = 8 * 1024 * 1024;
84
85#[derive(Clone, Copy, Debug)]
100struct AutoIncState {
101 column_id: u16,
102 next: i64,
103 seeded: bool,
104}
105
106type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
107
108fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
111 schema.auto_increment_column().map(|c| AutoIncState {
112 column_id: c.id,
113 next: 0,
114 seeded: false,
115 })
116}
117
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
131pub enum IndexBuildPolicy {
132 #[default]
134 Deferred,
135 Eager,
137}
138
139#[derive(Clone)]
140struct ReversePkSegment {
141 values: HashMap<RowId, Vec<u8>>,
142 removed: HashSet<RowId>,
143}
144
145#[derive(Clone)]
146struct ReversePkMap {
147 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
148 active: ReversePkSegment,
149}
150
151impl ReversePkMap {
152 fn new() -> Self {
153 Self {
154 frozen: Arc::new(Vec::new()),
155 active: ReversePkSegment {
156 values: HashMap::new(),
157 removed: HashSet::new(),
158 },
159 }
160 }
161
162 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
163 let mut map = Self::new();
164 map.active.values.extend(entries);
165 map
166 }
167
168 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
169 self.active.removed.remove(&row_id);
170 self.active.values.insert(row_id, key);
171 }
172
173 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
174 if let Some(key) = self.active.values.get(row_id) {
175 return Some(key);
176 }
177 if self.active.removed.contains(row_id) {
178 return None;
179 }
180 for segment in self.frozen.iter().rev() {
181 if let Some(key) = segment.values.get(row_id) {
182 return Some(key);
183 }
184 if segment.removed.contains(row_id) {
185 return None;
186 }
187 }
188 None
189 }
190
191 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
192 let previous = self.get(row_id).cloned();
193 self.active.values.remove(row_id);
194 self.active.removed.insert(*row_id);
195 previous
196 }
197
198 fn clear(&mut self) {
199 *self = Self::new();
200 }
201
202 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
203 let mut entries = HashMap::new();
204 for segment in self
205 .frozen
206 .iter()
207 .map(Arc::as_ref)
208 .chain(std::iter::once(&self.active))
209 {
210 for row_id in &segment.removed {
211 entries.remove(row_id);
212 }
213 entries.extend(
214 segment
215 .values
216 .iter()
217 .map(|(row_id, key)| (*row_id, key.clone())),
218 );
219 }
220 entries
221 }
222
223 fn seal(&mut self) {
224 if self.active.values.is_empty() && self.active.removed.is_empty() {
225 return;
226 }
227 let active = std::mem::replace(
228 &mut self.active,
229 ReversePkSegment {
230 values: HashMap::new(),
231 removed: HashSet::new(),
232 },
233 );
234 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
235 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
236 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
237 values: self.entries(),
238 removed: HashSet::new(),
239 })]);
240 }
241 }
242}
243
244#[derive(Clone)]
246pub struct Table {
247 dir: PathBuf,
248 table_id: u64,
249 name: String,
253 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
258 read_only: bool,
261 wal: WalSink,
262 memtable: Memtable,
263 mutable_run: MutableRun,
268 mutable_run_spill_bytes: u64,
270 compaction_zstd_level: i32,
273 allocator: RowIdAllocator,
274 epoch: Arc<EpochAuthority>,
275 persisted_epoch: u64,
278 data_generation: u64,
281 schema: Schema,
282 hot: HotIndex,
283 kek: Option<Arc<Kek>>,
286 column_keys: HashMap<u16, ([u8; 32], u8)>,
290 run_refs: Vec<RunRef>,
291 retiring: Vec<crate::manifest::RetiredRun>,
294 next_run_id: u64,
295 sync_byte_threshold: u64,
296 current_txn_id: u64,
301 bitmap: HashMap<u16, BitmapIndex>,
302 ann: HashMap<u16, AnnIndex>,
303 fm: HashMap<u16, FmIndex>,
304 sparse: HashMap<u16, SparseIndex>,
305 minhash: HashMap<u16, MinHashIndex>,
306 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
309 pk_by_row: ReversePkMap,
311 pinned: BTreeMap<Epoch, usize>,
314 pub(crate) live_count: u64,
317 reservoir: crate::reservoir::Reservoir,
320 reservoir_complete: bool,
328 had_deletes: bool,
332 agg_cache: Arc<HashMap<u64, CachedAgg>>,
336 global_idx_epoch: u64,
340 indexes_complete: bool,
345 index_build_policy: IndexBuildPolicy,
347 pk_by_row_complete: bool,
354 flushed_epoch: u64,
357 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
360 snapshots: Arc<crate::retention::SnapshotRegistry>,
363 commit_lock: Arc<parking_lot::Mutex<()>>,
365 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
368 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
378 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
387 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
389 pending_delete_rids: roaring::RoaringBitmap,
392 pending_put_cols: std::collections::HashSet<u16>,
395 pending_rows: Vec<Row>,
401 pending_rows_auto_inc: Vec<bool>,
402 pending_dels: Vec<RowId>,
405 pending_truncate: Option<Epoch>,
409 auto_inc: Option<AutoIncState>,
412 ttl: Option<TtlPolicy>,
415}
416
417const _: () = {
424 const fn assert_sync<T: ?Sized + Sync>() {}
425 assert_sync::<Table>();
426};
427
428enum CachedData {
434 Rows(Arc<Vec<Row>>),
435 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
436}
437
438impl CachedData {
439 fn approx_bytes(&self) -> u64 {
440 match self {
441 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
442 CachedData::Columns(c) => c
443 .iter()
444 .map(|(_, c)| c.approx_bytes())
445 .sum::<u64>()
446 .saturating_add(c.len() as u64 * 16),
447 }
448 }
449}
450
451struct CachedEntry {
455 data: CachedData,
456 footprint: roaring::RoaringBitmap,
457 condition_cols: Vec<u16>,
458}
459
460struct ResultCache {
471 entries: std::collections::HashMap<u64, CachedEntry>,
472 order: std::collections::VecDeque<u64>,
473 bytes: u64,
474 max_bytes: u64,
475 dir: Option<std::path::PathBuf>,
476 #[allow(dead_code)]
477 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
478}
479
480#[derive(serde::Serialize, serde::Deserialize)]
482struct SerializedEntry {
483 condition_cols: Vec<u16>,
484 footprint_bits: Vec<u32>,
485 data: SerializedData,
486}
487
488#[derive(serde::Serialize, serde::Deserialize)]
489enum SerializedData {
490 Rows(Vec<Row>),
491 Columns(Vec<(u16, columnar::NativeColumn)>),
492}
493
494impl SerializedEntry {
495 fn from_entry(entry: &CachedEntry) -> Self {
496 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
497 let data = match &entry.data {
498 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
499 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
500 };
501 Self {
502 condition_cols: entry.condition_cols.clone(),
503 footprint_bits,
504 data,
505 }
506 }
507
508 fn into_entry(self) -> Option<CachedEntry> {
509 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
510 let data = match self.data {
511 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
512 SerializedData::Columns(c) => {
513 if !c.iter().all(|(_, col)| col.validate()) {
516 return None;
517 }
518 CachedData::Columns(Arc::new(c))
519 }
520 };
521 Some(CachedEntry {
522 data,
523 footprint,
524 condition_cols: self.condition_cols,
525 })
526 }
527}
528
529impl ResultCache {
530 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
531
532 fn new() -> Self {
533 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
534 }
535
536 fn with_max_bytes(max_bytes: u64) -> Self {
537 Self {
538 entries: std::collections::HashMap::new(),
539 order: std::collections::VecDeque::new(),
540 bytes: 0,
541 max_bytes,
542 dir: None,
543 cache_dek: None,
544 }
545 }
546
547 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
548 let _ = std::fs::create_dir_all(&dir);
549 self.dir = Some(dir);
550 self
551 }
552
553 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
554 self.cache_dek = dek;
555 self
556 }
557
558 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
559 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
560 }
561
562 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
566 let Some(path) = self.disk_path(key) else {
567 return;
568 };
569 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
570 Ok(s) => s,
571 Err(_) => return,
572 };
573 let on_disk = if let Some(dek) = &self.cache_dek {
575 match self.encrypt_cache(&serialized, dek) {
576 Some(b) => b,
577 None => return,
578 }
579 } else {
580 serialized
581 };
582 let tmp = path.with_extension("tmp");
583 use std::io::Write;
584 let write = || -> std::io::Result<()> {
585 let mut f = std::fs::File::create(&tmp)?;
586 f.write_all(&on_disk)?;
587 f.flush()?;
588 Ok(())
589 };
590 if write().is_err() {
591 let _ = std::fs::remove_file(&tmp);
592 return;
593 }
594 let _ = std::fs::rename(&tmp, &path);
595 }
596
597 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
599 let path = self.disk_path(key)?;
600 let bytes = std::fs::read(&path).ok()?;
601 let plaintext = if let Some(dek) = &self.cache_dek {
602 self.decrypt_cache(&bytes, dek)?
603 } else {
604 bytes
605 };
606 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
607 serialized.into_entry()
608 }
609
610 fn remove_from_disk(&self, key: u64) {
612 if let Some(path) = self.disk_path(key) {
613 let _ = std::fs::remove_file(&path);
614 }
615 }
616
617 #[cfg(feature = "encryption")]
619 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
620 use crate::encryption::Cipher;
621 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
622 let mut nonce = [0u8; 12];
623 crate::encryption::fill_random(&mut nonce);
624 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
625 let mut out = Vec::with_capacity(12 + ct.len());
626 out.extend_from_slice(&nonce);
627 out.extend_from_slice(&ct);
628 Some(out)
629 }
630
631 #[cfg(not(feature = "encryption"))]
632 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
633 None
634 }
635
636 #[cfg(feature = "encryption")]
638 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
639 use crate::encryption::Cipher;
640 if bytes.len() < 28 {
641 return None;
642 }
643 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
644 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
645 let ct = &bytes[12..];
646 cipher.decrypt_page(&nonce, ct).ok()
647 }
648
649 #[cfg(not(feature = "encryption"))]
650 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
651 None
652 }
653
654 fn load_persistent(&mut self) {
657 let Some(dir) = self.dir.as_ref().cloned() else {
658 return;
659 };
660 let entries = match std::fs::read_dir(&dir) {
661 Ok(e) => e,
662 Err(_) => return,
663 };
664 for entry in entries.flatten() {
665 let path = entry.path();
666 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
668 let _ = std::fs::remove_file(&path);
669 continue;
670 }
671 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
672 continue;
673 }
674 let stem = match path.file_stem().and_then(|s| s.to_str()) {
675 Some(s) => s,
676 None => continue,
677 };
678 let key = match u64::from_str_radix(stem, 16) {
679 Ok(k) => k,
680 Err(_) => continue,
681 };
682 let bytes = match std::fs::read(&path) {
683 Ok(b) => b,
684 Err(_) => continue,
685 };
686 let plaintext = if let Some(dek) = &self.cache_dek {
688 match self.decrypt_cache(&bytes, dek) {
689 Some(p) => p,
690 None => {
691 let _ = std::fs::remove_file(&path);
692 continue;
693 }
694 }
695 } else {
696 bytes
697 };
698 match bincode::deserialize::<SerializedEntry>(&plaintext) {
699 Ok(serialized) => {
700 if let Some(entry) = serialized.into_entry() {
701 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
702 self.entries.insert(key, entry);
703 self.order.push_back(key);
704 } else {
705 let _ = std::fs::remove_file(&path);
706 }
707 }
708 Err(_) => {
709 let _ = std::fs::remove_file(&path);
710 }
711 }
712 }
713 self.evict();
714 }
715
716 fn set_max_bytes(&mut self, max_bytes: u64) {
717 self.max_bytes = max_bytes;
718 self.evict();
719 }
720
721 fn touch(&mut self, key: u64) {
723 self.order.retain(|k| *k != key);
724 self.order.push_back(key);
725 }
726
727 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
728 let res = self.entries.get(&key).and_then(|e| match &e.data {
729 CachedData::Rows(r) => Some(r.clone()),
730 CachedData::Columns(_) => None,
731 });
732 if res.is_some() {
733 self.touch(key);
734 return res;
735 }
736 if let Some(entry) = self.load_from_disk(key) {
738 let res = match &entry.data {
739 CachedData::Rows(r) => Some(r.clone()),
740 CachedData::Columns(_) => None,
741 };
742 if res.is_some() {
743 let approx = entry.data.approx_bytes();
744 self.bytes = self.bytes.saturating_add(approx);
745 self.entries.insert(key, entry);
746 self.order.push_back(key);
747 self.evict();
748 return res;
749 }
750 }
751 None
752 }
753
754 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
755 let res = self.entries.get(&key).and_then(|e| match &e.data {
756 CachedData::Columns(c) => Some(c.clone()),
757 CachedData::Rows(_) => None,
758 });
759 if res.is_some() {
760 self.touch(key);
761 return res;
762 }
763 if let Some(entry) = self.load_from_disk(key) {
765 let res = match &entry.data {
766 CachedData::Columns(c) => Some(c.clone()),
767 CachedData::Rows(_) => None,
768 };
769 if res.is_some() {
770 let approx = entry.data.approx_bytes();
771 self.bytes = self.bytes.saturating_add(approx);
772 self.entries.insert(key, entry);
773 self.order.push_back(key);
774 self.evict();
775 return res;
776 }
777 }
778 None
779 }
780
781 fn insert(&mut self, key: u64, entry: CachedEntry) {
782 let approx = entry.data.approx_bytes();
783 if self.entries.remove(&key).is_some() {
784 self.order.retain(|k| *k != key);
785 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
786 }
787 self.store_to_disk(key, &entry);
789 self.bytes = self.bytes.saturating_add(approx);
790 self.entries.insert(key, entry);
791 self.order.push_back(key);
792 self.evict();
793 }
794
795 fn invalidate(
804 &mut self,
805 delete_rids: &roaring::RoaringBitmap,
806 put_cols: &std::collections::HashSet<u16>,
807 ) {
808 if self.entries.is_empty() {
809 return;
810 }
811 let has_deletes = !delete_rids.is_empty();
812 let to_remove: std::collections::HashSet<u64> = self
813 .entries
814 .iter()
815 .filter(|(_, e)| {
816 let delete_hit = if e.footprint.is_empty() {
817 has_deletes
818 } else {
819 e.footprint.intersection_len(delete_rids) > 0
820 };
821 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
822 delete_hit || col_hit
823 })
824 .map(|(&k, _)| k)
825 .collect();
826 for key in &to_remove {
827 if let Some(e) = self.entries.remove(key) {
828 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
829 }
830 self.remove_from_disk(*key);
831 }
832 if !to_remove.is_empty() {
833 self.order.retain(|k| !to_remove.contains(k));
834 }
835 }
836
837 fn clear(&mut self) {
838 if let Some(dir) = &self.dir {
840 if let Ok(entries) = std::fs::read_dir(dir) {
841 for entry in entries.flatten() {
842 let path = entry.path();
843 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
844 let _ = std::fs::remove_file(&path);
845 }
846 }
847 }
848 }
849 self.entries.clear();
850 self.order.clear();
851 self.bytes = 0;
852 }
853
854 fn evict(&mut self) {
855 while self.bytes > self.max_bytes {
856 let Some(k) = self.order.pop_front() else {
857 break;
858 };
859 if let Some(e) = self.entries.remove(&k) {
860 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
861 self.remove_from_disk(k);
865 }
866 }
867 }
868}
869
870type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
877
878fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
879 let _ = kek;
880 #[cfg(feature = "encryption")]
881 {
882 if let Some(k) = kek {
883 return (
884 Some(k.derive_table_wal_key(_table_id)),
885 Some(k.derive_cache_key()),
886 );
887 }
888 }
889 (None, None)
890}
891
892#[cfg(feature = "encryption")]
894fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
895 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
896}
897
898#[cfg(not(feature = "encryption"))]
899fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
900 Box::new(crate::encryption::PlaintextCipher)
901}
902
903fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
904 let Some(kek) = kek else {
905 return HashMap::new();
906 };
907 #[cfg(feature = "encryption")]
908 {
909 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
910 schema
911 .columns
912 .iter()
913 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
914 .map(|c| {
915 let scheme = if schema
916 .indexes
917 .iter()
918 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
919 {
920 SCHEME_OPE_RANGE
921 } else {
922 SCHEME_HMAC_EQ
923 };
924 let key: [u8; 32] = *kek.derive_column_key(c.id);
925 (c.id, (key, scheme))
926 })
927 .collect()
928 }
929 #[cfg(not(feature = "encryption"))]
930 {
931 let _ = (kek, schema);
932 HashMap::new()
933 }
934}
935
936pub(crate) struct SharedCtx {
941 pub epoch: Arc<EpochAuthority>,
942 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
943 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
944 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
945 pub kek: Option<Arc<Kek>>,
946 pub commit_lock: Arc<parking_lot::Mutex<()>>,
952 pub shared: Option<SharedWalCtx>,
956 pub table_name: Option<String>,
959 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
962 pub read_only: bool,
964}
965
966#[derive(Clone)]
972pub(crate) struct SharedWalCtx {
973 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
974 pub group: Arc<GroupCommit>,
975 pub poisoned: Arc<AtomicBool>,
976 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
977 pub change_wake: tokio::sync::broadcast::Sender<()>,
978}
979
980enum WalSink {
983 Private(Wal),
984 Shared(SharedWalCtx),
985 ReadOnly,
986}
987
988impl Clone for WalSink {
989 fn clone(&self) -> Self {
990 match self {
991 Self::Shared(shared) => Self::Shared(shared.clone()),
992 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
993 }
994 }
995}
996
997impl SharedCtx {
998 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1002 let n_shards = if cache_dir.is_some() {
1006 1
1007 } else {
1008 crate::cache::CACHE_SHARDS
1009 };
1010 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1011 let page_cache = if let Some(d) = cache_dir {
1012 Arc::new(crate::cache::Sharded::new(1, || {
1013 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1014 }))
1015 } else {
1016 Arc::new(crate::cache::Sharded::new(n_shards, || {
1017 crate::cache::PageCache::new(per_shard)
1018 }))
1019 };
1020 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1021 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1022 crate::cache::CACHE_SHARDS,
1023 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1024 ));
1025 Self {
1026 epoch: Arc::new(EpochAuthority::new(0)),
1027 page_cache,
1028 decoded_cache,
1029 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1030 kek,
1031 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1032 shared: None,
1033 table_name: None,
1034 auth: None,
1035 read_only: false,
1036 }
1037 }
1038}
1039
1040fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1044 use crate::query::Condition;
1045 match c {
1046 Condition::Pk(_)
1048 | Condition::BitmapEq { .. }
1049 | Condition::BitmapIn { .. }
1050 | Condition::BytesPrefix { .. }
1051 | Condition::IsNull { .. }
1052 | Condition::IsNotNull { .. } => 0,
1053 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1055 1
1056 }
1057 Condition::FmContains { .. }
1059 | Condition::FmContainsAll { .. }
1060 | Condition::Ann { .. }
1061 | Condition::SparseMatch { .. } => 2,
1062 }
1063}
1064
1065impl Table {
1066 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1067 let dir = dir.as_ref().to_path_buf();
1068 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
1069 Self::create_in(&dir, schema, table_id, ctx)
1070 }
1071
1072 #[cfg(feature = "encryption")]
1083 pub fn create_encrypted(
1084 dir: impl AsRef<Path>,
1085 schema: Schema,
1086 table_id: u64,
1087 passphrase: &str,
1088 ) -> Result<Self> {
1089 let dir = dir.as_ref();
1090 std::fs::create_dir_all(dir.join(META_DIR))?;
1091 let salt = crate::encryption::random_salt();
1092 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
1093 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1094 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1095 Self::create_in(dir, schema, table_id, ctx)
1096 }
1097
1098 #[cfg(feature = "encryption")]
1103 pub fn create_with_key(
1104 dir: impl AsRef<Path>,
1105 schema: Schema,
1106 table_id: u64,
1107 key: &[u8],
1108 ) -> Result<Self> {
1109 let dir = dir.as_ref();
1110 std::fs::create_dir_all(dir.join(META_DIR))?;
1111 let salt = crate::encryption::random_salt();
1112 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
1113 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1114 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1115 Self::create_in(dir, schema, table_id, ctx)
1116 }
1117
1118 #[cfg(feature = "encryption")]
1120 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1121 let dir = dir.as_ref();
1122 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1123 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1124 MongrelError::NotFound(format!(
1125 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1126 salt_path
1127 ))
1128 })?;
1129 if salt_bytes.len() != crate::encryption::SALT_LEN {
1130 return Err(MongrelError::InvalidArgument(format!(
1131 "salt file is {} bytes, expected {}",
1132 salt_bytes.len(),
1133 crate::encryption::SALT_LEN
1134 )));
1135 }
1136 let mut salt = [0u8; crate::encryption::SALT_LEN];
1137 salt.copy_from_slice(&salt_bytes);
1138 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1139 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1140 Self::open_in(dir, ctx)
1141 }
1142
1143 pub(crate) fn create_in(
1144 dir: impl AsRef<Path>,
1145 schema: Schema,
1146 table_id: u64,
1147 ctx: SharedCtx,
1148 ) -> Result<Self> {
1149 schema.validate_auto_increment()?;
1150 schema.validate_defaults()?;
1151 schema.validate_ai()?;
1152 for index in &schema.indexes {
1153 index.validate_options()?;
1154 }
1155 let dir = dir.as_ref().to_path_buf();
1156 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1157 write_schema(&dir, &schema)?;
1158 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1159 let (wal, current_txn_id) = match ctx.shared.clone() {
1162 Some(s) => (WalSink::Shared(s), 0),
1163 None => {
1164 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1165 let mut w = if let Some(ref dk) = wal_dek {
1166 Wal::create_with_cipher(
1167 dir.join(WAL_DIR).join("seg-000000.wal"),
1168 Epoch(0),
1169 Some(make_cipher(dk)),
1170 0,
1171 )?
1172 } else {
1173 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1174 };
1175 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1176 (WalSink::Private(w), 1)
1177 }
1178 };
1179 let mut manifest = Manifest::new(table_id, schema.schema_id);
1180 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1185 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1186 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1187 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1188 let auto_inc = resolve_auto_inc(&schema);
1189 let rcache_dir = dir.join(RCACHE_DIR);
1190 Ok(Self {
1191 dir,
1192 table_id,
1193 name: ctx.table_name.unwrap_or_default(),
1194 auth: ctx.auth,
1195 read_only: ctx.read_only,
1196 wal,
1197 memtable: Memtable::new(),
1198 mutable_run: MutableRun::new(),
1199 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1200 compaction_zstd_level: 3,
1201 allocator: RowIdAllocator::new(0),
1202 epoch: ctx.epoch,
1203 persisted_epoch: 0,
1204 data_generation: 0,
1205 schema,
1206 hot: HotIndex::new(),
1207 kek: ctx.kek,
1208 column_keys,
1209 run_refs: Vec::new(),
1210 retiring: Vec::new(),
1211 next_run_id: 1,
1212 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1213 current_txn_id,
1214 bitmap,
1215 ann,
1216 fm,
1217 sparse,
1218 minhash,
1219 learned_range: Arc::new(HashMap::new()),
1220 pk_by_row: ReversePkMap::new(),
1221 pinned: BTreeMap::new(),
1222 live_count: 0,
1223 reservoir: crate::reservoir::Reservoir::default(),
1224 reservoir_complete: true,
1225 had_deletes: false,
1226 agg_cache: Arc::new(HashMap::new()),
1227 global_idx_epoch: 0,
1228 indexes_complete: true,
1229 index_build_policy: IndexBuildPolicy::default(),
1230 pk_by_row_complete: false,
1231 flushed_epoch: 0,
1232 page_cache: ctx.page_cache,
1233 decoded_cache: ctx.decoded_cache,
1234 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1235 snapshots: ctx.snapshots,
1236 commit_lock: ctx.commit_lock,
1237 result_cache: Arc::new(parking_lot::Mutex::new(
1238 ResultCache::new()
1239 .with_dir(rcache_dir)
1240 .with_cache_dek(cache_dek.clone()),
1241 )),
1242 pending_delete_rids: roaring::RoaringBitmap::new(),
1243 pending_put_cols: std::collections::HashSet::new(),
1244 pending_rows: Vec::new(),
1245 pending_rows_auto_inc: Vec::new(),
1246 pending_dels: Vec::new(),
1247 pending_truncate: None,
1248 wal_dek,
1249 auto_inc,
1250 ttl: None,
1251 })
1252 }
1253
1254 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1258 let dir = dir.as_ref();
1259 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1260 Self::open_in(dir, ctx)
1261 }
1262
1263 #[cfg(feature = "encryption")]
1266 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1267 let dir = dir.as_ref();
1268 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1269 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1270 MongrelError::NotFound(format!(
1271 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1272 salt_path
1273 ))
1274 })?;
1275 let salt_len = crate::encryption::SALT_LEN;
1276 if salt_bytes.len() != salt_len {
1277 return Err(MongrelError::InvalidArgument(format!(
1278 "encryption salt is {} bytes, expected {salt_len}",
1279 salt_bytes.len()
1280 )));
1281 }
1282 let mut salt = [0u8; 16];
1283 salt.copy_from_slice(&salt_bytes);
1284 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1285 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1286 let t = Self::open_in(dir, ctx)?;
1287 Ok(t)
1288 }
1289
1290 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1291 let dir = dir.as_ref().to_path_buf();
1292 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1293 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1294 let schema: Schema = read_schema(&dir)?;
1295 schema.validate_ai()?;
1296 for index in &schema.indexes {
1297 index.validate_options()?;
1298 }
1299 let replay_epoch = Epoch(manifest.current_epoch);
1300 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1301 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1305 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1306 None => {
1307 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1308 let replayed = match &active {
1310 Some(path) => {
1311 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1312 crate::wal::replay_with_cipher(path, cipher)?
1313 }
1314 None => Vec::new(),
1315 };
1316 let mut w = match &active {
1317 Some(path) => Wal::create_with_cipher(
1318 path,
1319 replay_epoch,
1320 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1321 0,
1322 )?,
1323 None => Wal::create_with_cipher(
1324 dir.join(WAL_DIR).join("seg-000000.wal"),
1325 replay_epoch,
1326 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1327 0,
1328 )?,
1329 };
1330 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1331 (WalSink::Private(w), replayed, 1)
1332 }
1333 };
1334
1335 let mut memtable = Memtable::new();
1336 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1337 let persisted_epoch = manifest.current_epoch;
1338 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1345 s.next = manifest.auto_inc_next;
1346 s.seeded = manifest.auto_inc_next > 0;
1347 s
1348 });
1349
1350 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1357 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1358 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1359 std::collections::BTreeMap::new();
1360 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1361 let mut saw_delete = false;
1362 for record in replayed {
1363 let txn_id = record.txn_id;
1364 match record.op {
1365 Op::Put { rows, .. } => {
1366 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1367 for row in &rows {
1368 allocator.advance_to(row.row_id);
1369 if let Some(ai) = auto_inc.as_mut() {
1370 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1371 if *n + 1 > ai.next {
1372 ai.next = *n + 1;
1373 }
1374 }
1375 }
1376 }
1377 staged_puts.entry(txn_id).or_default().extend(rows);
1378 }
1379 Op::Delete { row_ids, .. } => {
1380 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1381 }
1382 Op::TxnCommit { epoch, .. } => {
1383 let commit_epoch = Epoch(epoch);
1384 if let Some(puts) = staged_puts.remove(&txn_id) {
1385 if commit_epoch.0 > manifest.flushed_epoch {
1386 for row in &puts {
1387 memtable.upsert(row.clone());
1388 }
1389 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1390 }
1391 }
1392 if let Some(dels) = staged_deletes.remove(&txn_id) {
1393 saw_delete = true;
1394 if commit_epoch.0 > manifest.flushed_epoch {
1395 for rid in dels {
1396 memtable.tombstone(rid, commit_epoch);
1397 replayed_deletes.push((rid, commit_epoch));
1398 }
1399 }
1400 }
1401 }
1402 Op::TxnAbort => {
1403 staged_puts.remove(&txn_id);
1404 staged_deletes.remove(&txn_id);
1405 }
1406 Op::TruncateTable { .. }
1407 | Op::ExternalTableState { .. }
1408 | Op::Flush { .. }
1409 | Op::Ddl(_)
1410 | Op::BeforeImage { .. }
1411 | Op::CommitTimestamp { .. } => {}
1412 }
1413 }
1414
1415 let rcache_dir = dir.join(RCACHE_DIR);
1416 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1417 let mut db = Self {
1418 dir,
1419 table_id: manifest.table_id,
1420 name: ctx.table_name.unwrap_or_default(),
1421 auth: ctx.auth,
1422 read_only: ctx.read_only,
1423 wal,
1424 memtable,
1425 mutable_run: MutableRun::new(),
1426 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1427 compaction_zstd_level: 3,
1428 allocator,
1429 epoch: ctx.epoch,
1430 persisted_epoch,
1431 data_generation: persisted_epoch,
1432 schema,
1433 hot: HotIndex::new(),
1434 kek: ctx.kek,
1435 column_keys,
1436 run_refs: manifest.runs.clone(),
1437 retiring: manifest.retiring.clone(),
1438 next_run_id: manifest
1439 .runs
1440 .iter()
1441 .map(|r| r.run_id as u64 + 1)
1442 .max()
1443 .unwrap_or(1),
1444 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1445 current_txn_id,
1446 bitmap: HashMap::new(),
1447 ann: HashMap::new(),
1448 fm: HashMap::new(),
1449 sparse: HashMap::new(),
1450 minhash: HashMap::new(),
1451 learned_range: Arc::new(HashMap::new()),
1452 pk_by_row: ReversePkMap::new(),
1453 pinned: BTreeMap::new(),
1454 live_count: manifest.live_count,
1455 reservoir: crate::reservoir::Reservoir::default(),
1456 reservoir_complete: false,
1457 had_deletes: saw_delete
1458 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1459 != manifest.live_count,
1460 agg_cache: Arc::new(HashMap::new()),
1461 global_idx_epoch: manifest.global_idx_epoch,
1462 indexes_complete: true,
1463 index_build_policy: IndexBuildPolicy::default(),
1464 pk_by_row_complete: false,
1465 flushed_epoch: manifest.flushed_epoch,
1466 page_cache: ctx.page_cache,
1467 decoded_cache: ctx.decoded_cache,
1468 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1469 snapshots: ctx.snapshots,
1470 commit_lock: ctx.commit_lock,
1471 result_cache: Arc::new(parking_lot::Mutex::new(
1472 ResultCache::new()
1473 .with_dir(rcache_dir)
1474 .with_cache_dek(cache_dek.clone()),
1475 )),
1476 pending_delete_rids: roaring::RoaringBitmap::new(),
1477 pending_put_cols: std::collections::HashSet::new(),
1478 pending_rows: Vec::new(),
1479 pending_rows_auto_inc: Vec::new(),
1480 pending_dels: Vec::new(),
1481 pending_truncate: None,
1482 wal_dek,
1483 auto_inc,
1484 ttl: manifest.ttl,
1485 };
1486
1487 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1490
1491 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1496 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1497 c.epoch_built == manifest.global_idx_epoch
1498 && manifest.global_idx_epoch > 0
1499 && manifest
1500 .runs
1501 .iter()
1502 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1503 });
1504 if let Some(loaded) = checkpoint {
1505 if checkpoint_valid {
1506 db.hot = loaded.hot;
1507 db.bitmap = loaded.bitmap;
1508 db.ann = loaded.ann;
1509 db.fm = loaded.fm;
1510 db.sparse = loaded.sparse;
1511 db.minhash = loaded.minhash;
1512 db.learned_range = Arc::new(loaded.learned_range);
1513 }
1516 }
1517 if !checkpoint_valid {
1518 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1519 db.bitmap = bitmap;
1520 db.ann = ann;
1521 db.fm = fm;
1522 db.sparse = sparse;
1523 db.minhash = minhash;
1524 db.rebuild_indexes_from_runs()?;
1525 db.build_learned_ranges()?;
1526 }
1527
1528 for (epoch, group) in replayed_puts {
1533 let (losers, winner_pks) = db.partition_pk_winners(&group);
1534 for (key, &row_id) in &winner_pks {
1535 if let Some(old_rid) = db.hot.get(key) {
1536 if old_rid != row_id {
1537 db.tombstone_row(old_rid, epoch, false);
1538 }
1539 }
1540 }
1541 for &loser_rid in &losers {
1542 db.tombstone_row(loser_rid, epoch, false);
1543 }
1544 for (key, row_id) in winner_pks {
1545 db.insert_hot_pk(key, row_id);
1546 }
1547 if db.schema.primary_key().is_none() {
1548 for r in &group {
1549 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1550 }
1551 }
1552 for r in &group {
1553 if !losers.contains(&r.row_id) {
1554 db.index_row(r);
1555 }
1556 }
1557 }
1558 for (rid, epoch) in &replayed_deletes {
1562 db.remove_hot_for_row(*rid, *epoch);
1563 }
1564
1565 db.result_cache.lock().load_persistent();
1572 Ok(db)
1573 }
1574
1575 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1581 if self.reservoir_complete {
1582 return Ok(());
1583 }
1584 self.rebuild_reservoir()?;
1585 self.reservoir_complete = true;
1586 Ok(())
1587 }
1588
1589 fn rebuild_reservoir(&mut self) -> Result<()> {
1592 let snap = self.snapshot();
1593 let rows = self.visible_rows(snap)?;
1594 self.reservoir.reset();
1595 for r in rows {
1596 self.reservoir.offer(r.row_id.0);
1597 }
1598 Ok(())
1599 }
1600
1601 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1602 self.hot = HotIndex::new();
1603 self.pk_by_row.clear();
1604 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1605 self.bitmap = bitmap;
1606 self.ann = ann;
1607 self.fm = fm;
1608 self.sparse = sparse;
1609 self.minhash = minhash;
1610 let snapshot = Epoch(u64::MAX);
1611 let ttl_now = unix_nanos_now();
1612 for rr in self.run_refs.clone() {
1613 let mut reader = self.open_reader(rr.run_id)?;
1614 for row in reader.visible_rows(snapshot)? {
1615 if self.row_expired_at(&row, ttl_now) {
1616 continue;
1617 }
1618 let tok_row = self.tokenized_for_indexes(&row);
1619 index_into(
1620 &self.schema,
1621 &tok_row,
1622 &mut self.hot,
1623 &mut self.bitmap,
1624 &mut self.ann,
1625 &mut self.fm,
1626 &mut self.sparse,
1627 &mut self.minhash,
1628 );
1629 }
1630 }
1631 for row in self.mutable_run.visible_versions(snapshot) {
1632 if row.deleted {
1633 self.remove_hot_for_row(row.row_id, snapshot);
1634 } else if !self.row_expired_at(&row, ttl_now) {
1635 self.index_row(&row);
1636 }
1637 }
1638 for row in self.memtable.visible_versions(snapshot) {
1639 if row.deleted {
1640 self.remove_hot_for_row(row.row_id, snapshot);
1641 } else if !self.row_expired_at(&row, ttl_now) {
1642 self.index_row(&row);
1643 }
1644 }
1645 self.refresh_pk_by_row_from_hot();
1646 Ok(())
1647 }
1648
1649 fn refresh_pk_by_row_from_hot(&mut self) {
1650 self.pk_by_row_complete = true;
1651 if self.schema.primary_key().is_none() {
1652 self.pk_by_row.clear();
1653 return;
1654 }
1655 self.pk_by_row = ReversePkMap::from_entries(
1661 self.hot
1662 .entries()
1663 .into_iter()
1664 .map(|(key, row_id)| (row_id, key)),
1665 );
1666 }
1667
1668 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1669 if self.schema.primary_key().is_some() {
1670 self.pk_by_row.insert(row_id, key.clone());
1671 }
1672 self.hot.insert(key, row_id);
1673 }
1674
1675 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1679 self.learned_range = Arc::new(HashMap::new());
1680 if self.run_refs.len() != 1 {
1681 return Ok(());
1682 }
1683 let cols: Vec<(u16, usize)> = self
1684 .schema
1685 .indexes
1686 .iter()
1687 .filter(|i| i.kind == IndexKind::LearnedRange)
1688 .map(|i| {
1689 (
1690 i.column_id,
1691 i.options
1692 .learned_range
1693 .as_ref()
1694 .map(|options| options.epsilon)
1695 .unwrap_or(16),
1696 )
1697 })
1698 .collect();
1699 if cols.is_empty() {
1700 return Ok(());
1701 }
1702 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1703 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1704 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1705 _ => return Ok(()),
1706 };
1707 for (cid, epsilon) in cols {
1708 let ty = self
1709 .schema
1710 .columns
1711 .iter()
1712 .find(|c| c.id == cid)
1713 .map(|c| c.ty.clone())
1714 .unwrap_or(TypeId::Int64);
1715 match ty {
1716 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1717 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1718 let pairs: Vec<(i64, u64)> = data
1719 .iter()
1720 .zip(row_ids.iter())
1721 .map(|(v, r)| (*v, *r))
1722 .collect();
1723 Arc::make_mut(&mut self.learned_range).insert(
1724 cid,
1725 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
1726 );
1727 }
1728 }
1729 TypeId::Float64 => {
1730 if let columnar::NativeColumn::Float64 { data, .. } =
1731 reader.column_native(cid)?
1732 {
1733 let pairs: Vec<(f64, u64)> = data
1734 .iter()
1735 .zip(row_ids.iter())
1736 .map(|(v, r)| (*v, *r))
1737 .collect();
1738 Arc::make_mut(&mut self.learned_range).insert(
1739 cid,
1740 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
1741 );
1742 }
1743 }
1744 _ => {}
1745 }
1746 }
1747 Ok(())
1748 }
1749
1750 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1757 if self.indexes_complete {
1758 crate::trace::QueryTrace::record(|t| {
1759 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1760 });
1761 return Ok(());
1762 }
1763 crate::trace::QueryTrace::record(|t| {
1764 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1765 });
1766 self.rebuild_indexes_from_runs()?;
1767 self.build_learned_ranges()?;
1768 self.indexes_complete = true;
1769 let epoch = self.current_epoch();
1770 self.checkpoint_indexes(epoch);
1771 Ok(())
1772 }
1773
1774 fn pending_epoch(&self) -> Epoch {
1775 Epoch(self.epoch.visible().0 + 1)
1776 }
1777
1778 fn is_shared(&self) -> bool {
1781 matches!(self.wal, WalSink::Shared(_))
1782 }
1783
1784 fn ensure_txn_id(&mut self) -> u64 {
1788 if self.current_txn_id == 0 {
1789 let id = match &self.wal {
1790 WalSink::Shared(s) => {
1791 let mut g = s.txn_ids.lock();
1792 let v = *g;
1793 *g = g.wrapping_add(1);
1794 v
1795 }
1796 WalSink::Private(_) => 1,
1797 WalSink::ReadOnly => 1,
1798 };
1799 self.current_txn_id = id;
1800 }
1801 self.current_txn_id
1802 }
1803
1804 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1807 self.ensure_writable()?;
1808 let txn_id = self.ensure_txn_id();
1809 let table_id = self.table_id;
1810 match &mut self.wal {
1811 WalSink::Private(w) => {
1812 w.append_txn(txn_id, op)?;
1813 }
1814 WalSink::Shared(s) => {
1815 s.wal.lock().append(txn_id, table_id, op)?;
1816 }
1817 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
1818 }
1819 Ok(())
1820 }
1821
1822 fn ensure_writable(&self) -> Result<()> {
1823 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
1824 Err(MongrelError::ReadOnlyReplica)
1825 } else {
1826 Ok(())
1827 }
1828 }
1829
1830 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1841 match &self.auth {
1842 Some(checker) => checker.check(&self.name, perm),
1843 None => Ok(()),
1844 }
1845 }
1846 pub fn require_select(&self) -> Result<()> {
1851 self.require(crate::auth_state::RequiredPermission::Select)
1852 }
1853 fn require_insert(&self) -> Result<()> {
1854 self.require(crate::auth_state::RequiredPermission::Insert)
1855 }
1856 #[allow(dead_code)]
1860 fn require_update(&self) -> Result<()> {
1861 self.require(crate::auth_state::RequiredPermission::Update)
1862 }
1863 fn require_delete(&self) -> Result<()> {
1864 self.require(crate::auth_state::RequiredPermission::Delete)
1865 }
1866
1867 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1870 self.require_insert()?;
1871 Ok(self.put_returning(columns)?.0)
1872 }
1873
1874 pub fn put_returning(
1879 &mut self,
1880 mut columns: Vec<(u16, Value)>,
1881 ) -> Result<(RowId, Option<i64>)> {
1882 self.require_insert()?;
1883 let assigned = self.fill_auto_inc(&mut columns)?;
1884 self.apply_defaults(&mut columns)?;
1885 self.schema.validate_values(&columns)?;
1886 let row_id = if self.schema.clustered {
1891 self.derive_clustered_row_id(&columns)?
1892 } else {
1893 self.allocator.alloc()
1894 };
1895 let epoch = self.pending_epoch();
1896 let mut row = Row::new(row_id, epoch);
1897 for (col_id, val) in columns {
1898 row.columns.insert(col_id, val);
1899 }
1900 self.commit_rows(vec![row], assigned.is_some())?;
1901 Ok((row_id, assigned))
1902 }
1903
1904 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1907 self.require_insert()?;
1908 Ok(self
1909 .put_batch_returning(batch)?
1910 .into_iter()
1911 .map(|(r, _)| r)
1912 .collect())
1913 }
1914
1915 pub fn put_batch_returning(
1918 &mut self,
1919 batch: Vec<Vec<(u16, Value)>>,
1920 ) -> Result<Vec<(RowId, Option<i64>)>> {
1921 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1922 for mut cols in batch {
1923 let assigned = self.fill_auto_inc(&mut cols)?;
1924 self.apply_defaults(&mut cols)?;
1925 filled.push((cols, assigned));
1926 }
1927 for (cols, _) in &filled {
1928 self.schema.validate_values(cols)?;
1929 }
1930 let epoch = self.pending_epoch();
1931 let mut rows = Vec::with_capacity(filled.len());
1932 let mut ids = Vec::with_capacity(filled.len());
1933 for (cols, assigned) in filled {
1934 let row_id = if self.schema.clustered {
1935 self.derive_clustered_row_id(&cols)?
1936 } else {
1937 self.allocator.alloc()
1938 };
1939 let mut row = Row::new(row_id, epoch);
1940 for (c, v) in cols {
1941 row.columns.insert(c, v);
1942 }
1943 ids.push((row_id, assigned));
1944 rows.push(row);
1945 }
1946 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1947 self.commit_rows(rows, all_auto_generated)?;
1948 Ok(ids)
1949 }
1950
1951 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1957 self.ensure_writable()?;
1958 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1959 return Ok(None);
1960 };
1961 let pos = columns.iter().position(|(c, _)| *c == cid);
1962 let assigned = match pos {
1963 Some(i) => match &columns[i].1 {
1964 Value::Null => {
1965 let next = self.alloc_auto_inc_value()?;
1966 columns[i].1 = Value::Int64(next);
1967 Some(next)
1968 }
1969 Value::Int64(n) => {
1970 self.advance_auto_inc_past(*n)?;
1971 None
1972 }
1973 other => {
1974 return Err(MongrelError::InvalidArgument(format!(
1975 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1976 other
1977 )))
1978 }
1979 },
1980 None => {
1981 let next = self.alloc_auto_inc_value()?;
1982 columns.push((cid, Value::Int64(next)));
1983 Some(next)
1984 }
1985 };
1986 Ok(assigned)
1987 }
1988
1989 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1995 for col in &self.schema.columns {
1996 let Some(expr) = &col.default_value else {
1997 continue;
1998 };
1999 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2001 continue;
2002 }
2003 let pos = columns.iter().position(|(c, _)| *c == col.id);
2004 let needs_default = match pos {
2005 None => true,
2006 Some(i) => matches!(columns[i].1, Value::Null),
2007 };
2008 if !needs_default {
2009 continue;
2010 }
2011 let v = match expr {
2012 crate::schema::DefaultExpr::Static(v) => v.clone(),
2013 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2014 crate::schema::DefaultExpr::Uuid => {
2015 let mut buf = [0u8; 16];
2016 getrandom::getrandom(&mut buf)
2017 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2018 Value::Uuid(buf)
2019 }
2020 };
2021 match pos {
2022 None => columns.push((col.id, v)),
2023 Some(i) => columns[i].1 = v,
2024 }
2025 }
2026 Ok(())
2027 }
2028
2029 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2031 self.ensure_auto_inc_seeded()?;
2032 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2034 let v = ai.next;
2035 ai.next = ai.next.saturating_add(1);
2036 Ok(v)
2037 }
2038
2039 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2042 self.ensure_auto_inc_seeded()?;
2043 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2044 let floor = used.saturating_add(1).max(1);
2045 if ai.next < floor {
2046 ai.next = floor;
2047 }
2048 Ok(())
2049 }
2050
2051 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2056 let needs_seed = match self.auto_inc {
2057 Some(ai) => !ai.seeded,
2058 None => return Ok(()),
2059 };
2060 if !needs_seed {
2061 return Ok(());
2062 }
2063 if self.seed_empty_auto_inc() {
2064 return Ok(());
2065 }
2066 let cid = self
2067 .auto_inc
2068 .as_ref()
2069 .expect("auto-inc column present")
2070 .column_id;
2071 let max = self.scan_max_int64(cid)?;
2072 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2073 let floor = max.saturating_add(1).max(1);
2074 if ai.next < floor {
2075 ai.next = floor;
2076 }
2077 ai.seeded = true;
2078 Ok(())
2079 }
2080
2081 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2082 if n == 0 || self.auto_inc.is_none() {
2083 return Ok(None);
2084 }
2085 self.ensure_auto_inc_seeded()?;
2086 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2087 let start = ai.next;
2088 ai.next = ai.next.saturating_add(n as i64);
2089 Ok(Some(start))
2090 }
2091
2092 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2096 let mut max: i64 = 0;
2097 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2098 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2099 if *n > max {
2100 max = *n;
2101 }
2102 }
2103 }
2104 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2105 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2106 if *n > max {
2107 max = *n;
2108 }
2109 }
2110 }
2111 for rr in self.run_refs.clone() {
2112 let reader = self.open_reader(rr.run_id)?;
2113 if let Some(stats) = reader.column_page_stats(column_id) {
2114 for s in stats {
2115 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2116 if n > max {
2117 max = n;
2118 }
2119 }
2120 }
2121 } else if reader.has_column(column_id) {
2122 if let columnar::NativeColumn::Int64 { data, validity } =
2123 reader.column_native_shared(column_id)?
2124 {
2125 for (i, n) in data.iter().enumerate() {
2126 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2127 {
2128 max = *n;
2129 }
2130 }
2131 }
2132 }
2133 }
2134 Ok(max)
2135 }
2136
2137 fn seed_empty_auto_inc(&mut self) -> bool {
2138 let Some(ai) = self.auto_inc.as_mut() else {
2139 return false;
2140 };
2141 if ai.seeded || self.live_count != 0 {
2142 return false;
2143 }
2144 if ai.next < 1 {
2145 ai.next = 1;
2146 }
2147 ai.seeded = true;
2148 true
2149 }
2150
2151 fn advance_auto_inc_from_native_columns(
2152 &mut self,
2153 columns: &[(u16, columnar::NativeColumn)],
2154 n: usize,
2155 live_before: u64,
2156 ) -> Result<()> {
2157 let Some(ai) = self.auto_inc.as_mut() else {
2158 return Ok(());
2159 };
2160 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2161 return Ok(());
2162 };
2163 let columnar::NativeColumn::Int64 { data, validity } = col else {
2164 return Err(MongrelError::InvalidArgument(format!(
2165 "AUTO_INCREMENT column {} must be Int64",
2166 ai.column_id
2167 )));
2168 };
2169 let max = if native_int64_strictly_increasing(col, n) {
2170 data.get(n.saturating_sub(1)).copied()
2171 } else {
2172 data.iter()
2173 .take(n)
2174 .enumerate()
2175 .filter_map(|(i, v)| {
2176 if validity.is_empty() || columnar::validity_bit(validity, i) {
2177 Some(*v)
2178 } else {
2179 None
2180 }
2181 })
2182 .max()
2183 };
2184 if let Some(max) = max {
2185 let floor = max.saturating_add(1).max(1);
2186 if ai.next < floor {
2187 ai.next = floor;
2188 }
2189 if ai.seeded || live_before == 0 {
2190 ai.seeded = true;
2191 }
2192 }
2193 Ok(())
2194 }
2195
2196 fn fill_auto_inc_native_columns(
2197 &mut self,
2198 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2199 n: usize,
2200 ) -> Result<()> {
2201 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2202 return Ok(());
2203 };
2204 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2205 if let Some(start) = self.alloc_auto_inc_range(n)? {
2206 columns.push((
2207 cid,
2208 columnar::NativeColumn::Int64 {
2209 data: (start..start.saturating_add(n as i64)).collect(),
2210 validity: vec![0xFF; n.div_ceil(8)],
2211 },
2212 ));
2213 }
2214 return Ok(());
2215 };
2216
2217 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2218 return Err(MongrelError::InvalidArgument(format!(
2219 "AUTO_INCREMENT column {cid} must be Int64"
2220 )));
2221 };
2222 if data.len() < n {
2223 return Err(MongrelError::InvalidArgument(format!(
2224 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2225 data.len()
2226 )));
2227 }
2228 if columnar::all_non_null(validity, n) {
2229 return Ok(());
2230 }
2231 if validity.iter().all(|b| *b == 0) {
2232 if let Some(start) = self.alloc_auto_inc_range(n)? {
2233 for (i, slot) in data.iter_mut().take(n).enumerate() {
2234 *slot = start.saturating_add(i as i64);
2235 }
2236 *validity = vec![0xFF; n.div_ceil(8)];
2237 }
2238 return Ok(());
2239 }
2240
2241 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2242 for (i, slot) in data.iter_mut().enumerate().take(n) {
2243 if columnar::validity_bit(validity, i) {
2244 self.advance_auto_inc_past(*slot)?;
2245 } else {
2246 *slot = self.alloc_auto_inc_value()?;
2247 }
2248 }
2249 *validity = new_validity;
2250 Ok(())
2251 }
2252
2253 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2267 self.ensure_writable()?;
2268 if self.auto_inc.is_none() {
2269 return Ok(None);
2270 }
2271 Ok(Some(self.alloc_auto_inc_value()?))
2272 }
2273
2274 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2280 let payload = bincode::serialize(&rows)?;
2281 self.wal_append_data(Op::Put {
2282 table_id: self.table_id,
2283 rows: payload,
2284 })?;
2285 if self.is_shared() {
2286 self.pending_rows_auto_inc
2287 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2288 self.pending_rows.extend(rows);
2289 } else {
2290 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2291 }
2292 Ok(())
2293 }
2294
2295 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2300 self.apply_put_rows_inner(rows, true)
2301 }
2302
2303 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2304 if check_existing_pk {
2305 self.ensure_indexes_complete()?;
2306 }
2307 if rows.len() == 1 {
2311 let row = rows.into_iter().next().expect("len checked");
2312 return self.apply_put_row_single(row, check_existing_pk);
2313 }
2314 let pk_id = self.schema.primary_key().map(|c| c.id);
2331 let probe = match pk_id {
2332 Some(pid) => {
2333 check_existing_pk
2334 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2335 }
2336 None => false,
2337 };
2338 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2341 for r in rows {
2342 for &cid in r.columns.keys() {
2343 self.pending_put_cols.insert(cid);
2344 }
2345 match pk_id {
2346 Some(pid) if probe || maintain_pk_by_row => {
2347 if let Some(pk_val) = r.columns.get(&pid) {
2348 let key = self.index_lookup_key(pid, pk_val);
2349 if probe {
2350 if let Some(old_rid) = self.hot.get(&key) {
2351 if old_rid != r.row_id {
2352 self.tombstone_row(old_rid, r.committed_epoch, true);
2353 }
2354 }
2355 }
2356 if maintain_pk_by_row {
2357 self.pk_by_row.insert(r.row_id, key);
2358 }
2359 }
2360 }
2361 Some(_) => {}
2362 None => {
2363 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2364 }
2365 }
2366 self.index_row(&r);
2367 self.reservoir.offer(r.row_id.0);
2368 self.memtable.upsert(r);
2369 self.live_count = self.live_count.saturating_add(1);
2372 }
2373 self.data_generation = self.data_generation.wrapping_add(1);
2374 Ok(())
2375 }
2376
2377 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2381 for &cid in row.columns.keys() {
2382 self.pending_put_cols.insert(cid);
2383 }
2384 let epoch = row.committed_epoch;
2385 if let Some(pk_col) = self.schema.primary_key() {
2386 let pk_id = pk_col.id;
2387 if let Some(pk_val) = row.columns.get(&pk_id) {
2388 let maintain_pk_by_row = self.pk_by_row_complete;
2392 if check_existing_pk || maintain_pk_by_row {
2393 let key = self.index_lookup_key(pk_id, pk_val);
2394 if check_existing_pk {
2395 if let Some(old_rid) = self.hot.get(&key) {
2396 if old_rid != row.row_id {
2397 self.tombstone_row(old_rid, epoch, true);
2398 }
2399 }
2400 }
2401 if maintain_pk_by_row {
2402 self.pk_by_row.insert(row.row_id, key);
2403 }
2404 }
2405 }
2406 } else {
2407 self.hot
2408 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2409 }
2410 self.index_row(&row);
2411 self.reservoir.offer(row.row_id.0);
2412 self.memtable.upsert(row);
2413 self.live_count = self.live_count.saturating_add(1);
2414 self.data_generation = self.data_generation.wrapping_add(1);
2415 Ok(())
2416 }
2417
2418 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2421 self.allocator.alloc()
2422 }
2423
2424 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2430 let pk = self.schema.primary_key().ok_or_else(|| {
2431 MongrelError::Schema("clustered table requires a single-column primary key".into())
2432 })?;
2433 let pk_val = columns
2434 .iter()
2435 .find(|(id, _)| *id == pk.id)
2436 .map(|(_, v)| v)
2437 .ok_or_else(|| {
2438 MongrelError::Schema(format!(
2439 "clustered table missing primary key column {} ({})",
2440 pk.id, pk.name
2441 ))
2442 })?;
2443 let key_bytes = pk_val.encode_key();
2444 let mut hash: u64 = 0xcbf29ce484222325;
2446 for &b in &key_bytes {
2447 hash ^= b as u64;
2448 hash = hash.wrapping_mul(0x100000001b3);
2449 }
2450 Ok(RowId(hash.max(1)))
2453 }
2454
2455 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2463 self.ensure_indexes_complete()?;
2464 let n = rows.len();
2465 for r in rows {
2466 for &cid in r.columns.keys() {
2467 self.pending_put_cols.insert(cid);
2468 }
2469 }
2470 let (losers, winner_pks) = self.partition_pk_winners(rows);
2471 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2472 for (key, &row_id) in &winner_pks {
2474 if let Some(old_rid) = self.hot.get(key) {
2475 if old_rid != row_id {
2476 self.tombstone_row(old_rid, epoch, true);
2477 }
2478 }
2479 }
2480 for &loser_rid in &losers {
2483 self.tombstone_row(loser_rid, epoch, false);
2484 }
2485 for (key, row_id) in winner_pks {
2487 self.insert_hot_pk(key, row_id);
2488 }
2489 if self.schema.primary_key().is_none() {
2490 for r in rows {
2491 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2492 }
2493 }
2494 for r in rows {
2495 self.allocator.advance_to(r.row_id);
2496 if !losers.contains(&r.row_id) {
2497 self.index_row(r);
2498 }
2499 }
2500 for r in rows {
2501 if !losers.contains(&r.row_id) {
2502 self.reservoir.offer(r.row_id.0);
2503 }
2504 }
2505 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2506 self.data_generation = self.data_generation.wrapping_add(1);
2507 Ok(())
2508 }
2509
2510 pub(crate) fn recover_apply(
2515 &mut self,
2516 rows: Vec<Row>,
2517 deletes: Vec<(RowId, Epoch)>,
2518 ) -> Result<()> {
2519 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2523 std::collections::BTreeMap::new();
2524 for row in rows {
2525 self.allocator.advance_to(row.row_id);
2526 if let Some(ai) = self.auto_inc.as_mut() {
2531 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2532 if *n + 1 > ai.next {
2533 ai.next = *n + 1;
2534 }
2535 }
2536 }
2537 by_epoch.entry(row.committed_epoch).or_default().push(row);
2538 }
2539 for (epoch, group) in by_epoch {
2540 let (losers, winner_pks) = self.partition_pk_winners(&group);
2541 for (key, &row_id) in &winner_pks {
2543 if let Some(old_rid) = self.hot.get(key) {
2544 if old_rid != row_id {
2545 self.tombstone_row(old_rid, epoch, false);
2546 }
2547 }
2548 }
2549 for (key, row_id) in winner_pks {
2550 self.insert_hot_pk(key, row_id);
2551 }
2552 if self.schema.primary_key().is_none() {
2553 for r in &group {
2554 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2555 }
2556 }
2557 for r in &group {
2558 if !losers.contains(&r.row_id) {
2559 self.memtable.upsert(r.clone());
2560 self.index_row(r);
2561 }
2562 }
2563 }
2564 for (rid, epoch) in deletes {
2565 self.memtable.tombstone(rid, epoch);
2566 self.remove_hot_for_row(rid, epoch);
2567 }
2568 self.reservoir_complete = false;
2571 Ok(())
2572 }
2573
2574 pub(crate) fn flushed_epoch(&self) -> u64 {
2576 self.flushed_epoch
2577 }
2578
2579 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2580 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2581 }
2582
2583 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2585 self.schema.validate_values(cells)
2586 }
2587
2588 fn validate_columns_not_null(
2592 &self,
2593 columns: &[(u16, columnar::NativeColumn)],
2594 n: usize,
2595 ) -> Result<()> {
2596 let by_id: HashMap<u16, &columnar::NativeColumn> =
2597 columns.iter().map(|(id, c)| (*id, c)).collect();
2598 for col in &self.schema.columns {
2599 if !col.flags.contains(ColumnFlags::NULLABLE) {
2600 match by_id.get(&col.id) {
2601 None => {
2602 return Err(MongrelError::InvalidArgument(format!(
2603 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2604 col.name, col.id
2605 )));
2606 }
2607 Some(c) => {
2608 if c.null_count(n) != 0 {
2609 return Err(MongrelError::InvalidArgument(format!(
2610 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2611 col.name, col.id
2612 )));
2613 }
2614 }
2615 }
2616 }
2617 if let TypeId::Enum { variants } = &col.ty {
2618 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2619 if by_id.contains_key(&col.id) {
2620 return Err(MongrelError::InvalidArgument(format!(
2621 "column '{}' ({}) enum requires a bytes column",
2622 col.name, col.id
2623 )));
2624 }
2625 continue;
2626 };
2627 for index in 0..n {
2628 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2629 continue;
2630 };
2631 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2632 return Err(MongrelError::InvalidArgument(format!(
2633 "column '{}' ({}) enum value {:?} is not one of {:?}",
2634 col.name,
2635 col.id,
2636 String::from_utf8_lossy(value),
2637 variants
2638 )));
2639 }
2640 }
2641 }
2642 }
2643 Ok(())
2644 }
2645
2646 fn bulk_pk_winner_indices(
2651 &self,
2652 columns: &[(u16, columnar::NativeColumn)],
2653 n: usize,
2654 ) -> Option<Vec<usize>> {
2655 let pk_col = self.schema.primary_key()?;
2656 let pk_id = pk_col.id;
2657 let pk_ty = pk_col.ty.clone();
2658 let by_id: HashMap<u16, &columnar::NativeColumn> =
2659 columns.iter().map(|(id, c)| (*id, c)).collect();
2660 let pk_native = by_id.get(&pk_id)?;
2661 if native_int64_strictly_increasing(pk_native, n) {
2662 return None;
2663 }
2664 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2666 let mut null_pk_rows: Vec<usize> = Vec::new();
2667 for i in 0..n {
2668 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2669 Some(key) => {
2670 last.insert(key, i);
2671 }
2672 None => null_pk_rows.push(i),
2673 }
2674 }
2675 let mut winners: HashSet<usize> = last.values().copied().collect();
2676 for i in null_pk_rows {
2677 winners.insert(i);
2678 }
2679 Some((0..n).filter(|i| winners.contains(i)).collect())
2680 }
2681
2682 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2684 self.require_delete()?;
2685 let epoch = self.pending_epoch();
2686 self.wal_append_data(Op::Delete {
2687 table_id: self.table_id,
2688 row_ids: vec![row_id],
2689 })?;
2690 if self.is_shared() {
2691 self.pending_dels.push(row_id);
2692 } else {
2693 self.apply_delete(row_id, epoch);
2694 }
2695 Ok(())
2696 }
2697
2698 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2699 let pre = self.get(row_id, self.snapshot());
2700 self.delete(row_id)?;
2701 Ok(pre.map(|row| {
2702 let mut columns: Vec<_> = row.columns.into_iter().collect();
2703 columns.sort_by_key(|(id, _)| *id);
2704 OwnedRow { columns }
2705 }))
2706 }
2707
2708 pub fn truncate(&mut self) -> Result<()> {
2710 self.require_delete()?;
2711 let epoch = self.pending_epoch();
2712 self.wal_append_data(Op::TruncateTable {
2713 table_id: self.table_id,
2714 })?;
2715 self.pending_rows.clear();
2716 self.pending_rows_auto_inc.clear();
2717 self.pending_dels.clear();
2718 self.pending_truncate = Some(epoch);
2719 Ok(())
2720 }
2721
2722 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2724 for rr in std::mem::take(&mut self.run_refs) {
2725 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2726 }
2727 for r in std::mem::take(&mut self.retiring) {
2728 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2729 }
2730 self.memtable = Memtable::new();
2731 self.mutable_run = MutableRun::new();
2732 self.hot = HotIndex::new();
2733 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2734 self.bitmap = bitmap;
2735 self.ann = ann;
2736 self.fm = fm;
2737 self.sparse = sparse;
2738 self.minhash = minhash;
2739 self.learned_range = Arc::new(HashMap::new());
2740 self.pk_by_row.clear();
2741 self.pk_by_row_complete = false;
2742 self.live_count = 0;
2743 self.reservoir = crate::reservoir::Reservoir::default();
2744 self.reservoir_complete = true;
2745 self.had_deletes = true;
2746 self.agg_cache = Arc::new(HashMap::new());
2747 self.global_idx_epoch = 0;
2748 self.indexes_complete = true;
2749 self.pending_delete_rids.clear();
2750 self.pending_put_cols.clear();
2751 self.pending_rows.clear();
2752 self.pending_rows_auto_inc.clear();
2753 self.pending_dels.clear();
2754 self.clear_result_cache();
2755 self.invalidate_index_checkpoint();
2756 self.data_generation = self.data_generation.wrapping_add(1);
2757 Ok(())
2758 }
2759
2760 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2763 self.remove_hot_for_row(row_id, epoch);
2764 self.tombstone_row(row_id, epoch, true);
2765 self.data_generation = self.data_generation.wrapping_add(1);
2766 }
2767
2768 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2772 let tombstone = Row {
2773 row_id,
2774 committed_epoch: epoch,
2775 columns: std::collections::HashMap::new(),
2776 deleted: true,
2777 };
2778 self.memtable.upsert(tombstone);
2779 self.pk_by_row.remove(&row_id);
2780 if adjust_live_count {
2781 self.live_count = self.live_count.saturating_sub(1);
2782 }
2783 self.pending_delete_rids.insert(row_id.0 as u32);
2785 self.had_deletes = true;
2788 self.agg_cache = Arc::new(HashMap::new());
2789 }
2790
2791 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2795 let Some(pk_col) = self.schema.primary_key() else {
2796 return;
2797 };
2798 if self.pk_by_row_complete {
2801 if let Some(key) = self.pk_by_row.remove(&row_id) {
2802 if self.hot.get(&key) == Some(row_id) {
2803 self.hot.remove(&key);
2804 }
2805 }
2806 return;
2807 }
2808 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2827 if self.indexes_complete {
2828 let pk_val = self
2829 .memtable
2830 .get_version(row_id, lookup_epoch)
2831 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2832 .or_else(|| {
2833 self.mutable_run
2834 .get_version(row_id, lookup_epoch)
2835 .filter(|(_, r)| !r.deleted)
2836 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2837 })
2838 .or_else(|| {
2839 self.run_refs.iter().find_map(|rr| {
2840 let mut reader = self.open_reader(rr.run_id).ok()?;
2841 let (_, deleted, val) = reader
2842 .get_version_column(row_id, lookup_epoch, pk_col.id)
2843 .ok()??;
2844 if deleted {
2845 return None;
2846 }
2847 val
2848 })
2849 });
2850 if let Some(pk_val) = pk_val {
2851 let key = self.index_lookup_key(pk_col.id, &pk_val);
2852 if self.hot.get(&key) == Some(row_id) {
2853 self.hot.remove(&key);
2854 }
2855 return;
2856 }
2857 }
2858 self.refresh_pk_by_row_from_hot();
2863 if let Some(key) = self.pk_by_row.remove(&row_id) {
2864 if self.hot.get(&key) == Some(row_id) {
2865 self.hot.remove(&key);
2866 }
2867 }
2868 }
2869
2870 fn partition_pk_winners(
2875 &self,
2876 rows: &[Row],
2877 ) -> (
2878 std::collections::HashSet<RowId>,
2879 std::collections::HashMap<Vec<u8>, RowId>,
2880 ) {
2881 let mut losers = std::collections::HashSet::new();
2882 let Some(pk_col) = self.schema.primary_key() else {
2883 return (losers, std::collections::HashMap::new());
2884 };
2885 let pk_id = pk_col.id;
2886 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2887 std::collections::HashMap::new();
2888 for r in rows {
2889 let Some(pk_val) = r.columns.get(&pk_id) else {
2890 continue;
2891 };
2892 let key = self.index_lookup_key(pk_id, pk_val);
2893 if let Some(&old_rid) = winners.get(&key) {
2894 losers.insert(old_rid);
2895 }
2896 winners.insert(key, r.row_id);
2897 }
2898 (losers, winners)
2899 }
2900
2901 fn index_row(&mut self, row: &Row) {
2902 if row.deleted {
2903 return;
2904 }
2905 let any_predicate = self
2913 .schema
2914 .indexes
2915 .iter()
2916 .any(|idx| idx.predicate.is_some());
2917 if any_predicate {
2918 let columns_map: HashMap<u16, &Value> =
2919 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2920 let name_to_id: HashMap<&str, u16> = self
2921 .schema
2922 .columns
2923 .iter()
2924 .map(|c| (c.name.as_str(), c.id))
2925 .collect();
2926 for idx in &self.schema.indexes {
2927 if let Some(pred) = &idx.predicate {
2928 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2929 continue; }
2931 }
2932 index_into_single(
2934 idx,
2935 &self.schema,
2936 row,
2937 &mut self.hot,
2938 &mut self.bitmap,
2939 &mut self.ann,
2940 &mut self.fm,
2941 &mut self.sparse,
2942 &mut self.minhash,
2943 );
2944 }
2945 return;
2946 }
2947 if self.column_keys.is_empty() {
2951 index_into(
2952 &self.schema,
2953 row,
2954 &mut self.hot,
2955 &mut self.bitmap,
2956 &mut self.ann,
2957 &mut self.fm,
2958 &mut self.sparse,
2959 &mut self.minhash,
2960 );
2961 return;
2962 }
2963 let effective_row = self.tokenized_for_indexes(row);
2964 index_into(
2965 &self.schema,
2966 &effective_row,
2967 &mut self.hot,
2968 &mut self.bitmap,
2969 &mut self.ann,
2970 &mut self.fm,
2971 &mut self.sparse,
2972 &mut self.minhash,
2973 );
2974 }
2975
2976 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2982 if self.column_keys.is_empty() {
2983 return row.clone();
2984 }
2985 #[cfg(feature = "encryption")]
2986 {
2987 use crate::encryption::SCHEME_HMAC_EQ;
2988 let mut tok = row.clone();
2989 for (&cid, &(_, scheme)) in &self.column_keys {
2990 if scheme != SCHEME_HMAC_EQ {
2991 continue;
2992 }
2993 if let Some(v) = tok.columns.get(&cid).cloned() {
2994 if let Some(t) = self.tokenize_value(cid, &v) {
2995 tok.columns.insert(cid, t);
2996 }
2997 }
2998 }
2999 tok
3000 }
3001 #[cfg(not(feature = "encryption"))]
3002 {
3003 row.clone()
3004 }
3005 }
3006
3007 pub fn commit(&mut self) -> Result<Epoch> {
3012 self.ensure_writable()?;
3013 if self.is_shared() {
3014 self.commit_shared()
3015 } else {
3016 self.commit_private()
3017 }
3018 }
3019
3020 fn commit_private(&mut self) -> Result<Epoch> {
3022 let commit_lock = Arc::clone(&self.commit_lock);
3026 let _g = commit_lock.lock();
3027 let new_epoch = self.epoch.bump_assigned();
3028 let txn_id = self.current_txn_id;
3029 match &mut self.wal {
3033 WalSink::Private(w) => {
3034 w.append_txn(
3035 txn_id,
3036 Op::TxnCommit {
3037 epoch: new_epoch.0,
3038 added_runs: Vec::new(),
3039 },
3040 )?;
3041 w.sync()?;
3042 }
3043 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3044 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3045 }
3046 if let Some(epoch) = self.pending_truncate.take() {
3048 self.apply_truncate(epoch)?;
3049 }
3050 self.invalidate_pending_cache();
3051 self.persist_manifest(new_epoch)?;
3052 self.epoch.publish_in_order(new_epoch);
3056 self.current_txn_id += 1;
3057 self.data_generation = self.data_generation.wrapping_add(1);
3058 Ok(new_epoch)
3059 }
3060
3061 fn commit_shared(&mut self) -> Result<Epoch> {
3067 use std::sync::atomic::Ordering;
3068 let s = match &self.wal {
3069 WalSink::Shared(s) => s.clone(),
3070 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3071 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3072 };
3073 if s.poisoned.load(Ordering::Relaxed) {
3074 return Err(MongrelError::Other(
3075 "database poisoned by fsync error".into(),
3076 ));
3077 }
3078 let commit_lock = Arc::clone(&self.commit_lock);
3085 let _g = commit_lock.lock();
3086 let txn_id = self.ensure_txn_id();
3089 let (new_epoch, commit_seq) = {
3090 let mut wal = s.wal.lock();
3091 let new_epoch = self.epoch.bump_assigned();
3092 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
3093 (new_epoch, seq)
3094 };
3095 s.group
3096 .await_durable(&s.wal, commit_seq)
3097 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
3098
3099 if self.pending_truncate.take().is_some() {
3102 self.apply_truncate(new_epoch)?;
3103 }
3104 let mut rows = std::mem::take(&mut self.pending_rows);
3105 if !rows.is_empty() {
3106 for r in &mut rows {
3107 r.committed_epoch = new_epoch;
3108 }
3109 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3110 let all_auto_generated =
3111 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3112 self.apply_put_rows_inner(rows, !all_auto_generated)?;
3113 } else {
3114 self.pending_rows_auto_inc.clear();
3115 }
3116 let dels = std::mem::take(&mut self.pending_dels);
3117 for rid in dels {
3118 self.apply_delete(rid, new_epoch);
3119 }
3120
3121 self.invalidate_pending_cache();
3122 self.persist_manifest(new_epoch)?;
3123 self.epoch.publish_in_order(new_epoch);
3124 let _ = s.change_wake.send(());
3125 self.current_txn_id = 0;
3127 self.data_generation = self.data_generation.wrapping_add(1);
3128 Ok(new_epoch)
3129 }
3130
3131 pub fn flush(&mut self) -> Result<Epoch> {
3139 self.ensure_indexes_complete()?;
3140 let epoch = self.commit()?;
3141 let rows = self.memtable.drain_sorted();
3142 if !rows.is_empty() {
3143 self.mutable_run.insert_many(rows);
3144 }
3145 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3146 self.spill_mutable_run(epoch)?;
3147 self.mark_flushed(epoch)?;
3151 self.persist_manifest(epoch)?;
3152 self.build_learned_ranges()?;
3153 self.checkpoint_indexes(epoch);
3156 }
3157 Ok(epoch)
3160 }
3161
3162 pub fn force_flush(&mut self) -> Result<Epoch> {
3171 let saved = self.mutable_run_spill_bytes;
3172 self.mutable_run_spill_bytes = 1;
3173 let result = self.flush();
3174 self.mutable_run_spill_bytes = saved;
3175 result
3176 }
3177
3178 pub fn close(&mut self) -> Result<()> {
3185 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3186 self.force_flush()?;
3187 }
3188 Ok(())
3189 }
3190
3191 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3198 let op = Op::Flush {
3199 table_id: self.table_id,
3200 flushed_epoch: epoch.0,
3201 };
3202 match &mut self.wal {
3203 WalSink::Private(w) => {
3204 w.append_system(op)?;
3205 w.sync()?;
3206 }
3207 WalSink::Shared(s) => {
3208 s.wal.lock().append_system(op)?;
3213 }
3214 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3215 }
3216 self.flushed_epoch = epoch.0;
3217 if matches!(self.wal, WalSink::Private(_)) {
3218 self.rotate_wal(epoch)?;
3219 }
3220 Ok(())
3221 }
3222
3223 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3227 let rows = self.mutable_run.drain_sorted();
3228 if rows.is_empty() {
3229 return Ok(());
3230 }
3231 let run_id = self.next_run_id;
3232 self.next_run_id += 1;
3233 let path = self.run_path(run_id);
3234 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3235 if let Some(kek) = &self.kek {
3236 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3237 }
3238 let header = writer.write(&path, &rows)?;
3239 self.run_refs.push(RunRef {
3240 run_id: run_id as u128,
3241 level: 0,
3242 epoch_created: epoch.0,
3243 row_count: header.row_count,
3244 });
3245 Ok(())
3246 }
3247
3248 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3252 self.mutable_run_spill_bytes = bytes.max(1);
3253 }
3254
3255 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3259 self.compaction_zstd_level = level;
3260 }
3261
3262 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3266 self.result_cache.lock().set_max_bytes(max_bytes);
3267 }
3268
3269 pub(crate) fn clear_result_cache(&mut self) {
3273 self.result_cache.lock().clear();
3274 }
3275
3276 pub fn mutable_run_len(&self) -> usize {
3278 self.mutable_run.len()
3279 }
3280
3281 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3284 self.mutable_run.drain_sorted()
3285 }
3286
3287 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3292 let epoch = self.commit()?;
3293 let n = batch.len();
3294 if n == 0 {
3295 return Ok(epoch);
3296 }
3297 for row in &batch {
3298 self.schema.validate_values(row)?;
3299 }
3300 let live_before = self.live_count;
3301 self.spill_mutable_run(epoch)?;
3305 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3306 && self.indexes_complete
3307 && self.run_refs.is_empty()
3308 && self.memtable.is_empty()
3309 && self.mutable_run.is_empty();
3310 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3316 use rayon::prelude::*;
3317 self.schema
3318 .columns
3319 .par_iter()
3320 .map(|cdef| {
3321 (
3322 cdef.id,
3323 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3324 )
3325 })
3326 .collect::<Vec<_>>()
3327 };
3328 drop(batch);
3329 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3334 self.validate_columns_not_null(&user_columns, n)?;
3335 let winner_idx = self
3336 .bulk_pk_winner_indices(&user_columns, n)
3337 .filter(|idx| idx.len() != n);
3338 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3339 match winner_idx.as_deref() {
3340 Some(idx) => {
3341 let compacted = user_columns
3342 .iter()
3343 .map(|(id, c)| (*id, c.gather(idx)))
3344 .collect();
3345 (compacted, idx.len())
3346 }
3347 None => (std::mem::take(&mut user_columns), n),
3348 };
3349 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3350 let first = self.allocator.alloc_range(write_n as u64).0;
3351 for rid in first..first + write_n as u64 {
3352 self.reservoir.offer(rid);
3353 }
3354 let run_id = self.next_run_id;
3355 self.next_run_id += 1;
3356 let path = self.run_path(run_id);
3357 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3358 .clean(true)
3359 .with_lz4()
3360 .with_native_endian();
3361 if let Some(kek) = &self.kek {
3362 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3363 }
3364 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3365 self.run_refs.push(RunRef {
3366 run_id: run_id as u128,
3367 level: 0,
3368 epoch_created: epoch.0,
3369 row_count: header.row_count,
3370 });
3371 self.live_count = self.live_count.saturating_add(write_n as u64);
3372 if eager_index_build {
3373 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3374 self.index_columns_bulk(&write_columns, &row_ids);
3375 self.indexes_complete = true;
3376 self.build_learned_ranges()?;
3377 } else {
3378 self.indexes_complete = false;
3379 }
3380 self.mark_flushed(epoch)?;
3381 self.persist_manifest(epoch)?;
3382 if eager_index_build {
3383 self.checkpoint_indexes(epoch);
3384 }
3385 self.clear_result_cache();
3386 Ok(epoch)
3387 }
3388
3389 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3392 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3393 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3394 let segment_no = segment
3397 .file_stem()
3398 .and_then(|s| s.to_str())
3399 .and_then(|s| s.strip_prefix("seg-"))
3400 .and_then(|s| s.parse::<u64>().ok())
3401 .unwrap_or(0);
3402 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3403 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3404 wal.sync()?;
3405 self.wal = WalSink::Private(wal);
3406 Ok(())
3407 }
3408
3409 pub(crate) fn invalidate_pending_cache(&mut self) {
3414 self.result_cache
3415 .lock()
3416 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3417 self.pending_delete_rids.clear();
3418 self.pending_put_cols.clear();
3419 }
3420
3421 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3422 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3423 m.current_epoch = epoch.0;
3424 m.next_row_id = self.allocator.current().0;
3425 m.runs = self.run_refs.clone();
3426 m.live_count = self.live_count;
3427 m.global_idx_epoch = self.global_idx_epoch;
3428 m.flushed_epoch = self.flushed_epoch;
3429 m.retiring = self.retiring.clone();
3430 m.auto_inc_next = match self.auto_inc {
3434 Some(ai) if ai.seeded => ai.next,
3435 _ => 0,
3436 };
3437 m.ttl = self.ttl;
3438 let meta_dek = self.manifest_meta_dek();
3439 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3440 Ok(())
3441 }
3442
3443 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3449 if !self.indexes_complete {
3452 return;
3453 }
3454 let snap = global_idx::IndexSnapshot {
3455 hot: &self.hot,
3456 bitmap: &self.bitmap,
3457 ann: &self.ann,
3458 fm: &self.fm,
3459 sparse: &self.sparse,
3460 minhash: &self.minhash,
3461 learned_range: &self.learned_range,
3462 };
3463 let idx_dek = self.idx_dek();
3465 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3466 .is_ok()
3467 {
3468 self.global_idx_epoch = epoch.0;
3469 let _ = self.persist_manifest(epoch);
3470 }
3471 }
3472
3473 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3476 self.global_idx_epoch = 0;
3477 global_idx::remove(&self.dir);
3478 let _ = self.persist_manifest(self.epoch.visible());
3479 }
3480
3481 pub(crate) fn mark_indexes_incomplete(&mut self) {
3482 self.indexes_complete = false;
3483 self.invalidate_index_checkpoint();
3484 }
3485
3486 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3489 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3490 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3491 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3492 best = Some((epoch, row));
3493 }
3494 }
3495 for rr in &self.run_refs {
3496 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3497 continue;
3498 };
3499 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3500 continue;
3501 };
3502 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3503 best = Some((epoch, row));
3504 }
3505 }
3506 let now_nanos = unix_nanos_now();
3507 match best {
3508 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3509 Some((_, r)) => Some(r),
3510 None => None,
3511 }
3512 }
3513
3514 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3518 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3519 let mut fold = |row: Row| {
3520 best.entry(row.row_id.0)
3521 .and_modify(|e| {
3522 if row.committed_epoch > e.0 {
3523 *e = (row.committed_epoch, row.clone());
3524 }
3525 })
3526 .or_insert_with(|| (row.committed_epoch, row));
3527 };
3528 for row in self.memtable.visible_versions(snapshot.epoch) {
3529 fold(row);
3530 }
3531 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3532 fold(row);
3533 }
3534 for rr in &self.run_refs {
3535 let mut reader = self.open_reader(rr.run_id)?;
3536 for row in reader.visible_versions(snapshot.epoch)? {
3537 fold(row);
3538 }
3539 }
3540 let now_nanos = unix_nanos_now();
3541 let mut out: Vec<Row> = best
3542 .into_values()
3543 .filter_map(|(_, r)| {
3544 if r.deleted || self.row_expired_at(&r, now_nanos) {
3545 None
3546 } else {
3547 Some(r)
3548 }
3549 })
3550 .collect();
3551 out.sort_by_key(|r| r.row_id);
3552 Ok(out)
3553 }
3554
3555 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3562 if self.ttl.is_none()
3563 && self.memtable.is_empty()
3564 && self.mutable_run.is_empty()
3565 && self.run_refs.len() == 1
3566 {
3567 let rr = self.run_refs[0].clone();
3568 let mut reader = self.open_reader(rr.run_id)?;
3569 let idxs = reader.visible_indices(snapshot.epoch)?;
3570 let mut cols = Vec::with_capacity(self.schema.columns.len());
3571 for cdef in &self.schema.columns {
3572 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3573 }
3574 return Ok(cols);
3575 }
3576 let rows = self.visible_rows(snapshot)?;
3578 let mut cols: Vec<(u16, Vec<Value>)> = self
3579 .schema
3580 .columns
3581 .iter()
3582 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3583 .collect();
3584 for r in &rows {
3585 for (cid, vec) in cols.iter_mut() {
3586 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3587 }
3588 }
3589 Ok(cols)
3590 }
3591
3592 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3594 let row_id = self.hot.get(key)?;
3595 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3596 Some(row_id)
3597 } else {
3598 None
3599 }
3600 }
3601
3602 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3607 self.query_at_with_allowed(q, self.snapshot(), None)
3608 }
3609
3610 pub fn query_at_with_allowed(
3613 &mut self,
3614 q: &crate::query::Query,
3615 snapshot: Snapshot,
3616 allowed: Option<&std::collections::HashSet<RowId>>,
3617 ) -> Result<Vec<Row>> {
3618 self.query_at_with_allowed_after(q, snapshot, allowed, None)
3619 }
3620
3621 #[doc(hidden)]
3622 pub fn query_at_with_allowed_after(
3623 &mut self,
3624 q: &crate::query::Query,
3625 snapshot: Snapshot,
3626 allowed: Option<&std::collections::HashSet<RowId>>,
3627 after_row_id: Option<RowId>,
3628 ) -> Result<Vec<Row>> {
3629 self.require_select()?;
3630 self.ensure_indexes_complete()?;
3631 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3632 return Err(MongrelError::InvalidArgument(format!(
3633 "query exceeds {} conditions",
3634 crate::query::MAX_HARD_CONDITIONS
3635 )));
3636 }
3637 if let Some(limit) = q.limit {
3638 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
3639 return Err(MongrelError::InvalidArgument(format!(
3640 "query limit must be between 1 and {}",
3641 crate::query::MAX_FINAL_LIMIT
3642 )));
3643 }
3644 }
3645 if q.offset > crate::query::MAX_QUERY_OFFSET {
3646 return Err(MongrelError::InvalidArgument(format!(
3647 "query offset exceeds {}",
3648 crate::query::MAX_QUERY_OFFSET
3649 )));
3650 }
3651 self.query_conditions_at(
3652 &q.conditions,
3653 snapshot,
3654 allowed,
3655 q.limit,
3656 q.offset,
3657 after_row_id,
3658 )
3659 }
3660
3661 #[doc(hidden)]
3664 pub fn query_all_at(
3665 &mut self,
3666 conditions: &[crate::query::Condition],
3667 snapshot: Snapshot,
3668 ) -> Result<Vec<Row>> {
3669 self.require_select()?;
3670 self.ensure_indexes_complete()?;
3671 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3672 return Err(MongrelError::InvalidArgument(format!(
3673 "query exceeds {} conditions",
3674 crate::query::MAX_HARD_CONDITIONS
3675 )));
3676 }
3677 self.query_conditions_at(conditions, snapshot, None, None, 0, None)
3678 }
3679
3680 fn query_conditions_at(
3681 &self,
3682 conditions: &[crate::query::Condition],
3683 snapshot: Snapshot,
3684 allowed: Option<&std::collections::HashSet<RowId>>,
3685 limit: Option<usize>,
3686 offset: usize,
3687 after_row_id: Option<RowId>,
3688 ) -> Result<Vec<Row>> {
3689 crate::trace::QueryTrace::record(|t| {
3690 t.run_count = self.run_refs.len();
3691 t.memtable_rows = self.memtable.len();
3692 t.mutable_run_rows = self.mutable_run.len();
3693 });
3694 if conditions.is_empty() {
3698 crate::trace::QueryTrace::record(|t| {
3699 t.scan_mode = crate::trace::ScanMode::Materialized;
3700 t.row_materialized = true;
3701 });
3702 let mut rows = self.visible_rows(snapshot)?;
3703 if let Some(allowed) = allowed {
3704 rows.retain(|row| allowed.contains(&row.row_id));
3705 }
3706 if let Some(after_row_id) = after_row_id {
3707 rows.retain(|row| row.row_id > after_row_id);
3708 }
3709 rows.drain(..offset.min(rows.len()));
3710 if let Some(limit) = limit {
3711 rows.truncate(limit);
3712 }
3713 return Ok(rows);
3714 }
3715 crate::trace::QueryTrace::record(|t| {
3716 t.conditions_pushed = conditions.len();
3717 t.scan_mode = crate::trace::ScanMode::Materialized;
3718 t.row_materialized = true;
3719 });
3720 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
3727 ordered.sort_by_key(|c| condition_cost_rank(c));
3728 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3729 for c in &ordered {
3730 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
3731 let empty = s.is_empty();
3732 sets.push(s);
3733 if empty {
3734 break;
3735 }
3736 }
3737 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3738 if let Some(allowed) = allowed {
3739 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
3740 }
3741 if let Some(after_row_id) = after_row_id {
3742 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
3743 rids.drain(..first);
3744 }
3745 rids.drain(..offset.min(rids.len()));
3746 if let Some(limit) = limit {
3747 rids.truncate(limit);
3748 }
3749 self.rows_for_rids(&rids, snapshot)
3750 }
3751
3752 pub fn retrieve(
3754 &mut self,
3755 retriever: &crate::query::Retriever,
3756 ) -> Result<Vec<crate::query::RetrieverHit>> {
3757 self.retrieve_with_allowed(retriever, None)
3758 }
3759
3760 pub fn retrieve_at(
3761 &mut self,
3762 retriever: &crate::query::Retriever,
3763 snapshot: Snapshot,
3764 allowed: Option<&std::collections::HashSet<RowId>>,
3765 ) -> Result<Vec<crate::query::RetrieverHit>> {
3766 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
3767 }
3768
3769 pub fn retrieve_with_allowed(
3772 &mut self,
3773 retriever: &crate::query::Retriever,
3774 allowed: Option<&std::collections::HashSet<RowId>>,
3775 ) -> Result<Vec<crate::query::RetrieverHit>> {
3776 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
3777 }
3778
3779 pub fn retrieve_at_with_allowed(
3780 &mut self,
3781 retriever: &crate::query::Retriever,
3782 snapshot: Snapshot,
3783 allowed: Option<&std::collections::HashSet<RowId>>,
3784 ) -> Result<Vec<crate::query::RetrieverHit>> {
3785 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
3786 }
3787
3788 pub fn retrieve_at_with_allowed_and_context(
3789 &mut self,
3790 retriever: &crate::query::Retriever,
3791 snapshot: Snapshot,
3792 allowed: Option<&std::collections::HashSet<RowId>>,
3793 context: Option<&crate::query::AiExecutionContext>,
3794 ) -> Result<Vec<crate::query::RetrieverHit>> {
3795 self.require_select()?;
3796 self.ensure_indexes_complete()?;
3797 self.validate_retriever(retriever)?;
3798 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
3799 }
3800
3801 pub fn retrieve_at_with_candidate_authorization_and_context(
3802 &mut self,
3803 retriever: &crate::query::Retriever,
3804 snapshot: Snapshot,
3805 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3806 context: Option<&crate::query::AiExecutionContext>,
3807 ) -> Result<Vec<crate::query::RetrieverHit>> {
3808 self.require_select()?;
3809 self.ensure_indexes_complete()?;
3810 self.retrieve_at_with_candidate_authorization_on_generation(
3811 retriever,
3812 snapshot,
3813 authorization,
3814 context,
3815 )
3816 }
3817
3818 #[doc(hidden)]
3819 pub fn retrieve_at_with_candidate_authorization_on_generation(
3820 &self,
3821 retriever: &crate::query::Retriever,
3822 snapshot: Snapshot,
3823 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3824 context: Option<&crate::query::AiExecutionContext>,
3825 ) -> Result<Vec<crate::query::RetrieverHit>> {
3826 self.require_select()?;
3827 self.validate_retriever(retriever)?;
3828 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
3829 }
3830
3831 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3832 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
3833 let (column_id, k) = match retriever {
3834 Retriever::Ann {
3835 column_id,
3836 query,
3837 k,
3838 } => {
3839 let index = self.ann.get(column_id).ok_or_else(|| {
3840 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3841 })?;
3842 if query.len() != index.dim() {
3843 return Err(MongrelError::InvalidArgument(format!(
3844 "ANN query dimension must be {}, got {}",
3845 index.dim(),
3846 query.len()
3847 )));
3848 }
3849 if query.iter().any(|value| !value.is_finite()) {
3850 return Err(MongrelError::InvalidArgument(
3851 "ANN query values must be finite".into(),
3852 ));
3853 }
3854 (*column_id, *k)
3855 }
3856 Retriever::Sparse {
3857 column_id,
3858 query,
3859 k,
3860 } => {
3861 if !self.sparse.contains_key(column_id) {
3862 return Err(MongrelError::InvalidArgument(format!(
3863 "column {column_id} has no Sparse index"
3864 )));
3865 }
3866 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3867 return Err(MongrelError::InvalidArgument(
3868 "Sparse query must be non-empty with finite weights".into(),
3869 ));
3870 }
3871 if query.len() > MAX_SPARSE_TERMS {
3872 return Err(MongrelError::InvalidArgument(format!(
3873 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
3874 )));
3875 }
3876 (*column_id, *k)
3877 }
3878 Retriever::MinHash {
3879 column_id,
3880 members,
3881 k,
3882 } => {
3883 if !self.minhash.contains_key(column_id) {
3884 return Err(MongrelError::InvalidArgument(format!(
3885 "column {column_id} has no MinHash index"
3886 )));
3887 }
3888 if members.is_empty() {
3889 return Err(MongrelError::InvalidArgument(
3890 "MinHash members must not be empty".into(),
3891 ));
3892 }
3893 if members.len() > MAX_SET_MEMBERS {
3894 return Err(MongrelError::InvalidArgument(format!(
3895 "MinHash query exceeds {MAX_SET_MEMBERS} members"
3896 )));
3897 }
3898 let mut total_bytes = 0usize;
3899 for member in members {
3900 let bytes = member.encoded_len();
3901 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
3902 return Err(MongrelError::InvalidArgument(format!(
3903 "MinHash member exceeds {} bytes",
3904 crate::query::MAX_SET_MEMBER_BYTES
3905 )));
3906 }
3907 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
3908 MongrelError::InvalidArgument("MinHash input size overflow".into())
3909 })?;
3910 }
3911 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
3912 return Err(MongrelError::InvalidArgument(format!(
3913 "MinHash input exceeds {} bytes",
3914 crate::query::MAX_SET_INPUT_BYTES
3915 )));
3916 }
3917 (*column_id, *k)
3918 }
3919 };
3920 if k == 0 {
3921 return Err(MongrelError::InvalidArgument(
3922 "retriever k must be > 0".into(),
3923 ));
3924 }
3925 if k > MAX_RETRIEVER_K {
3926 return Err(MongrelError::InvalidArgument(format!(
3927 "retriever k exceeds {MAX_RETRIEVER_K}"
3928 )));
3929 }
3930 debug_assert!(self
3931 .schema
3932 .columns
3933 .iter()
3934 .any(|column| column.id == column_id));
3935 Ok(())
3936 }
3937
3938 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3939 use crate::query::Condition;
3940 match condition {
3941 Condition::Ann {
3942 column_id,
3943 query,
3944 k,
3945 } => self.validate_retriever(&crate::query::Retriever::Ann {
3946 column_id: *column_id,
3947 query: query.clone(),
3948 k: *k,
3949 }),
3950 Condition::SparseMatch {
3951 column_id,
3952 query,
3953 k,
3954 } => self.validate_retriever(&crate::query::Retriever::Sparse {
3955 column_id: *column_id,
3956 query: query.clone(),
3957 k: *k,
3958 }),
3959 Condition::MinHashSimilar {
3960 column_id,
3961 query,
3962 k,
3963 } => {
3964 if !self.minhash.contains_key(column_id) {
3965 return Err(MongrelError::InvalidArgument(format!(
3966 "column {column_id} has no MinHash index"
3967 )));
3968 }
3969 if query.is_empty() || *k == 0 {
3970 return Err(MongrelError::InvalidArgument(
3971 "MinHash query must be non-empty and k must be > 0".into(),
3972 ));
3973 }
3974 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
3975 {
3976 return Err(MongrelError::InvalidArgument(format!(
3977 "MinHash query must have <= {} members and k <= {}",
3978 crate::query::MAX_SET_MEMBERS,
3979 crate::query::MAX_RETRIEVER_K
3980 )));
3981 }
3982 Ok(())
3983 }
3984 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
3985 Err(MongrelError::InvalidArgument(format!(
3986 "bitmap IN exceeds {} values",
3987 crate::query::MAX_SET_MEMBERS
3988 )))
3989 }
3990 Condition::FmContainsAll { patterns, .. }
3991 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
3992 {
3993 Err(MongrelError::InvalidArgument(format!(
3994 "FM query exceeds {} patterns",
3995 crate::query::MAX_HARD_CONDITIONS
3996 )))
3997 }
3998 _ => Ok(()),
3999 }
4000 }
4001
4002 fn retrieve_filtered(
4003 &self,
4004 retriever: &crate::query::Retriever,
4005 snapshot: Snapshot,
4006 hard_filter: Option<&RowIdSet>,
4007 allowed: Option<&std::collections::HashSet<RowId>>,
4008 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4009 context: Option<&crate::query::AiExecutionContext>,
4010 ) -> Result<Vec<crate::query::RetrieverHit>> {
4011 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
4012 let started = std::time::Instant::now();
4013 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
4014 Retriever::Ann {
4015 column_id,
4016 query,
4017 k,
4018 } => {
4019 let Some(index) = self.ann.get(column_id) else {
4020 return Ok(Vec::new());
4021 };
4022 let cap = index.len();
4023 if cap == 0 {
4024 return Ok(Vec::new());
4025 }
4026 let mut breadth = (*k).max(1).min(cap);
4027 let mut eligibility = std::collections::HashMap::new();
4028 let mut filtered = loop {
4029 let mut seen = std::collections::HashSet::new();
4030 if let Some(context) = context {
4031 context.checkpoint()?;
4032 }
4033 let raw = index.search_with_context(query, breadth, context)?;
4034 let unchecked: Vec<_> = raw
4035 .iter()
4036 .map(|(row_id, _)| *row_id)
4037 .filter(|row_id| !eligibility.contains_key(row_id))
4038 .filter(|row_id| {
4039 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4040 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4041 })
4042 .collect();
4043 let eligible = self.eligible_and_authorized_candidate_ids(
4044 &unchecked,
4045 *column_id,
4046 snapshot,
4047 candidate_authorization,
4048 context,
4049 )?;
4050 for row_id in unchecked {
4051 eligibility.insert(row_id, eligible.contains(&row_id));
4052 }
4053 let filtered: Vec<_> = raw
4054 .into_iter()
4055 .filter(|(row_id, _)| {
4056 seen.insert(*row_id)
4057 && eligibility.get(row_id).copied().unwrap_or(false)
4058 })
4059 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
4060 .collect();
4061 if filtered.len() >= *k || breadth >= cap {
4062 break filtered;
4063 }
4064 breadth = breadth.saturating_mul(2).min(cap);
4065 };
4066 filtered.truncate(*k);
4067 filtered
4068 }
4069 Retriever::Sparse {
4070 column_id,
4071 query,
4072 k,
4073 } => self
4074 .sparse
4075 .get(column_id)
4076 .map(|index| -> Result<Vec<_>> {
4077 let mut breadth = (*k).max(1);
4078 let mut eligibility = std::collections::HashMap::new();
4079 loop {
4080 if let Some(context) = context {
4081 context.checkpoint()?;
4082 }
4083 let raw = index.search_with_context(query, breadth, context)?;
4084 let unchecked: Vec<_> = raw
4085 .iter()
4086 .map(|(row_id, _)| *row_id)
4087 .filter(|row_id| !eligibility.contains_key(row_id))
4088 .filter(|row_id| {
4089 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4090 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4091 })
4092 .collect();
4093 let eligible = self.eligible_and_authorized_candidate_ids(
4094 &unchecked,
4095 *column_id,
4096 snapshot,
4097 candidate_authorization,
4098 context,
4099 )?;
4100 for row_id in unchecked {
4101 eligibility.insert(row_id, eligible.contains(&row_id));
4102 }
4103 let filtered: Vec<_> = raw
4104 .iter()
4105 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4106 .take(*k)
4107 .map(|(row_id, score)| {
4108 (*row_id, RetrieverScore::SparseDotProduct(*score))
4109 })
4110 .collect();
4111 if filtered.len() >= *k || raw.len() < breadth {
4112 break Ok(filtered);
4113 }
4114 let next = breadth.saturating_mul(2);
4115 if next == breadth {
4116 break Ok(filtered);
4117 }
4118 breadth = next;
4119 }
4120 })
4121 .transpose()?
4122 .unwrap_or_default(),
4123 Retriever::MinHash {
4124 column_id,
4125 members,
4126 k,
4127 } => self
4128 .minhash
4129 .get(column_id)
4130 .map(|index| -> Result<Vec<_>> {
4131 let mut hashes = Vec::with_capacity(members.len());
4132 for member in members {
4133 if let Some(context) = context {
4134 context.consume(crate::query::work_units(
4135 member.encoded_len(),
4136 crate::query::PARSE_WORK_QUANTUM,
4137 ))?;
4138 }
4139 hashes.push(member.hash_v1());
4140 }
4141 let mut breadth = (*k).max(1);
4142 let mut eligibility = std::collections::HashMap::new();
4143 loop {
4144 if let Some(context) = context {
4145 context.checkpoint()?;
4146 }
4147 let raw = index.search_with_context(&hashes, breadth, context)?;
4148 let unchecked: Vec<_> = raw
4149 .iter()
4150 .map(|(row_id, _)| *row_id)
4151 .filter(|row_id| !eligibility.contains_key(row_id))
4152 .filter(|row_id| {
4153 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4154 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4155 })
4156 .collect();
4157 let eligible = self.eligible_and_authorized_candidate_ids(
4158 &unchecked,
4159 *column_id,
4160 snapshot,
4161 candidate_authorization,
4162 context,
4163 )?;
4164 for row_id in unchecked {
4165 eligibility.insert(row_id, eligible.contains(&row_id));
4166 }
4167 let filtered: Vec<_> = raw
4168 .iter()
4169 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4170 .take(*k)
4171 .map(|(row_id, score)| {
4172 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
4173 })
4174 .collect();
4175 if filtered.len() >= *k || raw.len() < breadth {
4176 break Ok(filtered);
4177 }
4178 let next = breadth.saturating_mul(2);
4179 if next == breadth {
4180 break Ok(filtered);
4181 }
4182 breadth = next;
4183 }
4184 })
4185 .transpose()?
4186 .unwrap_or_default(),
4187 };
4188 let elapsed = started.elapsed().as_nanos() as u64;
4189 crate::trace::QueryTrace::record(|trace| {
4190 match retriever {
4191 Retriever::Ann { .. } => {
4192 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
4193 }
4194 Retriever::Sparse { .. } => {
4195 trace.sparse_candidate_nanos =
4196 trace.sparse_candidate_nanos.saturating_add(elapsed)
4197 }
4198 Retriever::MinHash { .. } => {
4199 trace.minhash_candidate_nanos =
4200 trace.minhash_candidate_nanos.saturating_add(elapsed)
4201 }
4202 }
4203 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
4204 });
4205 Ok(scored
4206 .into_iter()
4207 .enumerate()
4208 .map(|(rank, (row_id, score))| RetrieverHit {
4209 row_id,
4210 rank: rank + 1,
4211 score,
4212 })
4213 .collect())
4214 }
4215
4216 fn eligible_candidate_ids(
4217 &self,
4218 candidates: &[RowId],
4219 _column_id: u16,
4220 snapshot: Snapshot,
4221 context: Option<&crate::query::AiExecutionContext>,
4222 ) -> Result<std::collections::HashSet<RowId>> {
4223 if !self.had_deletes
4224 && self.ttl.is_none()
4225 && self.pending_put_cols.is_empty()
4226 && snapshot.epoch == self.snapshot().epoch
4227 {
4228 return Ok(candidates.iter().copied().collect());
4229 }
4230 let mut readers: Vec<_> = self
4231 .run_refs
4232 .iter()
4233 .map(|run| self.open_reader(run.run_id))
4234 .collect::<Result<_>>()?;
4235 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4236 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
4237 for &row_id in candidates {
4238 if let Some(context) = context {
4239 context.consume(1)?;
4240 }
4241 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4242 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4243 let overlay = match (mem, mutable) {
4244 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
4245 (Some(value), None) | (None, Some(value)) => Some(value),
4246 (None, None) => None,
4247 };
4248 if let Some((_, row)) = overlay {
4249 if !row.deleted && !self.row_expired_at(&row, now) {
4250 eligible.insert(row_id);
4251 }
4252 continue;
4253 }
4254 let mut best: Option<(Epoch, bool, usize)> = None;
4255 for (index, reader) in readers.iter_mut().enumerate() {
4256 if let Some((epoch, deleted)) =
4257 reader.get_version_visibility(row_id, snapshot.epoch)?
4258 {
4259 if best
4260 .as_ref()
4261 .map(|(best_epoch, ..)| epoch > *best_epoch)
4262 .unwrap_or(true)
4263 {
4264 best = Some((epoch, deleted, index));
4265 }
4266 }
4267 }
4268 let Some((_, false, reader_index)) = best else {
4269 continue;
4270 };
4271 if let Some(ttl) = self.ttl {
4272 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4273 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4274 {
4275 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4276 continue;
4277 }
4278 }
4279 }
4280 eligible.insert(row_id);
4281 }
4282 Ok(eligible)
4283 }
4284
4285 fn eligible_and_authorized_candidate_ids(
4286 &self,
4287 candidates: &[RowId],
4288 column_id: u16,
4289 snapshot: Snapshot,
4290 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4291 context: Option<&crate::query::AiExecutionContext>,
4292 ) -> Result<std::collections::HashSet<RowId>> {
4293 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
4294 let Some(authorization) = authorization else {
4295 return Ok(eligible);
4296 };
4297 let candidates: Vec<_> = eligible.into_iter().collect();
4298 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
4299 }
4300
4301 fn policy_allowed_candidate_ids(
4302 &self,
4303 candidates: &[RowId],
4304 snapshot: Snapshot,
4305 authorization: &crate::security::CandidateAuthorization<'_>,
4306 context: Option<&crate::query::AiExecutionContext>,
4307 ) -> Result<std::collections::HashSet<RowId>> {
4308 let started = std::time::Instant::now();
4309 if candidates.is_empty()
4310 || authorization.principal.is_admin
4311 || !authorization.security.rls_enabled(authorization.table)
4312 {
4313 return Ok(candidates.iter().copied().collect());
4314 }
4315 if let Some(context) = context {
4316 context.checkpoint()?;
4317 }
4318 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
4319 let mut rows: std::collections::HashMap<RowId, Row> = candidates
4320 .iter()
4321 .map(|row_id| {
4322 (
4323 *row_id,
4324 Row {
4325 row_id: *row_id,
4326 committed_epoch: snapshot.epoch,
4327 columns: std::collections::HashMap::new(),
4328 deleted: false,
4329 },
4330 )
4331 })
4332 .collect();
4333 let columns = authorization
4334 .security
4335 .select_policy_columns(authorization.table, authorization.principal);
4336 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4337 let mut decoded = 0usize;
4338 for column_id in &columns {
4339 if let Some(context) = context {
4340 context.checkpoint()?;
4341 }
4342 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4343 &row_ids, *column_id, snapshot, query_now, context,
4344 )? {
4345 if let Some(row) = rows.get_mut(&row_id) {
4346 row.columns.insert(*column_id, value);
4347 decoded = decoded.saturating_add(1);
4348 }
4349 }
4350 }
4351 if let Some(context) = context {
4352 context.consume(candidates.len().saturating_add(decoded))?;
4353 }
4354 let allowed = rows
4355 .into_values()
4356 .filter_map(|row| {
4357 authorization
4358 .security
4359 .row_allowed(
4360 authorization.table,
4361 crate::security::PolicyCommand::Select,
4362 &row,
4363 authorization.principal,
4364 false,
4365 )
4366 .then_some(row.row_id)
4367 })
4368 .collect();
4369 crate::trace::QueryTrace::record(|trace| {
4370 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
4371 trace.rls_policy_columns_decoded =
4372 trace.rls_policy_columns_decoded.saturating_add(decoded);
4373 trace.authorization_nanos = trace
4374 .authorization_nanos
4375 .saturating_add(started.elapsed().as_nanos() as u64);
4376 });
4377 Ok(allowed)
4378 }
4379
4380 pub fn search(
4382 &mut self,
4383 request: &crate::query::SearchRequest,
4384 ) -> Result<Vec<crate::query::SearchHit>> {
4385 self.search_with_allowed(request, None)
4386 }
4387
4388 pub fn search_at(
4389 &mut self,
4390 request: &crate::query::SearchRequest,
4391 snapshot: Snapshot,
4392 authorized: Option<&std::collections::HashSet<RowId>>,
4393 ) -> Result<Vec<crate::query::SearchHit>> {
4394 self.search_at_with_allowed(request, snapshot, authorized)
4395 }
4396
4397 pub fn search_with_allowed(
4398 &mut self,
4399 request: &crate::query::SearchRequest,
4400 authorized: Option<&std::collections::HashSet<RowId>>,
4401 ) -> Result<Vec<crate::query::SearchHit>> {
4402 self.search_at_with_allowed(request, self.snapshot(), authorized)
4403 }
4404
4405 pub fn search_at_with_allowed(
4406 &mut self,
4407 request: &crate::query::SearchRequest,
4408 snapshot: Snapshot,
4409 authorized: Option<&std::collections::HashSet<RowId>>,
4410 ) -> Result<Vec<crate::query::SearchHit>> {
4411 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
4412 }
4413
4414 pub fn search_at_with_allowed_and_context(
4415 &mut self,
4416 request: &crate::query::SearchRequest,
4417 snapshot: Snapshot,
4418 authorized: Option<&std::collections::HashSet<RowId>>,
4419 context: Option<&crate::query::AiExecutionContext>,
4420 ) -> Result<Vec<crate::query::SearchHit>> {
4421 self.ensure_indexes_complete()?;
4422 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
4423 }
4424
4425 pub fn search_at_with_candidate_authorization_and_context(
4426 &mut self,
4427 request: &crate::query::SearchRequest,
4428 snapshot: Snapshot,
4429 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4430 context: Option<&crate::query::AiExecutionContext>,
4431 ) -> Result<Vec<crate::query::SearchHit>> {
4432 self.ensure_indexes_complete()?;
4433 self.search_at_with_filters_and_context(
4434 request,
4435 snapshot,
4436 None,
4437 authorization,
4438 context,
4439 None,
4440 )
4441 }
4442
4443 #[doc(hidden)]
4444 pub fn search_at_with_candidate_authorization_on_generation(
4445 &self,
4446 request: &crate::query::SearchRequest,
4447 snapshot: Snapshot,
4448 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4449 context: Option<&crate::query::AiExecutionContext>,
4450 ) -> Result<Vec<crate::query::SearchHit>> {
4451 self.search_at_with_filters_and_context(
4452 request,
4453 snapshot,
4454 None,
4455 authorization,
4456 context,
4457 None,
4458 )
4459 }
4460
4461 #[doc(hidden)]
4462 pub fn search_at_with_candidate_authorization_on_generation_after(
4463 &self,
4464 request: &crate::query::SearchRequest,
4465 snapshot: Snapshot,
4466 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4467 context: Option<&crate::query::AiExecutionContext>,
4468 after: Option<crate::query::SearchAfter>,
4469 ) -> Result<Vec<crate::query::SearchHit>> {
4470 self.search_at_with_filters_and_context(
4471 request,
4472 snapshot,
4473 None,
4474 authorization,
4475 context,
4476 after,
4477 )
4478 }
4479
4480 fn search_at_with_filters_and_context(
4481 &self,
4482 request: &crate::query::SearchRequest,
4483 snapshot: Snapshot,
4484 authorized: Option<&std::collections::HashSet<RowId>>,
4485 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4486 context: Option<&crate::query::AiExecutionContext>,
4487 after: Option<crate::query::SearchAfter>,
4488 ) -> Result<Vec<crate::query::SearchHit>> {
4489 use crate::query::{
4490 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
4491 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
4492 };
4493 let total_started = std::time::Instant::now();
4494 self.require_select()?;
4495 if request.limit == 0 {
4496 return Err(MongrelError::InvalidArgument(
4497 "search limit must be > 0".into(),
4498 ));
4499 }
4500 if request.limit > MAX_FINAL_LIMIT {
4501 return Err(MongrelError::InvalidArgument(format!(
4502 "search limit exceeds {MAX_FINAL_LIMIT}"
4503 )));
4504 }
4505 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
4506 return Err(MongrelError::InvalidArgument(
4507 "search-after score must be finite".into(),
4508 ));
4509 }
4510 if request.retrievers.is_empty() {
4511 return Err(MongrelError::InvalidArgument(
4512 "search requires at least one retriever".into(),
4513 ));
4514 }
4515 if request.retrievers.len() > MAX_RETRIEVERS {
4516 return Err(MongrelError::InvalidArgument(format!(
4517 "search exceeds {MAX_RETRIEVERS} retrievers"
4518 )));
4519 }
4520 if request.must.len() > MAX_HARD_CONDITIONS {
4521 return Err(MongrelError::InvalidArgument(format!(
4522 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
4523 )));
4524 }
4525 for condition in &request.must {
4526 self.validate_condition(condition)?;
4527 }
4528 if request.must.iter().any(|condition| {
4529 matches!(
4530 condition,
4531 Condition::Ann { .. }
4532 | Condition::SparseMatch { .. }
4533 | Condition::MinHashSimilar { .. }
4534 )
4535 }) {
4536 return Err(MongrelError::InvalidArgument(
4537 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
4538 .into(),
4539 ));
4540 }
4541 let mut names = std::collections::HashSet::new();
4542 for named in &request.retrievers {
4543 if named.name.is_empty()
4544 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
4545 || !names.insert(named.name.as_str())
4546 {
4547 return Err(MongrelError::InvalidArgument(format!(
4548 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
4549 crate::query::MAX_RETRIEVER_NAME_BYTES
4550 )));
4551 }
4552 if !named.weight.is_finite()
4553 || named.weight < 0.0
4554 || named.weight > MAX_RETRIEVER_WEIGHT
4555 {
4556 return Err(MongrelError::InvalidArgument(format!(
4557 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4558 )));
4559 }
4560 self.validate_retriever(&named.retriever)?;
4561 }
4562 let projection = request
4563 .projection
4564 .clone()
4565 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
4566 if projection.len() > MAX_PROJECTION_COLUMNS {
4567 return Err(MongrelError::InvalidArgument(format!(
4568 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
4569 )));
4570 }
4571 for column_id in &projection {
4572 if !self
4573 .schema
4574 .columns
4575 .iter()
4576 .any(|column| column.id == *column_id)
4577 {
4578 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
4579 }
4580 }
4581 if let Some(crate::query::Rerank::ExactVector {
4582 embedding_column,
4583 query,
4584 candidate_limit,
4585 weight,
4586 ..
4587 }) = &request.rerank
4588 {
4589 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
4590 {
4591 return Err(MongrelError::InvalidArgument(format!(
4592 "rerank candidate_limit must be between search limit and {}",
4593 crate::query::MAX_RETRIEVER_K
4594 )));
4595 }
4596 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
4597 return Err(MongrelError::InvalidArgument(format!(
4598 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4599 )));
4600 }
4601 let column = self
4602 .schema
4603 .columns
4604 .iter()
4605 .find(|column| column.id == *embedding_column)
4606 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
4607 let crate::schema::TypeId::Embedding { dim } = column.ty else {
4608 return Err(MongrelError::InvalidArgument(format!(
4609 "rerank column {embedding_column} is not an embedding"
4610 )));
4611 };
4612 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
4613 return Err(MongrelError::InvalidArgument(format!(
4614 "rerank query must contain {dim} finite values"
4615 )));
4616 }
4617 }
4618
4619 let hard_filter_started = std::time::Instant::now();
4620 let hard_filter = if request.must.is_empty() {
4621 None
4622 } else {
4623 let mut sets = Vec::with_capacity(request.must.len());
4624 for condition in &request.must {
4625 if let Some(context) = context {
4626 context.checkpoint()?;
4627 }
4628 sets.push(self.resolve_condition(condition, snapshot)?);
4629 }
4630 Some(RowIdSet::intersect_many(sets))
4631 };
4632 crate::trace::QueryTrace::record(|trace| {
4633 trace.hard_filter_nanos = trace
4634 .hard_filter_nanos
4635 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
4636 });
4637 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
4638 return Ok(Vec::new());
4639 }
4640
4641 let constant = match request.fusion {
4642 Fusion::ReciprocalRank { constant } => constant,
4643 };
4644 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
4645 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
4646 let mut fusion_nanos = 0u64;
4647 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
4648 std::collections::HashMap::new();
4649 for named in retrievers {
4650 if named.weight == 0.0 {
4651 continue;
4652 }
4653 if let Some(context) = context {
4654 context.checkpoint()?;
4655 }
4656 let hits = self.retrieve_filtered(
4657 &named.retriever,
4658 snapshot,
4659 hard_filter.as_ref(),
4660 authorized,
4661 candidate_authorization,
4662 context,
4663 )?;
4664 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
4665 let fusion_started = std::time::Instant::now();
4666 for hit in hits {
4667 if let Some(context) = context {
4668 context.consume(1)?;
4669 }
4670 let contribution = named.weight / (constant as f64 + hit.rank as f64);
4671 if !contribution.is_finite() {
4672 return Err(MongrelError::InvalidArgument(
4673 "retriever contribution must be finite".into(),
4674 ));
4675 }
4676 let max_fused_candidates = context.map_or(
4677 crate::query::MAX_FUSED_CANDIDATES,
4678 crate::query::AiExecutionContext::max_fused_candidates,
4679 );
4680 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
4681 return Err(MongrelError::WorkBudgetExceeded);
4682 }
4683 let entry = fused.entry(hit.row_id).or_default();
4684 entry.0 += contribution;
4685 if !entry.0.is_finite() {
4686 return Err(MongrelError::InvalidArgument(
4687 "fused score must be finite".into(),
4688 ));
4689 }
4690 entry.1.push(ComponentScore {
4691 retriever_name: retriever_name.clone(),
4692 rank: hit.rank,
4693 raw_score: hit.score,
4694 contribution,
4695 });
4696 }
4697 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
4698 }
4699 let union_size = fused.len();
4700 let mut ranked: Vec<_> = fused
4701 .into_iter()
4702 .map(|(row_id, (fused_score, components))| {
4703 (row_id, fused_score, components, None, fused_score)
4704 })
4705 .collect();
4706 let order = |(a_row, _, _, _, a_score): &(
4707 RowId,
4708 f64,
4709 Vec<ComponentScore>,
4710 Option<f32>,
4711 f64,
4712 ),
4713 (b_row, _, _, _, b_score): &(
4714 RowId,
4715 f64,
4716 Vec<ComponentScore>,
4717 Option<f32>,
4718 f64,
4719 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
4720 if let Some(crate::query::Rerank::ExactVector {
4721 embedding_column,
4722 query,
4723 metric,
4724 candidate_limit,
4725 weight,
4726 }) = &request.rerank
4727 {
4728 let fused_order = |(a_row, a_score, ..): &(
4729 RowId,
4730 f64,
4731 Vec<ComponentScore>,
4732 Option<f32>,
4733 f64,
4734 ),
4735 (b_row, b_score, ..): &(
4736 RowId,
4737 f64,
4738 Vec<ComponentScore>,
4739 Option<f32>,
4740 f64,
4741 )| {
4742 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
4743 };
4744 let selection_started = std::time::Instant::now();
4745 if let Some(context) = context {
4746 context.consume(ranked.len())?;
4747 }
4748 if ranked.len() > *candidate_limit {
4749 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
4750 ranked.truncate(*candidate_limit);
4751 }
4752 ranked.sort_by(fused_order);
4753 fusion_nanos =
4754 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4755 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4756 if let Some(context) = context {
4757 context.consume(row_ids.len())?;
4758 }
4759 let query_now =
4760 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4761 let gather_started = std::time::Instant::now();
4762 let vectors = self.values_for_rids_batch_at_with_context(
4763 &row_ids,
4764 *embedding_column,
4765 snapshot,
4766 query_now,
4767 context,
4768 )?;
4769 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
4770 let vector_work =
4771 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
4772 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
4773 if let Some(context) = context {
4774 context.consume(vector_work)?;
4775 }
4776 query
4777 .iter()
4778 .map(|value| f64::from(*value).powi(2))
4779 .sum::<f64>()
4780 .sqrt()
4781 } else {
4782 0.0
4783 };
4784 let score_started = std::time::Instant::now();
4785 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
4786 for (row_id, value) in vectors {
4787 let Value::Embedding(vector) = value else {
4788 continue;
4789 };
4790 let score = match metric {
4791 crate::query::VectorMetric::DotProduct => {
4792 if let Some(context) = context {
4793 context.consume(vector_work)?;
4794 }
4795 query
4796 .iter()
4797 .zip(&vector)
4798 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4799 .sum::<f64>()
4800 }
4801 crate::query::VectorMetric::Cosine => {
4802 if let Some(context) = context {
4803 context.consume(vector_work.saturating_mul(2))?;
4804 }
4805 let dot = query
4806 .iter()
4807 .zip(&vector)
4808 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4809 .sum::<f64>();
4810 let norm = vector
4811 .iter()
4812 .map(|value| f64::from(*value).powi(2))
4813 .sum::<f64>()
4814 .sqrt();
4815 if query_norm == 0.0 || norm == 0.0 {
4816 0.0
4817 } else {
4818 dot / (query_norm * norm)
4819 }
4820 }
4821 crate::query::VectorMetric::Euclidean => {
4822 if let Some(context) = context {
4823 context.consume(vector_work)?;
4824 }
4825 query
4826 .iter()
4827 .zip(&vector)
4828 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
4829 .sum::<f64>()
4830 .sqrt()
4831 }
4832 };
4833 if !score.is_finite() {
4834 return Err(MongrelError::InvalidArgument(
4835 "exact rerank score must be finite".into(),
4836 ));
4837 }
4838 scores.insert(row_id, score as f32);
4839 }
4840 let mut reranked = Vec::with_capacity(ranked.len());
4841 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
4842 let Some(score) = scores.get(&row_id).copied() else {
4843 continue;
4844 };
4845 let ordering_score = match metric {
4846 crate::query::VectorMetric::Euclidean => -f64::from(score),
4847 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
4848 f64::from(score)
4849 }
4850 };
4851 let final_score = fused_score + *weight * ordering_score;
4852 if !final_score.is_finite() {
4853 return Err(MongrelError::InvalidArgument(
4854 "final rerank score must be finite".into(),
4855 ));
4856 }
4857 reranked.push((row_id, fused_score, components, Some(score), final_score));
4858 }
4859 ranked = reranked;
4860 ranked.sort_by(order);
4861 crate::trace::QueryTrace::record(|trace| {
4862 trace.exact_vector_gather_nanos =
4863 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
4864 trace.exact_vector_score_nanos = trace
4865 .exact_vector_score_nanos
4866 .saturating_add(score_started.elapsed().as_nanos() as u64);
4867 });
4868 }
4869 if let Some(after) = after {
4870 ranked.retain(|(row_id, _, _, _, final_score)| {
4871 final_score.total_cmp(&after.final_score).is_lt()
4872 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
4873 });
4874 }
4875 let projection_started = std::time::Instant::now();
4876 let sentinel = projection
4877 .first()
4878 .copied()
4879 .or_else(|| self.schema.columns.first().map(|column| column.id));
4880 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4881 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
4882 let mut projection_rows = 0usize;
4883 let mut projection_cells = 0usize;
4884 while out.len() < request.limit && !ranked.is_empty() {
4885 if let Some(context) = context {
4886 context.checkpoint()?;
4887 context.consume(ranked.len())?;
4888 }
4889 let needed = request.limit - out.len();
4890 let window_size = ranked
4891 .len()
4892 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
4893 let selection_started = std::time::Instant::now();
4894 let mut remainder = if ranked.len() > window_size {
4895 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
4896 ranked.split_off(window_size)
4897 } else {
4898 Vec::new()
4899 };
4900 ranked.sort_by(order);
4901 fusion_nanos =
4902 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4903 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4904 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
4905 if let Some(context) = context {
4906 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
4907 }
4908 projection_rows = projection_rows.saturating_add(row_ids.len());
4909 projection_cells =
4910 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
4911 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
4912 std::collections::HashMap::new();
4913 if let Some(column_id) = sentinel {
4914 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4915 &row_ids, column_id, snapshot, query_now, context,
4916 )? {
4917 cells.entry(row_id).or_default().insert(column_id, value);
4918 }
4919 }
4920 for &column_id in &projection {
4921 if Some(column_id) == sentinel {
4922 continue;
4923 }
4924 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4925 &row_ids, column_id, snapshot, query_now, context,
4926 )? {
4927 cells.entry(row_id).or_default().insert(column_id, value);
4928 }
4929 }
4930 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
4931 ranked.drain(..)
4932 {
4933 let Some(row_cells) = cells.remove(&row_id) else {
4934 continue;
4935 };
4936 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4937 let final_rank = out.len() + 1;
4938 out.push(SearchHit {
4939 row_id,
4940 cells: projection
4941 .iter()
4942 .filter_map(|column_id| {
4943 row_cells
4944 .get(column_id)
4945 .cloned()
4946 .map(|value| (*column_id, value))
4947 })
4948 .collect(),
4949 components,
4950 fused_score,
4951 exact_rerank_score,
4952 final_score,
4953 final_rank,
4954 });
4955 if out.len() == request.limit {
4956 break;
4957 }
4958 }
4959 ranked.append(&mut remainder);
4960 }
4961 crate::trace::QueryTrace::record(|trace| {
4962 trace.union_size = union_size;
4963 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
4964 trace.projection_nanos = trace
4965 .projection_nanos
4966 .saturating_add(projection_started.elapsed().as_nanos() as u64);
4967 trace.total_nanos = trace
4968 .total_nanos
4969 .saturating_add(total_started.elapsed().as_nanos() as u64);
4970 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
4971 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
4972 if let Some(context) = context {
4973 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
4974 }
4975 });
4976 Ok(out)
4977 }
4978
4979 pub fn set_similarity(
4982 &mut self,
4983 request: &crate::query::SetSimilarityRequest,
4984 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4985 self.set_similarity_with_allowed(request, None)
4986 }
4987
4988 pub fn set_similarity_at(
4989 &mut self,
4990 request: &crate::query::SetSimilarityRequest,
4991 snapshot: Snapshot,
4992 allowed: Option<&std::collections::HashSet<RowId>>,
4993 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
4994 self.set_similarity_explained_at(request, snapshot, allowed)
4995 .map(|(hits, _)| hits)
4996 }
4997
4998 pub fn ann_rerank(
5000 &mut self,
5001 request: &crate::query::AnnRerankRequest,
5002 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5003 self.ann_rerank_with_allowed(request, None)
5004 }
5005
5006 pub fn ann_rerank_with_allowed(
5007 &mut self,
5008 request: &crate::query::AnnRerankRequest,
5009 allowed: Option<&std::collections::HashSet<RowId>>,
5010 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5011 self.ann_rerank_at(request, self.snapshot(), allowed)
5012 }
5013
5014 pub fn ann_rerank_at(
5015 &mut self,
5016 request: &crate::query::AnnRerankRequest,
5017 snapshot: Snapshot,
5018 allowed: Option<&std::collections::HashSet<RowId>>,
5019 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5020 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
5021 }
5022
5023 pub fn ann_rerank_at_with_context(
5024 &mut self,
5025 request: &crate::query::AnnRerankRequest,
5026 snapshot: Snapshot,
5027 allowed: Option<&std::collections::HashSet<RowId>>,
5028 context: Option<&crate::query::AiExecutionContext>,
5029 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5030 self.ensure_indexes_complete()?;
5031 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
5032 }
5033
5034 pub fn ann_rerank_at_with_candidate_authorization_and_context(
5035 &mut self,
5036 request: &crate::query::AnnRerankRequest,
5037 snapshot: Snapshot,
5038 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5039 context: Option<&crate::query::AiExecutionContext>,
5040 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5041 self.ensure_indexes_complete()?;
5042 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
5043 }
5044
5045 #[doc(hidden)]
5046 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
5047 &self,
5048 request: &crate::query::AnnRerankRequest,
5049 snapshot: Snapshot,
5050 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5051 context: Option<&crate::query::AiExecutionContext>,
5052 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5053 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
5054 }
5055
5056 fn ann_rerank_at_with_filters_and_context(
5057 &self,
5058 request: &crate::query::AnnRerankRequest,
5059 snapshot: Snapshot,
5060 allowed: Option<&std::collections::HashSet<RowId>>,
5061 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5062 context: Option<&crate::query::AiExecutionContext>,
5063 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5064 use crate::query::{
5065 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
5066 };
5067 if request.candidate_k == 0 || request.limit == 0 {
5068 return Err(MongrelError::InvalidArgument(
5069 "candidate_k and limit must be > 0".into(),
5070 ));
5071 }
5072 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
5073 return Err(MongrelError::InvalidArgument(format!(
5074 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
5075 )));
5076 }
5077 let retriever = Retriever::Ann {
5078 column_id: request.column_id,
5079 query: request.query.clone(),
5080 k: request.candidate_k,
5081 };
5082 self.require_select()?;
5083 self.validate_retriever(&retriever)?;
5084 let hits = self.retrieve_filtered(
5085 &retriever,
5086 snapshot,
5087 None,
5088 allowed,
5089 candidate_authorization,
5090 context,
5091 )?;
5092 let distances: std::collections::HashMap<_, _> = hits
5093 .iter()
5094 .filter_map(|hit| match hit.score {
5095 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
5096 _ => None,
5097 })
5098 .collect();
5099 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
5100 if let Some(context) = context {
5101 context.consume(row_ids.len())?;
5102 }
5103 let gather_started = std::time::Instant::now();
5104 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5105 let values = self.values_for_rids_batch_at_with_context(
5106 &row_ids,
5107 request.column_id,
5108 snapshot,
5109 query_now,
5110 context,
5111 )?;
5112 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5113 let score_started = std::time::Instant::now();
5114 let vector_work =
5115 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
5116 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
5117 if let Some(context) = context {
5118 context.consume(vector_work)?;
5119 }
5120 request
5121 .query
5122 .iter()
5123 .map(|value| f64::from(*value).powi(2))
5124 .sum::<f64>()
5125 .sqrt()
5126 } else {
5127 0.0
5128 };
5129 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
5130 for (row_id, value) in values {
5131 let Value::Embedding(vector) = value else {
5132 continue;
5133 };
5134 let exact_score = match request.metric {
5135 VectorMetric::DotProduct => {
5136 if let Some(context) = context {
5137 context.consume(vector_work)?;
5138 }
5139 request
5140 .query
5141 .iter()
5142 .zip(&vector)
5143 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5144 .sum::<f64>()
5145 }
5146 VectorMetric::Cosine => {
5147 if let Some(context) = context {
5148 context.consume(vector_work.saturating_mul(2))?;
5149 }
5150 let dot = request
5151 .query
5152 .iter()
5153 .zip(&vector)
5154 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5155 .sum::<f64>();
5156 let norm = vector
5157 .iter()
5158 .map(|value| f64::from(*value).powi(2))
5159 .sum::<f64>()
5160 .sqrt();
5161 if query_norm == 0.0 || norm == 0.0 {
5162 0.0
5163 } else {
5164 dot / (query_norm * norm)
5165 }
5166 }
5167 VectorMetric::Euclidean => {
5168 if let Some(context) = context {
5169 context.consume(vector_work)?;
5170 }
5171 request
5172 .query
5173 .iter()
5174 .zip(&vector)
5175 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
5176 .sum::<f64>()
5177 .sqrt()
5178 }
5179 };
5180 let exact_score = exact_score as f32;
5181 if !exact_score.is_finite() {
5182 return Err(MongrelError::InvalidArgument(
5183 "exact ANN score must be finite".into(),
5184 ));
5185 }
5186 reranked.push(AnnRerankHit {
5187 row_id,
5188 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
5189 exact_score,
5190 });
5191 }
5192 reranked.sort_by(|left, right| {
5193 let score = match request.metric {
5194 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
5195 VectorMetric::Cosine | VectorMetric::DotProduct => {
5196 right.exact_score.total_cmp(&left.exact_score)
5197 }
5198 };
5199 score.then_with(|| left.row_id.cmp(&right.row_id))
5200 });
5201 reranked.truncate(request.limit);
5202 crate::trace::QueryTrace::record(|trace| {
5203 trace.exact_vector_gather_nanos =
5204 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
5205 trace.exact_vector_score_nanos = trace
5206 .exact_vector_score_nanos
5207 .saturating_add(score_started.elapsed().as_nanos() as u64);
5208 });
5209 Ok(reranked)
5210 }
5211
5212 pub fn set_similarity_with_allowed(
5213 &mut self,
5214 request: &crate::query::SetSimilarityRequest,
5215 allowed: Option<&std::collections::HashSet<RowId>>,
5216 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5217 self.set_similarity_explained_at(request, self.snapshot(), allowed)
5218 .map(|(hits, _)| hits)
5219 }
5220
5221 pub fn set_similarity_explained(
5222 &mut self,
5223 request: &crate::query::SetSimilarityRequest,
5224 ) -> Result<(
5225 Vec<crate::query::SetSimilarityHit>,
5226 crate::query::SetSimilarityTrace,
5227 )> {
5228 self.set_similarity_explained_at(request, self.snapshot(), None)
5229 }
5230
5231 fn set_similarity_explained_at(
5232 &mut self,
5233 request: &crate::query::SetSimilarityRequest,
5234 snapshot: Snapshot,
5235 allowed: Option<&std::collections::HashSet<RowId>>,
5236 ) -> Result<(
5237 Vec<crate::query::SetSimilarityHit>,
5238 crate::query::SetSimilarityTrace,
5239 )> {
5240 self.ensure_indexes_complete()?;
5241 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
5242 }
5243
5244 pub fn set_similarity_at_with_context(
5245 &mut self,
5246 request: &crate::query::SetSimilarityRequest,
5247 snapshot: Snapshot,
5248 allowed: Option<&std::collections::HashSet<RowId>>,
5249 context: Option<&crate::query::AiExecutionContext>,
5250 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5251 self.ensure_indexes_complete()?;
5252 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
5253 .map(|(hits, _)| hits)
5254 }
5255
5256 pub fn set_similarity_at_with_candidate_authorization_and_context(
5257 &mut self,
5258 request: &crate::query::SetSimilarityRequest,
5259 snapshot: Snapshot,
5260 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5261 context: Option<&crate::query::AiExecutionContext>,
5262 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5263 self.ensure_indexes_complete()?;
5264 self.set_similarity_explained_at_with_context(
5265 request,
5266 snapshot,
5267 None,
5268 authorization,
5269 context,
5270 )
5271 .map(|(hits, _)| hits)
5272 }
5273
5274 #[doc(hidden)]
5275 pub fn set_similarity_at_with_candidate_authorization_on_generation(
5276 &self,
5277 request: &crate::query::SetSimilarityRequest,
5278 snapshot: Snapshot,
5279 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5280 context: Option<&crate::query::AiExecutionContext>,
5281 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5282 self.set_similarity_explained_at_with_context(
5283 request,
5284 snapshot,
5285 None,
5286 authorization,
5287 context,
5288 )
5289 .map(|(hits, _)| hits)
5290 }
5291
5292 fn set_similarity_explained_at_with_context(
5293 &self,
5294 request: &crate::query::SetSimilarityRequest,
5295 snapshot: Snapshot,
5296 allowed: Option<&std::collections::HashSet<RowId>>,
5297 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5298 context: Option<&crate::query::AiExecutionContext>,
5299 ) -> Result<(
5300 Vec<crate::query::SetSimilarityHit>,
5301 crate::query::SetSimilarityTrace,
5302 )> {
5303 use crate::query::{
5304 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
5305 MAX_SET_MEMBERS,
5306 };
5307 let mut trace = crate::query::SetSimilarityTrace::default();
5308 if request.members.is_empty() {
5309 return Ok((Vec::new(), trace));
5310 }
5311 if request.candidate_k == 0 || request.limit == 0 {
5312 return Err(MongrelError::InvalidArgument(
5313 "candidate_k and limit must be > 0".into(),
5314 ));
5315 }
5316 if request.candidate_k > MAX_RETRIEVER_K
5317 || request.limit > MAX_FINAL_LIMIT
5318 || request.members.len() > MAX_SET_MEMBERS
5319 {
5320 return Err(MongrelError::InvalidArgument(format!(
5321 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
5322 )));
5323 }
5324 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
5325 return Err(MongrelError::InvalidArgument(
5326 "min_jaccard must be finite and between 0 and 1".into(),
5327 ));
5328 }
5329 let started = std::time::Instant::now();
5330 let retriever = Retriever::MinHash {
5331 column_id: request.column_id,
5332 members: request.members.clone(),
5333 k: request.candidate_k,
5334 };
5335 self.require_select()?;
5336 self.validate_retriever(&retriever)?;
5337 let hits = self.retrieve_filtered(
5338 &retriever,
5339 snapshot,
5340 None,
5341 allowed,
5342 candidate_authorization,
5343 context,
5344 )?;
5345 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
5346 trace.candidate_count = hits.len();
5347 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
5348 if let Some(context) = context {
5349 context.consume(row_ids.len())?;
5350 }
5351 let started = std::time::Instant::now();
5352 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5353 let values = self.values_for_rids_batch_at_with_context(
5354 &row_ids,
5355 request.column_id,
5356 snapshot,
5357 query_now,
5358 context,
5359 )?;
5360 trace.gather_us = started.elapsed().as_micros() as u64;
5361 if let Some(context) = context {
5362 context.consume(request.members.len())?;
5363 }
5364 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
5365 let estimates: std::collections::HashMap<_, _> = hits
5366 .into_iter()
5367 .filter_map(|hit| match hit.score {
5368 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
5369 _ => None,
5370 })
5371 .collect();
5372 let started = std::time::Instant::now();
5373 let mut parsed = Vec::with_capacity(values.len());
5374 for (row_id, value) in values {
5375 let Value::Bytes(bytes) = value else {
5376 continue;
5377 };
5378 if let Some(context) = context {
5379 context.consume(crate::query::work_units(
5380 bytes.len(),
5381 crate::query::PARSE_WORK_QUANTUM,
5382 ))?;
5383 }
5384 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
5385 continue;
5386 };
5387 if let Some(context) = context {
5388 context.consume(members.len())?;
5389 }
5390 let stored = members
5391 .into_iter()
5392 .filter_map(|member| match member {
5393 serde_json::Value::String(value) => {
5394 Some(crate::query::SetMember::String(value))
5395 }
5396 serde_json::Value::Number(value) => {
5397 Some(crate::query::SetMember::Number(value))
5398 }
5399 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
5400 _ => None,
5401 })
5402 .collect::<std::collections::HashSet<_>>();
5403 parsed.push((row_id, stored));
5404 }
5405 trace.parse_us = started.elapsed().as_micros() as u64;
5406 trace.verified_count = parsed.len();
5407 let started = std::time::Instant::now();
5408 let mut exact = Vec::new();
5409 for (row_id, stored) in parsed {
5410 if let Some(context) = context {
5411 context.consume(query.len().saturating_add(stored.len()))?;
5412 }
5413 let union = query.union(&stored).count();
5414 let score = if union == 0 {
5415 1.0
5416 } else {
5417 query.intersection(&stored).count() as f32 / union as f32
5418 };
5419 if score >= request.min_jaccard {
5420 exact.push(SetSimilarityHit {
5421 row_id,
5422 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
5423 exact_jaccard: score,
5424 });
5425 }
5426 }
5427 exact.sort_by(|a, b| {
5428 b.exact_jaccard
5429 .total_cmp(&a.exact_jaccard)
5430 .then_with(|| a.row_id.cmp(&b.row_id))
5431 });
5432 exact.truncate(request.limit);
5433 trace.score_us = started.elapsed().as_micros() as u64;
5434 crate::trace::QueryTrace::record(|query_trace| {
5435 query_trace.exact_set_gather_nanos = query_trace
5436 .exact_set_gather_nanos
5437 .saturating_add(trace.gather_us.saturating_mul(1_000));
5438 query_trace.exact_set_parse_nanos = query_trace
5439 .exact_set_parse_nanos
5440 .saturating_add(trace.parse_us.saturating_mul(1_000));
5441 query_trace.exact_set_score_nanos = query_trace
5442 .exact_set_score_nanos
5443 .saturating_add(trace.score_us.saturating_mul(1_000));
5444 });
5445 Ok((exact, trace))
5446 }
5447
5448 fn values_for_rids_batch_at(
5450 &self,
5451 row_ids: &[u64],
5452 column_id: u16,
5453 snapshot: Snapshot,
5454 now: i64,
5455 ) -> Result<Vec<(RowId, Value)>> {
5456 if self.ttl.is_none()
5457 && self.memtable.is_empty()
5458 && self.mutable_run.is_empty()
5459 && self.run_refs.len() == 1
5460 {
5461 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5462 if row_ids.len().saturating_mul(24) < reader.row_count() {
5467 let mut values = Vec::with_capacity(row_ids.len());
5468 for &raw_row_id in row_ids {
5469 let row_id = RowId(raw_row_id);
5470 if let Some((_, false, Some(value))) =
5471 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5472 {
5473 values.push((row_id, value));
5474 }
5475 }
5476 return Ok(values);
5477 }
5478 let (positions, visible_row_ids) =
5479 reader.visible_positions_with_rids(snapshot.epoch)?;
5480 let requested: Vec<(RowId, usize)> = row_ids
5481 .iter()
5482 .filter_map(|raw| {
5483 visible_row_ids
5484 .binary_search(&(*raw as i64))
5485 .ok()
5486 .map(|index| (RowId(*raw), positions[index]))
5487 })
5488 .collect();
5489 let values = reader.gather_column(
5490 column_id,
5491 &requested
5492 .iter()
5493 .map(|(_, position)| *position)
5494 .collect::<Vec<_>>(),
5495 )?;
5496 return Ok(requested
5497 .into_iter()
5498 .zip(values)
5499 .map(|((row_id, _), value)| (row_id, value))
5500 .collect());
5501 }
5502 self.values_for_rids_at(row_ids, column_id, snapshot, now)
5503 }
5504
5505 fn values_for_rids_batch_at_with_context(
5506 &self,
5507 row_ids: &[u64],
5508 column_id: u16,
5509 snapshot: Snapshot,
5510 now: i64,
5511 context: Option<&crate::query::AiExecutionContext>,
5512 ) -> Result<Vec<(RowId, Value)>> {
5513 let Some(context) = context else {
5514 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
5515 };
5516 let mut values = Vec::with_capacity(row_ids.len());
5517 for chunk in row_ids.chunks(256) {
5518 context.checkpoint()?;
5519 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
5520 }
5521 Ok(values)
5522 }
5523
5524 fn values_for_rids_at(
5526 &self,
5527 row_ids: &[u64],
5528 column_id: u16,
5529 snapshot: Snapshot,
5530 now: i64,
5531 ) -> Result<Vec<(RowId, Value)>> {
5532 let mut readers: Vec<_> = self
5533 .run_refs
5534 .iter()
5535 .map(|run| self.open_reader(run.run_id))
5536 .collect::<Result<_>>()?;
5537 let mut out = Vec::with_capacity(row_ids.len());
5538 for &raw_row_id in row_ids {
5539 let row_id = RowId(raw_row_id);
5540 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5541 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5542 let overlay = match (mem, mutable) {
5543 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
5544 (a_epoch, a)
5545 } else {
5546 (b_epoch, b)
5547 }),
5548 (Some(value), None) | (None, Some(value)) => Some(value),
5549 (None, None) => None,
5550 };
5551 if let Some((_, row)) = overlay {
5552 if !row.deleted && !self.row_expired_at(&row, now) {
5553 if let Some(value) = row.columns.get(&column_id) {
5554 out.push((row_id, value.clone()));
5555 }
5556 }
5557 continue;
5558 }
5559
5560 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
5561 for (index, reader) in readers.iter_mut().enumerate() {
5562 if let Some((epoch, deleted, value)) =
5563 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5564 {
5565 if best
5566 .as_ref()
5567 .map(|(best_epoch, ..)| epoch > *best_epoch)
5568 .unwrap_or(true)
5569 {
5570 best = Some((epoch, deleted, value, index));
5571 }
5572 }
5573 }
5574 let Some((_, false, Some(value), reader_index)) = best else {
5575 continue;
5576 };
5577 if let Some(ttl) = self.ttl {
5578 if ttl.column_id != column_id {
5579 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5580 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5581 {
5582 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5583 continue;
5584 }
5585 }
5586 } else if let Value::Int64(timestamp) = value {
5587 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5588 continue;
5589 }
5590 }
5591 }
5592 out.push((row_id, value));
5593 }
5594 Ok(out)
5595 }
5596
5597 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
5602 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now())
5603 }
5604
5605 pub fn rows_for_rids_with_context(
5606 &self,
5607 rids: &[u64],
5608 snapshot: Snapshot,
5609 context: &crate::query::AiExecutionContext,
5610 ) -> Result<Vec<Row>> {
5611 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
5612 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos())
5613 }
5614
5615 fn rows_for_rids_at_time(
5616 &self,
5617 rids: &[u64],
5618 snapshot: Snapshot,
5619 ttl_now: i64,
5620 ) -> Result<Vec<Row>> {
5621 use std::collections::HashMap;
5622 let mut rows = Vec::with_capacity(rids.len());
5623 let tier_size = self.memtable.len() + self.mutable_run.len();
5640 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
5641 if rids.len().saturating_mul(24) < tier_size {
5642 for &rid in rids {
5643 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
5644 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
5645 let newest = match (mem, mrun) {
5646 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
5647 (Some((_, mr)), None) => Some(mr),
5648 (None, Some((_, rr))) => Some(rr),
5649 (None, None) => None,
5650 };
5651 if let Some(row) = newest {
5652 overlay.insert(rid, row);
5653 }
5654 }
5655 } else {
5656 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
5657 overlay
5658 .entry(row.row_id.0)
5659 .and_modify(|e| {
5660 if row.committed_epoch > e.committed_epoch {
5661 *e = row.clone();
5662 }
5663 })
5664 .or_insert(row);
5665 };
5666 for row in self.memtable.visible_versions(snapshot.epoch) {
5667 fold_newest(row, &mut overlay);
5668 }
5669 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5670 fold_newest(row, &mut overlay);
5671 }
5672 }
5673 if self.run_refs.len() == 1 {
5674 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5675 if rids.len().saturating_mul(24) < reader.row_count() {
5683 for &rid in rids {
5684 if let Some(r) = overlay.get(&rid) {
5685 if !r.deleted {
5686 rows.push(r.clone());
5687 }
5688 continue;
5689 }
5690 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
5691 if !row.deleted {
5692 rows.push(row);
5693 }
5694 }
5695 }
5696 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5697 return Ok(rows);
5698 }
5699 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
5708 enum Src {
5711 Overlay,
5712 Run,
5713 }
5714 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
5715 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
5716 for rid in rids {
5717 if overlay.contains_key(rid) {
5718 plan.push(Src::Overlay);
5719 continue;
5720 }
5721 match vis_rids.binary_search(&(*rid as i64)) {
5722 Ok(i) => {
5723 plan.push(Src::Run);
5724 fetch.push(positions[i]);
5725 }
5726 Err(_) => { }
5727 }
5728 }
5729 let fetched = reader.materialize_batch(&fetch)?;
5730 let mut fetched_iter = fetched.into_iter();
5731 for (rid, src) in rids.iter().zip(plan) {
5732 match src {
5733 Src::Overlay => {
5734 if let Some(r) = overlay.get(rid) {
5735 if !r.deleted {
5736 rows.push(r.clone());
5737 }
5738 }
5739 }
5740 Src::Run => {
5741 if let Some(row) = fetched_iter.next() {
5742 if !row.deleted {
5743 rows.push(row);
5744 }
5745 }
5746 }
5747 }
5748 }
5749 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5750 return Ok(rows);
5751 }
5752 let mut readers: Vec<_> = self
5756 .run_refs
5757 .iter()
5758 .map(|rr| self.open_reader(rr.run_id))
5759 .collect::<Result<Vec<_>>>()?;
5760 for rid in rids {
5761 if let Some(r) = overlay.get(rid) {
5762 if !r.deleted {
5763 rows.push(r.clone());
5764 }
5765 continue;
5766 }
5767 let mut best: Option<(Epoch, Row)> = None;
5768 for reader in readers.iter_mut() {
5769 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
5770 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
5771 best = Some((epoch, row));
5772 }
5773 }
5774 }
5775 if let Some((_, r)) = best {
5776 if !r.deleted {
5777 rows.push(r);
5778 }
5779 }
5780 }
5781 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5782 Ok(rows)
5783 }
5784
5785 pub fn indexes_complete(&self) -> bool {
5795 self.indexes_complete
5796 }
5797
5798 pub fn index_build_policy(&self) -> IndexBuildPolicy {
5800 self.index_build_policy
5801 }
5802
5803 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
5807 self.index_build_policy = policy;
5808 }
5809
5810 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
5815 if !self.indexes_complete {
5819 return None;
5820 }
5821 let b = self.bitmap.get(&column_id)?;
5822 let result: Vec<Vec<u8>> = b
5823 .keys()
5824 .into_iter()
5825 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
5826 .collect();
5827 Some(result)
5828 }
5829
5830 pub fn fk_join_row_ids(
5831 &self,
5832 fk_column_id: u16,
5833 pk_values: &[Vec<u8>],
5834 fk_conditions: &[crate::query::Condition],
5835 snapshot: Snapshot,
5836 ) -> Result<Vec<u64>> {
5837 let Some(b) = self.bitmap.get(&fk_column_id) else {
5838 return Ok(Vec::new());
5839 };
5840 let mut join_set = {
5841 let mut acc = roaring::RoaringBitmap::new();
5842 for v in pk_values {
5843 acc |= b.get(v);
5844 }
5845 RowIdSet::from_roaring(acc)
5846 };
5847 if !fk_conditions.is_empty() {
5848 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5849 sets.push(join_set);
5850 for c in fk_conditions {
5851 sets.push(self.resolve_condition(c, snapshot)?);
5852 }
5853 join_set = RowIdSet::intersect_many(sets);
5854 }
5855 Ok(join_set.into_sorted_vec())
5856 }
5857
5858 pub fn fk_join_count(
5864 &self,
5865 fk_column_id: u16,
5866 pk_values: &[Vec<u8>],
5867 fk_conditions: &[crate::query::Condition],
5868 snapshot: Snapshot,
5869 ) -> Result<u64> {
5870 let Some(b) = self.bitmap.get(&fk_column_id) else {
5871 return Ok(0);
5872 };
5873 let mut acc = roaring::RoaringBitmap::new();
5874 for v in pk_values {
5875 acc |= b.get(v);
5876 }
5877 if fk_conditions.is_empty() {
5878 return Ok(acc.len());
5879 }
5880 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5881 sets.push(RowIdSet::from_roaring(acc));
5882 for c in fk_conditions {
5883 sets.push(self.resolve_condition(c, snapshot)?);
5884 }
5885 Ok(RowIdSet::intersect_many(sets).len() as u64)
5886 }
5887
5888 fn resolve_condition(
5893 &self,
5894 c: &crate::query::Condition,
5895 snapshot: Snapshot,
5896 ) -> Result<RowIdSet> {
5897 self.resolve_condition_with_allowed(c, snapshot, None)
5898 }
5899
5900 fn resolve_condition_with_allowed(
5901 &self,
5902 c: &crate::query::Condition,
5903 snapshot: Snapshot,
5904 allowed: Option<&std::collections::HashSet<RowId>>,
5905 ) -> Result<RowIdSet> {
5906 use crate::query::Condition;
5907 self.validate_condition(c)?;
5908 Ok(match c {
5909 Condition::Pk(key) => {
5910 let lookup = self
5911 .schema
5912 .primary_key()
5913 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5914 .unwrap_or_else(|| key.clone());
5915 self.hot
5916 .get(&lookup)
5917 .map(|r| RowIdSet::one(r.0))
5918 .unwrap_or_else(RowIdSet::empty)
5919 }
5920 Condition::BitmapEq { column_id, value } => {
5921 let lookup = self.index_lookup_key_bytes(*column_id, value);
5922 self.bitmap
5923 .get(column_id)
5924 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5925 .unwrap_or_else(RowIdSet::empty)
5926 }
5927 Condition::BitmapIn { column_id, values } => {
5928 let bm = self.bitmap.get(column_id);
5929 let mut acc = roaring::RoaringBitmap::new();
5930 if let Some(b) = bm {
5931 for v in values {
5932 let lookup = self.index_lookup_key_bytes(*column_id, v);
5933 acc |= b.get(&lookup);
5934 }
5935 }
5936 RowIdSet::from_roaring(acc)
5937 }
5938 Condition::BytesPrefix { column_id, prefix } => {
5939 if let Some(b) = self.bitmap.get(column_id) {
5944 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5945 let mut acc = roaring::RoaringBitmap::new();
5946 for key in b.keys() {
5947 if key.starts_with(&lookup_prefix) {
5948 acc |= b.get(&key);
5949 }
5950 }
5951 RowIdSet::from_roaring(acc)
5952 } else {
5953 RowIdSet::empty()
5954 }
5955 }
5956 Condition::FmContains { column_id, pattern } => self
5957 .fm
5958 .get(column_id)
5959 .map(|f| {
5960 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
5961 })
5962 .unwrap_or_else(RowIdSet::empty),
5963 Condition::FmContainsAll {
5964 column_id,
5965 patterns,
5966 } => {
5967 if let Some(f) = self.fm.get(column_id) {
5970 let sets: Vec<RowIdSet> = patterns
5971 .iter()
5972 .map(|pat| {
5973 RowIdSet::from_unsorted(
5974 f.locate(pat).into_iter().map(|r| r.0).collect(),
5975 )
5976 })
5977 .collect();
5978 RowIdSet::intersect_many(sets)
5979 } else {
5980 RowIdSet::empty()
5981 }
5982 }
5983 Condition::Ann {
5984 column_id,
5985 query,
5986 k,
5987 } => RowIdSet::from_unsorted(
5988 self.retrieve_filtered(
5989 &crate::query::Retriever::Ann {
5990 column_id: *column_id,
5991 query: query.clone(),
5992 k: *k,
5993 },
5994 snapshot,
5995 None,
5996 allowed,
5997 None,
5998 None,
5999 )?
6000 .into_iter()
6001 .map(|hit| hit.row_id.0)
6002 .collect(),
6003 ),
6004 Condition::SparseMatch {
6005 column_id,
6006 query,
6007 k,
6008 } => RowIdSet::from_unsorted(
6009 self.retrieve_filtered(
6010 &crate::query::Retriever::Sparse {
6011 column_id: *column_id,
6012 query: query.clone(),
6013 k: *k,
6014 },
6015 snapshot,
6016 None,
6017 allowed,
6018 None,
6019 None,
6020 )?
6021 .into_iter()
6022 .map(|hit| hit.row_id.0)
6023 .collect(),
6024 ),
6025 Condition::MinHashSimilar {
6026 column_id,
6027 query,
6028 k,
6029 } => match self.minhash.get(column_id) {
6030 Some(index) => {
6031 let candidates = index.candidate_row_ids(query);
6032 let eligible =
6033 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
6034 RowIdSet::from_unsorted(
6035 index
6036 .search_filtered(query, *k, |row_id| {
6037 eligible.contains(&row_id)
6038 && allowed.map_or(true, |allowed| allowed.contains(&row_id))
6039 })
6040 .into_iter()
6041 .map(|(row_id, _)| row_id.0)
6042 .collect(),
6043 )
6044 }
6045 None => RowIdSet::empty(),
6046 },
6047 Condition::Range { column_id, lo, hi } => {
6048 let mut set = if let Some(li) = self.learned_range.get(column_id) {
6057 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
6058 } else if self.run_refs.len() == 1 {
6059 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6060 r.range_row_id_set_i64(*column_id, *lo, *hi)?
6061 } else {
6062 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
6063 };
6064 set.remove_many(self.overlay_rid_set(snapshot));
6065 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
6066 set
6067 }
6068 Condition::RangeF64 {
6069 column_id,
6070 lo,
6071 lo_inclusive,
6072 hi,
6073 hi_inclusive,
6074 } => {
6075 let mut set = if let Some(li) = self.learned_range.get(column_id) {
6078 RowIdSet::from_unsorted(
6079 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
6080 .into_iter()
6081 .collect(),
6082 )
6083 } else if self.run_refs.len() == 1 {
6084 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6085 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
6086 } else {
6087 return self.range_scan_f64(
6088 *column_id,
6089 *lo,
6090 *lo_inclusive,
6091 *hi,
6092 *hi_inclusive,
6093 snapshot,
6094 );
6095 };
6096 set.remove_many(self.overlay_rid_set(snapshot));
6097 self.range_scan_overlay_f64(
6098 &mut set,
6099 *column_id,
6100 *lo,
6101 *lo_inclusive,
6102 *hi,
6103 *hi_inclusive,
6104 snapshot,
6105 );
6106 set
6107 }
6108 Condition::IsNull { column_id } => {
6109 let mut set = if self.run_refs.len() == 1 {
6110 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6111 r.null_row_id_set(*column_id, true)?
6112 } else {
6113 return self.null_scan(*column_id, true, snapshot);
6114 };
6115 set.remove_many(self.overlay_rid_set(snapshot));
6116 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
6117 set
6118 }
6119 Condition::IsNotNull { column_id } => {
6120 let mut set = if self.run_refs.len() == 1 {
6121 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6122 r.null_row_id_set(*column_id, false)?
6123 } else {
6124 return self.null_scan(*column_id, false, snapshot);
6125 };
6126 set.remove_many(self.overlay_rid_set(snapshot));
6127 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
6128 set
6129 }
6130 })
6131 }
6132
6133 fn range_scan_i64(
6141 &self,
6142 column_id: u16,
6143 lo: i64,
6144 hi: i64,
6145 snapshot: Snapshot,
6146 ) -> Result<RowIdSet> {
6147 let mut row_ids = Vec::new();
6148 let overlay_rids = self.overlay_rid_set(snapshot);
6149 for rr in &self.run_refs {
6150 let mut reader = self.open_reader(rr.run_id)?;
6151 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
6152 for rid in matched {
6153 if !overlay_rids.contains(&rid) {
6154 row_ids.push(rid);
6155 }
6156 }
6157 }
6158 let mut s = RowIdSet::from_unsorted(row_ids);
6159 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
6160 Ok(s)
6161 }
6162
6163 fn range_scan_f64(
6166 &self,
6167 column_id: u16,
6168 lo: f64,
6169 lo_inclusive: bool,
6170 hi: f64,
6171 hi_inclusive: bool,
6172 snapshot: Snapshot,
6173 ) -> Result<RowIdSet> {
6174 let mut row_ids = Vec::new();
6175 let overlay_rids = self.overlay_rid_set(snapshot);
6176 for rr in &self.run_refs {
6177 let mut reader = self.open_reader(rr.run_id)?;
6178 let matched = reader.range_row_ids_visible_f64(
6179 column_id,
6180 lo,
6181 lo_inclusive,
6182 hi,
6183 hi_inclusive,
6184 snapshot.epoch,
6185 )?;
6186 for rid in matched {
6187 if !overlay_rids.contains(&rid) {
6188 row_ids.push(rid);
6189 }
6190 }
6191 }
6192 let mut s = RowIdSet::from_unsorted(row_ids);
6193 self.range_scan_overlay_f64(
6194 &mut s,
6195 column_id,
6196 lo,
6197 lo_inclusive,
6198 hi,
6199 hi_inclusive,
6200 snapshot,
6201 );
6202 Ok(s)
6203 }
6204
6205 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
6207 let mut s = HashSet::new();
6208 for row in self.memtable.visible_versions(snapshot.epoch) {
6209 s.insert(row.row_id.0);
6210 }
6211 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6212 s.insert(row.row_id.0);
6213 }
6214 s
6215 }
6216
6217 fn range_scan_overlay_i64(
6218 &self,
6219 s: &mut RowIdSet,
6220 column_id: u16,
6221 lo: i64,
6222 hi: i64,
6223 snapshot: Snapshot,
6224 ) {
6225 let mut newest: HashMap<u64, &Row> = HashMap::new();
6230 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6231 let memtable = self.memtable.visible_versions(snapshot.epoch);
6232 for r in &mutable {
6233 newest.entry(r.row_id.0).or_insert(r);
6234 }
6235 for r in &memtable {
6236 newest.insert(r.row_id.0, r);
6237 }
6238 for row in newest.values() {
6239 if !row.deleted {
6240 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
6241 if *v >= lo && *v <= hi {
6242 s.insert(row.row_id.0);
6243 }
6244 }
6245 }
6246 }
6247 }
6248
6249 #[allow(clippy::too_many_arguments)]
6250 fn range_scan_overlay_f64(
6251 &self,
6252 s: &mut RowIdSet,
6253 column_id: u16,
6254 lo: f64,
6255 lo_inclusive: bool,
6256 hi: f64,
6257 hi_inclusive: bool,
6258 snapshot: Snapshot,
6259 ) {
6260 let mut newest: HashMap<u64, &Row> = HashMap::new();
6263 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6264 let memtable = self.memtable.visible_versions(snapshot.epoch);
6265 for r in &mutable {
6266 newest.entry(r.row_id.0).or_insert(r);
6267 }
6268 for r in &memtable {
6269 newest.insert(r.row_id.0, r);
6270 }
6271 for row in newest.values() {
6272 if !row.deleted {
6273 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
6274 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
6275 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
6276 if ok_lo && ok_hi {
6277 s.insert(row.row_id.0);
6278 }
6279 }
6280 }
6281 }
6282 }
6283
6284 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
6287 let mut row_ids = Vec::new();
6288 let overlay_rids = self.overlay_rid_set(snapshot);
6289 for rr in &self.run_refs {
6290 let mut reader = self.open_reader(rr.run_id)?;
6291 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
6292 for rid in matched {
6293 if !overlay_rids.contains(&rid) {
6294 row_ids.push(rid);
6295 }
6296 }
6297 }
6298 let mut s = RowIdSet::from_unsorted(row_ids);
6299 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
6300 Ok(s)
6301 }
6302
6303 fn null_scan_overlay(
6307 &self,
6308 s: &mut RowIdSet,
6309 column_id: u16,
6310 want_nulls: bool,
6311 snapshot: Snapshot,
6312 ) {
6313 let mut newest: HashMap<u64, &Row> = HashMap::new();
6314 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6315 let memtable = self.memtable.visible_versions(snapshot.epoch);
6316 for r in &mutable {
6317 newest.entry(r.row_id.0).or_insert(r);
6318 }
6319 for r in &memtable {
6320 newest.insert(r.row_id.0, r);
6321 }
6322 for row in newest.values() {
6323 if row.deleted {
6324 continue;
6325 }
6326 let is_null = !row.columns.contains_key(&column_id)
6327 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
6328 if is_null == want_nulls {
6329 s.insert(row.row_id.0);
6330 }
6331 }
6332 }
6333
6334 pub fn snapshot(&self) -> Snapshot {
6335 Snapshot::at(self.epoch.visible())
6336 }
6337
6338 pub fn data_generation(&self) -> u64 {
6340 self.data_generation
6341 }
6342
6343 pub(crate) fn table_id(&self) -> u64 {
6344 self.table_id
6345 }
6346
6347 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
6348 self.ensure_indexes_complete()?;
6349 self.memtable.seal();
6350 self.mutable_run.seal();
6351 self.hot.seal();
6352 for index in self.bitmap.values_mut() {
6353 index.seal();
6354 }
6355 for index in self.ann.values_mut() {
6356 index.seal();
6357 }
6358 for index in self.fm.values_mut() {
6359 index.seal();
6360 }
6361 for index in self.sparse.values_mut() {
6362 index.seal();
6363 }
6364 for index in self.minhash.values_mut() {
6365 index.seal();
6366 }
6367 self.pk_by_row.seal();
6368 let mut generation = self.clone();
6369 generation.read_only = true;
6370 generation.wal = WalSink::ReadOnly;
6371 generation.pending_delete_rids.clear();
6372 generation.pending_put_cols.clear();
6373 generation.pending_rows.clear();
6374 generation.pending_rows_auto_inc.clear();
6375 generation.pending_dels.clear();
6376 generation.pending_truncate = None;
6377 generation.agg_cache = Arc::new(HashMap::new());
6378 Ok(generation)
6379 }
6380
6381 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
6382 (std::mem::size_of::<Self>() as u64)
6383 .saturating_add(self.memtable.approx_bytes())
6384 .saturating_add(self.mutable_run.approx_bytes())
6385 .saturating_add(self.live_count.saturating_mul(64))
6386 }
6387
6388 pub fn pin_snapshot(&mut self) -> Snapshot {
6391 let e = self.epoch.visible();
6392 *self.pinned.entry(e).or_insert(0) += 1;
6393 Snapshot::at(e)
6394 }
6395
6396 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
6398 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
6399 *count -= 1;
6400 if *count == 0 {
6401 self.pinned.remove(&snap.epoch);
6402 }
6403 }
6404 }
6405
6406 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
6416 let local = self.pinned.keys().next().copied();
6417 let global = self.snapshots.min_pinned();
6418 let history = self.snapshots.history_floor(self.current_epoch());
6419 [local, global, history].into_iter().flatten().min()
6420 }
6421
6422 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
6426 self.ensure_writable()?;
6427 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
6428 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
6429 }
6430
6431 pub fn clear_ttl(&mut self) -> Result<()> {
6432 self.ensure_writable()?;
6433 self.apply_ttl_policy_at(None, self.current_epoch())
6434 }
6435
6436 pub fn ttl(&self) -> Option<TtlPolicy> {
6437 self.ttl
6438 }
6439
6440 pub(crate) fn prepare_ttl_policy(
6441 &self,
6442 column_name: &str,
6443 duration_nanos: u64,
6444 ) -> Result<TtlPolicy> {
6445 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
6446 return Err(MongrelError::InvalidArgument(
6447 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
6448 ));
6449 }
6450 let column = self
6451 .schema
6452 .columns
6453 .iter()
6454 .find(|column| column.name == column_name)
6455 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
6456 if column.ty != TypeId::TimestampNanos {
6457 return Err(MongrelError::Schema(format!(
6458 "TTL column {column_name} must be TimestampNanos, is {:?}",
6459 column.ty
6460 )));
6461 }
6462 Ok(TtlPolicy {
6463 column_id: column.id,
6464 duration_nanos,
6465 })
6466 }
6467
6468 pub(crate) fn apply_ttl_policy_at(
6469 &mut self,
6470 policy: Option<TtlPolicy>,
6471 epoch: Epoch,
6472 ) -> Result<()> {
6473 if let Some(policy) = policy {
6474 let column = self
6475 .schema
6476 .columns
6477 .iter()
6478 .find(|column| column.id == policy.column_id)
6479 .ok_or_else(|| {
6480 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
6481 })?;
6482 if column.ty != TypeId::TimestampNanos
6483 || policy.duration_nanos == 0
6484 || policy.duration_nanos > i64::MAX as u64
6485 {
6486 return Err(MongrelError::Schema("invalid TTL policy".into()));
6487 }
6488 }
6489 self.ttl = policy;
6490 self.agg_cache = Arc::new(HashMap::new());
6491 self.clear_result_cache();
6492 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6493 self.persist_manifest(epoch)
6494 }
6495
6496 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
6497 let Some(policy) = self.ttl else {
6498 return false;
6499 };
6500 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
6501 return false;
6502 };
6503 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
6504 }
6505
6506 pub fn current_epoch(&self) -> Epoch {
6507 self.epoch.visible()
6508 }
6509
6510 pub fn memtable_len(&self) -> usize {
6511 self.memtable.len()
6512 }
6513
6514 pub fn count(&self) -> u64 {
6517 if self.ttl.is_none()
6518 && self.pending_put_cols.is_empty()
6519 && self.pending_delete_rids.is_empty()
6520 && self.pending_rows.is_empty()
6521 && self.pending_dels.is_empty()
6522 && self.pending_truncate.is_none()
6523 {
6524 self.live_count
6525 } else {
6526 self.visible_rows(self.snapshot())
6527 .map(|rows| rows.len() as u64)
6528 .unwrap_or(self.live_count)
6529 }
6530 }
6531
6532 pub fn count_conditions(
6536 &mut self,
6537 conditions: &[crate::query::Condition],
6538 snapshot: Snapshot,
6539 ) -> Result<Option<u64>> {
6540 use crate::query::Condition;
6541 if self.ttl.is_some() {
6542 if conditions.is_empty() {
6543 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
6544 }
6545 let mut sets = Vec::with_capacity(conditions.len());
6546 for condition in conditions {
6547 sets.push(self.resolve_condition(condition, snapshot)?);
6548 }
6549 let survivors = RowIdSet::intersect_many(sets);
6550 let rows = self.visible_rows(snapshot)?;
6551 return Ok(Some(
6552 rows.into_iter()
6553 .filter(|row| survivors.contains(row.row_id.0))
6554 .count() as u64,
6555 ));
6556 }
6557 if conditions.is_empty() {
6558 return Ok(Some(self.count()));
6559 }
6560 let served = |c: &Condition| {
6561 matches!(
6562 c,
6563 Condition::Pk(_)
6564 | Condition::BitmapEq { .. }
6565 | Condition::BitmapIn { .. }
6566 | Condition::BytesPrefix { .. }
6567 | Condition::FmContains { .. }
6568 | Condition::FmContainsAll { .. }
6569 | Condition::Ann { .. }
6570 | Condition::Range { .. }
6571 | Condition::RangeF64 { .. }
6572 | Condition::SparseMatch { .. }
6573 | Condition::MinHashSimilar { .. }
6574 | Condition::IsNull { .. }
6575 | Condition::IsNotNull { .. }
6576 )
6577 };
6578 if !conditions.iter().all(served) {
6579 return Ok(None);
6580 }
6581 self.ensure_indexes_complete()?;
6582 if !self.pending_put_cols.is_empty()
6583 || !self.pending_delete_rids.is_empty()
6584 || !self.pending_rows.is_empty()
6585 || !self.pending_dels.is_empty()
6586 || self.pending_truncate.is_some()
6587 {
6588 let mut sets = Vec::with_capacity(conditions.len());
6589 for condition in conditions {
6590 sets.push(self.resolve_condition(condition, snapshot)?);
6591 }
6592 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6593 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
6594 }
6595 let mut sets = Vec::with_capacity(conditions.len());
6596 for condition in conditions {
6597 sets.push(self.resolve_condition(condition, snapshot)?);
6598 }
6599 let mut rids = RowIdSet::intersect_many(sets);
6600 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6610 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
6611 }
6612 let count = rids.len() as u64;
6613 crate::trace::QueryTrace::record(|t| {
6614 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
6615 t.survivor_count = Some(count as usize);
6616 t.conditions_pushed = conditions.len();
6617 });
6618 Ok(Some(count))
6619 }
6620
6621 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
6626 let mut out = Vec::new();
6627 for row in self.memtable.visible_versions(snapshot.epoch) {
6628 if row.deleted {
6629 out.push(row.row_id.0);
6630 }
6631 }
6632 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6633 if row.deleted {
6634 out.push(row.row_id.0);
6635 }
6636 }
6637 out
6638 }
6639
6640 pub fn bulk_load_columns(
6649 &mut self,
6650 user_columns: Vec<(u16, columnar::NativeColumn)>,
6651 ) -> Result<Epoch> {
6652 self.bulk_load_columns_with(user_columns, 3, false, true)
6653 }
6654
6655 pub fn bulk_load_fast(
6662 &mut self,
6663 user_columns: Vec<(u16, columnar::NativeColumn)>,
6664 ) -> Result<Epoch> {
6665 self.bulk_load_columns_with(user_columns, -1, true, false)
6666 }
6667
6668 fn bulk_load_columns_with(
6669 &mut self,
6670 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
6671 zstd_level: i32,
6672 force_plain: bool,
6673 lz4: bool,
6674 ) -> Result<Epoch> {
6675 let epoch = self.commit()?;
6676 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
6677 if n == 0 {
6678 return Ok(epoch);
6679 }
6680 let live_before = self.live_count;
6681 self.spill_mutable_run(epoch)?;
6683 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6684 && self.indexes_complete
6685 && self.run_refs.is_empty()
6686 && self.memtable.is_empty()
6687 && self.mutable_run.is_empty();
6688 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6691 self.validate_columns_not_null(&user_columns, n)?;
6692 let winner_idx = self
6693 .bulk_pk_winner_indices(&user_columns, n)
6694 .filter(|idx| idx.len() != n);
6695 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6696 match winner_idx.as_deref() {
6697 Some(idx) => {
6698 let compacted = user_columns
6699 .iter()
6700 .map(|(id, c)| (*id, c.gather(idx)))
6701 .collect();
6702 (compacted, idx.len())
6703 }
6704 None => (user_columns, n),
6705 };
6706 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6707 let first = self.allocator.alloc_range(write_n as u64).0;
6708 for rid in first..first + write_n as u64 {
6709 self.reservoir.offer(rid);
6710 }
6711 let run_id = self.next_run_id;
6712 self.next_run_id += 1;
6713 let path = self.run_path(run_id);
6714 let mut writer =
6715 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
6716 if force_plain {
6717 writer = writer.with_plain();
6718 } else if lz4 {
6719 writer = writer.with_lz4();
6722 } else {
6723 writer = writer.with_zstd_level(zstd_level);
6724 }
6725 if let Some(kek) = &self.kek {
6726 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6727 }
6728 let header = writer.write_native(&path, &write_columns, write_n, first)?;
6729 self.run_refs.push(RunRef {
6730 run_id: run_id as u128,
6731 level: 0,
6732 epoch_created: epoch.0,
6733 row_count: header.row_count,
6734 });
6735 self.live_count = self.live_count.saturating_add(write_n as u64);
6736 if eager_index_build {
6737 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6738 self.index_columns_bulk(&write_columns, &row_ids);
6739 self.indexes_complete = true;
6740 self.build_learned_ranges()?;
6741 } else {
6742 self.indexes_complete = false;
6746 }
6747 self.mark_flushed(epoch)?;
6748 self.persist_manifest(epoch)?;
6749 if eager_index_build {
6750 self.checkpoint_indexes(epoch);
6751 }
6752 self.clear_result_cache();
6753 self.data_generation = self.data_generation.wrapping_add(1);
6754 Ok(epoch)
6755 }
6756
6757 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
6775 let n = row_ids.len();
6776 if n == 0 {
6777 return;
6778 }
6779 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
6780 columns.iter().map(|(id, c)| (*id, c)).collect();
6781 let ty_of: std::collections::HashMap<u16, TypeId> = self
6782 .schema
6783 .columns
6784 .iter()
6785 .map(|c| (c.id, c.ty.clone()))
6786 .collect();
6787 let pk_id = self.schema.primary_key().map(|c| c.id);
6788
6789 for (i, &rid) in row_ids.iter().enumerate() {
6790 let row_id = RowId(rid);
6791 if let Some(pid) = pk_id {
6792 if let Some(col) = by_id.get(&pid) {
6793 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
6794 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
6795 self.insert_hot_pk(key, row_id);
6796 }
6797 }
6798 }
6799 for idef in &self.schema.indexes {
6800 let Some(col) = by_id.get(&idef.column_id) else {
6801 continue;
6802 };
6803 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
6804 match idef.kind {
6805 IndexKind::Bitmap => {
6806 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
6807 if let Some(key) =
6808 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
6809 {
6810 b.insert(key, row_id);
6811 }
6812 }
6813 }
6814 IndexKind::FmIndex => {
6815 if let Some(f) = self.fm.get_mut(&idef.column_id) {
6816 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6817 f.insert(bytes.to_vec(), row_id);
6818 }
6819 }
6820 }
6821 IndexKind::Sparse => {
6822 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
6823 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6824 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
6825 s.insert(&terms, row_id);
6826 }
6827 }
6828 }
6829 }
6830 IndexKind::MinHash => {
6831 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
6832 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6833 let tokens = crate::index::token_hashes_from_bytes(bytes);
6834 mh.insert(&tokens, row_id);
6835 }
6836 }
6837 }
6838 _ => {}
6839 }
6840 }
6841 }
6842 }
6843
6844 pub fn visible_columns_native(
6849 &self,
6850 snapshot: Snapshot,
6851 projection: Option<&[u16]>,
6852 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
6853 let wanted: Vec<u16> = match projection {
6854 Some(p) => p.to_vec(),
6855 None => self.schema.columns.iter().map(|c| c.id).collect(),
6856 };
6857 if self.ttl.is_none()
6858 && self.memtable.is_empty()
6859 && self.mutable_run.is_empty()
6860 && self.run_refs.len() == 1
6861 {
6862 let rr = self.run_refs[0].clone();
6863 let mut reader = self.open_reader(rr.run_id)?;
6864 let idxs = reader.visible_indices_native(snapshot.epoch)?;
6865 let all_visible = idxs.len() == reader.row_count();
6866 if reader.has_mmap() {
6872 use rayon::prelude::*;
6873 let valid: Vec<u16> = wanted
6876 .iter()
6877 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
6878 .copied()
6879 .collect();
6880 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
6882 .par_iter()
6883 .filter_map(|cid| {
6884 reader
6885 .column_native_shared(*cid)
6886 .ok()
6887 .map(|col| (*cid, col))
6888 })
6889 .collect();
6890 let cols = decoded
6891 .into_iter()
6892 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
6893 .collect();
6894 return Ok(cols);
6895 }
6896 let mut cols = Vec::with_capacity(wanted.len());
6897 for cid in &wanted {
6898 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
6899 Some(c) => c,
6900 None => continue,
6901 };
6902 let col = reader.column_native(cdef.id)?;
6903 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
6904 }
6905 return Ok(cols);
6906 }
6907 let vcols = self.visible_columns(snapshot)?;
6908 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
6909 let out: Vec<(u16, columnar::NativeColumn)> = vcols
6910 .into_iter()
6911 .filter(|(id, _)| want_set.contains(id))
6912 .map(|(id, vals)| {
6913 let ty = self
6914 .schema
6915 .columns
6916 .iter()
6917 .find(|c| c.id == id)
6918 .map(|c| c.ty.clone())
6919 .unwrap_or(TypeId::Bytes);
6920 (id, columnar::values_to_native(ty, &vals))
6921 })
6922 .collect();
6923 Ok(out)
6924 }
6925
6926 pub fn run_count(&self) -> usize {
6927 self.run_refs.len()
6928 }
6929
6930 pub fn memtable_is_empty(&self) -> bool {
6932 self.memtable.is_empty()
6933 }
6934
6935 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
6939 self.page_cache.stats()
6940 }
6941
6942 pub fn reset_page_cache_stats(&self) {
6944 self.page_cache.reset_stats();
6945 }
6946
6947 pub fn run_ids(&self) -> Vec<u128> {
6950 self.run_refs.iter().map(|r| r.run_id).collect()
6951 }
6952
6953 pub fn single_run_is_clean(&self) -> bool {
6957 if self.ttl.is_some() || self.run_refs.len() != 1 {
6958 return false;
6959 }
6960 self.open_reader(self.run_refs[0].run_id)
6961 .map(|r| r.is_clean())
6962 .unwrap_or(false)
6963 }
6964
6965 fn resolve_footprint(
6972 &self,
6973 conditions: &[crate::query::Condition],
6974 snapshot: Snapshot,
6975 ) -> roaring::RoaringBitmap {
6976 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6977 return roaring::RoaringBitmap::new();
6978 }
6979 if self.run_refs.is_empty() {
6980 return roaring::RoaringBitmap::new();
6981 }
6982 if self.run_refs.len() == 1 {
6984 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
6985 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
6986 return rids.to_roaring_lossy();
6987 }
6988 }
6989 }
6990 roaring::RoaringBitmap::new()
6991 }
6992
6993 pub fn query_columns_native_cached(
7004 &mut self,
7005 conditions: &[crate::query::Condition],
7006 projection: Option<&[u16]>,
7007 snapshot: Snapshot,
7008 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7009 if self.ttl.is_some() {
7012 return self.query_columns_native(conditions, projection, snapshot);
7013 }
7014 if conditions.is_empty() {
7015 return self.query_columns_native(conditions, projection, snapshot);
7016 }
7017 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
7021 if let Some(hit) = self.result_cache.lock().get_columns(key) {
7022 crate::trace::QueryTrace::record(|t| {
7023 t.result_cache_hit = true;
7024 t.scan_mode = crate::trace::ScanMode::NativePushdown;
7025 });
7026 return Ok(Some((*hit).clone()));
7027 }
7028 let res = self.query_columns_native(conditions, projection, snapshot)?;
7029 if let Some(cols) = &res {
7030 let footprint = self.resolve_footprint(conditions, snapshot);
7031 let condition_cols = crate::query::condition_columns(conditions);
7032 self.result_cache.lock().insert(
7033 key,
7034 CachedEntry {
7035 data: CachedData::Columns(Arc::new(cols.clone())),
7036 footprint,
7037 condition_cols,
7038 },
7039 );
7040 }
7041 Ok(res)
7042 }
7043
7044 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
7049 if self.ttl.is_some() {
7050 return self.query(q);
7051 }
7052 if q.conditions.is_empty() {
7053 return self.query(q);
7054 }
7055 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
7056 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
7057 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
7058 if let Some(hit) = self.result_cache.lock().get_rows(key) {
7059 crate::trace::QueryTrace::record(|t| {
7060 t.result_cache_hit = true;
7061 t.scan_mode = crate::trace::ScanMode::Materialized;
7062 });
7063 return Ok((*hit).clone());
7064 }
7065 let rows = self.query(q)?;
7066 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
7067 let condition_cols = crate::query::condition_columns(&q.conditions);
7068 self.result_cache.lock().insert(
7069 key,
7070 CachedEntry {
7071 data: CachedData::Rows(Arc::new(rows.clone())),
7072 footprint,
7073 condition_cols,
7074 },
7075 );
7076 Ok(rows)
7077 }
7078
7079 #[allow(clippy::type_complexity)]
7094 pub fn query_columns_native_traced(
7095 &mut self,
7096 conditions: &[crate::query::Condition],
7097 projection: Option<&[u16]>,
7098 snapshot: Snapshot,
7099 ) -> Result<(
7100 Option<Vec<(u16, columnar::NativeColumn)>>,
7101 crate::trace::QueryTrace,
7102 )> {
7103 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7104 self.query_columns_native(conditions, projection, snapshot)
7105 });
7106 Ok((result?, trace))
7107 }
7108
7109 #[allow(clippy::type_complexity)]
7112 pub fn query_columns_native_cached_traced(
7113 &mut self,
7114 conditions: &[crate::query::Condition],
7115 projection: Option<&[u16]>,
7116 snapshot: Snapshot,
7117 ) -> Result<(
7118 Option<Vec<(u16, columnar::NativeColumn)>>,
7119 crate::trace::QueryTrace,
7120 )> {
7121 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7122 self.query_columns_native_cached(conditions, projection, snapshot)
7123 });
7124 Ok((result?, trace))
7125 }
7126
7127 pub fn native_page_cursor_traced(
7129 &self,
7130 snapshot: Snapshot,
7131 projection: Vec<(u16, TypeId)>,
7132 conditions: &[crate::query::Condition],
7133 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
7134 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7135 self.native_page_cursor(snapshot, projection, conditions)
7136 });
7137 Ok((result?, trace))
7138 }
7139
7140 pub fn native_multi_run_cursor_traced(
7142 &self,
7143 snapshot: Snapshot,
7144 projection: Vec<(u16, TypeId)>,
7145 conditions: &[crate::query::Condition],
7146 ) -> Result<(
7147 Option<crate::cursor::MultiRunCursor>,
7148 crate::trace::QueryTrace,
7149 )> {
7150 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7151 self.native_multi_run_cursor(snapshot, projection, conditions)
7152 });
7153 Ok((result?, trace))
7154 }
7155
7156 pub fn count_conditions_traced(
7158 &mut self,
7159 conditions: &[crate::query::Condition],
7160 snapshot: Snapshot,
7161 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
7162 let (result, trace) =
7163 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
7164 Ok((result?, trace))
7165 }
7166
7167 pub fn query_traced(
7169 &mut self,
7170 q: &crate::query::Query,
7171 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
7172 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
7173 Ok((result?, trace))
7174 }
7175
7176 pub fn query_columns_native(
7181 &mut self,
7182 conditions: &[crate::query::Condition],
7183 projection: Option<&[u16]>,
7184 snapshot: Snapshot,
7185 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7186 use crate::query::Condition;
7187 if self.ttl.is_some() {
7190 return Ok(None);
7191 }
7192 if conditions.is_empty() {
7193 return Ok(None);
7194 }
7195 self.ensure_indexes_complete()?;
7196
7197 let served = |c: &Condition| {
7202 matches!(
7203 c,
7204 Condition::Pk(_)
7205 | Condition::BitmapEq { .. }
7206 | Condition::BitmapIn { .. }
7207 | Condition::BytesPrefix { .. }
7208 | Condition::FmContains { .. }
7209 | Condition::FmContainsAll { .. }
7210 | Condition::Ann { .. }
7211 | Condition::Range { .. }
7212 | Condition::RangeF64 { .. }
7213 | Condition::SparseMatch { .. }
7214 | Condition::MinHashSimilar { .. }
7215 | Condition::IsNull { .. }
7216 | Condition::IsNotNull { .. }
7217 )
7218 };
7219 if !conditions.iter().all(served) {
7220 return Ok(None);
7221 }
7222 let fast_path =
7223 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
7224 crate::trace::QueryTrace::record(|t| {
7225 t.run_count = self.run_refs.len();
7226 t.memtable_rows = self.memtable.len();
7227 t.mutable_run_rows = self.mutable_run.len();
7228 t.conditions_pushed = conditions.len();
7229 t.learned_range_used = conditions.iter().any(|c| match c {
7230 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
7231 self.learned_range.contains_key(column_id)
7232 }
7233 _ => false,
7234 });
7235 });
7236 let col_ids: Vec<u16> = projection
7238 .map(|p| p.to_vec())
7239 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
7240 let proj_pairs: Vec<(u16, TypeId)> = col_ids
7241 .iter()
7242 .map(|&cid| {
7243 let ty = self
7244 .schema
7245 .columns
7246 .iter()
7247 .find(|c| c.id == cid)
7248 .map(|c| c.ty.clone())
7249 .unwrap_or(TypeId::Bytes);
7250 (cid, ty)
7251 })
7252 .collect();
7253
7254 if fast_path {
7260 let needs_column = conditions.iter().any(|c| match c {
7263 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
7264 Condition::RangeF64 { column_id, .. } => {
7265 !self.learned_range.contains_key(column_id)
7266 }
7267 _ => false,
7268 });
7269 let mut reader_opt: Option<RunReader> = if needs_column {
7270 Some(self.open_reader(self.run_refs[0].run_id)?)
7271 } else {
7272 None
7273 };
7274 let mut sets: Vec<RowIdSet> = Vec::new();
7275 for c in conditions {
7276 let s = match c {
7277 Condition::Range { column_id, lo, hi }
7278 if !self.learned_range.contains_key(column_id) =>
7279 {
7280 if reader_opt.is_none() {
7281 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7282 }
7283 reader_opt
7284 .as_mut()
7285 .expect("reader opened for range")
7286 .range_row_id_set_i64(*column_id, *lo, *hi)?
7287 }
7288 Condition::RangeF64 {
7289 column_id,
7290 lo,
7291 lo_inclusive,
7292 hi,
7293 hi_inclusive,
7294 } if !self.learned_range.contains_key(column_id) => {
7295 if reader_opt.is_none() {
7296 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7297 }
7298 reader_opt
7299 .as_mut()
7300 .expect("reader opened for range")
7301 .range_row_id_set_f64(
7302 *column_id,
7303 *lo,
7304 *lo_inclusive,
7305 *hi,
7306 *hi_inclusive,
7307 )?
7308 }
7309 _ => self.resolve_condition(c, snapshot)?,
7310 };
7311 sets.push(s);
7312 }
7313 let candidates = RowIdSet::intersect_many(sets);
7314 crate::trace::QueryTrace::record(|t| {
7315 t.survivor_count = Some(candidates.len());
7316 });
7317 if candidates.is_empty() {
7318 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
7319 .iter()
7320 .map(|&id| {
7321 (
7322 id,
7323 columnar::null_native(
7324 proj_pairs
7325 .iter()
7326 .find(|(c, _)| c == &id)
7327 .map(|(_, t)| t.clone())
7328 .unwrap_or(TypeId::Bytes),
7329 0,
7330 ),
7331 )
7332 })
7333 .collect();
7334 return Ok(Some(cols));
7335 }
7336 let mut reader = match reader_opt.take() {
7337 Some(r) => r,
7338 None => self.open_reader(self.run_refs[0].run_id)?,
7339 };
7340 let candidate_ids = candidates.into_sorted_vec();
7341 let (positions, fast_rid) = if let Some(positions) =
7342 reader.positions_for_row_ids_fast(&candidate_ids)
7343 {
7344 (positions, true)
7345 } else {
7346 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
7347 match col {
7348 columnar::NativeColumn::Int64 { data, .. } => {
7349 let mut p: Vec<usize> = candidate_ids
7350 .iter()
7351 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
7352 .collect();
7353 p.sort_unstable();
7354 (p, false)
7355 }
7356 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
7357 }
7358 };
7359 crate::trace::QueryTrace::record(|t| {
7360 t.scan_mode = crate::trace::ScanMode::NativePushdown;
7361 t.fast_row_id_map = fast_rid;
7362 });
7363 let mut cols = Vec::with_capacity(col_ids.len());
7364 for cid in &col_ids {
7365 let col = reader.column_native(*cid)?;
7366 cols.push((*cid, col.gather(&positions)));
7367 }
7368 return Ok(Some(cols));
7369 }
7370
7371 if !self.run_refs.is_empty() {
7384 use crate::cursor::{drain_cursor_to_columns, Cursor};
7385 let remaining: usize;
7386 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
7387 let c = self
7388 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
7389 .expect("single-run cursor should build when run_refs.len() == 1");
7390 remaining = c.remaining_rows();
7391 Box::new(c)
7392 } else {
7393 let c = self
7394 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
7395 .expect("multi-run cursor should build when run_refs.len() >= 1");
7396 remaining = c.remaining_rows();
7397 Box::new(c)
7398 };
7399 crate::trace::QueryTrace::record(|t| {
7400 if t.survivor_count.is_none() {
7401 t.survivor_count = Some(remaining);
7402 }
7403 });
7404 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
7405 return Ok(Some(cols));
7406 }
7407
7408 crate::trace::QueryTrace::record(|t| {
7413 t.scan_mode = crate::trace::ScanMode::Materialized;
7414 t.row_materialized = true;
7415 });
7416 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7417 for c in conditions {
7418 sets.push(self.resolve_condition(c, snapshot)?);
7419 }
7420 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7421 let rows = self.rows_for_rids(&rids, snapshot)?;
7422 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
7423 for (cid, ty) in &proj_pairs {
7424 let vals: Vec<Value> = rows
7425 .iter()
7426 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
7427 .collect();
7428 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
7429 }
7430 Ok(Some(cols))
7431 }
7432
7433 pub fn native_page_cursor(
7448 &self,
7449 snapshot: Snapshot,
7450 projection: Vec<(u16, TypeId)>,
7451 conditions: &[crate::query::Condition],
7452 ) -> Result<Option<NativePageCursor>> {
7453 use crate::cursor::build_page_plans;
7454 if self.ttl.is_some() {
7455 return Ok(None);
7456 }
7457 if !conditions.is_empty() && !self.indexes_complete {
7460 return Ok(None);
7461 }
7462 if self.run_refs.len() != 1 {
7463 return Ok(None);
7464 }
7465 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7466 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7467
7468 let overlay_rids: HashSet<u64> = {
7471 let mut s = HashSet::new();
7472 for row in self.memtable.visible_versions(snapshot.epoch) {
7473 s.insert(row.row_id.0);
7474 }
7475 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7476 s.insert(row.row_id.0);
7477 }
7478 s
7479 };
7480
7481 let survivors = if conditions.is_empty() {
7485 None
7486 } else {
7487 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
7488 };
7489
7490 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
7497 survivors.clone()
7498 } else if let Some(s) = &survivors {
7499 let mut run_set = s.clone();
7500 run_set.remove_many(overlay_rids.iter().copied());
7501 Some(run_set)
7502 } else {
7503 Some(RowIdSet::from_unsorted(
7504 rids.iter()
7505 .map(|&r| r as u64)
7506 .filter(|r| !overlay_rids.contains(r))
7507 .collect(),
7508 ))
7509 };
7510
7511 let overlay_rows = if overlay_rids.is_empty() {
7512 Vec::new()
7513 } else {
7514 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7515 self.overlay_visible_rows(snapshot, bound)
7516 };
7517
7518 let plans = if positions.is_empty() {
7520 Vec::new()
7521 } else {
7522 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
7523 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
7524 };
7525
7526 let overlay = if overlay_rows.is_empty() {
7528 None
7529 } else {
7530 let filtered =
7531 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7532 if filtered.is_empty() {
7533 None
7534 } else {
7535 Some(self.materialize_overlay(&filtered, &projection))
7536 }
7537 };
7538
7539 let overlay_row_count = overlay
7540 .as_ref()
7541 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7542 .unwrap_or(0);
7543 crate::trace::QueryTrace::record(|t| {
7544 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
7545 t.run_count = self.run_refs.len();
7546 t.memtable_rows = self.memtable.len();
7547 t.mutable_run_rows = self.mutable_run.len();
7548 t.overlay_rows = overlay_row_count;
7549 t.conditions_pushed = conditions.len();
7550 t.pages_decoded = plans
7551 .iter()
7552 .map(|p| p.positions.len())
7553 .sum::<usize>()
7554 .min(1);
7555 });
7556
7557 Ok(Some(NativePageCursor::new_with_overlay(
7558 reader, projection, plans, overlay,
7559 )))
7560 }
7561 #[allow(clippy::type_complexity)]
7571 pub fn native_multi_run_cursor(
7572 &self,
7573 snapshot: Snapshot,
7574 projection: Vec<(u16, TypeId)>,
7575 conditions: &[crate::query::Condition],
7576 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
7577 use crate::cursor::{MultiRunCursor, RunStream};
7578 use crate::sorted_run::SYS_ROW_ID;
7579 use std::collections::{BinaryHeap, HashMap, HashSet};
7580 if self.ttl.is_some() {
7581 return Ok(None);
7582 }
7583 if !conditions.is_empty() && !self.indexes_complete {
7586 return Ok(None);
7587 }
7588 if self.run_refs.is_empty() {
7589 return Ok(None);
7590 }
7591
7592 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
7594 Vec::with_capacity(self.run_refs.len());
7595 for rr in &self.run_refs {
7596 let mut reader = self.open_reader(rr.run_id)?;
7597 let (rids, eps, del) = reader.system_columns_native()?;
7598 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
7599 run_meta.push((reader, rids, eps, del, page_rows));
7600 }
7601
7602 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
7606 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
7607 for i in 0..rids.len() {
7608 let rid = rids[i] as u64;
7609 let e = eps[i] as u64;
7610 if e > snapshot.epoch.0 {
7611 continue;
7612 }
7613 let is_del = del[i] != 0;
7614 best.entry(rid)
7615 .and_modify(|cur| {
7616 if e > cur.0 {
7617 *cur = (e, run_idx, i, is_del);
7618 }
7619 })
7620 .or_insert((e, run_idx, i, is_del));
7621 }
7622 }
7623
7624 let overlay_rids: HashSet<u64> = {
7626 let mut s = HashSet::new();
7627 for row in self.memtable.visible_versions(snapshot.epoch) {
7628 s.insert(row.row_id.0);
7629 }
7630 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7631 s.insert(row.row_id.0);
7632 }
7633 s
7634 };
7635
7636 let survivors: Option<RowIdSet> = if conditions.is_empty() {
7638 None
7639 } else {
7640 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7641 for c in conditions {
7642 sets.push(self.resolve_condition(c, snapshot)?);
7643 }
7644 Some(RowIdSet::intersect_many(sets))
7645 };
7646
7647 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
7651 for (rid, (_, run_idx, pos, deleted)) in &best {
7652 if *deleted {
7653 continue;
7654 }
7655 if overlay_rids.contains(rid) {
7656 continue;
7657 }
7658 if let Some(s) = &survivors {
7659 if !s.contains(*rid) {
7660 continue;
7661 }
7662 }
7663 per_run[*run_idx].push((*rid, *pos));
7664 }
7665 for v in per_run.iter_mut() {
7666 v.sort_unstable_by_key(|&(rid, _)| rid);
7667 }
7668
7669 let mut streams = Vec::with_capacity(run_meta.len());
7671 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
7672 let mut total = 0usize;
7673 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
7674 let mut starts = Vec::with_capacity(page_rows.len());
7675 let mut acc = 0usize;
7676 for &r in &page_rows {
7677 starts.push(acc);
7678 acc += r;
7679 }
7680 let mut survivors_vec: Vec<(u64, usize, usize)> =
7681 Vec::with_capacity(per_run[run_idx].len());
7682 for &(rid, pos) in &per_run[run_idx] {
7683 let page_seq = match starts.partition_point(|&s| s <= pos) {
7684 0 => continue,
7685 p => p - 1,
7686 };
7687 let within = pos - starts[page_seq];
7688 survivors_vec.push((rid, page_seq, within));
7689 }
7690 total += survivors_vec.len();
7691 if let Some(&(rid, _, _)) = survivors_vec.first() {
7692 heap.push(std::cmp::Reverse((rid, run_idx)));
7693 }
7694 streams.push(RunStream::new(reader, survivors_vec, page_rows));
7695 }
7696
7697 let overlay_rows = if overlay_rids.is_empty() {
7699 Vec::new()
7700 } else {
7701 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7702 self.overlay_visible_rows(snapshot, bound)
7703 };
7704 let overlay = if overlay_rows.is_empty() {
7705 None
7706 } else {
7707 let filtered =
7708 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7709 if filtered.is_empty() {
7710 None
7711 } else {
7712 Some(self.materialize_overlay(&filtered, &projection))
7713 }
7714 };
7715
7716 let overlay_row_count = overlay
7717 .as_ref()
7718 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7719 .unwrap_or(0);
7720 crate::trace::QueryTrace::record(|t| {
7721 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
7722 t.run_count = self.run_refs.len();
7723 t.memtable_rows = self.memtable.len();
7724 t.mutable_run_rows = self.mutable_run.len();
7725 t.overlay_rows = overlay_row_count;
7726 t.conditions_pushed = conditions.len();
7727 t.survivor_count = Some(total);
7728 });
7729
7730 Ok(Some(MultiRunCursor::new(
7731 streams, projection, heap, total, overlay,
7732 )))
7733 }
7734
7735 fn overlay_materialization_bound<'a>(
7747 conditions: &[crate::query::Condition],
7748 survivors: &'a Option<RowIdSet>,
7749 ) -> Option<&'a RowIdSet> {
7750 use crate::query::Condition;
7751 let has_range = conditions
7752 .iter()
7753 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7754 if has_range {
7755 None
7756 } else {
7757 survivors.as_ref()
7758 }
7759 }
7760
7761 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
7773 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
7774 let mut fold = |row: Row| {
7775 if let Some(b) = bound {
7776 if !b.contains(row.row_id.0) {
7777 return;
7778 }
7779 }
7780 best.entry(row.row_id.0)
7781 .and_modify(|(be, br)| {
7782 if row.committed_epoch > *be {
7783 *be = row.committed_epoch;
7784 *br = row.clone();
7785 }
7786 })
7787 .or_insert_with(|| (row.committed_epoch, row));
7788 };
7789 for row in self.memtable.visible_versions(snapshot.epoch) {
7790 fold(row);
7791 }
7792 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7793 fold(row);
7794 }
7795 let mut out: Vec<Row> = best
7796 .into_values()
7797 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
7798 .collect();
7799 out.sort_by_key(|r| r.row_id);
7800 out
7801 }
7802
7803 fn filter_overlay_rows(
7811 &self,
7812 rows: Vec<Row>,
7813 conditions: &[crate::query::Condition],
7814 survivors: Option<&RowIdSet>,
7815 snapshot: Snapshot,
7816 ) -> Result<Vec<Row>> {
7817 if conditions.is_empty() {
7818 return Ok(rows);
7819 }
7820 use crate::query::Condition;
7821 let all_index_served = !conditions
7825 .iter()
7826 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7827 if all_index_served {
7828 return Ok(rows
7829 .into_iter()
7830 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
7831 .collect());
7832 }
7833 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7836 for c in conditions {
7837 let s = match c {
7838 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
7839 _ => self.resolve_condition(c, snapshot)?,
7840 };
7841 per_cond_sets.push(s);
7842 }
7843 Ok(rows
7844 .into_iter()
7845 .filter(|row| {
7846 conditions.iter().enumerate().all(|(i, c)| match c {
7847 Condition::Range { column_id, lo, hi } => {
7848 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
7849 }
7850 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
7851 match row.columns.get(column_id) {
7852 Some(Value::Float64(v)) => {
7853 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
7854 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
7855 lo_ok && hi_ok
7856 }
7857 _ => false,
7858 }
7859 }
7860 _ => per_cond_sets[i].contains(row.row_id.0),
7861 })
7862 })
7863 .collect())
7864 }
7865
7866 fn materialize_overlay(
7869 &self,
7870 rows: &[Row],
7871 projection: &[(u16, TypeId)],
7872 ) -> Vec<columnar::NativeColumn> {
7873 if projection.is_empty() {
7874 return vec![columnar::null_native(TypeId::Int64, rows.len())];
7875 }
7876 let mut cols = Vec::with_capacity(projection.len());
7877 for (cid, ty) in projection {
7878 let vals: Vec<Value> = rows
7879 .iter()
7880 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
7881 .collect();
7882 cols.push(columnar::values_to_native(ty.clone(), &vals));
7883 }
7884 cols
7885 }
7886
7887 fn resolve_survivor_rids(
7892 &self,
7893 conditions: &[crate::query::Condition],
7894 reader: &mut RunReader,
7895 snapshot: Snapshot,
7896 ) -> Result<RowIdSet> {
7897 use crate::query::Condition;
7898 let mut sets: Vec<RowIdSet> = Vec::new();
7899 for c in conditions {
7900 self.validate_condition(c)?;
7901 let s: RowIdSet = match c {
7902 Condition::Pk(key) => {
7903 let lookup = self
7904 .schema
7905 .primary_key()
7906 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7907 .unwrap_or_else(|| key.clone());
7908 self.hot
7909 .get(&lookup)
7910 .map(|r| RowIdSet::one(r.0))
7911 .unwrap_or_else(RowIdSet::empty)
7912 }
7913 Condition::BitmapEq { column_id, value } => {
7914 let lookup = self.index_lookup_key_bytes(*column_id, value);
7915 self.bitmap
7916 .get(column_id)
7917 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7918 .unwrap_or_else(RowIdSet::empty)
7919 }
7920 Condition::BitmapIn { column_id, values } => {
7921 let bm = self.bitmap.get(column_id);
7922 let mut acc = roaring::RoaringBitmap::new();
7923 if let Some(b) = bm {
7924 for v in values {
7925 let lookup = self.index_lookup_key_bytes(*column_id, v);
7926 acc |= b.get(&lookup);
7927 }
7928 }
7929 RowIdSet::from_roaring(acc)
7930 }
7931 Condition::BytesPrefix { column_id, prefix } => {
7932 if let Some(b) = self.bitmap.get(column_id) {
7933 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7934 let mut acc = roaring::RoaringBitmap::new();
7935 for key in b.keys() {
7936 if key.starts_with(&lookup_prefix) {
7937 acc |= b.get(&key);
7938 }
7939 }
7940 RowIdSet::from_roaring(acc)
7941 } else {
7942 RowIdSet::empty()
7943 }
7944 }
7945 Condition::FmContains { column_id, pattern } => self
7946 .fm
7947 .get(column_id)
7948 .map(|f| {
7949 RowIdSet::from_unsorted(
7950 f.locate(pattern).into_iter().map(|r| r.0).collect(),
7951 )
7952 })
7953 .unwrap_or_else(RowIdSet::empty),
7954 Condition::FmContainsAll {
7955 column_id,
7956 patterns,
7957 } => {
7958 if let Some(f) = self.fm.get(column_id) {
7959 let sets: Vec<RowIdSet> = patterns
7960 .iter()
7961 .map(|pat| {
7962 RowIdSet::from_unsorted(
7963 f.locate(pat).into_iter().map(|r| r.0).collect(),
7964 )
7965 })
7966 .collect();
7967 RowIdSet::intersect_many(sets)
7968 } else {
7969 RowIdSet::empty()
7970 }
7971 }
7972 Condition::Ann {
7973 column_id,
7974 query,
7975 k,
7976 } => RowIdSet::from_unsorted(
7977 self.retrieve_filtered(
7978 &crate::query::Retriever::Ann {
7979 column_id: *column_id,
7980 query: query.clone(),
7981 k: *k,
7982 },
7983 snapshot,
7984 None,
7985 None,
7986 None,
7987 None,
7988 )?
7989 .into_iter()
7990 .map(|hit| hit.row_id.0)
7991 .collect(),
7992 ),
7993 Condition::SparseMatch {
7994 column_id,
7995 query,
7996 k,
7997 } => RowIdSet::from_unsorted(
7998 self.retrieve_filtered(
7999 &crate::query::Retriever::Sparse {
8000 column_id: *column_id,
8001 query: query.clone(),
8002 k: *k,
8003 },
8004 snapshot,
8005 None,
8006 None,
8007 None,
8008 None,
8009 )?
8010 .into_iter()
8011 .map(|hit| hit.row_id.0)
8012 .collect(),
8013 ),
8014 Condition::MinHashSimilar {
8015 column_id,
8016 query,
8017 k,
8018 } => match self.minhash.get(column_id) {
8019 Some(index) => {
8020 let candidates = index.candidate_row_ids(query);
8021 let eligible =
8022 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
8023 RowIdSet::from_unsorted(
8024 index
8025 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
8026 .into_iter()
8027 .map(|(row_id, _)| row_id.0)
8028 .collect(),
8029 )
8030 }
8031 None => RowIdSet::empty(),
8032 },
8033 Condition::Range { column_id, lo, hi } => {
8034 if let Some(li) = self.learned_range.get(column_id) {
8035 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
8036 } else {
8037 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
8038 }
8039 }
8040 Condition::RangeF64 {
8041 column_id,
8042 lo,
8043 lo_inclusive,
8044 hi,
8045 hi_inclusive,
8046 } => {
8047 if let Some(li) = self.learned_range.get(column_id) {
8048 RowIdSet::from_unsorted(
8049 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
8050 .into_iter()
8051 .collect(),
8052 )
8053 } else {
8054 reader.range_row_id_set_f64(
8055 *column_id,
8056 *lo,
8057 *lo_inclusive,
8058 *hi,
8059 *hi_inclusive,
8060 )?
8061 }
8062 }
8063 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
8064 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
8065 };
8066 sets.push(s);
8067 }
8068 Ok(RowIdSet::intersect_many(sets))
8069 }
8070
8071 pub fn scan_cursor(
8092 &self,
8093 snapshot: Snapshot,
8094 projection: Vec<(u16, TypeId)>,
8095 conditions: &[crate::query::Condition],
8096 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
8097 if self.ttl.is_some() {
8098 return Ok(None);
8099 }
8100 if !conditions.is_empty() && !self.indexes_complete {
8106 return Ok(None);
8107 }
8108 if self.run_refs.len() == 1 {
8109 Ok(self
8110 .native_page_cursor(snapshot, projection, conditions)?
8111 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
8112 } else {
8113 Ok(self
8114 .native_multi_run_cursor(snapshot, projection, conditions)?
8115 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
8116 }
8117 }
8118
8119 pub fn aggregate_native(
8133 &self,
8134 snapshot: Snapshot,
8135 column: Option<u16>,
8136 conditions: &[crate::query::Condition],
8137 agg: NativeAgg,
8138 ) -> Result<Option<NativeAggResult>> {
8139 if self.ttl.is_some() {
8140 return Ok(None);
8141 }
8142 if self.run_refs.len() == 1 && conditions.is_empty() {
8144 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
8145 return Ok(Some(res));
8146 }
8147 }
8148 if matches!(agg, NativeAgg::Count) && column.is_none() {
8150 return Ok(self
8151 .scan_cursor(snapshot, Vec::new(), conditions)?
8152 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
8153 }
8154 let cid = match column {
8157 Some(c) => c,
8158 None => return Ok(None),
8159 };
8160 let ty = self.column_type(cid);
8161 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
8162 else {
8163 return Ok(None);
8164 };
8165 match ty {
8166 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8167 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
8168 Ok(Some(pack_int(agg, count, sum, mn, mx)))
8169 }
8170 TypeId::Float64 => {
8171 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
8172 Ok(Some(pack_float(agg, count, sum, mn, mx)))
8173 }
8174 _ => Ok(None),
8175 }
8176 }
8177
8178 fn aggregate_from_stats(
8186 &self,
8187 snapshot: Snapshot,
8188 column: Option<u16>,
8189 agg: NativeAgg,
8190 ) -> Result<Option<NativeAggResult>> {
8191 let cid = match (agg, column) {
8192 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
8193 _ => return Ok(None), };
8195 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
8196 return Ok(None);
8197 };
8198 let Some(cs) = stats.get(&cid) else {
8199 return Ok(None);
8200 };
8201 match agg {
8202 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
8204 self.live_count.saturating_sub(cs.null_count),
8205 ))),
8206 NativeAgg::Min | NativeAgg::Max => {
8207 let bound = if agg == NativeAgg::Min {
8208 &cs.min
8209 } else {
8210 &cs.max
8211 };
8212 match bound {
8213 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
8214 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
8215 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
8220 None => Ok(None),
8221 }
8222 }
8223 _ => Ok(None),
8224 }
8225 }
8226
8227 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
8236 if self.ttl.is_some() {
8237 return Ok(None);
8238 }
8239 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
8240 return Ok(None);
8241 }
8242 self.ensure_indexes_complete()?;
8245 let reader = self.open_reader(self.run_refs[0].run_id)?;
8246 if self.live_count != reader.row_count() as u64 {
8247 return Ok(None);
8248 }
8249 let Some(bm) = self.bitmap.get(&column_id) else {
8250 return Ok(None); };
8252 let mut distinct = bm.value_count() as u64;
8253 if !bm.get(&Value::Null.encode_key()).is_empty() {
8256 distinct = distinct.saturating_sub(1);
8257 }
8258 Ok(Some(distinct))
8259 }
8260
8261 pub fn aggregate_incremental(
8273 &mut self,
8274 cache_key: u64,
8275 conditions: &[crate::query::Condition],
8276 column: Option<u16>,
8277 agg: NativeAgg,
8278 ) -> Result<IncrementalAggResult> {
8279 let snap = self.snapshot();
8280 let cur_wm = self.allocator.current().0;
8281 let cur_epoch = snap.epoch.0;
8282 let incremental_ok = self.ttl.is_none()
8289 && !self.had_deletes
8290 && self.memtable.is_empty()
8291 && self.mutable_run.is_empty();
8292
8293 if incremental_ok {
8296 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
8297 if cached.epoch == cur_epoch {
8298 return Ok(IncrementalAggResult {
8299 state: cached.state,
8300 incremental: true,
8301 delta_rows: 0,
8302 });
8303 }
8304 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
8305 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
8306 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
8307 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8308 let delta_state = agg_state_from_rows(
8309 &delta_rows,
8310 conditions,
8311 &index_sets,
8312 column,
8313 agg,
8314 &self.schema,
8315 )?;
8316 let merged = cached.state.merge(delta_state);
8317 let delta_n = delta_rids.len() as u64;
8318 Arc::make_mut(&mut self.agg_cache).insert(
8319 cache_key,
8320 CachedAgg {
8321 state: merged.clone(),
8322 watermark: cur_wm,
8323 epoch: cur_epoch,
8324 },
8325 );
8326 return Ok(IncrementalAggResult {
8327 state: merged,
8328 incremental: true,
8329 delta_rows: delta_n,
8330 });
8331 }
8332 }
8333 }
8334
8335 let cursor_ok =
8340 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8341 let state = if cursor_ok && agg != NativeAgg::Avg {
8342 match self.aggregate_native(snap, column, conditions, agg)? {
8343 Some(result) => {
8344 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
8345 }
8346 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
8347 }
8348 } else {
8349 self.agg_state_full_scan(conditions, column, agg, snap)?
8350 };
8351 if incremental_ok {
8353 Arc::make_mut(&mut self.agg_cache).insert(
8354 cache_key,
8355 CachedAgg {
8356 state: state.clone(),
8357 watermark: cur_wm,
8358 epoch: cur_epoch,
8359 },
8360 );
8361 }
8362 Ok(IncrementalAggResult {
8363 state,
8364 incremental: false,
8365 delta_rows: 0,
8366 })
8367 }
8368
8369 fn agg_state_full_scan(
8372 &self,
8373 conditions: &[crate::query::Condition],
8374 column: Option<u16>,
8375 agg: NativeAgg,
8376 snap: Snapshot,
8377 ) -> Result<AggState> {
8378 let rows = self.visible_rows(snap)?;
8379 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8380 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
8381 }
8382
8383 fn resolve_index_conditions(
8386 &self,
8387 conditions: &[crate::query::Condition],
8388 snapshot: Snapshot,
8389 ) -> Result<Vec<RowIdSet>> {
8390 use crate::query::Condition;
8391 let mut sets = Vec::new();
8392 for c in conditions {
8393 if matches!(
8394 c,
8395 Condition::Ann { .. }
8396 | Condition::SparseMatch { .. }
8397 | Condition::MinHashSimilar { .. }
8398 ) {
8399 sets.push(self.resolve_condition(c, snapshot)?);
8400 }
8401 }
8402 Ok(sets)
8403 }
8404
8405 fn column_type(&self, cid: u16) -> TypeId {
8406 self.schema
8407 .columns
8408 .iter()
8409 .find(|c| c.id == cid)
8410 .map(|c| c.ty.clone())
8411 .unwrap_or(TypeId::Bytes)
8412 }
8413
8414 pub fn approx_aggregate(
8423 &mut self,
8424 conditions: &[crate::query::Condition],
8425 column: Option<u16>,
8426 agg: ApproxAgg,
8427 z: f64,
8428 ) -> Result<Option<ApproxResult>> {
8429 use crate::query::Condition;
8430 self.ensure_reservoir_complete()?;
8431 let snapshot = self.snapshot();
8432 let n_pop = self.count();
8433 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
8434 if sample_rids.is_empty() {
8435 return Ok(None);
8436 }
8437 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
8439 let s = live_sample.len();
8440 if s == 0 {
8441 return Ok(None);
8442 }
8443
8444 let mut index_sets: Vec<RowIdSet> = Vec::new();
8447 for c in conditions {
8448 if matches!(
8449 c,
8450 Condition::Ann { .. }
8451 | Condition::SparseMatch { .. }
8452 | Condition::MinHashSimilar { .. }
8453 ) {
8454 index_sets.push(self.resolve_condition(c, snapshot)?);
8455 }
8456 }
8457
8458 let cid = match (agg, column) {
8460 (ApproxAgg::Count, _) => None,
8461 (_, Some(c)) => Some(c),
8462 _ => return Ok(None),
8463 };
8464 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
8465 for r in &live_sample {
8466 if !conditions
8468 .iter()
8469 .all(|c| condition_matches_row(c, r, &self.schema))
8470 {
8471 continue;
8472 }
8473 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
8475 continue;
8476 }
8477 if let Some(cid) = cid {
8478 if let Some(v) = as_f64(r.columns.get(&cid)) {
8479 passing_vals.push(v);
8480 } } else {
8482 passing_vals.push(0.0); }
8484 }
8485 let m = passing_vals.len();
8486
8487 let (point, half) = match agg {
8488 ApproxAgg::Count => {
8489 let p = m as f64 / s as f64;
8491 let point = n_pop as f64 * p;
8492 let var = if s > 1 {
8493 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
8494 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8495 } else {
8496 0.0
8497 };
8498 (point, z * var.sqrt())
8499 }
8500 ApproxAgg::Sum => {
8501 let y: Vec<f64> = live_sample
8503 .iter()
8504 .map(|r| {
8505 let passes_row = conditions
8506 .iter()
8507 .all(|c| condition_matches_row(c, r, &self.schema))
8508 && index_sets.iter().all(|set| set.contains(r.row_id.0));
8509 if passes_row {
8510 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
8511 } else {
8512 0.0
8513 }
8514 })
8515 .collect();
8516 let mean_y = y.iter().sum::<f64>() / s as f64;
8517 let point = n_pop as f64 * mean_y;
8518 let var = if s > 1 {
8519 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
8520 let var_y = ss / (s - 1) as f64;
8521 n_pop as f64 * n_pop as f64 * var_y / s as f64
8522 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8523 } else {
8524 0.0
8525 };
8526 (point, z * var.sqrt())
8527 }
8528 ApproxAgg::Avg => {
8529 if m == 0 {
8530 return Ok(Some(ApproxResult {
8531 point: 0.0,
8532 ci_low: 0.0,
8533 ci_high: 0.0,
8534 n_population: n_pop,
8535 n_sample_live: s,
8536 n_passing: 0,
8537 }));
8538 }
8539 let mean = passing_vals.iter().sum::<f64>() / m as f64;
8540 let half = if m > 1 {
8541 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
8542 let sd = (ss / (m - 1) as f64).sqrt();
8543 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
8544 z * sd / (m as f64).sqrt() * fpc.sqrt()
8545 } else {
8546 0.0
8547 };
8548 (mean, half)
8549 }
8550 };
8551
8552 Ok(Some(ApproxResult {
8553 point,
8554 ci_low: point - half,
8555 ci_high: point + half,
8556 n_population: n_pop,
8557 n_sample_live: s,
8558 n_passing: m,
8559 }))
8560 }
8561
8562 pub fn exact_column_stats(
8570 &self,
8571 _snapshot: Snapshot,
8572 projection: &[u16],
8573 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
8574 if self.ttl.is_some()
8575 || !(self.memtable.is_empty()
8576 && self.mutable_run.is_empty()
8577 && self.run_refs.len() == 1)
8578 {
8579 return Ok(None);
8580 }
8581 let reader = self.open_reader(self.run_refs[0].run_id)?;
8582 if self.live_count != reader.row_count() as u64 {
8583 return Ok(None);
8584 }
8585 let mut out = HashMap::new();
8586 for &cid in projection {
8587 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
8588 Some(c) => c,
8589 None => continue,
8590 };
8591 let Some(stats) = reader.column_page_stats(cid) else {
8593 out.insert(
8594 cid,
8595 ColumnStat {
8596 min: None,
8597 max: None,
8598 null_count: self.live_count,
8599 },
8600 );
8601 continue;
8602 };
8603 let stat = match cdef.ty {
8604 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8605 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
8606 min: mn.map(Value::Int64),
8607 max: mx.map(Value::Int64),
8608 null_count: n,
8609 })
8610 }
8611 TypeId::Float64 => {
8612 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
8613 min: mn.map(Value::Float64),
8614 max: mx.map(Value::Float64),
8615 null_count: n,
8616 })
8617 }
8618 _ => None,
8619 };
8620 if let Some(s) = stat {
8621 out.insert(cid, s);
8622 }
8623 }
8624 Ok(Some(out))
8625 }
8626
8627 pub fn dir(&self) -> &Path {
8628 &self.dir
8629 }
8630
8631 pub fn schema(&self) -> &Schema {
8632 &self.schema
8633 }
8634
8635 pub(crate) fn set_catalog_name(&mut self, name: String) {
8636 self.name = name;
8637 }
8638
8639 pub(crate) fn prepare_alter_column(
8640 &mut self,
8641 column_name: &str,
8642 change: &AlterColumn,
8643 ) -> Result<ColumnDef> {
8644 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
8645 return Err(MongrelError::InvalidArgument(
8646 "ALTER COLUMN requires committing staged writes first".into(),
8647 ));
8648 }
8649 let old = self
8650 .schema
8651 .columns
8652 .iter()
8653 .find(|c| c.name == column_name)
8654 .cloned()
8655 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
8656 let mut next = old.clone();
8657
8658 if let Some(name) = &change.name {
8659 let trimmed = name.trim();
8660 if trimmed.is_empty() {
8661 return Err(MongrelError::InvalidArgument(
8662 "ALTER COLUMN name must not be empty".into(),
8663 ));
8664 }
8665 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
8666 return Err(MongrelError::Schema(format!(
8667 "column {trimmed} already exists"
8668 )));
8669 }
8670 next.name = trimmed.to_string();
8671 }
8672
8673 if let Some(ty) = &change.ty {
8674 next.ty = ty.clone();
8675 }
8676 if let Some(flags) = change.flags {
8677 validate_alter_column_flags(old.flags, flags)?;
8678 next.flags = flags;
8679 }
8680
8681 if let Some(default_change) = &change.default_value {
8682 next.default_value = default_change.clone();
8683 }
8684
8685 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
8686 if old.flags.contains(ColumnFlags::NULLABLE)
8687 && !next.flags.contains(ColumnFlags::NULLABLE)
8688 && self.column_has_nulls(old.id)?
8689 {
8690 return Err(MongrelError::InvalidArgument(format!(
8691 "column '{}' contains NULL values",
8692 old.name
8693 )));
8694 }
8695 Ok(next)
8696 }
8697
8698 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
8699 let idx = self
8700 .schema
8701 .columns
8702 .iter()
8703 .position(|c| c.id == column.id)
8704 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
8705 if self.schema.columns[idx] == column {
8706 return Ok(());
8707 }
8708 self.schema.columns[idx] = column;
8709 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8710 self.schema.validate_auto_increment()?;
8711 self.schema.validate_defaults()?;
8712 self.auto_inc = resolve_auto_inc(&self.schema);
8713 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
8714 write_schema(&self.dir, &self.schema)?;
8715 self.clear_result_cache();
8716 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8717 self.persist_manifest(self.current_epoch())?;
8718 Ok(())
8719 }
8720
8721 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
8722 self.ensure_writable()?;
8723 let column = self.prepare_alter_column(column_name, &change)?;
8724 self.apply_altered_column(column.clone())?;
8725 Ok(column)
8726 }
8727
8728 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
8729 if self.live_count == 0 {
8730 return Ok(false);
8731 }
8732 let snap = self.snapshot();
8733 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
8734 Ok(columns
8735 .first()
8736 .map(|(_, col)| col.null_count(col.len()) != 0)
8737 .unwrap_or(true))
8738 }
8739
8740 fn has_stored_versions(&self) -> bool {
8741 !self.memtable.is_empty()
8742 || !self.mutable_run.is_empty()
8743 || self.run_refs.iter().any(|r| r.row_count > 0)
8744 || !self.retiring.is_empty()
8745 }
8746
8747 pub fn add_column(
8752 &mut self,
8753 name: &str,
8754 ty: TypeId,
8755 flags: ColumnFlags,
8756 default_value: Option<crate::schema::DefaultExpr>,
8757 ) -> Result<u16> {
8758 self.add_column_with_id(name, ty, flags, default_value, None)
8759 }
8760
8761 pub fn add_column_with_id(
8762 &mut self,
8763 name: &str,
8764 ty: TypeId,
8765 flags: ColumnFlags,
8766 default_value: Option<crate::schema::DefaultExpr>,
8767 requested_id: Option<u16>,
8768 ) -> Result<u16> {
8769 self.ensure_writable()?;
8770 if self.schema.columns.iter().any(|c| c.name == name) {
8771 return Err(MongrelError::Schema(format!(
8772 "column {name} already exists"
8773 )));
8774 }
8775 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
8776 if self.schema.columns.iter().any(|c| c.id == id) {
8777 return Err(MongrelError::Schema(format!(
8778 "column id {id} already exists"
8779 )));
8780 }
8781 id
8782 } else {
8783 self.schema
8784 .columns
8785 .iter()
8786 .map(|c| c.id)
8787 .max()
8788 .unwrap_or(0)
8789 .checked_add(1)
8790 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
8791 };
8792 self.schema.columns.push(ColumnDef {
8793 id,
8794 name: name.to_string(),
8795 ty,
8796 flags,
8797 default_value,
8798 });
8799 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8800 self.schema.validate_auto_increment()?;
8801 self.schema.validate_defaults()?;
8802 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
8803 self.auto_inc = resolve_auto_inc(&self.schema);
8804 }
8805 write_schema(&self.dir, &self.schema)?;
8806 self.clear_result_cache();
8807 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8809 self.persist_manifest(self.current_epoch())?;
8810 Ok(id)
8811 }
8812
8813 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
8822 self.ensure_writable()?;
8823 let cid = self
8824 .schema
8825 .columns
8826 .iter()
8827 .find(|c| c.name == column_name)
8828 .map(|c| c.id)
8829 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
8830 let ty = self
8831 .schema
8832 .columns
8833 .iter()
8834 .find(|c| c.id == cid)
8835 .map(|c| c.ty.clone())
8836 .unwrap_or(TypeId::Int64);
8837 if !matches!(
8838 ty,
8839 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
8840 ) {
8841 return Err(MongrelError::Schema(format!(
8842 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
8843 )));
8844 }
8845 if self
8846 .schema
8847 .indexes
8848 .iter()
8849 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
8850 {
8851 return Ok(()); }
8853 self.schema.indexes.push(IndexDef {
8854 name: format!("{}_learned_range", column_name),
8855 column_id: cid,
8856 kind: IndexKind::LearnedRange,
8857 predicate: None,
8858 options: Default::default(),
8859 });
8860 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8861 write_schema(&self.dir, &self.schema)?;
8862 self.build_learned_ranges()?;
8863 Ok(())
8864 }
8865
8866 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
8869 self.sync_byte_threshold = threshold;
8870 if let WalSink::Private(w) = &mut self.wal {
8871 w.set_sync_byte_threshold(threshold);
8872 }
8873 }
8874
8875 pub fn page_cache_flush(&self) {
8879 self.page_cache.flush_to_disk();
8880 }
8881
8882 pub fn page_cache_len(&self) -> usize {
8884 self.page_cache.len()
8885 }
8886
8887 pub fn decoded_cache_len(&self) -> usize {
8890 self.decoded_cache.len()
8891 }
8892
8893 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
8896 self.memtable.drain_sorted()
8897 }
8898
8899 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
8900 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
8901 }
8902
8903 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
8904 self.run_refs.iter().map(|run| run.run_id)
8905 }
8906
8907 pub(crate) fn table_dir(&self) -> &Path {
8908 &self.dir
8909 }
8910
8911 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
8912 &self.schema
8913 }
8914
8915 pub(crate) fn alloc_run_id(&mut self) -> u64 {
8916 let id = self.next_run_id;
8917 self.next_run_id += 1;
8918 id
8919 }
8920
8921 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
8922 self.run_refs.push(run_ref);
8923 }
8924
8925 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
8935 self.retiring.push(crate::manifest::RetiredRun {
8936 run_id,
8937 retire_epoch,
8938 });
8939 }
8940
8941 pub(crate) fn reap_retiring(
8945 &mut self,
8946 min_active: Epoch,
8947 backup_pinned: &std::collections::HashSet<u128>,
8948 ) -> Result<usize> {
8949 if self.retiring.is_empty() {
8950 return Ok(0);
8951 }
8952 let mut reaped = 0;
8953 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
8954 for r in std::mem::take(&mut self.retiring) {
8960 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
8961 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
8962 reaped += 1;
8963 } else {
8964 kept.push(r);
8965 }
8966 }
8967 self.retiring = kept;
8968 if reaped > 0 {
8969 self.persist_manifest(self.current_epoch())?;
8970 }
8971 Ok(reaped)
8972 }
8973
8974 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
8975 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
8976 return false;
8977 }
8978 self.live_count = self.live_count.saturating_add(run_ref.row_count);
8979 self.run_refs.push(run_ref);
8980 self.indexes_complete = false;
8981 true
8982 }
8983
8984 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
8985 self.kek.as_ref()
8986 }
8987
8988 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
8989 let mut reader = RunReader::open_with_cache(
8990 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
8991 self.schema.clone(),
8992 self.kek.clone(),
8993 Some(self.page_cache.clone()),
8994 Some(self.decoded_cache.clone()),
8995 self.table_id,
8996 Some(&self.verified_runs),
8997 )?;
8998 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
9002 reader.set_uniform_epoch(Epoch(rr.epoch_created));
9003 }
9004 Ok(reader)
9005 }
9006
9007 pub(crate) fn run_refs(&self) -> &[RunRef] {
9008 &self.run_refs
9009 }
9010
9011 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
9012 self.retiring.iter().map(|run| run.run_id)
9013 }
9014
9015 pub(crate) fn runs_dir(&self) -> PathBuf {
9016 self.dir.join(RUNS_DIR)
9017 }
9018
9019 pub(crate) fn wal_dir(&self) -> PathBuf {
9020 self.dir.join(WAL_DIR)
9021 }
9022
9023 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
9024 self.run_refs = refs;
9025 }
9026
9027 pub(crate) fn next_run_id(&self) -> u64 {
9028 self.next_run_id
9029 }
9030
9031 pub(crate) fn compaction_zstd_level(&self) -> i32 {
9032 self.compaction_zstd_level
9033 }
9034
9035 pub(crate) fn bump_next_run_id(&mut self) {
9036 self.next_run_id += 1;
9037 }
9038
9039 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
9040 self.kek.clone()
9041 }
9042
9043 #[cfg(feature = "encryption")]
9047 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
9048 self.kek.as_ref().map(|k| k.derive_idx_key())
9049 }
9050
9051 #[cfg(not(feature = "encryption"))]
9052 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
9053 None
9054 }
9055
9056 #[cfg(feature = "encryption")]
9060 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
9061 self.kek.as_ref().map(|k| *k.derive_meta_key())
9062 }
9063
9064 #[cfg(not(feature = "encryption"))]
9065 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
9066 None
9067 }
9068
9069 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
9072 self.column_keys
9073 .iter()
9074 .map(|(&id, &(_, scheme))| (id, scheme))
9075 .collect()
9076 }
9077
9078 #[cfg(feature = "encryption")]
9083 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
9084 self.tokenize_value_enc(column_id, v)
9085 }
9086
9087 #[cfg(feature = "encryption")]
9088 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
9089 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
9090 let (key, scheme) = self.column_keys.get(&column_id)?;
9091 let token: Vec<u8> = match (*scheme, v) {
9092 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
9093 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
9094 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
9095 _ => hmac_token(key, &v.encode_key()).to_vec(),
9096 };
9097 Some(Value::Bytes(token))
9098 }
9099
9100 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
9102 self.index_lookup_key_bytes(column_id, &v.encode_key())
9103 }
9104
9105 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
9108 #[cfg(feature = "encryption")]
9109 {
9110 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
9111 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
9112 if *scheme == SCHEME_HMAC_EQ {
9113 return hmac_token(key, encoded).to_vec();
9114 }
9115 }
9116 }
9117 let _ = column_id;
9118 encoded.to_vec()
9119 }
9120}
9121
9122fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
9123 let columnar::NativeColumn::Int64 { data, validity } = col else {
9124 return false;
9125 };
9126 if data.len() < n || !columnar::all_non_null(validity, n) {
9127 return false;
9128 }
9129 data.iter()
9130 .take(n)
9131 .zip(data.iter().skip(1))
9132 .all(|(a, b)| a < b)
9133}
9134
9135#[derive(Debug, Clone)]
9139pub struct ColumnStat {
9140 pub min: Option<Value>,
9141 pub max: Option<Value>,
9142 pub null_count: u64,
9143}
9144
9145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9147pub enum NativeAgg {
9148 Count,
9149 Sum,
9150 Min,
9151 Max,
9152 Avg,
9153}
9154
9155#[derive(Debug, Clone, PartialEq)]
9157pub enum NativeAggResult {
9158 Count(u64),
9159 Int(i64),
9160 Float(f64),
9161 Null,
9163}
9164
9165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9167pub enum ApproxAgg {
9168 Count,
9169 Sum,
9170 Avg,
9171}
9172
9173#[derive(Debug, Clone)]
9177pub struct ApproxResult {
9178 pub point: f64,
9180 pub ci_low: f64,
9182 pub ci_high: f64,
9184 pub n_population: u64,
9186 pub n_sample_live: usize,
9188 pub n_passing: usize,
9190}
9191
9192#[derive(Debug, Clone, PartialEq)]
9197pub enum AggState {
9198 Count(u64),
9200 SumI {
9202 sum: i128,
9203 count: u64,
9204 },
9205 SumF {
9207 sum: f64,
9208 count: u64,
9209 },
9210 AvgI {
9212 sum: i128,
9213 count: u64,
9214 },
9215 AvgF {
9217 sum: f64,
9218 count: u64,
9219 },
9220 MinI(i64),
9222 MaxI(i64),
9223 MinF(f64),
9225 MaxF(f64),
9226 Empty,
9228}
9229
9230impl AggState {
9231 pub fn merge(self, other: AggState) -> AggState {
9233 use AggState::*;
9234 match (self, other) {
9235 (Empty, x) | (x, Empty) => x,
9236 (Count(a), Count(b)) => Count(a + b),
9237 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
9238 sum: sa + sb,
9239 count: ca + cb,
9240 },
9241 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
9242 sum: sa + sb,
9243 count: ca + cb,
9244 },
9245 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
9246 sum: sa + sb,
9247 count: ca + cb,
9248 },
9249 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
9250 sum: sa + sb,
9251 count: ca + cb,
9252 },
9253 (MinI(a), MinI(b)) => MinI(a.min(b)),
9254 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
9255 (MinF(a), MinF(b)) => MinF(a.min(b)),
9256 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
9257 _ => Empty, }
9259 }
9260
9261 pub fn point(&self) -> Option<f64> {
9263 match self {
9264 AggState::Count(n) => Some(*n as f64),
9265 AggState::SumI { sum, .. } => Some(*sum as f64),
9266 AggState::SumF { sum, .. } => Some(*sum),
9267 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
9268 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
9269 AggState::MinI(n) => Some(*n as f64),
9270 AggState::MaxI(n) => Some(*n as f64),
9271 AggState::MinF(n) => Some(*n),
9272 AggState::MaxF(n) => Some(*n),
9273 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
9274 }
9275 }
9276
9277 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
9281 let is_float = matches!(ty, Some(TypeId::Float64));
9282 match (agg, result) {
9283 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
9284 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
9285 sum: x as i128,
9286 count: 1, },
9288 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
9289 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
9290 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
9291 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
9292 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
9293 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
9294 (NativeAgg::Count, _) => AggState::Empty,
9295 (_, NativeAggResult::Null) => AggState::Empty,
9296 _ => {
9297 let _ = is_float;
9298 AggState::Empty
9299 }
9300 }
9301 }
9302}
9303
9304#[derive(Debug, Clone)]
9307pub struct CachedAgg {
9308 pub state: AggState,
9309 pub watermark: u64,
9310 pub epoch: u64,
9311}
9312
9313#[derive(Debug, Clone)]
9315pub struct IncrementalAggResult {
9316 pub state: AggState,
9318 pub incremental: bool,
9321 pub delta_rows: u64,
9323}
9324
9325fn agg_state_from_rows(
9329 rows: &[Row],
9330 conditions: &[crate::query::Condition],
9331 index_sets: &[RowIdSet],
9332 column: Option<u16>,
9333 agg: NativeAgg,
9334 schema: &Schema,
9335) -> Result<AggState> {
9336 let mut count: u64 = 0;
9337 let mut sum_i: i128 = 0;
9338 let mut sum_f: f64 = 0.0;
9339 let mut mn_i: i64 = i64::MAX;
9340 let mut mx_i: i64 = i64::MIN;
9341 let mut mn_f: f64 = f64::INFINITY;
9342 let mut mx_f: f64 = f64::NEG_INFINITY;
9343 let mut saw_int = false;
9344 let mut saw_float = false;
9345 for r in rows {
9346 if !conditions
9347 .iter()
9348 .all(|c| condition_matches_row(c, r, schema))
9349 {
9350 continue;
9351 }
9352 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
9353 continue;
9354 }
9355 match agg {
9356 NativeAgg::Count => match column {
9357 None => count += 1,
9359 Some(cid) => match r.columns.get(&cid) {
9362 None | Some(Value::Null) => {}
9363 Some(_) => count += 1,
9364 },
9365 },
9366 _ => match column.and_then(|cid| r.columns.get(&cid)) {
9367 Some(Value::Int64(n)) => {
9368 count += 1;
9369 sum_i += *n as i128;
9370 mn_i = mn_i.min(*n);
9371 mx_i = mx_i.max(*n);
9372 saw_int = true;
9373 }
9374 Some(Value::Float64(f)) => {
9375 count += 1;
9376 sum_f += f;
9377 mn_f = mn_f.min(*f);
9378 mx_f = mx_f.max(*f);
9379 saw_float = true;
9380 }
9381 _ => {}
9382 },
9383 }
9384 }
9385 Ok(match agg {
9386 NativeAgg::Count => {
9387 if count == 0 {
9388 AggState::Empty
9389 } else {
9390 AggState::Count(count)
9391 }
9392 }
9393 NativeAgg::Sum => {
9394 if count == 0 {
9395 AggState::Empty
9396 } else if saw_int {
9397 AggState::SumI { sum: sum_i, count }
9398 } else {
9399 AggState::SumF { sum: sum_f, count }
9400 }
9401 }
9402 NativeAgg::Avg => {
9403 if count == 0 {
9404 AggState::Empty
9405 } else if saw_int {
9406 AggState::AvgI { sum: sum_i, count }
9407 } else {
9408 AggState::AvgF { sum: sum_f, count }
9409 }
9410 }
9411 NativeAgg::Min => {
9412 if !saw_int && !saw_float {
9413 AggState::Empty
9414 } else if saw_int {
9415 AggState::MinI(mn_i)
9416 } else {
9417 AggState::MinF(mn_f)
9418 }
9419 }
9420 NativeAgg::Max => {
9421 if !saw_int && !saw_float {
9422 AggState::Empty
9423 } else if saw_int {
9424 AggState::MaxI(mx_i)
9425 } else {
9426 AggState::MaxF(mx_f)
9427 }
9428 }
9429 })
9430}
9431
9432fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
9436 use crate::query::Condition;
9437 match c {
9438 Condition::Pk(key) => match schema.primary_key() {
9439 Some(pk) => row
9440 .columns
9441 .get(&pk.id)
9442 .map(|v| v.encode_key() == *key)
9443 .unwrap_or(false),
9444 None => false,
9445 },
9446 Condition::BitmapEq { column_id, value } => row
9447 .columns
9448 .get(column_id)
9449 .map(|v| v.encode_key() == *value)
9450 .unwrap_or(false),
9451 Condition::BitmapIn { column_id, values } => {
9452 let key = row.columns.get(column_id).map(|v| v.encode_key());
9453 match key {
9454 Some(k) => values.contains(&k),
9455 None => false,
9456 }
9457 }
9458 Condition::BytesPrefix { column_id, prefix } => row
9459 .columns
9460 .get(column_id)
9461 .map(|v| v.encode_key().starts_with(prefix))
9462 .unwrap_or(false),
9463 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
9464 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
9465 _ => false,
9466 },
9467 Condition::RangeF64 {
9468 column_id,
9469 lo,
9470 lo_inclusive,
9471 hi,
9472 hi_inclusive,
9473 } => match row.columns.get(column_id) {
9474 Some(Value::Float64(n)) => {
9475 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
9476 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
9477 lo_ok && hi_ok
9478 }
9479 _ => false,
9480 },
9481 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
9482 Some(Value::Bytes(b)) => {
9483 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
9484 }
9485 _ => false,
9486 },
9487 Condition::FmContainsAll {
9488 column_id,
9489 patterns,
9490 } => match row.columns.get(column_id) {
9491 Some(Value::Bytes(b)) => patterns
9492 .iter()
9493 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
9494 _ => false,
9495 },
9496 Condition::Ann { .. }
9497 | Condition::SparseMatch { .. }
9498 | Condition::MinHashSimilar { .. } => true,
9499 Condition::IsNull { column_id } => {
9500 matches!(row.columns.get(column_id), Some(Value::Null) | None)
9501 }
9502 Condition::IsNotNull { column_id } => {
9503 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
9504 }
9505 }
9506}
9507
9508fn as_f64(v: Option<&Value>) -> Option<f64> {
9510 match v {
9511 Some(Value::Int64(n)) => Some(*n as f64),
9512 Some(Value::Float64(f)) => Some(*f),
9513 _ => None,
9514 }
9515}
9516
9517fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
9521 let mut count: u64 = 0;
9522 let mut sum: i128 = 0;
9523 let mut mn: i64 = i64::MAX;
9524 let mut mx: i64 = i64::MIN;
9525 while let Some(cols) = cursor.next_batch()? {
9526 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
9527 if crate::columnar::all_non_null(validity, data.len()) {
9528 count += data.len() as u64;
9530 sum += data.iter().map(|&v| v as i128).sum::<i128>();
9531 mn = mn.min(*data.iter().min().unwrap_or(&mn));
9532 mx = mx.max(*data.iter().max().unwrap_or(&mx));
9533 } else {
9534 for (i, &v) in data.iter().enumerate() {
9535 if crate::columnar::validity_bit(validity, i) {
9536 count += 1;
9537 sum += v as i128;
9538 mn = mn.min(v);
9539 mx = mx.max(v);
9540 }
9541 }
9542 }
9543 }
9544 }
9545 Ok((count, sum, mn, mx))
9546}
9547
9548fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
9550 let mut count: u64 = 0;
9551 let mut sum: f64 = 0.0;
9552 let mut mn: f64 = f64::INFINITY;
9553 let mut mx: f64 = f64::NEG_INFINITY;
9554 while let Some(cols) = cursor.next_batch()? {
9555 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
9556 if crate::columnar::all_non_null(validity, data.len()) {
9557 count += data.len() as u64;
9558 sum += data.iter().sum::<f64>();
9559 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
9560 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
9561 } else {
9562 for (i, &v) in data.iter().enumerate() {
9563 if crate::columnar::validity_bit(validity, i) {
9564 count += 1;
9565 sum += v;
9566 mn = mn.min(v);
9567 mx = mx.max(v);
9568 }
9569 }
9570 }
9571 }
9572 }
9573 Ok((count, sum, mn, mx))
9574}
9575
9576fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
9577 if count == 0 && !matches!(agg, NativeAgg::Count) {
9578 return NativeAggResult::Null;
9579 }
9580 match agg {
9581 NativeAgg::Count => NativeAggResult::Count(count),
9582 NativeAgg::Sum => match sum.try_into() {
9585 Ok(v) => NativeAggResult::Int(v),
9586 Err(_) => NativeAggResult::Null,
9587 },
9588 NativeAgg::Min => NativeAggResult::Int(mn),
9589 NativeAgg::Max => NativeAggResult::Int(mx),
9590 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
9591 }
9592}
9593
9594fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
9595 if count == 0 && !matches!(agg, NativeAgg::Count) {
9596 return NativeAggResult::Null;
9597 }
9598 match agg {
9599 NativeAgg::Count => NativeAggResult::Count(count),
9600 NativeAgg::Sum => NativeAggResult::Float(sum),
9601 NativeAgg::Min => NativeAggResult::Float(mn),
9602 NativeAgg::Max => NativeAggResult::Float(mx),
9603 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
9604 }
9605}
9606
9607fn agg_int(
9610 stats: &[crate::page::PageStat],
9611 decode: fn(Option<&[u8]>) -> Option<i64>,
9612) -> Option<(Option<i64>, Option<i64>, u64)> {
9613 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
9614 let mut any = false;
9615 for s in stats {
9616 if let Some(v) = decode(s.min.as_deref()) {
9617 mn = mn.min(v);
9618 any = true;
9619 }
9620 if let Some(v) = decode(s.max.as_deref()) {
9621 mx = mx.max(v);
9622 any = true;
9623 }
9624 nulls += s.null_count;
9625 }
9626 any.then_some((Some(mn), Some(mx), nulls))
9627}
9628
9629fn agg_float(
9631 stats: &[crate::page::PageStat],
9632 decode: fn(Option<&[u8]>) -> Option<f64>,
9633) -> Option<(Option<f64>, Option<f64>, u64)> {
9634 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
9635 let mut any = false;
9636 for s in stats {
9637 if let Some(v) = decode(s.min.as_deref()) {
9638 mn = mn.min(v);
9639 any = true;
9640 }
9641 if let Some(v) = decode(s.max.as_deref()) {
9642 mx = mx.max(v);
9643 any = true;
9644 }
9645 nulls += s.null_count;
9646 }
9647 any.then_some((Some(mn), Some(mx), nulls))
9648}
9649
9650type SecondaryIndexes = (
9652 HashMap<u16, BitmapIndex>,
9653 HashMap<u16, AnnIndex>,
9654 HashMap<u16, FmIndex>,
9655 HashMap<u16, SparseIndex>,
9656 HashMap<u16, MinHashIndex>,
9657);
9658
9659fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
9660 let mut bitmap = HashMap::new();
9661 let mut ann = HashMap::new();
9662 let mut fm = HashMap::new();
9663 let mut sparse = HashMap::new();
9664 let mut minhash = HashMap::new();
9665 for idef in &schema.indexes {
9666 match idef.kind {
9667 IndexKind::Bitmap => {
9668 bitmap.insert(idef.column_id, BitmapIndex::new());
9669 }
9670 IndexKind::Ann => {
9671 let dim = schema
9672 .columns
9673 .iter()
9674 .find(|c| c.id == idef.column_id)
9675 .and_then(|c| match c.ty {
9676 TypeId::Embedding { dim } => Some(dim as usize),
9677 _ => None,
9678 })
9679 .unwrap_or(0);
9680 let options = idef.options.ann.clone().unwrap_or_default();
9681 ann.insert(
9682 idef.column_id,
9683 AnnIndex::with_options(
9684 dim,
9685 options.m,
9686 options.ef_construction,
9687 options.ef_search,
9688 ),
9689 );
9690 }
9691 IndexKind::FmIndex => {
9692 fm.insert(idef.column_id, FmIndex::new());
9693 }
9694 IndexKind::Sparse => {
9695 sparse.insert(idef.column_id, SparseIndex::new());
9696 }
9697 IndexKind::MinHash => {
9698 let options = idef.options.minhash.clone().unwrap_or_default();
9699 minhash.insert(
9700 idef.column_id,
9701 MinHashIndex::with_options(options.permutations, options.bands),
9702 );
9703 }
9704 _ => {}
9705 }
9706 }
9707 (bitmap, ann, fm, sparse, minhash)
9708}
9709
9710const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
9711 | ColumnFlags::AUTO_INCREMENT
9712 | ColumnFlags::ENCRYPTED
9713 | ColumnFlags::ENCRYPTED_INDEXABLE
9714 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
9715
9716fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
9717 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
9718 return Err(MongrelError::Schema(
9719 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
9720 ));
9721 }
9722 Ok(())
9723}
9724
9725fn validate_alter_column_type(
9726 schema: &Schema,
9727 old: &ColumnDef,
9728 next: &ColumnDef,
9729 has_stored_versions: bool,
9730) -> Result<()> {
9731 if old.ty == next.ty {
9732 return Ok(());
9733 }
9734 if schema.indexes.iter().any(|i| i.column_id == old.id) {
9735 return Err(MongrelError::Schema(format!(
9736 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
9737 old.name
9738 )));
9739 }
9740 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
9741 return Ok(());
9742 }
9743 Err(MongrelError::Schema(format!(
9744 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
9745 old.ty, next.ty
9746 )))
9747}
9748
9749fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
9750 matches!(
9751 (old, new),
9752 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
9753 )
9754}
9755
9756fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
9762 let mut prev: Option<i64> = None;
9763 for r in rows {
9764 match r.columns.get(&pk_id) {
9765 Some(Value::Int64(v)) => {
9766 if prev.is_some_and(|p| p >= *v) {
9767 return false;
9768 }
9769 prev = Some(*v);
9770 }
9771 _ => return false,
9772 }
9773 }
9774 true
9775}
9776
9777#[allow(clippy::too_many_arguments)]
9778fn index_into(
9779 schema: &Schema,
9780 row: &Row,
9781 hot: &mut HotIndex,
9782 bitmap: &mut HashMap<u16, BitmapIndex>,
9783 ann: &mut HashMap<u16, AnnIndex>,
9784 fm: &mut HashMap<u16, FmIndex>,
9785 sparse: &mut HashMap<u16, SparseIndex>,
9786 minhash: &mut HashMap<u16, MinHashIndex>,
9787) {
9788 for idef in &schema.indexes {
9789 let Some(val) = row.columns.get(&idef.column_id) else {
9790 continue;
9791 };
9792 match idef.kind {
9793 IndexKind::Bitmap => {
9794 if let Some(b) = bitmap.get_mut(&idef.column_id) {
9795 b.insert(val.encode_key(), row.row_id);
9796 }
9797 }
9798 IndexKind::Ann => {
9799 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
9800 a.insert_validated(v, row.row_id);
9801 }
9802 }
9803 IndexKind::FmIndex => {
9804 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
9805 f.insert(b.clone(), row.row_id);
9806 }
9807 }
9808 IndexKind::Sparse => {
9809 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
9810 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
9813 s.insert(&terms, row.row_id);
9814 }
9815 }
9816 }
9817 IndexKind::MinHash => {
9818 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
9819 let tokens = crate::index::token_hashes_from_bytes(b);
9822 mh.insert(&tokens, row.row_id);
9823 }
9824 }
9825 _ => {}
9826 }
9827 }
9828 if let Some(pk_col) = schema.primary_key() {
9829 if let Some(pk_val) = row.columns.get(&pk_col.id) {
9830 hot.insert(pk_val.encode_key(), row.row_id);
9831 }
9832 }
9833}
9834
9835#[allow(clippy::too_many_arguments)]
9838fn index_into_single(
9839 idef: &IndexDef,
9840 _schema: &Schema,
9841 row: &Row,
9842 _hot: &mut HotIndex,
9843 bitmap: &mut HashMap<u16, BitmapIndex>,
9844 ann: &mut HashMap<u16, AnnIndex>,
9845 fm: &mut HashMap<u16, FmIndex>,
9846 sparse: &mut HashMap<u16, SparseIndex>,
9847 minhash: &mut HashMap<u16, MinHashIndex>,
9848) {
9849 let Some(val) = row.columns.get(&idef.column_id) else {
9850 return;
9851 };
9852 match idef.kind {
9853 IndexKind::Bitmap => {
9854 if let Some(b) = bitmap.get_mut(&idef.column_id) {
9855 b.insert(val.encode_key(), row.row_id);
9856 }
9857 }
9858 IndexKind::Ann => {
9859 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
9860 a.insert_validated(v, row.row_id);
9861 }
9862 }
9863 IndexKind::FmIndex => {
9864 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
9865 f.insert(b.clone(), row.row_id);
9866 }
9867 }
9868 IndexKind::Sparse => {
9869 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
9870 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
9871 s.insert(&terms, row.row_id);
9872 }
9873 }
9874 }
9875 IndexKind::MinHash => {
9876 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
9877 let tokens = crate::index::token_hashes_from_bytes(b);
9878 mh.insert(&tokens, row.row_id);
9879 }
9880 }
9881 _ => {}
9882 }
9883}
9884
9885fn eval_partial_predicate(
9891 pred: &str,
9892 columns_map: &HashMap<u16, &Value>,
9893 name_to_id: &HashMap<&str, u16>,
9894) -> bool {
9895 let lower = pred.trim().to_ascii_lowercase();
9896 if let Some(rest) = lower.strip_suffix(" is not null") {
9898 let col_name = rest.trim();
9899 if let Some(col_id) = name_to_id.get(col_name) {
9900 return columns_map
9901 .get(col_id)
9902 .is_some_and(|v| !matches!(v, Value::Null));
9903 }
9904 }
9905 if let Some(rest) = lower.strip_suffix(" is null") {
9907 let col_name = rest.trim();
9908 if let Some(col_id) = name_to_id.get(col_name) {
9909 return columns_map
9910 .get(col_id)
9911 .map_or(true, |v| matches!(v, Value::Null));
9912 }
9913 }
9914 true
9917}
9918
9919#[allow(dead_code)]
9925fn bulk_index_key(
9926 column_keys: &HashMap<u16, ([u8; 32], u8)>,
9927 column_id: u16,
9928 ty: TypeId,
9929 col: &columnar::NativeColumn,
9930 i: usize,
9931) -> Option<Vec<u8>> {
9932 let encoded = columnar::encode_key_native(ty, col, i)?;
9933 #[cfg(feature = "encryption")]
9934 {
9935 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
9936 if let Some((key, scheme)) = column_keys.get(&column_id) {
9937 return Some(match (*scheme, col) {
9938 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
9939 (_, columnar::NativeColumn::Int64 { data, .. }) => {
9940 ope_token_i64(key, data[i]).to_vec()
9941 }
9942 (_, columnar::NativeColumn::Float64 { data, .. }) => {
9943 ope_token_f64(key, data[i]).to_vec()
9944 }
9945 _ => hmac_token(key, &encoded).to_vec(),
9946 });
9947 }
9948 }
9949 #[cfg(not(feature = "encryption"))]
9950 {
9951 let _ = (column_id, column_keys, col);
9952 }
9953 Some(encoded)
9954}
9955
9956pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
9957 let json = serde_json::to_string_pretty(schema)
9958 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
9959 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
9960 Ok(())
9961}
9962
9963fn read_schema(dir: &Path) -> Result<Schema> {
9964 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
9965 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
9966}
9967
9968fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
9969 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
9970}
9971
9972fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
9973 let n = list_wal_numbers(wal_dir)?;
9974 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
9975}
9976
9977fn next_wal_number(wal_dir: &Path) -> Result<u32> {
9978 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
9979}
9980
9981fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
9982 let _ = std::fs::create_dir_all(wal_dir);
9983 let mut max_n = None;
9984 for entry in std::fs::read_dir(wal_dir)? {
9985 let entry = entry?;
9986 let fname = entry.file_name();
9987 let Some(s) = fname.to_str() else {
9988 continue;
9989 };
9990 let Some(stripped) = s.strip_prefix("seg-") else {
9991 continue;
9992 };
9993 let Some(stripped) = stripped.strip_suffix(".wal") else {
9994 continue;
9995 };
9996 if let Ok(n) = stripped.parse::<u32>() {
9997 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
9998 }
9999 }
10000 Ok(max_n)
10001}