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 ann_candidate_cap(
65 index_len: usize,
66 context: Option<&crate::query::AiExecutionContext>,
67) -> usize {
68 index_len
69 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
70 .min(context.map_or(
71 crate::query::MAX_RAW_INDEX_CANDIDATES,
72 crate::query::AiExecutionContext::max_fused_candidates,
73 ))
74}
75
76#[cfg(test)]
77mod ann_candidate_cap_tests {
78 use super::*;
79
80 #[test]
81 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
82 assert_eq!(
83 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
84 crate::query::MAX_RAW_INDEX_CANDIDATES,
85 );
86 let context = crate::query::AiExecutionContext::with_limits(
87 std::time::Duration::from_secs(1),
88 usize::MAX,
89 17,
90 );
91 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
92 }
93}
94
95fn civil_from_days(z: i64) -> (i64, u32, u32) {
96 let z = z + 719_468;
97 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
98 let doe = z - era * 146_097;
99 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
100 let y = yoe + era * 400;
101 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
102 let mp = (5 * doy + 2) / 153;
103 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
104 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
105 (if m <= 2 { y + 1 } else { y }, m, d)
106}
107
108const 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;
115
116#[derive(Clone, Copy, Debug)]
131struct AutoIncState {
132 column_id: u16,
133 next: i64,
134 seeded: bool,
135}
136
137type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
138
139fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
142 schema.auto_increment_column().map(|c| AutoIncState {
143 column_id: c.id,
144 next: 0,
145 seeded: false,
146 })
147}
148
149#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
162pub enum IndexBuildPolicy {
163 #[default]
165 Deferred,
166 Eager,
168}
169
170#[derive(Clone)]
171struct ReversePkSegment {
172 values: HashMap<RowId, Vec<u8>>,
173 removed: HashSet<RowId>,
174}
175
176#[derive(Clone)]
177struct ReversePkMap {
178 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
179 active: ReversePkSegment,
180}
181
182impl ReversePkMap {
183 fn new() -> Self {
184 Self {
185 frozen: Arc::new(Vec::new()),
186 active: ReversePkSegment {
187 values: HashMap::new(),
188 removed: HashSet::new(),
189 },
190 }
191 }
192
193 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
194 let mut map = Self::new();
195 map.active.values.extend(entries);
196 map
197 }
198
199 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
200 self.active.removed.remove(&row_id);
201 self.active.values.insert(row_id, key);
202 }
203
204 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
205 if let Some(key) = self.active.values.get(row_id) {
206 return Some(key);
207 }
208 if self.active.removed.contains(row_id) {
209 return None;
210 }
211 for segment in self.frozen.iter().rev() {
212 if let Some(key) = segment.values.get(row_id) {
213 return Some(key);
214 }
215 if segment.removed.contains(row_id) {
216 return None;
217 }
218 }
219 None
220 }
221
222 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
223 let previous = self.get(row_id).cloned();
224 self.active.values.remove(row_id);
225 self.active.removed.insert(*row_id);
226 previous
227 }
228
229 fn clear(&mut self) {
230 *self = Self::new();
231 }
232
233 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
234 let mut entries = HashMap::new();
235 for segment in self
236 .frozen
237 .iter()
238 .map(Arc::as_ref)
239 .chain(std::iter::once(&self.active))
240 {
241 for row_id in &segment.removed {
242 entries.remove(row_id);
243 }
244 entries.extend(
245 segment
246 .values
247 .iter()
248 .map(|(row_id, key)| (*row_id, key.clone())),
249 );
250 }
251 entries
252 }
253
254 fn seal(&mut self) {
255 if self.active.values.is_empty() && self.active.removed.is_empty() {
256 return;
257 }
258 let active = std::mem::replace(
259 &mut self.active,
260 ReversePkSegment {
261 values: HashMap::new(),
262 removed: HashSet::new(),
263 },
264 );
265 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
266 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
267 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
268 values: self.entries(),
269 removed: HashSet::new(),
270 })]);
271 }
272 }
273}
274
275#[derive(Clone)]
277pub struct Table {
278 dir: PathBuf,
279 table_id: u64,
280 name: String,
284 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
289 read_only: bool,
292 wal: WalSink,
293 memtable: Memtable,
294 mutable_run: MutableRun,
299 mutable_run_spill_bytes: u64,
301 compaction_zstd_level: i32,
304 allocator: RowIdAllocator,
305 epoch: Arc<EpochAuthority>,
306 persisted_epoch: u64,
309 data_generation: u64,
312 schema: Schema,
313 hot: HotIndex,
314 kek: Option<Arc<Kek>>,
317 column_keys: HashMap<u16, ([u8; 32], u8)>,
321 run_refs: Vec<RunRef>,
322 retiring: Vec<crate::manifest::RetiredRun>,
325 next_run_id: u64,
326 sync_byte_threshold: u64,
327 current_txn_id: u64,
332 bitmap: HashMap<u16, BitmapIndex>,
333 ann: HashMap<u16, AnnIndex>,
334 fm: HashMap<u16, FmIndex>,
335 sparse: HashMap<u16, SparseIndex>,
336 minhash: HashMap<u16, MinHashIndex>,
337 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
340 pk_by_row: ReversePkMap,
342 pinned: BTreeMap<Epoch, usize>,
345 pub(crate) live_count: u64,
348 reservoir: crate::reservoir::Reservoir,
351 reservoir_complete: bool,
359 had_deletes: bool,
363 agg_cache: Arc<HashMap<u64, CachedAgg>>,
367 global_idx_epoch: u64,
371 indexes_complete: bool,
376 index_build_policy: IndexBuildPolicy,
378 pk_by_row_complete: bool,
385 flushed_epoch: u64,
388 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
391 snapshots: Arc<crate::retention::SnapshotRegistry>,
394 commit_lock: Arc<parking_lot::Mutex<()>>,
396 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
399 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
409 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
418 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
420 pending_delete_rids: roaring::RoaringBitmap,
423 pending_put_cols: std::collections::HashSet<u16>,
426 pending_rows: Vec<Row>,
432 pending_rows_auto_inc: Vec<bool>,
433 pending_dels: Vec<RowId>,
436 pending_truncate: Option<Epoch>,
440 auto_inc: Option<AutoIncState>,
443 ttl: Option<TtlPolicy>,
446}
447
448const _: () = {
455 const fn assert_sync<T: ?Sized + Sync>() {}
456 assert_sync::<Table>();
457};
458
459enum CachedData {
465 Rows(Arc<Vec<Row>>),
466 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
467}
468
469impl CachedData {
470 fn approx_bytes(&self) -> u64 {
471 match self {
472 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
473 CachedData::Columns(c) => c
474 .iter()
475 .map(|(_, c)| c.approx_bytes())
476 .sum::<u64>()
477 .saturating_add(c.len() as u64 * 16),
478 }
479 }
480}
481
482struct CachedEntry {
486 data: CachedData,
487 footprint: roaring::RoaringBitmap,
488 condition_cols: Vec<u16>,
489}
490
491struct ResultCache {
502 entries: std::collections::HashMap<u64, CachedEntry>,
503 order: std::collections::VecDeque<u64>,
504 bytes: u64,
505 max_bytes: u64,
506 dir: Option<std::path::PathBuf>,
507 #[allow(dead_code)]
508 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
509}
510
511#[derive(serde::Serialize, serde::Deserialize)]
513struct SerializedEntry {
514 condition_cols: Vec<u16>,
515 footprint_bits: Vec<u32>,
516 data: SerializedData,
517}
518
519#[derive(serde::Serialize, serde::Deserialize)]
520enum SerializedData {
521 Rows(Vec<Row>),
522 Columns(Vec<(u16, columnar::NativeColumn)>),
523}
524
525impl SerializedEntry {
526 fn from_entry(entry: &CachedEntry) -> Self {
527 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
528 let data = match &entry.data {
529 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
530 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
531 };
532 Self {
533 condition_cols: entry.condition_cols.clone(),
534 footprint_bits,
535 data,
536 }
537 }
538
539 fn into_entry(self) -> Option<CachedEntry> {
540 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
541 let data = match self.data {
542 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
543 SerializedData::Columns(c) => {
544 if !c.iter().all(|(_, col)| col.validate()) {
547 return None;
548 }
549 CachedData::Columns(Arc::new(c))
550 }
551 };
552 Some(CachedEntry {
553 data,
554 footprint,
555 condition_cols: self.condition_cols,
556 })
557 }
558}
559
560impl ResultCache {
561 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
562
563 fn new() -> Self {
564 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
565 }
566
567 fn with_max_bytes(max_bytes: u64) -> Self {
568 Self {
569 entries: std::collections::HashMap::new(),
570 order: std::collections::VecDeque::new(),
571 bytes: 0,
572 max_bytes,
573 dir: None,
574 cache_dek: None,
575 }
576 }
577
578 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
579 let _ = std::fs::create_dir_all(&dir);
580 self.dir = Some(dir);
581 self
582 }
583
584 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
585 self.cache_dek = dek;
586 self
587 }
588
589 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
590 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
591 }
592
593 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
597 let Some(path) = self.disk_path(key) else {
598 return;
599 };
600 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
601 Ok(s) => s,
602 Err(_) => return,
603 };
604 let on_disk = if let Some(dek) = &self.cache_dek {
606 match self.encrypt_cache(&serialized, dek) {
607 Some(b) => b,
608 None => return,
609 }
610 } else {
611 serialized
612 };
613 let tmp = path.with_extension("tmp");
614 use std::io::Write;
615 let write = || -> std::io::Result<()> {
616 let mut f = std::fs::File::create(&tmp)?;
617 f.write_all(&on_disk)?;
618 f.flush()?;
619 Ok(())
620 };
621 if write().is_err() {
622 let _ = std::fs::remove_file(&tmp);
623 return;
624 }
625 let _ = std::fs::rename(&tmp, &path);
626 }
627
628 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
630 let path = self.disk_path(key)?;
631 let bytes = std::fs::read(&path).ok()?;
632 let plaintext = if let Some(dek) = &self.cache_dek {
633 self.decrypt_cache(&bytes, dek)?
634 } else {
635 bytes
636 };
637 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
638 serialized.into_entry()
639 }
640
641 fn remove_from_disk(&self, key: u64) {
643 if let Some(path) = self.disk_path(key) {
644 let _ = std::fs::remove_file(&path);
645 }
646 }
647
648 #[cfg(feature = "encryption")]
650 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
651 use crate::encryption::Cipher;
652 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
653 let mut nonce = [0u8; 12];
654 crate::encryption::fill_random(&mut nonce);
655 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
656 let mut out = Vec::with_capacity(12 + ct.len());
657 out.extend_from_slice(&nonce);
658 out.extend_from_slice(&ct);
659 Some(out)
660 }
661
662 #[cfg(not(feature = "encryption"))]
663 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
664 None
665 }
666
667 #[cfg(feature = "encryption")]
669 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
670 use crate::encryption::Cipher;
671 if bytes.len() < 28 {
672 return None;
673 }
674 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
675 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
676 let ct = &bytes[12..];
677 cipher.decrypt_page(&nonce, ct).ok()
678 }
679
680 #[cfg(not(feature = "encryption"))]
681 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
682 None
683 }
684
685 fn load_persistent(&mut self) {
688 let Some(dir) = self.dir.as_ref().cloned() else {
689 return;
690 };
691 let entries = match std::fs::read_dir(&dir) {
692 Ok(e) => e,
693 Err(_) => return,
694 };
695 for entry in entries.flatten() {
696 let path = entry.path();
697 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
699 let _ = std::fs::remove_file(&path);
700 continue;
701 }
702 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
703 continue;
704 }
705 let stem = match path.file_stem().and_then(|s| s.to_str()) {
706 Some(s) => s,
707 None => continue,
708 };
709 let key = match u64::from_str_radix(stem, 16) {
710 Ok(k) => k,
711 Err(_) => continue,
712 };
713 let bytes = match std::fs::read(&path) {
714 Ok(b) => b,
715 Err(_) => continue,
716 };
717 let plaintext = if let Some(dek) = &self.cache_dek {
719 match self.decrypt_cache(&bytes, dek) {
720 Some(p) => p,
721 None => {
722 let _ = std::fs::remove_file(&path);
723 continue;
724 }
725 }
726 } else {
727 bytes
728 };
729 match bincode::deserialize::<SerializedEntry>(&plaintext) {
730 Ok(serialized) => {
731 if let Some(entry) = serialized.into_entry() {
732 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
733 self.entries.insert(key, entry);
734 self.order.push_back(key);
735 } else {
736 let _ = std::fs::remove_file(&path);
737 }
738 }
739 Err(_) => {
740 let _ = std::fs::remove_file(&path);
741 }
742 }
743 }
744 self.evict();
745 }
746
747 fn set_max_bytes(&mut self, max_bytes: u64) {
748 self.max_bytes = max_bytes;
749 self.evict();
750 }
751
752 fn touch(&mut self, key: u64) {
754 self.order.retain(|k| *k != key);
755 self.order.push_back(key);
756 }
757
758 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
759 let res = self.entries.get(&key).and_then(|e| match &e.data {
760 CachedData::Rows(r) => Some(r.clone()),
761 CachedData::Columns(_) => None,
762 });
763 if res.is_some() {
764 self.touch(key);
765 return res;
766 }
767 if let Some(entry) = self.load_from_disk(key) {
769 let res = match &entry.data {
770 CachedData::Rows(r) => Some(r.clone()),
771 CachedData::Columns(_) => None,
772 };
773 if res.is_some() {
774 let approx = entry.data.approx_bytes();
775 self.bytes = self.bytes.saturating_add(approx);
776 self.entries.insert(key, entry);
777 self.order.push_back(key);
778 self.evict();
779 return res;
780 }
781 }
782 None
783 }
784
785 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
786 let res = self.entries.get(&key).and_then(|e| match &e.data {
787 CachedData::Columns(c) => Some(c.clone()),
788 CachedData::Rows(_) => None,
789 });
790 if res.is_some() {
791 self.touch(key);
792 return res;
793 }
794 if let Some(entry) = self.load_from_disk(key) {
796 let res = match &entry.data {
797 CachedData::Columns(c) => Some(c.clone()),
798 CachedData::Rows(_) => None,
799 };
800 if res.is_some() {
801 let approx = entry.data.approx_bytes();
802 self.bytes = self.bytes.saturating_add(approx);
803 self.entries.insert(key, entry);
804 self.order.push_back(key);
805 self.evict();
806 return res;
807 }
808 }
809 None
810 }
811
812 fn insert(&mut self, key: u64, entry: CachedEntry) {
813 let approx = entry.data.approx_bytes();
814 if self.entries.remove(&key).is_some() {
815 self.order.retain(|k| *k != key);
816 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
817 }
818 self.store_to_disk(key, &entry);
820 self.bytes = self.bytes.saturating_add(approx);
821 self.entries.insert(key, entry);
822 self.order.push_back(key);
823 self.evict();
824 }
825
826 fn invalidate(
835 &mut self,
836 delete_rids: &roaring::RoaringBitmap,
837 put_cols: &std::collections::HashSet<u16>,
838 ) {
839 if self.entries.is_empty() {
840 return;
841 }
842 let has_deletes = !delete_rids.is_empty();
843 let to_remove: std::collections::HashSet<u64> = self
844 .entries
845 .iter()
846 .filter(|(_, e)| {
847 let delete_hit = if e.footprint.is_empty() {
848 has_deletes
849 } else {
850 e.footprint.intersection_len(delete_rids) > 0
851 };
852 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
853 delete_hit || col_hit
854 })
855 .map(|(&k, _)| k)
856 .collect();
857 for key in &to_remove {
858 if let Some(e) = self.entries.remove(key) {
859 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
860 }
861 self.remove_from_disk(*key);
862 }
863 if !to_remove.is_empty() {
864 self.order.retain(|k| !to_remove.contains(k));
865 }
866 }
867
868 fn clear(&mut self) {
869 if let Some(dir) = &self.dir {
871 if let Ok(entries) = std::fs::read_dir(dir) {
872 for entry in entries.flatten() {
873 let path = entry.path();
874 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
875 let _ = std::fs::remove_file(&path);
876 }
877 }
878 }
879 }
880 self.entries.clear();
881 self.order.clear();
882 self.bytes = 0;
883 }
884
885 fn evict(&mut self) {
886 while self.bytes > self.max_bytes {
887 let Some(k) = self.order.pop_front() else {
888 break;
889 };
890 if let Some(e) = self.entries.remove(&k) {
891 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
892 self.remove_from_disk(k);
896 }
897 }
898 }
899}
900
901type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
908
909fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
910 let _ = kek;
911 #[cfg(feature = "encryption")]
912 {
913 if let Some(k) = kek {
914 return (
915 Some(k.derive_table_wal_key(_table_id)),
916 Some(k.derive_cache_key()),
917 );
918 }
919 }
920 (None, None)
921}
922
923#[cfg(feature = "encryption")]
925fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
926 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
927}
928
929#[cfg(not(feature = "encryption"))]
930fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
931 Box::new(crate::encryption::PlaintextCipher)
932}
933
934fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
935 let Some(kek) = kek else {
936 return HashMap::new();
937 };
938 #[cfg(feature = "encryption")]
939 {
940 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
941 schema
942 .columns
943 .iter()
944 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
945 .map(|c| {
946 let scheme = if schema
947 .indexes
948 .iter()
949 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
950 {
951 SCHEME_OPE_RANGE
952 } else {
953 SCHEME_HMAC_EQ
954 };
955 let key: [u8; 32] = *kek.derive_column_key(c.id);
956 (c.id, (key, scheme))
957 })
958 .collect()
959 }
960 #[cfg(not(feature = "encryption"))]
961 {
962 let _ = (kek, schema);
963 HashMap::new()
964 }
965}
966
967pub(crate) struct SharedCtx {
972 pub epoch: Arc<EpochAuthority>,
973 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
974 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
975 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
976 pub kek: Option<Arc<Kek>>,
977 pub commit_lock: Arc<parking_lot::Mutex<()>>,
983 pub shared: Option<SharedWalCtx>,
987 pub table_name: Option<String>,
990 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
993 pub read_only: bool,
995}
996
997#[derive(Clone)]
1003pub(crate) struct SharedWalCtx {
1004 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1005 pub group: Arc<GroupCommit>,
1006 pub poisoned: Arc<AtomicBool>,
1007 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1008 pub change_wake: tokio::sync::broadcast::Sender<()>,
1009}
1010
1011enum WalSink {
1014 Private(Wal),
1015 Shared(SharedWalCtx),
1016 ReadOnly,
1017}
1018
1019impl Clone for WalSink {
1020 fn clone(&self) -> Self {
1021 match self {
1022 Self::Shared(shared) => Self::Shared(shared.clone()),
1023 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1024 }
1025 }
1026}
1027
1028impl SharedCtx {
1029 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1033 let n_shards = if cache_dir.is_some() {
1037 1
1038 } else {
1039 crate::cache::CACHE_SHARDS
1040 };
1041 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1042 let page_cache = if let Some(d) = cache_dir {
1043 Arc::new(crate::cache::Sharded::new(1, || {
1044 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1045 }))
1046 } else {
1047 Arc::new(crate::cache::Sharded::new(n_shards, || {
1048 crate::cache::PageCache::new(per_shard)
1049 }))
1050 };
1051 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1052 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1053 crate::cache::CACHE_SHARDS,
1054 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1055 ));
1056 Self {
1057 epoch: Arc::new(EpochAuthority::new(0)),
1058 page_cache,
1059 decoded_cache,
1060 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1061 kek,
1062 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1063 shared: None,
1064 table_name: None,
1065 auth: None,
1066 read_only: false,
1067 }
1068 }
1069}
1070
1071fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1075 use crate::query::Condition;
1076 match c {
1077 Condition::Pk(_)
1079 | Condition::BitmapEq { .. }
1080 | Condition::BitmapIn { .. }
1081 | Condition::BytesPrefix { .. }
1082 | Condition::IsNull { .. }
1083 | Condition::IsNotNull { .. } => 0,
1084 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1086 1
1087 }
1088 Condition::FmContains { .. }
1090 | Condition::FmContainsAll { .. }
1091 | Condition::Ann { .. }
1092 | Condition::SparseMatch { .. } => 2,
1093 }
1094}
1095
1096impl Table {
1097 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1098 let dir = dir.as_ref().to_path_buf();
1099 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
1100 Self::create_in(&dir, schema, table_id, ctx)
1101 }
1102
1103 #[cfg(feature = "encryption")]
1114 pub fn create_encrypted(
1115 dir: impl AsRef<Path>,
1116 schema: Schema,
1117 table_id: u64,
1118 passphrase: &str,
1119 ) -> Result<Self> {
1120 let dir = dir.as_ref();
1121 std::fs::create_dir_all(dir.join(META_DIR))?;
1122 let salt = crate::encryption::random_salt();
1123 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
1124 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1125 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1126 Self::create_in(dir, schema, table_id, ctx)
1127 }
1128
1129 #[cfg(feature = "encryption")]
1134 pub fn create_with_key(
1135 dir: impl AsRef<Path>,
1136 schema: Schema,
1137 table_id: u64,
1138 key: &[u8],
1139 ) -> Result<Self> {
1140 let dir = dir.as_ref();
1141 std::fs::create_dir_all(dir.join(META_DIR))?;
1142 let salt = crate::encryption::random_salt();
1143 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
1144 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1145 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1146 Self::create_in(dir, schema, table_id, ctx)
1147 }
1148
1149 #[cfg(feature = "encryption")]
1151 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1152 let dir = dir.as_ref();
1153 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1154 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1155 MongrelError::NotFound(format!(
1156 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1157 salt_path
1158 ))
1159 })?;
1160 if salt_bytes.len() != crate::encryption::SALT_LEN {
1161 return Err(MongrelError::InvalidArgument(format!(
1162 "salt file is {} bytes, expected {}",
1163 salt_bytes.len(),
1164 crate::encryption::SALT_LEN
1165 )));
1166 }
1167 let mut salt = [0u8; crate::encryption::SALT_LEN];
1168 salt.copy_from_slice(&salt_bytes);
1169 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1170 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1171 Self::open_in(dir, ctx)
1172 }
1173
1174 pub(crate) fn create_in(
1175 dir: impl AsRef<Path>,
1176 schema: Schema,
1177 table_id: u64,
1178 ctx: SharedCtx,
1179 ) -> Result<Self> {
1180 schema.validate_auto_increment()?;
1181 schema.validate_defaults()?;
1182 schema.validate_ai()?;
1183 for index in &schema.indexes {
1184 index.validate_options()?;
1185 }
1186 let dir = dir.as_ref().to_path_buf();
1187 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1188 write_schema(&dir, &schema)?;
1189 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1190 let (wal, current_txn_id) = match ctx.shared.clone() {
1193 Some(s) => (WalSink::Shared(s), 0),
1194 None => {
1195 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1196 let mut w = if let Some(ref dk) = wal_dek {
1197 Wal::create_with_cipher(
1198 dir.join(WAL_DIR).join("seg-000000.wal"),
1199 Epoch(0),
1200 Some(make_cipher(dk)),
1201 0,
1202 )?
1203 } else {
1204 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1205 };
1206 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1207 (WalSink::Private(w), 1)
1208 }
1209 };
1210 let mut manifest = Manifest::new(table_id, schema.schema_id);
1211 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1216 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1217 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1218 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1219 let auto_inc = resolve_auto_inc(&schema);
1220 let rcache_dir = dir.join(RCACHE_DIR);
1221 Ok(Self {
1222 dir,
1223 table_id,
1224 name: ctx.table_name.unwrap_or_default(),
1225 auth: ctx.auth,
1226 read_only: ctx.read_only,
1227 wal,
1228 memtable: Memtable::new(),
1229 mutable_run: MutableRun::new(),
1230 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1231 compaction_zstd_level: 3,
1232 allocator: RowIdAllocator::new(0),
1233 epoch: ctx.epoch,
1234 persisted_epoch: 0,
1235 data_generation: 0,
1236 schema,
1237 hot: HotIndex::new(),
1238 kek: ctx.kek,
1239 column_keys,
1240 run_refs: Vec::new(),
1241 retiring: Vec::new(),
1242 next_run_id: 1,
1243 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1244 current_txn_id,
1245 bitmap,
1246 ann,
1247 fm,
1248 sparse,
1249 minhash,
1250 learned_range: Arc::new(HashMap::new()),
1251 pk_by_row: ReversePkMap::new(),
1252 pinned: BTreeMap::new(),
1253 live_count: 0,
1254 reservoir: crate::reservoir::Reservoir::default(),
1255 reservoir_complete: true,
1256 had_deletes: false,
1257 agg_cache: Arc::new(HashMap::new()),
1258 global_idx_epoch: 0,
1259 indexes_complete: true,
1260 index_build_policy: IndexBuildPolicy::default(),
1261 pk_by_row_complete: false,
1262 flushed_epoch: 0,
1263 page_cache: ctx.page_cache,
1264 decoded_cache: ctx.decoded_cache,
1265 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1266 snapshots: ctx.snapshots,
1267 commit_lock: ctx.commit_lock,
1268 result_cache: Arc::new(parking_lot::Mutex::new(
1269 ResultCache::new()
1270 .with_dir(rcache_dir)
1271 .with_cache_dek(cache_dek.clone()),
1272 )),
1273 pending_delete_rids: roaring::RoaringBitmap::new(),
1274 pending_put_cols: std::collections::HashSet::new(),
1275 pending_rows: Vec::new(),
1276 pending_rows_auto_inc: Vec::new(),
1277 pending_dels: Vec::new(),
1278 pending_truncate: None,
1279 wal_dek,
1280 auto_inc,
1281 ttl: None,
1282 })
1283 }
1284
1285 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1289 let dir = dir.as_ref();
1290 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1291 Self::open_in(dir, ctx)
1292 }
1293
1294 #[cfg(feature = "encryption")]
1297 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1298 let dir = dir.as_ref();
1299 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1300 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1301 MongrelError::NotFound(format!(
1302 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1303 salt_path
1304 ))
1305 })?;
1306 let salt_len = crate::encryption::SALT_LEN;
1307 if salt_bytes.len() != salt_len {
1308 return Err(MongrelError::InvalidArgument(format!(
1309 "encryption salt is {} bytes, expected {salt_len}",
1310 salt_bytes.len()
1311 )));
1312 }
1313 let mut salt = [0u8; 16];
1314 salt.copy_from_slice(&salt_bytes);
1315 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1316 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1317 let t = Self::open_in(dir, ctx)?;
1318 Ok(t)
1319 }
1320
1321 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1322 let dir = dir.as_ref().to_path_buf();
1323 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1324 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1325 let schema: Schema = read_schema(&dir)?;
1326 schema.validate_ai()?;
1327 for index in &schema.indexes {
1328 index.validate_options()?;
1329 }
1330 let replay_epoch = Epoch(manifest.current_epoch);
1331 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1332 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1336 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1337 None => {
1338 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1339 let replayed = match &active {
1341 Some(path) => {
1342 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1343 crate::wal::replay_with_cipher(path, cipher)?
1344 }
1345 None => Vec::new(),
1346 };
1347 let mut w = match &active {
1348 Some(path) => Wal::create_with_cipher(
1349 path,
1350 replay_epoch,
1351 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1352 0,
1353 )?,
1354 None => Wal::create_with_cipher(
1355 dir.join(WAL_DIR).join("seg-000000.wal"),
1356 replay_epoch,
1357 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1358 0,
1359 )?,
1360 };
1361 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1362 (WalSink::Private(w), replayed, 1)
1363 }
1364 };
1365
1366 let mut memtable = Memtable::new();
1367 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1368 let persisted_epoch = manifest.current_epoch;
1369 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1376 s.next = manifest.auto_inc_next;
1377 s.seeded = manifest.auto_inc_next > 0;
1378 s
1379 });
1380
1381 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1388 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1389 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1390 std::collections::BTreeMap::new();
1391 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1392 let mut saw_delete = false;
1393 for record in replayed {
1394 let txn_id = record.txn_id;
1395 match record.op {
1396 Op::Put { rows, .. } => {
1397 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1398 for row in &rows {
1399 allocator.advance_to(row.row_id);
1400 if let Some(ai) = auto_inc.as_mut() {
1401 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1402 if *n + 1 > ai.next {
1403 ai.next = *n + 1;
1404 }
1405 }
1406 }
1407 }
1408 staged_puts.entry(txn_id).or_default().extend(rows);
1409 }
1410 Op::Delete { row_ids, .. } => {
1411 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1412 }
1413 Op::TxnCommit { epoch, .. } => {
1414 let commit_epoch = Epoch(epoch);
1415 if let Some(puts) = staged_puts.remove(&txn_id) {
1416 if commit_epoch.0 > manifest.flushed_epoch {
1417 for row in &puts {
1418 memtable.upsert(row.clone());
1419 }
1420 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1421 }
1422 }
1423 if let Some(dels) = staged_deletes.remove(&txn_id) {
1424 saw_delete = true;
1425 if commit_epoch.0 > manifest.flushed_epoch {
1426 for rid in dels {
1427 memtable.tombstone(rid, commit_epoch);
1428 replayed_deletes.push((rid, commit_epoch));
1429 }
1430 }
1431 }
1432 }
1433 Op::TxnAbort => {
1434 staged_puts.remove(&txn_id);
1435 staged_deletes.remove(&txn_id);
1436 }
1437 Op::TruncateTable { .. }
1438 | Op::ExternalTableState { .. }
1439 | Op::Flush { .. }
1440 | Op::Ddl(_)
1441 | Op::BeforeImage { .. }
1442 | Op::CommitTimestamp { .. } => {}
1443 }
1444 }
1445
1446 let rcache_dir = dir.join(RCACHE_DIR);
1447 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1448 let mut db = Self {
1449 dir,
1450 table_id: manifest.table_id,
1451 name: ctx.table_name.unwrap_or_default(),
1452 auth: ctx.auth,
1453 read_only: ctx.read_only,
1454 wal,
1455 memtable,
1456 mutable_run: MutableRun::new(),
1457 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1458 compaction_zstd_level: 3,
1459 allocator,
1460 epoch: ctx.epoch,
1461 persisted_epoch,
1462 data_generation: persisted_epoch,
1463 schema,
1464 hot: HotIndex::new(),
1465 kek: ctx.kek,
1466 column_keys,
1467 run_refs: manifest.runs.clone(),
1468 retiring: manifest.retiring.clone(),
1469 next_run_id: manifest
1470 .runs
1471 .iter()
1472 .map(|r| r.run_id as u64 + 1)
1473 .max()
1474 .unwrap_or(1),
1475 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1476 current_txn_id,
1477 bitmap: HashMap::new(),
1478 ann: HashMap::new(),
1479 fm: HashMap::new(),
1480 sparse: HashMap::new(),
1481 minhash: HashMap::new(),
1482 learned_range: Arc::new(HashMap::new()),
1483 pk_by_row: ReversePkMap::new(),
1484 pinned: BTreeMap::new(),
1485 live_count: manifest.live_count,
1486 reservoir: crate::reservoir::Reservoir::default(),
1487 reservoir_complete: false,
1488 had_deletes: saw_delete
1489 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1490 != manifest.live_count,
1491 agg_cache: Arc::new(HashMap::new()),
1492 global_idx_epoch: manifest.global_idx_epoch,
1493 indexes_complete: true,
1494 index_build_policy: IndexBuildPolicy::default(),
1495 pk_by_row_complete: false,
1496 flushed_epoch: manifest.flushed_epoch,
1497 page_cache: ctx.page_cache,
1498 decoded_cache: ctx.decoded_cache,
1499 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1500 snapshots: ctx.snapshots,
1501 commit_lock: ctx.commit_lock,
1502 result_cache: Arc::new(parking_lot::Mutex::new(
1503 ResultCache::new()
1504 .with_dir(rcache_dir)
1505 .with_cache_dek(cache_dek.clone()),
1506 )),
1507 pending_delete_rids: roaring::RoaringBitmap::new(),
1508 pending_put_cols: std::collections::HashSet::new(),
1509 pending_rows: Vec::new(),
1510 pending_rows_auto_inc: Vec::new(),
1511 pending_dels: Vec::new(),
1512 pending_truncate: None,
1513 wal_dek,
1514 auto_inc,
1515 ttl: manifest.ttl,
1516 };
1517
1518 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1521
1522 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1527 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1528 c.epoch_built == manifest.global_idx_epoch
1529 && manifest.global_idx_epoch > 0
1530 && manifest
1531 .runs
1532 .iter()
1533 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1534 });
1535 if let Some(loaded) = checkpoint {
1536 if checkpoint_valid {
1537 db.hot = loaded.hot;
1538 db.bitmap = loaded.bitmap;
1539 db.ann = loaded.ann;
1540 db.fm = loaded.fm;
1541 db.sparse = loaded.sparse;
1542 db.minhash = loaded.minhash;
1543 db.learned_range = Arc::new(loaded.learned_range);
1544 }
1547 }
1548 if !checkpoint_valid {
1549 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1550 db.bitmap = bitmap;
1551 db.ann = ann;
1552 db.fm = fm;
1553 db.sparse = sparse;
1554 db.minhash = minhash;
1555 db.rebuild_indexes_from_runs()?;
1556 db.build_learned_ranges()?;
1557 }
1558
1559 for (epoch, group) in replayed_puts {
1564 let (losers, winner_pks) = db.partition_pk_winners(&group);
1565 for (key, &row_id) in &winner_pks {
1566 if let Some(old_rid) = db.hot.get(key) {
1567 if old_rid != row_id {
1568 db.tombstone_row(old_rid, epoch, false);
1569 }
1570 }
1571 }
1572 for &loser_rid in &losers {
1573 db.tombstone_row(loser_rid, epoch, false);
1574 }
1575 for (key, row_id) in winner_pks {
1576 db.insert_hot_pk(key, row_id);
1577 }
1578 if db.schema.primary_key().is_none() {
1579 for r in &group {
1580 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1581 }
1582 }
1583 for r in &group {
1584 if !losers.contains(&r.row_id) {
1585 db.index_row(r);
1586 }
1587 }
1588 }
1589 for (rid, epoch) in &replayed_deletes {
1593 db.remove_hot_for_row(*rid, *epoch);
1594 }
1595
1596 db.result_cache.lock().load_persistent();
1603 Ok(db)
1604 }
1605
1606 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1612 if self.reservoir_complete {
1613 return Ok(());
1614 }
1615 self.rebuild_reservoir()?;
1616 self.reservoir_complete = true;
1617 Ok(())
1618 }
1619
1620 fn rebuild_reservoir(&mut self) -> Result<()> {
1623 let snap = self.snapshot();
1624 let rows = self.visible_rows(snap)?;
1625 self.reservoir.reset();
1626 for r in rows {
1627 self.reservoir.offer(r.row_id.0);
1628 }
1629 Ok(())
1630 }
1631
1632 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1633 self.hot = HotIndex::new();
1634 self.pk_by_row.clear();
1635 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1636 self.bitmap = bitmap;
1637 self.ann = ann;
1638 self.fm = fm;
1639 self.sparse = sparse;
1640 self.minhash = minhash;
1641 let snapshot = Epoch(u64::MAX);
1642 let ttl_now = unix_nanos_now();
1643 for rr in self.run_refs.clone() {
1644 let mut reader = self.open_reader(rr.run_id)?;
1645 for row in reader.visible_rows(snapshot)? {
1646 if self.row_expired_at(&row, ttl_now) {
1647 continue;
1648 }
1649 let tok_row = self.tokenized_for_indexes(&row);
1650 index_into(
1651 &self.schema,
1652 &tok_row,
1653 &mut self.hot,
1654 &mut self.bitmap,
1655 &mut self.ann,
1656 &mut self.fm,
1657 &mut self.sparse,
1658 &mut self.minhash,
1659 );
1660 }
1661 }
1662 for row in self.mutable_run.visible_versions(snapshot) {
1663 if row.deleted {
1664 self.remove_hot_for_row(row.row_id, snapshot);
1665 } else if !self.row_expired_at(&row, ttl_now) {
1666 self.index_row(&row);
1667 }
1668 }
1669 for row in self.memtable.visible_versions(snapshot) {
1670 if row.deleted {
1671 self.remove_hot_for_row(row.row_id, snapshot);
1672 } else if !self.row_expired_at(&row, ttl_now) {
1673 self.index_row(&row);
1674 }
1675 }
1676 self.refresh_pk_by_row_from_hot();
1677 Ok(())
1678 }
1679
1680 fn refresh_pk_by_row_from_hot(&mut self) {
1681 self.pk_by_row_complete = true;
1682 if self.schema.primary_key().is_none() {
1683 self.pk_by_row.clear();
1684 return;
1685 }
1686 self.pk_by_row = ReversePkMap::from_entries(
1692 self.hot
1693 .entries()
1694 .into_iter()
1695 .map(|(key, row_id)| (row_id, key)),
1696 );
1697 }
1698
1699 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1700 if self.schema.primary_key().is_some() {
1701 self.pk_by_row.insert(row_id, key.clone());
1702 }
1703 self.hot.insert(key, row_id);
1704 }
1705
1706 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1710 self.learned_range = Arc::new(HashMap::new());
1711 if self.run_refs.len() != 1 {
1712 return Ok(());
1713 }
1714 let cols: Vec<(u16, usize)> = self
1715 .schema
1716 .indexes
1717 .iter()
1718 .filter(|i| i.kind == IndexKind::LearnedRange)
1719 .map(|i| {
1720 (
1721 i.column_id,
1722 i.options
1723 .learned_range
1724 .as_ref()
1725 .map(|options| options.epsilon)
1726 .unwrap_or(16),
1727 )
1728 })
1729 .collect();
1730 if cols.is_empty() {
1731 return Ok(());
1732 }
1733 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1734 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1735 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1736 _ => return Ok(()),
1737 };
1738 for (cid, epsilon) in cols {
1739 let ty = self
1740 .schema
1741 .columns
1742 .iter()
1743 .find(|c| c.id == cid)
1744 .map(|c| c.ty.clone())
1745 .unwrap_or(TypeId::Int64);
1746 match ty {
1747 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1748 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1749 let pairs: Vec<(i64, u64)> = data
1750 .iter()
1751 .zip(row_ids.iter())
1752 .map(|(v, r)| (*v, *r))
1753 .collect();
1754 Arc::make_mut(&mut self.learned_range).insert(
1755 cid,
1756 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
1757 );
1758 }
1759 }
1760 TypeId::Float64 => {
1761 if let columnar::NativeColumn::Float64 { data, .. } =
1762 reader.column_native(cid)?
1763 {
1764 let pairs: Vec<(f64, u64)> = data
1765 .iter()
1766 .zip(row_ids.iter())
1767 .map(|(v, r)| (*v, *r))
1768 .collect();
1769 Arc::make_mut(&mut self.learned_range).insert(
1770 cid,
1771 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
1772 );
1773 }
1774 }
1775 _ => {}
1776 }
1777 }
1778 Ok(())
1779 }
1780
1781 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1788 if self.indexes_complete {
1789 crate::trace::QueryTrace::record(|t| {
1790 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1791 });
1792 return Ok(());
1793 }
1794 crate::trace::QueryTrace::record(|t| {
1795 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1796 });
1797 self.rebuild_indexes_from_runs()?;
1798 self.build_learned_ranges()?;
1799 self.indexes_complete = true;
1800 let epoch = self.current_epoch();
1801 self.checkpoint_indexes(epoch);
1802 Ok(())
1803 }
1804
1805 fn pending_epoch(&self) -> Epoch {
1806 Epoch(self.epoch.visible().0 + 1)
1807 }
1808
1809 fn is_shared(&self) -> bool {
1812 matches!(self.wal, WalSink::Shared(_))
1813 }
1814
1815 fn ensure_txn_id(&mut self) -> u64 {
1819 if self.current_txn_id == 0 {
1820 let id = match &self.wal {
1821 WalSink::Shared(s) => {
1822 let mut g = s.txn_ids.lock();
1823 let v = *g;
1824 *g = g.wrapping_add(1);
1825 v
1826 }
1827 WalSink::Private(_) => 1,
1828 WalSink::ReadOnly => 1,
1829 };
1830 self.current_txn_id = id;
1831 }
1832 self.current_txn_id
1833 }
1834
1835 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1838 self.ensure_writable()?;
1839 let txn_id = self.ensure_txn_id();
1840 let table_id = self.table_id;
1841 match &mut self.wal {
1842 WalSink::Private(w) => {
1843 w.append_txn(txn_id, op)?;
1844 }
1845 WalSink::Shared(s) => {
1846 s.wal.lock().append(txn_id, table_id, op)?;
1847 }
1848 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
1849 }
1850 Ok(())
1851 }
1852
1853 fn ensure_writable(&self) -> Result<()> {
1854 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
1855 Err(MongrelError::ReadOnlyReplica)
1856 } else {
1857 Ok(())
1858 }
1859 }
1860
1861 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1872 match &self.auth {
1873 Some(checker) => checker.check(&self.name, perm),
1874 None => Ok(()),
1875 }
1876 }
1877 pub fn require_select(&self) -> Result<()> {
1882 self.require(crate::auth_state::RequiredPermission::Select)
1883 }
1884 fn require_insert(&self) -> Result<()> {
1885 self.require(crate::auth_state::RequiredPermission::Insert)
1886 }
1887 #[allow(dead_code)]
1891 fn require_update(&self) -> Result<()> {
1892 self.require(crate::auth_state::RequiredPermission::Update)
1893 }
1894 fn require_delete(&self) -> Result<()> {
1895 self.require(crate::auth_state::RequiredPermission::Delete)
1896 }
1897
1898 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1901 self.require_insert()?;
1902 Ok(self.put_returning(columns)?.0)
1903 }
1904
1905 pub fn put_returning(
1910 &mut self,
1911 mut columns: Vec<(u16, Value)>,
1912 ) -> Result<(RowId, Option<i64>)> {
1913 self.require_insert()?;
1914 let assigned = self.fill_auto_inc(&mut columns)?;
1915 self.apply_defaults(&mut columns)?;
1916 self.schema.validate_values(&columns)?;
1917 let row_id = if self.schema.clustered {
1922 self.derive_clustered_row_id(&columns)?
1923 } else {
1924 self.allocator.alloc()
1925 };
1926 let epoch = self.pending_epoch();
1927 let mut row = Row::new(row_id, epoch);
1928 for (col_id, val) in columns {
1929 row.columns.insert(col_id, val);
1930 }
1931 self.commit_rows(vec![row], assigned.is_some())?;
1932 Ok((row_id, assigned))
1933 }
1934
1935 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1938 self.require_insert()?;
1939 Ok(self
1940 .put_batch_returning(batch)?
1941 .into_iter()
1942 .map(|(r, _)| r)
1943 .collect())
1944 }
1945
1946 pub fn put_batch_returning(
1949 &mut self,
1950 batch: Vec<Vec<(u16, Value)>>,
1951 ) -> Result<Vec<(RowId, Option<i64>)>> {
1952 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1953 for mut cols in batch {
1954 let assigned = self.fill_auto_inc(&mut cols)?;
1955 self.apply_defaults(&mut cols)?;
1956 filled.push((cols, assigned));
1957 }
1958 for (cols, _) in &filled {
1959 self.schema.validate_values(cols)?;
1960 }
1961 let epoch = self.pending_epoch();
1962 let mut rows = Vec::with_capacity(filled.len());
1963 let mut ids = Vec::with_capacity(filled.len());
1964 for (cols, assigned) in filled {
1965 let row_id = if self.schema.clustered {
1966 self.derive_clustered_row_id(&cols)?
1967 } else {
1968 self.allocator.alloc()
1969 };
1970 let mut row = Row::new(row_id, epoch);
1971 for (c, v) in cols {
1972 row.columns.insert(c, v);
1973 }
1974 ids.push((row_id, assigned));
1975 rows.push(row);
1976 }
1977 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1978 self.commit_rows(rows, all_auto_generated)?;
1979 Ok(ids)
1980 }
1981
1982 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1988 self.ensure_writable()?;
1989 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1990 return Ok(None);
1991 };
1992 let pos = columns.iter().position(|(c, _)| *c == cid);
1993 let assigned = match pos {
1994 Some(i) => match &columns[i].1 {
1995 Value::Null => {
1996 let next = self.alloc_auto_inc_value()?;
1997 columns[i].1 = Value::Int64(next);
1998 Some(next)
1999 }
2000 Value::Int64(n) => {
2001 self.advance_auto_inc_past(*n)?;
2002 None
2003 }
2004 other => {
2005 return Err(MongrelError::InvalidArgument(format!(
2006 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2007 other
2008 )))
2009 }
2010 },
2011 None => {
2012 let next = self.alloc_auto_inc_value()?;
2013 columns.push((cid, Value::Int64(next)));
2014 Some(next)
2015 }
2016 };
2017 Ok(assigned)
2018 }
2019
2020 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2026 for col in &self.schema.columns {
2027 let Some(expr) = &col.default_value else {
2028 continue;
2029 };
2030 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2032 continue;
2033 }
2034 let pos = columns.iter().position(|(c, _)| *c == col.id);
2035 let needs_default = match pos {
2036 None => true,
2037 Some(i) => matches!(columns[i].1, Value::Null),
2038 };
2039 if !needs_default {
2040 continue;
2041 }
2042 let v = match expr {
2043 crate::schema::DefaultExpr::Static(v) => v.clone(),
2044 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2045 crate::schema::DefaultExpr::Uuid => {
2046 let mut buf = [0u8; 16];
2047 getrandom::getrandom(&mut buf)
2048 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2049 Value::Uuid(buf)
2050 }
2051 };
2052 match pos {
2053 None => columns.push((col.id, v)),
2054 Some(i) => columns[i].1 = v,
2055 }
2056 }
2057 Ok(())
2058 }
2059
2060 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2062 self.ensure_auto_inc_seeded()?;
2063 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2065 let v = ai.next;
2066 ai.next = ai.next.saturating_add(1);
2067 Ok(v)
2068 }
2069
2070 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2073 self.ensure_auto_inc_seeded()?;
2074 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2075 let floor = used.saturating_add(1).max(1);
2076 if ai.next < floor {
2077 ai.next = floor;
2078 }
2079 Ok(())
2080 }
2081
2082 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2087 let needs_seed = match self.auto_inc {
2088 Some(ai) => !ai.seeded,
2089 None => return Ok(()),
2090 };
2091 if !needs_seed {
2092 return Ok(());
2093 }
2094 if self.seed_empty_auto_inc() {
2095 return Ok(());
2096 }
2097 let cid = self
2098 .auto_inc
2099 .as_ref()
2100 .expect("auto-inc column present")
2101 .column_id;
2102 let max = self.scan_max_int64(cid)?;
2103 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2104 let floor = max.saturating_add(1).max(1);
2105 if ai.next < floor {
2106 ai.next = floor;
2107 }
2108 ai.seeded = true;
2109 Ok(())
2110 }
2111
2112 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2113 if n == 0 || self.auto_inc.is_none() {
2114 return Ok(None);
2115 }
2116 self.ensure_auto_inc_seeded()?;
2117 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2118 let start = ai.next;
2119 ai.next = ai.next.saturating_add(n as i64);
2120 Ok(Some(start))
2121 }
2122
2123 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2127 let mut max: i64 = 0;
2128 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2129 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2130 if *n > max {
2131 max = *n;
2132 }
2133 }
2134 }
2135 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2136 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2137 if *n > max {
2138 max = *n;
2139 }
2140 }
2141 }
2142 for rr in self.run_refs.clone() {
2143 let reader = self.open_reader(rr.run_id)?;
2144 if let Some(stats) = reader.column_page_stats(column_id) {
2145 for s in stats {
2146 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2147 if n > max {
2148 max = n;
2149 }
2150 }
2151 }
2152 } else if reader.has_column(column_id) {
2153 if let columnar::NativeColumn::Int64 { data, validity } =
2154 reader.column_native_shared(column_id)?
2155 {
2156 for (i, n) in data.iter().enumerate() {
2157 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2158 {
2159 max = *n;
2160 }
2161 }
2162 }
2163 }
2164 }
2165 Ok(max)
2166 }
2167
2168 fn seed_empty_auto_inc(&mut self) -> bool {
2169 let Some(ai) = self.auto_inc.as_mut() else {
2170 return false;
2171 };
2172 if ai.seeded || self.live_count != 0 {
2173 return false;
2174 }
2175 if ai.next < 1 {
2176 ai.next = 1;
2177 }
2178 ai.seeded = true;
2179 true
2180 }
2181
2182 fn advance_auto_inc_from_native_columns(
2183 &mut self,
2184 columns: &[(u16, columnar::NativeColumn)],
2185 n: usize,
2186 live_before: u64,
2187 ) -> Result<()> {
2188 let Some(ai) = self.auto_inc.as_mut() else {
2189 return Ok(());
2190 };
2191 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2192 return Ok(());
2193 };
2194 let columnar::NativeColumn::Int64 { data, validity } = col else {
2195 return Err(MongrelError::InvalidArgument(format!(
2196 "AUTO_INCREMENT column {} must be Int64",
2197 ai.column_id
2198 )));
2199 };
2200 let max = if native_int64_strictly_increasing(col, n) {
2201 data.get(n.saturating_sub(1)).copied()
2202 } else {
2203 data.iter()
2204 .take(n)
2205 .enumerate()
2206 .filter_map(|(i, v)| {
2207 if validity.is_empty() || columnar::validity_bit(validity, i) {
2208 Some(*v)
2209 } else {
2210 None
2211 }
2212 })
2213 .max()
2214 };
2215 if let Some(max) = max {
2216 let floor = max.saturating_add(1).max(1);
2217 if ai.next < floor {
2218 ai.next = floor;
2219 }
2220 if ai.seeded || live_before == 0 {
2221 ai.seeded = true;
2222 }
2223 }
2224 Ok(())
2225 }
2226
2227 fn fill_auto_inc_native_columns(
2228 &mut self,
2229 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2230 n: usize,
2231 ) -> Result<()> {
2232 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2233 return Ok(());
2234 };
2235 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2236 if let Some(start) = self.alloc_auto_inc_range(n)? {
2237 columns.push((
2238 cid,
2239 columnar::NativeColumn::Int64 {
2240 data: (start..start.saturating_add(n as i64)).collect(),
2241 validity: vec![0xFF; n.div_ceil(8)],
2242 },
2243 ));
2244 }
2245 return Ok(());
2246 };
2247
2248 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2249 return Err(MongrelError::InvalidArgument(format!(
2250 "AUTO_INCREMENT column {cid} must be Int64"
2251 )));
2252 };
2253 if data.len() < n {
2254 return Err(MongrelError::InvalidArgument(format!(
2255 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2256 data.len()
2257 )));
2258 }
2259 if columnar::all_non_null(validity, n) {
2260 return Ok(());
2261 }
2262 if validity.iter().all(|b| *b == 0) {
2263 if let Some(start) = self.alloc_auto_inc_range(n)? {
2264 for (i, slot) in data.iter_mut().take(n).enumerate() {
2265 *slot = start.saturating_add(i as i64);
2266 }
2267 *validity = vec![0xFF; n.div_ceil(8)];
2268 }
2269 return Ok(());
2270 }
2271
2272 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2273 for (i, slot) in data.iter_mut().enumerate().take(n) {
2274 if columnar::validity_bit(validity, i) {
2275 self.advance_auto_inc_past(*slot)?;
2276 } else {
2277 *slot = self.alloc_auto_inc_value()?;
2278 }
2279 }
2280 *validity = new_validity;
2281 Ok(())
2282 }
2283
2284 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2298 self.ensure_writable()?;
2299 if self.auto_inc.is_none() {
2300 return Ok(None);
2301 }
2302 Ok(Some(self.alloc_auto_inc_value()?))
2303 }
2304
2305 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2311 let payload = bincode::serialize(&rows)?;
2312 self.wal_append_data(Op::Put {
2313 table_id: self.table_id,
2314 rows: payload,
2315 })?;
2316 if self.is_shared() {
2317 self.pending_rows_auto_inc
2318 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2319 self.pending_rows.extend(rows);
2320 } else {
2321 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2322 }
2323 Ok(())
2324 }
2325
2326 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2331 self.apply_put_rows_inner(rows, true)
2332 }
2333
2334 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2335 if check_existing_pk {
2336 self.ensure_indexes_complete()?;
2337 }
2338 if rows.len() == 1 {
2342 let row = rows.into_iter().next().expect("len checked");
2343 return self.apply_put_row_single(row, check_existing_pk);
2344 }
2345 let pk_id = self.schema.primary_key().map(|c| c.id);
2362 let probe = match pk_id {
2363 Some(pid) => {
2364 check_existing_pk
2365 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2366 }
2367 None => false,
2368 };
2369 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2372 for r in rows {
2373 for &cid in r.columns.keys() {
2374 self.pending_put_cols.insert(cid);
2375 }
2376 match pk_id {
2377 Some(pid) if probe || maintain_pk_by_row => {
2378 if let Some(pk_val) = r.columns.get(&pid) {
2379 let key = self.index_lookup_key(pid, pk_val);
2380 if probe {
2381 if let Some(old_rid) = self.hot.get(&key) {
2382 if old_rid != r.row_id {
2383 self.tombstone_row(old_rid, r.committed_epoch, true);
2384 }
2385 }
2386 }
2387 if maintain_pk_by_row {
2388 self.pk_by_row.insert(r.row_id, key);
2389 }
2390 }
2391 }
2392 Some(_) => {}
2393 None => {
2394 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2395 }
2396 }
2397 self.index_row(&r);
2398 self.reservoir.offer(r.row_id.0);
2399 self.memtable.upsert(r);
2400 self.live_count = self.live_count.saturating_add(1);
2403 }
2404 self.data_generation = self.data_generation.wrapping_add(1);
2405 Ok(())
2406 }
2407
2408 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2412 for &cid in row.columns.keys() {
2413 self.pending_put_cols.insert(cid);
2414 }
2415 let epoch = row.committed_epoch;
2416 if let Some(pk_col) = self.schema.primary_key() {
2417 let pk_id = pk_col.id;
2418 if let Some(pk_val) = row.columns.get(&pk_id) {
2419 let maintain_pk_by_row = self.pk_by_row_complete;
2423 if check_existing_pk || maintain_pk_by_row {
2424 let key = self.index_lookup_key(pk_id, pk_val);
2425 if check_existing_pk {
2426 if let Some(old_rid) = self.hot.get(&key) {
2427 if old_rid != row.row_id {
2428 self.tombstone_row(old_rid, epoch, true);
2429 }
2430 }
2431 }
2432 if maintain_pk_by_row {
2433 self.pk_by_row.insert(row.row_id, key);
2434 }
2435 }
2436 }
2437 } else {
2438 self.hot
2439 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2440 }
2441 self.index_row(&row);
2442 self.reservoir.offer(row.row_id.0);
2443 self.memtable.upsert(row);
2444 self.live_count = self.live_count.saturating_add(1);
2445 self.data_generation = self.data_generation.wrapping_add(1);
2446 Ok(())
2447 }
2448
2449 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2452 self.allocator.alloc()
2453 }
2454
2455 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2461 let pk = self.schema.primary_key().ok_or_else(|| {
2462 MongrelError::Schema("clustered table requires a single-column primary key".into())
2463 })?;
2464 let pk_val = columns
2465 .iter()
2466 .find(|(id, _)| *id == pk.id)
2467 .map(|(_, v)| v)
2468 .ok_or_else(|| {
2469 MongrelError::Schema(format!(
2470 "clustered table missing primary key column {} ({})",
2471 pk.id, pk.name
2472 ))
2473 })?;
2474 let key_bytes = pk_val.encode_key();
2475 let mut hash: u64 = 0xcbf29ce484222325;
2477 for &b in &key_bytes {
2478 hash ^= b as u64;
2479 hash = hash.wrapping_mul(0x100000001b3);
2480 }
2481 Ok(RowId(hash.max(1)))
2484 }
2485
2486 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2494 self.ensure_indexes_complete()?;
2495 let n = rows.len();
2496 for r in rows {
2497 for &cid in r.columns.keys() {
2498 self.pending_put_cols.insert(cid);
2499 }
2500 }
2501 let (losers, winner_pks) = self.partition_pk_winners(rows);
2502 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2503 for (key, &row_id) in &winner_pks {
2505 if let Some(old_rid) = self.hot.get(key) {
2506 if old_rid != row_id {
2507 self.tombstone_row(old_rid, epoch, true);
2508 }
2509 }
2510 }
2511 for &loser_rid in &losers {
2514 self.tombstone_row(loser_rid, epoch, false);
2515 }
2516 for (key, row_id) in winner_pks {
2518 self.insert_hot_pk(key, row_id);
2519 }
2520 if self.schema.primary_key().is_none() {
2521 for r in rows {
2522 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2523 }
2524 }
2525 for r in rows {
2526 self.allocator.advance_to(r.row_id);
2527 if !losers.contains(&r.row_id) {
2528 self.index_row(r);
2529 }
2530 }
2531 for r in rows {
2532 if !losers.contains(&r.row_id) {
2533 self.reservoir.offer(r.row_id.0);
2534 }
2535 }
2536 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2537 self.data_generation = self.data_generation.wrapping_add(1);
2538 Ok(())
2539 }
2540
2541 pub(crate) fn recover_apply(
2546 &mut self,
2547 rows: Vec<Row>,
2548 deletes: Vec<(RowId, Epoch)>,
2549 ) -> Result<()> {
2550 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2554 std::collections::BTreeMap::new();
2555 for row in rows {
2556 self.allocator.advance_to(row.row_id);
2557 if let Some(ai) = self.auto_inc.as_mut() {
2562 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2563 if *n + 1 > ai.next {
2564 ai.next = *n + 1;
2565 }
2566 }
2567 }
2568 by_epoch.entry(row.committed_epoch).or_default().push(row);
2569 }
2570 for (epoch, group) in by_epoch {
2571 let (losers, winner_pks) = self.partition_pk_winners(&group);
2572 for (key, &row_id) in &winner_pks {
2574 if let Some(old_rid) = self.hot.get(key) {
2575 if old_rid != row_id {
2576 self.tombstone_row(old_rid, epoch, false);
2577 }
2578 }
2579 }
2580 for (key, row_id) in winner_pks {
2581 self.insert_hot_pk(key, row_id);
2582 }
2583 if self.schema.primary_key().is_none() {
2584 for r in &group {
2585 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2586 }
2587 }
2588 for r in &group {
2589 if !losers.contains(&r.row_id) {
2590 self.memtable.upsert(r.clone());
2591 self.index_row(r);
2592 }
2593 }
2594 }
2595 for (rid, epoch) in deletes {
2596 self.memtable.tombstone(rid, epoch);
2597 self.remove_hot_for_row(rid, epoch);
2598 }
2599 self.reservoir_complete = false;
2602 Ok(())
2603 }
2604
2605 pub(crate) fn flushed_epoch(&self) -> u64 {
2607 self.flushed_epoch
2608 }
2609
2610 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
2611 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
2612 }
2613
2614 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2616 self.schema.validate_values(cells)
2617 }
2618
2619 fn validate_columns_not_null(
2623 &self,
2624 columns: &[(u16, columnar::NativeColumn)],
2625 n: usize,
2626 ) -> Result<()> {
2627 let by_id: HashMap<u16, &columnar::NativeColumn> =
2628 columns.iter().map(|(id, c)| (*id, c)).collect();
2629 for col in &self.schema.columns {
2630 if !col.flags.contains(ColumnFlags::NULLABLE) {
2631 match by_id.get(&col.id) {
2632 None => {
2633 return Err(MongrelError::InvalidArgument(format!(
2634 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2635 col.name, col.id
2636 )));
2637 }
2638 Some(c) => {
2639 if c.null_count(n) != 0 {
2640 return Err(MongrelError::InvalidArgument(format!(
2641 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2642 col.name, col.id
2643 )));
2644 }
2645 }
2646 }
2647 }
2648 if let TypeId::Enum { variants } = &col.ty {
2649 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
2650 if by_id.contains_key(&col.id) {
2651 return Err(MongrelError::InvalidArgument(format!(
2652 "column '{}' ({}) enum requires a bytes column",
2653 col.name, col.id
2654 )));
2655 }
2656 continue;
2657 };
2658 for index in 0..n {
2659 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
2660 continue;
2661 };
2662 if !variants.iter().any(|variant| variant.as_bytes() == value) {
2663 return Err(MongrelError::InvalidArgument(format!(
2664 "column '{}' ({}) enum value {:?} is not one of {:?}",
2665 col.name,
2666 col.id,
2667 String::from_utf8_lossy(value),
2668 variants
2669 )));
2670 }
2671 }
2672 }
2673 }
2674 Ok(())
2675 }
2676
2677 fn bulk_pk_winner_indices(
2682 &self,
2683 columns: &[(u16, columnar::NativeColumn)],
2684 n: usize,
2685 ) -> Option<Vec<usize>> {
2686 let pk_col = self.schema.primary_key()?;
2687 let pk_id = pk_col.id;
2688 let pk_ty = pk_col.ty.clone();
2689 let by_id: HashMap<u16, &columnar::NativeColumn> =
2690 columns.iter().map(|(id, c)| (*id, c)).collect();
2691 let pk_native = by_id.get(&pk_id)?;
2692 if native_int64_strictly_increasing(pk_native, n) {
2693 return None;
2694 }
2695 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2697 let mut null_pk_rows: Vec<usize> = Vec::new();
2698 for i in 0..n {
2699 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2700 Some(key) => {
2701 last.insert(key, i);
2702 }
2703 None => null_pk_rows.push(i),
2704 }
2705 }
2706 let mut winners: HashSet<usize> = last.values().copied().collect();
2707 for i in null_pk_rows {
2708 winners.insert(i);
2709 }
2710 Some((0..n).filter(|i| winners.contains(i)).collect())
2711 }
2712
2713 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2715 self.require_delete()?;
2716 let epoch = self.pending_epoch();
2717 self.wal_append_data(Op::Delete {
2718 table_id: self.table_id,
2719 row_ids: vec![row_id],
2720 })?;
2721 if self.is_shared() {
2722 self.pending_dels.push(row_id);
2723 } else {
2724 self.apply_delete(row_id, epoch);
2725 }
2726 Ok(())
2727 }
2728
2729 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2730 let pre = self.get(row_id, self.snapshot());
2731 self.delete(row_id)?;
2732 Ok(pre.map(|row| {
2733 let mut columns: Vec<_> = row.columns.into_iter().collect();
2734 columns.sort_by_key(|(id, _)| *id);
2735 OwnedRow { columns }
2736 }))
2737 }
2738
2739 pub fn truncate(&mut self) -> Result<()> {
2741 self.require_delete()?;
2742 let epoch = self.pending_epoch();
2743 self.wal_append_data(Op::TruncateTable {
2744 table_id: self.table_id,
2745 })?;
2746 self.pending_rows.clear();
2747 self.pending_rows_auto_inc.clear();
2748 self.pending_dels.clear();
2749 self.pending_truncate = Some(epoch);
2750 Ok(())
2751 }
2752
2753 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2755 for rr in std::mem::take(&mut self.run_refs) {
2756 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2757 }
2758 for r in std::mem::take(&mut self.retiring) {
2759 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2760 }
2761 self.memtable = Memtable::new();
2762 self.mutable_run = MutableRun::new();
2763 self.hot = HotIndex::new();
2764 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2765 self.bitmap = bitmap;
2766 self.ann = ann;
2767 self.fm = fm;
2768 self.sparse = sparse;
2769 self.minhash = minhash;
2770 self.learned_range = Arc::new(HashMap::new());
2771 self.pk_by_row.clear();
2772 self.pk_by_row_complete = false;
2773 self.live_count = 0;
2774 self.reservoir = crate::reservoir::Reservoir::default();
2775 self.reservoir_complete = true;
2776 self.had_deletes = true;
2777 self.agg_cache = Arc::new(HashMap::new());
2778 self.global_idx_epoch = 0;
2779 self.indexes_complete = true;
2780 self.pending_delete_rids.clear();
2781 self.pending_put_cols.clear();
2782 self.pending_rows.clear();
2783 self.pending_rows_auto_inc.clear();
2784 self.pending_dels.clear();
2785 self.clear_result_cache();
2786 self.invalidate_index_checkpoint();
2787 self.data_generation = self.data_generation.wrapping_add(1);
2788 Ok(())
2789 }
2790
2791 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2794 self.remove_hot_for_row(row_id, epoch);
2795 self.tombstone_row(row_id, epoch, true);
2796 self.data_generation = self.data_generation.wrapping_add(1);
2797 }
2798
2799 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2803 let tombstone = Row {
2804 row_id,
2805 committed_epoch: epoch,
2806 columns: std::collections::HashMap::new(),
2807 deleted: true,
2808 };
2809 self.memtable.upsert(tombstone);
2810 self.pk_by_row.remove(&row_id);
2811 if adjust_live_count {
2812 self.live_count = self.live_count.saturating_sub(1);
2813 }
2814 self.pending_delete_rids.insert(row_id.0 as u32);
2816 self.had_deletes = true;
2819 self.agg_cache = Arc::new(HashMap::new());
2820 }
2821
2822 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2826 let Some(pk_col) = self.schema.primary_key() else {
2827 return;
2828 };
2829 if self.pk_by_row_complete {
2832 if let Some(key) = self.pk_by_row.remove(&row_id) {
2833 if self.hot.get(&key) == Some(row_id) {
2834 self.hot.remove(&key);
2835 }
2836 }
2837 return;
2838 }
2839 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2858 if self.indexes_complete {
2859 let pk_val = self
2860 .memtable
2861 .get_version(row_id, lookup_epoch)
2862 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2863 .or_else(|| {
2864 self.mutable_run
2865 .get_version(row_id, lookup_epoch)
2866 .filter(|(_, r)| !r.deleted)
2867 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2868 })
2869 .or_else(|| {
2870 self.run_refs.iter().find_map(|rr| {
2871 let mut reader = self.open_reader(rr.run_id).ok()?;
2872 let (_, deleted, val) = reader
2873 .get_version_column(row_id, lookup_epoch, pk_col.id)
2874 .ok()??;
2875 if deleted {
2876 return None;
2877 }
2878 val
2879 })
2880 });
2881 if let Some(pk_val) = pk_val {
2882 let key = self.index_lookup_key(pk_col.id, &pk_val);
2883 if self.hot.get(&key) == Some(row_id) {
2884 self.hot.remove(&key);
2885 }
2886 return;
2887 }
2888 }
2889 self.refresh_pk_by_row_from_hot();
2894 if let Some(key) = self.pk_by_row.remove(&row_id) {
2895 if self.hot.get(&key) == Some(row_id) {
2896 self.hot.remove(&key);
2897 }
2898 }
2899 }
2900
2901 fn partition_pk_winners(
2906 &self,
2907 rows: &[Row],
2908 ) -> (
2909 std::collections::HashSet<RowId>,
2910 std::collections::HashMap<Vec<u8>, RowId>,
2911 ) {
2912 let mut losers = std::collections::HashSet::new();
2913 let Some(pk_col) = self.schema.primary_key() else {
2914 return (losers, std::collections::HashMap::new());
2915 };
2916 let pk_id = pk_col.id;
2917 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2918 std::collections::HashMap::new();
2919 for r in rows {
2920 let Some(pk_val) = r.columns.get(&pk_id) else {
2921 continue;
2922 };
2923 let key = self.index_lookup_key(pk_id, pk_val);
2924 if let Some(&old_rid) = winners.get(&key) {
2925 losers.insert(old_rid);
2926 }
2927 winners.insert(key, r.row_id);
2928 }
2929 (losers, winners)
2930 }
2931
2932 fn index_row(&mut self, row: &Row) {
2933 if row.deleted {
2934 return;
2935 }
2936 let any_predicate = self
2944 .schema
2945 .indexes
2946 .iter()
2947 .any(|idx| idx.predicate.is_some());
2948 if any_predicate {
2949 let columns_map: HashMap<u16, &Value> =
2950 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2951 let name_to_id: HashMap<&str, u16> = self
2952 .schema
2953 .columns
2954 .iter()
2955 .map(|c| (c.name.as_str(), c.id))
2956 .collect();
2957 for idx in &self.schema.indexes {
2958 if let Some(pred) = &idx.predicate {
2959 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2960 continue; }
2962 }
2963 index_into_single(
2965 idx,
2966 &self.schema,
2967 row,
2968 &mut self.hot,
2969 &mut self.bitmap,
2970 &mut self.ann,
2971 &mut self.fm,
2972 &mut self.sparse,
2973 &mut self.minhash,
2974 );
2975 }
2976 return;
2977 }
2978 if self.column_keys.is_empty() {
2982 index_into(
2983 &self.schema,
2984 row,
2985 &mut self.hot,
2986 &mut self.bitmap,
2987 &mut self.ann,
2988 &mut self.fm,
2989 &mut self.sparse,
2990 &mut self.minhash,
2991 );
2992 return;
2993 }
2994 let effective_row = self.tokenized_for_indexes(row);
2995 index_into(
2996 &self.schema,
2997 &effective_row,
2998 &mut self.hot,
2999 &mut self.bitmap,
3000 &mut self.ann,
3001 &mut self.fm,
3002 &mut self.sparse,
3003 &mut self.minhash,
3004 );
3005 }
3006
3007 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3013 if self.column_keys.is_empty() {
3014 return row.clone();
3015 }
3016 #[cfg(feature = "encryption")]
3017 {
3018 use crate::encryption::SCHEME_HMAC_EQ;
3019 let mut tok = row.clone();
3020 for (&cid, &(_, scheme)) in &self.column_keys {
3021 if scheme != SCHEME_HMAC_EQ {
3022 continue;
3023 }
3024 if let Some(v) = tok.columns.get(&cid).cloned() {
3025 if let Some(t) = self.tokenize_value(cid, &v) {
3026 tok.columns.insert(cid, t);
3027 }
3028 }
3029 }
3030 tok
3031 }
3032 #[cfg(not(feature = "encryption"))]
3033 {
3034 row.clone()
3035 }
3036 }
3037
3038 pub fn commit(&mut self) -> Result<Epoch> {
3043 self.ensure_writable()?;
3044 if self.is_shared() {
3045 self.commit_shared()
3046 } else {
3047 self.commit_private()
3048 }
3049 }
3050
3051 fn commit_private(&mut self) -> Result<Epoch> {
3053 let commit_lock = Arc::clone(&self.commit_lock);
3057 let _g = commit_lock.lock();
3058 let new_epoch = self.epoch.bump_assigned();
3059 let txn_id = self.current_txn_id;
3060 match &mut self.wal {
3064 WalSink::Private(w) => {
3065 w.append_txn(
3066 txn_id,
3067 Op::TxnCommit {
3068 epoch: new_epoch.0,
3069 added_runs: Vec::new(),
3070 },
3071 )?;
3072 w.sync()?;
3073 }
3074 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3075 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3076 }
3077 if let Some(epoch) = self.pending_truncate.take() {
3079 self.apply_truncate(epoch)?;
3080 }
3081 self.invalidate_pending_cache();
3082 self.persist_manifest(new_epoch)?;
3083 self.epoch.publish_in_order(new_epoch);
3087 self.current_txn_id += 1;
3088 self.data_generation = self.data_generation.wrapping_add(1);
3089 Ok(new_epoch)
3090 }
3091
3092 fn commit_shared(&mut self) -> Result<Epoch> {
3098 use std::sync::atomic::Ordering;
3099 let s = match &self.wal {
3100 WalSink::Shared(s) => s.clone(),
3101 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3102 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3103 };
3104 if s.poisoned.load(Ordering::Relaxed) {
3105 return Err(MongrelError::Other(
3106 "database poisoned by fsync error".into(),
3107 ));
3108 }
3109 let commit_lock = Arc::clone(&self.commit_lock);
3116 let _g = commit_lock.lock();
3117 let txn_id = self.ensure_txn_id();
3120 let (new_epoch, commit_seq) = {
3121 let mut wal = s.wal.lock();
3122 let new_epoch = self.epoch.bump_assigned();
3123 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
3124 (new_epoch, seq)
3125 };
3126 s.group
3127 .await_durable(&s.wal, commit_seq)
3128 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
3129
3130 if self.pending_truncate.take().is_some() {
3133 self.apply_truncate(new_epoch)?;
3134 }
3135 let mut rows = std::mem::take(&mut self.pending_rows);
3136 if !rows.is_empty() {
3137 for r in &mut rows {
3138 r.committed_epoch = new_epoch;
3139 }
3140 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3141 let all_auto_generated =
3142 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3143 self.apply_put_rows_inner(rows, !all_auto_generated)?;
3144 } else {
3145 self.pending_rows_auto_inc.clear();
3146 }
3147 let dels = std::mem::take(&mut self.pending_dels);
3148 for rid in dels {
3149 self.apply_delete(rid, new_epoch);
3150 }
3151
3152 self.invalidate_pending_cache();
3153 self.persist_manifest(new_epoch)?;
3154 self.epoch.publish_in_order(new_epoch);
3155 let _ = s.change_wake.send(());
3156 self.current_txn_id = 0;
3158 self.data_generation = self.data_generation.wrapping_add(1);
3159 Ok(new_epoch)
3160 }
3161
3162 pub fn flush(&mut self) -> Result<Epoch> {
3170 self.ensure_indexes_complete()?;
3171 let epoch = self.commit()?;
3172 let rows = self.memtable.drain_sorted();
3173 if !rows.is_empty() {
3174 self.mutable_run.insert_many(rows);
3175 }
3176 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3177 self.spill_mutable_run(epoch)?;
3178 self.mark_flushed(epoch)?;
3182 self.persist_manifest(epoch)?;
3183 self.build_learned_ranges()?;
3184 self.checkpoint_indexes(epoch);
3187 }
3188 Ok(epoch)
3191 }
3192
3193 pub fn force_flush(&mut self) -> Result<Epoch> {
3202 let saved = self.mutable_run_spill_bytes;
3203 self.mutable_run_spill_bytes = 1;
3204 let result = self.flush();
3205 self.mutable_run_spill_bytes = saved;
3206 result
3207 }
3208
3209 pub fn close(&mut self) -> Result<()> {
3216 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3217 self.force_flush()?;
3218 }
3219 Ok(())
3220 }
3221
3222 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3229 let op = Op::Flush {
3230 table_id: self.table_id,
3231 flushed_epoch: epoch.0,
3232 };
3233 match &mut self.wal {
3234 WalSink::Private(w) => {
3235 w.append_system(op)?;
3236 w.sync()?;
3237 }
3238 WalSink::Shared(s) => {
3239 s.wal.lock().append_system(op)?;
3244 }
3245 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3246 }
3247 self.flushed_epoch = epoch.0;
3248 if matches!(self.wal, WalSink::Private(_)) {
3249 self.rotate_wal(epoch)?;
3250 }
3251 Ok(())
3252 }
3253
3254 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3258 let rows = self.mutable_run.drain_sorted();
3259 if rows.is_empty() {
3260 return Ok(());
3261 }
3262 let run_id = self.next_run_id;
3263 self.next_run_id += 1;
3264 let path = self.run_path(run_id);
3265 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3266 if let Some(kek) = &self.kek {
3267 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3268 }
3269 let header = writer.write(&path, &rows)?;
3270 self.run_refs.push(RunRef {
3271 run_id: run_id as u128,
3272 level: 0,
3273 epoch_created: epoch.0,
3274 row_count: header.row_count,
3275 });
3276 Ok(())
3277 }
3278
3279 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3283 self.mutable_run_spill_bytes = bytes.max(1);
3284 }
3285
3286 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3290 self.compaction_zstd_level = level;
3291 }
3292
3293 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3297 self.result_cache.lock().set_max_bytes(max_bytes);
3298 }
3299
3300 pub(crate) fn clear_result_cache(&mut self) {
3304 self.result_cache.lock().clear();
3305 }
3306
3307 pub fn mutable_run_len(&self) -> usize {
3309 self.mutable_run.len()
3310 }
3311
3312 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3315 self.mutable_run.drain_sorted()
3316 }
3317
3318 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3323 let epoch = self.commit()?;
3324 let n = batch.len();
3325 if n == 0 {
3326 return Ok(epoch);
3327 }
3328 for row in &batch {
3329 self.schema.validate_values(row)?;
3330 }
3331 let live_before = self.live_count;
3332 self.spill_mutable_run(epoch)?;
3336 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3337 && self.indexes_complete
3338 && self.run_refs.is_empty()
3339 && self.memtable.is_empty()
3340 && self.mutable_run.is_empty();
3341 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3347 use rayon::prelude::*;
3348 self.schema
3349 .columns
3350 .par_iter()
3351 .map(|cdef| {
3352 (
3353 cdef.id,
3354 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3355 )
3356 })
3357 .collect::<Vec<_>>()
3358 };
3359 drop(batch);
3360 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3365 self.validate_columns_not_null(&user_columns, n)?;
3366 let winner_idx = self
3367 .bulk_pk_winner_indices(&user_columns, n)
3368 .filter(|idx| idx.len() != n);
3369 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3370 match winner_idx.as_deref() {
3371 Some(idx) => {
3372 let compacted = user_columns
3373 .iter()
3374 .map(|(id, c)| (*id, c.gather(idx)))
3375 .collect();
3376 (compacted, idx.len())
3377 }
3378 None => (std::mem::take(&mut user_columns), n),
3379 };
3380 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3381 let first = self.allocator.alloc_range(write_n as u64).0;
3382 for rid in first..first + write_n as u64 {
3383 self.reservoir.offer(rid);
3384 }
3385 let run_id = self.next_run_id;
3386 self.next_run_id += 1;
3387 let path = self.run_path(run_id);
3388 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3389 .clean(true)
3390 .with_lz4()
3391 .with_native_endian();
3392 if let Some(kek) = &self.kek {
3393 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3394 }
3395 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3396 self.run_refs.push(RunRef {
3397 run_id: run_id as u128,
3398 level: 0,
3399 epoch_created: epoch.0,
3400 row_count: header.row_count,
3401 });
3402 self.live_count = self.live_count.saturating_add(write_n as u64);
3403 if eager_index_build {
3404 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3405 self.index_columns_bulk(&write_columns, &row_ids);
3406 self.indexes_complete = true;
3407 self.build_learned_ranges()?;
3408 } else {
3409 self.indexes_complete = false;
3410 }
3411 self.mark_flushed(epoch)?;
3412 self.persist_manifest(epoch)?;
3413 if eager_index_build {
3414 self.checkpoint_indexes(epoch);
3415 }
3416 self.clear_result_cache();
3417 Ok(epoch)
3418 }
3419
3420 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3423 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3424 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3425 let segment_no = segment
3428 .file_stem()
3429 .and_then(|s| s.to_str())
3430 .and_then(|s| s.strip_prefix("seg-"))
3431 .and_then(|s| s.parse::<u64>().ok())
3432 .unwrap_or(0);
3433 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3434 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3435 wal.sync()?;
3436 self.wal = WalSink::Private(wal);
3437 Ok(())
3438 }
3439
3440 pub(crate) fn invalidate_pending_cache(&mut self) {
3445 self.result_cache
3446 .lock()
3447 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3448 self.pending_delete_rids.clear();
3449 self.pending_put_cols.clear();
3450 }
3451
3452 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3453 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3454 m.current_epoch = epoch.0;
3455 m.next_row_id = self.allocator.current().0;
3456 m.runs = self.run_refs.clone();
3457 m.live_count = self.live_count;
3458 m.global_idx_epoch = self.global_idx_epoch;
3459 m.flushed_epoch = self.flushed_epoch;
3460 m.retiring = self.retiring.clone();
3461 m.auto_inc_next = match self.auto_inc {
3465 Some(ai) if ai.seeded => ai.next,
3466 _ => 0,
3467 };
3468 m.ttl = self.ttl;
3469 let meta_dek = self.manifest_meta_dek();
3470 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3471 Ok(())
3472 }
3473
3474 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3480 if !self.indexes_complete {
3483 return;
3484 }
3485 let snap = global_idx::IndexSnapshot {
3486 hot: &self.hot,
3487 bitmap: &self.bitmap,
3488 ann: &self.ann,
3489 fm: &self.fm,
3490 sparse: &self.sparse,
3491 minhash: &self.minhash,
3492 learned_range: &self.learned_range,
3493 };
3494 let idx_dek = self.idx_dek();
3496 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3497 .is_ok()
3498 {
3499 self.global_idx_epoch = epoch.0;
3500 let _ = self.persist_manifest(epoch);
3501 }
3502 }
3503
3504 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3507 self.global_idx_epoch = 0;
3508 global_idx::remove(&self.dir);
3509 let _ = self.persist_manifest(self.epoch.visible());
3510 }
3511
3512 pub(crate) fn mark_indexes_incomplete(&mut self) {
3513 self.indexes_complete = false;
3514 self.invalidate_index_checkpoint();
3515 }
3516
3517 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3520 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3521 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3522 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3523 best = Some((epoch, row));
3524 }
3525 }
3526 for rr in &self.run_refs {
3527 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3528 continue;
3529 };
3530 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3531 continue;
3532 };
3533 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3534 best = Some((epoch, row));
3535 }
3536 }
3537 let now_nanos = unix_nanos_now();
3538 match best {
3539 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
3540 Some((_, r)) => Some(r),
3541 None => None,
3542 }
3543 }
3544
3545 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3549 self.visible_rows_at_time(snapshot, unix_nanos_now())
3550 }
3551
3552 #[doc(hidden)]
3553 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
3554 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3555 let mut fold = |row: Row| {
3556 best.entry(row.row_id.0)
3557 .and_modify(|e| {
3558 if row.committed_epoch > e.0 {
3559 *e = (row.committed_epoch, row.clone());
3560 }
3561 })
3562 .or_insert_with(|| (row.committed_epoch, row));
3563 };
3564 for row in self.memtable.visible_versions(snapshot.epoch) {
3565 fold(row);
3566 }
3567 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3568 fold(row);
3569 }
3570 for rr in &self.run_refs {
3571 let mut reader = self.open_reader(rr.run_id)?;
3572 for row in reader.visible_versions(snapshot.epoch)? {
3573 fold(row);
3574 }
3575 }
3576 let mut out: Vec<Row> = best
3577 .into_values()
3578 .filter_map(|(_, r)| {
3579 if r.deleted || self.row_expired_at(&r, now_nanos) {
3580 None
3581 } else {
3582 Some(r)
3583 }
3584 })
3585 .collect();
3586 out.sort_by_key(|r| r.row_id);
3587 Ok(out)
3588 }
3589
3590 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3597 if self.ttl.is_none()
3598 && self.memtable.is_empty()
3599 && self.mutable_run.is_empty()
3600 && self.run_refs.len() == 1
3601 {
3602 let rr = self.run_refs[0].clone();
3603 let mut reader = self.open_reader(rr.run_id)?;
3604 let idxs = reader.visible_indices(snapshot.epoch)?;
3605 let mut cols = Vec::with_capacity(self.schema.columns.len());
3606 for cdef in &self.schema.columns {
3607 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3608 }
3609 return Ok(cols);
3610 }
3611 let rows = self.visible_rows(snapshot)?;
3613 let mut cols: Vec<(u16, Vec<Value>)> = self
3614 .schema
3615 .columns
3616 .iter()
3617 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3618 .collect();
3619 for r in &rows {
3620 for (cid, vec) in cols.iter_mut() {
3621 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3622 }
3623 }
3624 Ok(cols)
3625 }
3626
3627 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3629 let row_id = self.hot.get(key)?;
3630 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
3631 Some(row_id)
3632 } else {
3633 None
3634 }
3635 }
3636
3637 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3642 self.query_at_with_allowed(q, self.snapshot(), None)
3643 }
3644
3645 pub fn query_at_with_allowed(
3648 &mut self,
3649 q: &crate::query::Query,
3650 snapshot: Snapshot,
3651 allowed: Option<&std::collections::HashSet<RowId>>,
3652 ) -> Result<Vec<Row>> {
3653 self.query_at_with_allowed_after(q, snapshot, allowed, None)
3654 }
3655
3656 #[doc(hidden)]
3657 pub fn query_at_with_allowed_after(
3658 &mut self,
3659 q: &crate::query::Query,
3660 snapshot: Snapshot,
3661 allowed: Option<&std::collections::HashSet<RowId>>,
3662 after_row_id: Option<RowId>,
3663 ) -> Result<Vec<Row>> {
3664 self.query_at_with_allowed_after_at_time(
3665 q,
3666 snapshot,
3667 allowed,
3668 after_row_id,
3669 unix_nanos_now(),
3670 )
3671 }
3672
3673 #[doc(hidden)]
3674 pub fn query_at_with_allowed_after_at_time(
3675 &mut self,
3676 q: &crate::query::Query,
3677 snapshot: Snapshot,
3678 allowed: Option<&std::collections::HashSet<RowId>>,
3679 after_row_id: Option<RowId>,
3680 query_time_nanos: i64,
3681 ) -> Result<Vec<Row>> {
3682 self.require_select()?;
3683 self.ensure_indexes_complete()?;
3684 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3685 return Err(MongrelError::InvalidArgument(format!(
3686 "query exceeds {} conditions",
3687 crate::query::MAX_HARD_CONDITIONS
3688 )));
3689 }
3690 if let Some(limit) = q.limit {
3691 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
3692 return Err(MongrelError::InvalidArgument(format!(
3693 "query limit must be between 1 and {}",
3694 crate::query::MAX_FINAL_LIMIT
3695 )));
3696 }
3697 }
3698 if q.offset > crate::query::MAX_QUERY_OFFSET {
3699 return Err(MongrelError::InvalidArgument(format!(
3700 "query offset exceeds {}",
3701 crate::query::MAX_QUERY_OFFSET
3702 )));
3703 }
3704 self.query_conditions_at(
3705 &q.conditions,
3706 snapshot,
3707 allowed,
3708 q.limit,
3709 q.offset,
3710 after_row_id,
3711 query_time_nanos,
3712 )
3713 }
3714
3715 #[doc(hidden)]
3718 pub fn query_all_at(
3719 &mut self,
3720 conditions: &[crate::query::Condition],
3721 snapshot: Snapshot,
3722 ) -> Result<Vec<Row>> {
3723 self.require_select()?;
3724 self.ensure_indexes_complete()?;
3725 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
3726 return Err(MongrelError::InvalidArgument(format!(
3727 "query exceeds {} conditions",
3728 crate::query::MAX_HARD_CONDITIONS
3729 )));
3730 }
3731 self.query_conditions_at(conditions, snapshot, None, None, 0, None, unix_nanos_now())
3732 }
3733
3734 #[allow(clippy::too_many_arguments)]
3735 fn query_conditions_at(
3736 &self,
3737 conditions: &[crate::query::Condition],
3738 snapshot: Snapshot,
3739 allowed: Option<&std::collections::HashSet<RowId>>,
3740 limit: Option<usize>,
3741 offset: usize,
3742 after_row_id: Option<RowId>,
3743 query_time_nanos: i64,
3744 ) -> Result<Vec<Row>> {
3745 crate::trace::QueryTrace::record(|t| {
3746 t.run_count = self.run_refs.len();
3747 t.memtable_rows = self.memtable.len();
3748 t.mutable_run_rows = self.mutable_run.len();
3749 });
3750 if conditions.is_empty() {
3754 crate::trace::QueryTrace::record(|t| {
3755 t.scan_mode = crate::trace::ScanMode::Materialized;
3756 t.row_materialized = true;
3757 });
3758 let mut rows = self.visible_rows_at_time(snapshot, query_time_nanos)?;
3759 if let Some(allowed) = allowed {
3760 rows.retain(|row| allowed.contains(&row.row_id));
3761 }
3762 if let Some(after_row_id) = after_row_id {
3763 rows.retain(|row| row.row_id > after_row_id);
3764 }
3765 rows.drain(..offset.min(rows.len()));
3766 if let Some(limit) = limit {
3767 rows.truncate(limit);
3768 }
3769 return Ok(rows);
3770 }
3771 crate::trace::QueryTrace::record(|t| {
3772 t.conditions_pushed = conditions.len();
3773 t.scan_mode = crate::trace::ScanMode::Materialized;
3774 t.row_materialized = true;
3775 });
3776 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
3783 ordered.sort_by_key(|c| condition_cost_rank(c));
3784 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3785 for c in &ordered {
3786 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
3787 let empty = s.is_empty();
3788 sets.push(s);
3789 if empty {
3790 break;
3791 }
3792 }
3793 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3794 if let Some(allowed) = allowed {
3795 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
3796 }
3797 if let Some(after_row_id) = after_row_id {
3798 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
3799 rids.drain(..first);
3800 }
3801 rids.drain(..offset.min(rids.len()));
3802 if let Some(limit) = limit {
3803 rids.truncate(limit);
3804 }
3805 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos)
3806 }
3807
3808 pub fn retrieve(
3810 &mut self,
3811 retriever: &crate::query::Retriever,
3812 ) -> Result<Vec<crate::query::RetrieverHit>> {
3813 self.retrieve_with_allowed(retriever, None)
3814 }
3815
3816 pub fn retrieve_at(
3817 &mut self,
3818 retriever: &crate::query::Retriever,
3819 snapshot: Snapshot,
3820 allowed: Option<&std::collections::HashSet<RowId>>,
3821 ) -> Result<Vec<crate::query::RetrieverHit>> {
3822 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
3823 }
3824
3825 pub fn retrieve_with_allowed(
3828 &mut self,
3829 retriever: &crate::query::Retriever,
3830 allowed: Option<&std::collections::HashSet<RowId>>,
3831 ) -> Result<Vec<crate::query::RetrieverHit>> {
3832 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
3833 }
3834
3835 pub fn retrieve_at_with_allowed(
3836 &mut self,
3837 retriever: &crate::query::Retriever,
3838 snapshot: Snapshot,
3839 allowed: Option<&std::collections::HashSet<RowId>>,
3840 ) -> Result<Vec<crate::query::RetrieverHit>> {
3841 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
3842 }
3843
3844 pub fn retrieve_at_with_allowed_and_context(
3845 &mut self,
3846 retriever: &crate::query::Retriever,
3847 snapshot: Snapshot,
3848 allowed: Option<&std::collections::HashSet<RowId>>,
3849 context: Option<&crate::query::AiExecutionContext>,
3850 ) -> Result<Vec<crate::query::RetrieverHit>> {
3851 self.require_select()?;
3852 self.ensure_indexes_complete()?;
3853 self.validate_retriever(retriever)?;
3854 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
3855 }
3856
3857 pub fn retrieve_at_with_candidate_authorization_and_context(
3858 &mut self,
3859 retriever: &crate::query::Retriever,
3860 snapshot: Snapshot,
3861 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3862 context: Option<&crate::query::AiExecutionContext>,
3863 ) -> Result<Vec<crate::query::RetrieverHit>> {
3864 self.require_select()?;
3865 self.ensure_indexes_complete()?;
3866 self.retrieve_at_with_candidate_authorization_on_generation(
3867 retriever,
3868 snapshot,
3869 authorization,
3870 context,
3871 )
3872 }
3873
3874 #[doc(hidden)]
3875 pub fn retrieve_at_with_candidate_authorization_on_generation(
3876 &self,
3877 retriever: &crate::query::Retriever,
3878 snapshot: Snapshot,
3879 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
3880 context: Option<&crate::query::AiExecutionContext>,
3881 ) -> Result<Vec<crate::query::RetrieverHit>> {
3882 self.require_select()?;
3883 self.validate_retriever(retriever)?;
3884 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
3885 }
3886
3887 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
3888 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
3889 let (column_id, k) = match retriever {
3890 Retriever::Ann {
3891 column_id,
3892 query,
3893 k,
3894 } => {
3895 let index = self.ann.get(column_id).ok_or_else(|| {
3896 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
3897 })?;
3898 if query.len() != index.dim() {
3899 return Err(MongrelError::InvalidArgument(format!(
3900 "ANN query dimension must be {}, got {}",
3901 index.dim(),
3902 query.len()
3903 )));
3904 }
3905 if query.iter().any(|value| !value.is_finite()) {
3906 return Err(MongrelError::InvalidArgument(
3907 "ANN query values must be finite".into(),
3908 ));
3909 }
3910 (*column_id, *k)
3911 }
3912 Retriever::Sparse {
3913 column_id,
3914 query,
3915 k,
3916 } => {
3917 if !self.sparse.contains_key(column_id) {
3918 return Err(MongrelError::InvalidArgument(format!(
3919 "column {column_id} has no Sparse index"
3920 )));
3921 }
3922 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
3923 return Err(MongrelError::InvalidArgument(
3924 "Sparse query must be non-empty with finite weights".into(),
3925 ));
3926 }
3927 if query.len() > MAX_SPARSE_TERMS {
3928 return Err(MongrelError::InvalidArgument(format!(
3929 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
3930 )));
3931 }
3932 (*column_id, *k)
3933 }
3934 Retriever::MinHash {
3935 column_id,
3936 members,
3937 k,
3938 } => {
3939 if !self.minhash.contains_key(column_id) {
3940 return Err(MongrelError::InvalidArgument(format!(
3941 "column {column_id} has no MinHash index"
3942 )));
3943 }
3944 if members.is_empty() {
3945 return Err(MongrelError::InvalidArgument(
3946 "MinHash members must not be empty".into(),
3947 ));
3948 }
3949 if members.len() > MAX_SET_MEMBERS {
3950 return Err(MongrelError::InvalidArgument(format!(
3951 "MinHash query exceeds {MAX_SET_MEMBERS} members"
3952 )));
3953 }
3954 let mut total_bytes = 0usize;
3955 for member in members {
3956 let bytes = member.encoded_len();
3957 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
3958 return Err(MongrelError::InvalidArgument(format!(
3959 "MinHash member exceeds {} bytes",
3960 crate::query::MAX_SET_MEMBER_BYTES
3961 )));
3962 }
3963 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
3964 MongrelError::InvalidArgument("MinHash input size overflow".into())
3965 })?;
3966 }
3967 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
3968 return Err(MongrelError::InvalidArgument(format!(
3969 "MinHash input exceeds {} bytes",
3970 crate::query::MAX_SET_INPUT_BYTES
3971 )));
3972 }
3973 (*column_id, *k)
3974 }
3975 };
3976 if k == 0 {
3977 return Err(MongrelError::InvalidArgument(
3978 "retriever k must be > 0".into(),
3979 ));
3980 }
3981 if k > MAX_RETRIEVER_K {
3982 return Err(MongrelError::InvalidArgument(format!(
3983 "retriever k exceeds {MAX_RETRIEVER_K}"
3984 )));
3985 }
3986 debug_assert!(self
3987 .schema
3988 .columns
3989 .iter()
3990 .any(|column| column.id == column_id));
3991 Ok(())
3992 }
3993
3994 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
3995 use crate::query::Condition;
3996 match condition {
3997 Condition::Ann {
3998 column_id,
3999 query,
4000 k,
4001 } => self.validate_retriever(&crate::query::Retriever::Ann {
4002 column_id: *column_id,
4003 query: query.clone(),
4004 k: *k,
4005 }),
4006 Condition::SparseMatch {
4007 column_id,
4008 query,
4009 k,
4010 } => self.validate_retriever(&crate::query::Retriever::Sparse {
4011 column_id: *column_id,
4012 query: query.clone(),
4013 k: *k,
4014 }),
4015 Condition::MinHashSimilar {
4016 column_id,
4017 query,
4018 k,
4019 } => {
4020 if !self.minhash.contains_key(column_id) {
4021 return Err(MongrelError::InvalidArgument(format!(
4022 "column {column_id} has no MinHash index"
4023 )));
4024 }
4025 if query.is_empty() || *k == 0 {
4026 return Err(MongrelError::InvalidArgument(
4027 "MinHash query must be non-empty and k must be > 0".into(),
4028 ));
4029 }
4030 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
4031 {
4032 return Err(MongrelError::InvalidArgument(format!(
4033 "MinHash query must have <= {} members and k <= {}",
4034 crate::query::MAX_SET_MEMBERS,
4035 crate::query::MAX_RETRIEVER_K
4036 )));
4037 }
4038 Ok(())
4039 }
4040 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
4041 Err(MongrelError::InvalidArgument(format!(
4042 "bitmap IN exceeds {} values",
4043 crate::query::MAX_SET_MEMBERS
4044 )))
4045 }
4046 Condition::FmContainsAll { patterns, .. }
4047 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
4048 {
4049 Err(MongrelError::InvalidArgument(format!(
4050 "FM query exceeds {} patterns",
4051 crate::query::MAX_HARD_CONDITIONS
4052 )))
4053 }
4054 _ => Ok(()),
4055 }
4056 }
4057
4058 fn retrieve_filtered(
4059 &self,
4060 retriever: &crate::query::Retriever,
4061 snapshot: Snapshot,
4062 hard_filter: Option<&RowIdSet>,
4063 allowed: Option<&std::collections::HashSet<RowId>>,
4064 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4065 context: Option<&crate::query::AiExecutionContext>,
4066 ) -> Result<Vec<crate::query::RetrieverHit>> {
4067 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
4068 let started = std::time::Instant::now();
4069 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
4070 Retriever::Ann {
4071 column_id,
4072 query,
4073 k,
4074 } => {
4075 let Some(index) = self.ann.get(column_id) else {
4076 return Ok(Vec::new());
4077 };
4078 let cap = ann_candidate_cap(index.len(), context);
4079 if cap == 0 {
4080 return Ok(Vec::new());
4081 }
4082 let mut breadth = (*k).max(1).min(cap);
4083 let mut eligibility = std::collections::HashMap::new();
4084 let mut filtered = loop {
4085 let mut seen = std::collections::HashSet::new();
4086 if let Some(context) = context {
4087 context.checkpoint()?;
4088 }
4089 let raw = index.search_with_context(query, breadth, context)?;
4090 let unchecked: Vec<_> = raw
4091 .iter()
4092 .map(|(row_id, _)| *row_id)
4093 .filter(|row_id| !eligibility.contains_key(row_id))
4094 .filter(|row_id| {
4095 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4096 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4097 })
4098 .collect();
4099 let eligible = self.eligible_and_authorized_candidate_ids(
4100 &unchecked,
4101 *column_id,
4102 snapshot,
4103 candidate_authorization,
4104 context,
4105 )?;
4106 for row_id in unchecked {
4107 eligibility.insert(row_id, eligible.contains(&row_id));
4108 }
4109 let filtered: Vec<_> = raw
4110 .into_iter()
4111 .filter(|(row_id, _)| {
4112 seen.insert(*row_id)
4113 && eligibility.get(row_id).copied().unwrap_or(false)
4114 })
4115 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
4116 .collect();
4117 if filtered.len() >= *k || breadth >= cap {
4118 if filtered.len() < *k && index.len() > cap && breadth >= cap {
4119 crate::trace::QueryTrace::record(|trace| {
4120 trace.ann_candidate_cap_hit = true;
4121 });
4122 }
4123 break filtered;
4124 }
4125 breadth = breadth.saturating_mul(2).min(cap);
4126 };
4127 filtered.truncate(*k);
4128 filtered
4129 }
4130 Retriever::Sparse {
4131 column_id,
4132 query,
4133 k,
4134 } => self
4135 .sparse
4136 .get(column_id)
4137 .map(|index| -> Result<Vec<_>> {
4138 let mut breadth = (*k).max(1);
4139 let mut eligibility = std::collections::HashMap::new();
4140 loop {
4141 if let Some(context) = context {
4142 context.checkpoint()?;
4143 }
4144 let raw = index.search_with_context(query, breadth, context)?;
4145 let unchecked: Vec<_> = raw
4146 .iter()
4147 .map(|(row_id, _)| *row_id)
4148 .filter(|row_id| !eligibility.contains_key(row_id))
4149 .filter(|row_id| {
4150 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4151 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4152 })
4153 .collect();
4154 let eligible = self.eligible_and_authorized_candidate_ids(
4155 &unchecked,
4156 *column_id,
4157 snapshot,
4158 candidate_authorization,
4159 context,
4160 )?;
4161 for row_id in unchecked {
4162 eligibility.insert(row_id, eligible.contains(&row_id));
4163 }
4164 let filtered: Vec<_> = raw
4165 .iter()
4166 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4167 .take(*k)
4168 .map(|(row_id, score)| {
4169 (*row_id, RetrieverScore::SparseDotProduct(*score))
4170 })
4171 .collect();
4172 if filtered.len() >= *k || raw.len() < breadth {
4173 break Ok(filtered);
4174 }
4175 let next = breadth.saturating_mul(2);
4176 if next == breadth {
4177 break Ok(filtered);
4178 }
4179 breadth = next;
4180 }
4181 })
4182 .transpose()?
4183 .unwrap_or_default(),
4184 Retriever::MinHash {
4185 column_id,
4186 members,
4187 k,
4188 } => self
4189 .minhash
4190 .get(column_id)
4191 .map(|index| -> Result<Vec<_>> {
4192 let mut hashes = Vec::with_capacity(members.len());
4193 for member in members {
4194 if let Some(context) = context {
4195 context.consume(crate::query::work_units(
4196 member.encoded_len(),
4197 crate::query::PARSE_WORK_QUANTUM,
4198 ))?;
4199 }
4200 hashes.push(member.hash_v1());
4201 }
4202 let mut breadth = (*k).max(1);
4203 let mut eligibility = std::collections::HashMap::new();
4204 loop {
4205 if let Some(context) = context {
4206 context.checkpoint()?;
4207 }
4208 let raw = index.search_with_context(&hashes, breadth, context)?;
4209 let unchecked: Vec<_> = raw
4210 .iter()
4211 .map(|(row_id, _)| *row_id)
4212 .filter(|row_id| !eligibility.contains_key(row_id))
4213 .filter(|row_id| {
4214 hard_filter.map_or(true, |filter| filter.contains(row_id.0))
4215 && allowed.map_or(true, |allowed| allowed.contains(row_id))
4216 })
4217 .collect();
4218 let eligible = self.eligible_and_authorized_candidate_ids(
4219 &unchecked,
4220 *column_id,
4221 snapshot,
4222 candidate_authorization,
4223 context,
4224 )?;
4225 for row_id in unchecked {
4226 eligibility.insert(row_id, eligible.contains(&row_id));
4227 }
4228 let filtered: Vec<_> = raw
4229 .iter()
4230 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
4231 .take(*k)
4232 .map(|(row_id, score)| {
4233 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
4234 })
4235 .collect();
4236 if filtered.len() >= *k || raw.len() < breadth {
4237 break Ok(filtered);
4238 }
4239 let next = breadth.saturating_mul(2);
4240 if next == breadth {
4241 break Ok(filtered);
4242 }
4243 breadth = next;
4244 }
4245 })
4246 .transpose()?
4247 .unwrap_or_default(),
4248 };
4249 let elapsed = started.elapsed().as_nanos() as u64;
4250 crate::trace::QueryTrace::record(|trace| {
4251 match retriever {
4252 Retriever::Ann { .. } => {
4253 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
4254 }
4255 Retriever::Sparse { .. } => {
4256 trace.sparse_candidate_nanos =
4257 trace.sparse_candidate_nanos.saturating_add(elapsed)
4258 }
4259 Retriever::MinHash { .. } => {
4260 trace.minhash_candidate_nanos =
4261 trace.minhash_candidate_nanos.saturating_add(elapsed)
4262 }
4263 }
4264 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
4265 });
4266 Ok(scored
4267 .into_iter()
4268 .enumerate()
4269 .map(|(rank, (row_id, score))| RetrieverHit {
4270 row_id,
4271 rank: rank + 1,
4272 score,
4273 })
4274 .collect())
4275 }
4276
4277 fn eligible_candidate_ids(
4278 &self,
4279 candidates: &[RowId],
4280 _column_id: u16,
4281 snapshot: Snapshot,
4282 context: Option<&crate::query::AiExecutionContext>,
4283 ) -> Result<std::collections::HashSet<RowId>> {
4284 if !self.had_deletes
4285 && self.ttl.is_none()
4286 && self.pending_put_cols.is_empty()
4287 && snapshot.epoch == self.snapshot().epoch
4288 {
4289 return Ok(candidates.iter().copied().collect());
4290 }
4291 let mut readers: Vec<_> = self
4292 .run_refs
4293 .iter()
4294 .map(|run| self.open_reader(run.run_id))
4295 .collect::<Result<_>>()?;
4296 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4297 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
4298 for &row_id in candidates {
4299 if let Some(context) = context {
4300 context.consume(1)?;
4301 }
4302 let mem = self.memtable.get_version(row_id, snapshot.epoch);
4303 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
4304 let overlay = match (mem, mutable) {
4305 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
4306 (Some(value), None) | (None, Some(value)) => Some(value),
4307 (None, None) => None,
4308 };
4309 if let Some((_, row)) = overlay {
4310 if !row.deleted && !self.row_expired_at(&row, now) {
4311 eligible.insert(row_id);
4312 }
4313 continue;
4314 }
4315 let mut best: Option<(Epoch, bool, usize)> = None;
4316 for (index, reader) in readers.iter_mut().enumerate() {
4317 if let Some((epoch, deleted)) =
4318 reader.get_version_visibility(row_id, snapshot.epoch)?
4319 {
4320 if best
4321 .as_ref()
4322 .map(|(best_epoch, ..)| epoch > *best_epoch)
4323 .unwrap_or(true)
4324 {
4325 best = Some((epoch, deleted, index));
4326 }
4327 }
4328 }
4329 let Some((_, false, reader_index)) = best else {
4330 continue;
4331 };
4332 if let Some(ttl) = self.ttl {
4333 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
4334 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
4335 {
4336 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
4337 continue;
4338 }
4339 }
4340 }
4341 eligible.insert(row_id);
4342 }
4343 Ok(eligible)
4344 }
4345
4346 fn eligible_and_authorized_candidate_ids(
4347 &self,
4348 candidates: &[RowId],
4349 column_id: u16,
4350 snapshot: Snapshot,
4351 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4352 context: Option<&crate::query::AiExecutionContext>,
4353 ) -> Result<std::collections::HashSet<RowId>> {
4354 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
4355 let Some(authorization) = authorization else {
4356 return Ok(eligible);
4357 };
4358 let candidates: Vec<_> = eligible.into_iter().collect();
4359 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
4360 }
4361
4362 fn policy_allowed_candidate_ids(
4363 &self,
4364 candidates: &[RowId],
4365 snapshot: Snapshot,
4366 authorization: &crate::security::CandidateAuthorization<'_>,
4367 context: Option<&crate::query::AiExecutionContext>,
4368 ) -> Result<std::collections::HashSet<RowId>> {
4369 let started = std::time::Instant::now();
4370 if candidates.is_empty()
4371 || authorization.principal.is_admin
4372 || !authorization.security.rls_enabled(authorization.table)
4373 {
4374 return Ok(candidates.iter().copied().collect());
4375 }
4376 if let Some(context) = context {
4377 context.checkpoint()?;
4378 }
4379 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
4380 let mut rows: std::collections::HashMap<RowId, Row> = candidates
4381 .iter()
4382 .map(|row_id| {
4383 (
4384 *row_id,
4385 Row {
4386 row_id: *row_id,
4387 committed_epoch: snapshot.epoch,
4388 columns: std::collections::HashMap::new(),
4389 deleted: false,
4390 },
4391 )
4392 })
4393 .collect();
4394 let columns = authorization
4395 .security
4396 .select_policy_columns(authorization.table, authorization.principal);
4397 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4398 let mut decoded = 0usize;
4399 for column_id in &columns {
4400 if let Some(context) = context {
4401 context.checkpoint()?;
4402 }
4403 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4404 &row_ids, *column_id, snapshot, query_now, context,
4405 )? {
4406 if let Some(row) = rows.get_mut(&row_id) {
4407 row.columns.insert(*column_id, value);
4408 decoded = decoded.saturating_add(1);
4409 }
4410 }
4411 }
4412 if let Some(context) = context {
4413 context.consume(candidates.len().saturating_add(decoded))?;
4414 }
4415 let allowed = rows
4416 .into_values()
4417 .filter_map(|row| {
4418 authorization
4419 .security
4420 .row_allowed(
4421 authorization.table,
4422 crate::security::PolicyCommand::Select,
4423 &row,
4424 authorization.principal,
4425 false,
4426 )
4427 .then_some(row.row_id)
4428 })
4429 .collect();
4430 crate::trace::QueryTrace::record(|trace| {
4431 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
4432 trace.rls_policy_columns_decoded =
4433 trace.rls_policy_columns_decoded.saturating_add(decoded);
4434 trace.authorization_nanos = trace
4435 .authorization_nanos
4436 .saturating_add(started.elapsed().as_nanos() as u64);
4437 });
4438 Ok(allowed)
4439 }
4440
4441 pub fn search(
4443 &mut self,
4444 request: &crate::query::SearchRequest,
4445 ) -> Result<Vec<crate::query::SearchHit>> {
4446 self.search_with_allowed(request, None)
4447 }
4448
4449 pub fn search_at(
4450 &mut self,
4451 request: &crate::query::SearchRequest,
4452 snapshot: Snapshot,
4453 authorized: Option<&std::collections::HashSet<RowId>>,
4454 ) -> Result<Vec<crate::query::SearchHit>> {
4455 self.search_at_with_allowed(request, snapshot, authorized)
4456 }
4457
4458 pub fn search_with_allowed(
4459 &mut self,
4460 request: &crate::query::SearchRequest,
4461 authorized: Option<&std::collections::HashSet<RowId>>,
4462 ) -> Result<Vec<crate::query::SearchHit>> {
4463 self.search_at_with_allowed(request, self.snapshot(), authorized)
4464 }
4465
4466 pub fn search_at_with_allowed(
4467 &mut self,
4468 request: &crate::query::SearchRequest,
4469 snapshot: Snapshot,
4470 authorized: Option<&std::collections::HashSet<RowId>>,
4471 ) -> Result<Vec<crate::query::SearchHit>> {
4472 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
4473 }
4474
4475 pub fn search_at_with_allowed_and_context(
4476 &mut self,
4477 request: &crate::query::SearchRequest,
4478 snapshot: Snapshot,
4479 authorized: Option<&std::collections::HashSet<RowId>>,
4480 context: Option<&crate::query::AiExecutionContext>,
4481 ) -> Result<Vec<crate::query::SearchHit>> {
4482 self.ensure_indexes_complete()?;
4483 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
4484 }
4485
4486 pub fn search_at_with_candidate_authorization_and_context(
4487 &mut self,
4488 request: &crate::query::SearchRequest,
4489 snapshot: Snapshot,
4490 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4491 context: Option<&crate::query::AiExecutionContext>,
4492 ) -> Result<Vec<crate::query::SearchHit>> {
4493 self.ensure_indexes_complete()?;
4494 self.search_at_with_filters_and_context(
4495 request,
4496 snapshot,
4497 None,
4498 authorization,
4499 context,
4500 None,
4501 )
4502 }
4503
4504 #[doc(hidden)]
4505 pub fn search_at_with_candidate_authorization_on_generation(
4506 &self,
4507 request: &crate::query::SearchRequest,
4508 snapshot: Snapshot,
4509 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4510 context: Option<&crate::query::AiExecutionContext>,
4511 ) -> Result<Vec<crate::query::SearchHit>> {
4512 self.search_at_with_filters_and_context(
4513 request,
4514 snapshot,
4515 None,
4516 authorization,
4517 context,
4518 None,
4519 )
4520 }
4521
4522 #[doc(hidden)]
4523 pub fn search_at_with_candidate_authorization_on_generation_after(
4524 &self,
4525 request: &crate::query::SearchRequest,
4526 snapshot: Snapshot,
4527 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4528 context: Option<&crate::query::AiExecutionContext>,
4529 after: Option<crate::query::SearchAfter>,
4530 ) -> Result<Vec<crate::query::SearchHit>> {
4531 self.search_at_with_filters_and_context(
4532 request,
4533 snapshot,
4534 None,
4535 authorization,
4536 context,
4537 after,
4538 )
4539 }
4540
4541 fn search_at_with_filters_and_context(
4542 &self,
4543 request: &crate::query::SearchRequest,
4544 snapshot: Snapshot,
4545 authorized: Option<&std::collections::HashSet<RowId>>,
4546 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4547 context: Option<&crate::query::AiExecutionContext>,
4548 after: Option<crate::query::SearchAfter>,
4549 ) -> Result<Vec<crate::query::SearchHit>> {
4550 use crate::query::{
4551 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
4552 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
4553 };
4554 let total_started = std::time::Instant::now();
4555 let rank_offset = after.map_or(0, |after| after.returned_count);
4556 self.require_select()?;
4557 if request.limit == 0 {
4558 return Err(MongrelError::InvalidArgument(
4559 "search limit must be > 0".into(),
4560 ));
4561 }
4562 if request.limit > MAX_FINAL_LIMIT {
4563 return Err(MongrelError::InvalidArgument(format!(
4564 "search limit exceeds {MAX_FINAL_LIMIT}"
4565 )));
4566 }
4567 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
4568 return Err(MongrelError::InvalidArgument(
4569 "search-after score must be finite".into(),
4570 ));
4571 }
4572 if request.retrievers.is_empty() {
4573 return Err(MongrelError::InvalidArgument(
4574 "search requires at least one retriever".into(),
4575 ));
4576 }
4577 if request.retrievers.len() > MAX_RETRIEVERS {
4578 return Err(MongrelError::InvalidArgument(format!(
4579 "search exceeds {MAX_RETRIEVERS} retrievers"
4580 )));
4581 }
4582 if request.must.len() > MAX_HARD_CONDITIONS {
4583 return Err(MongrelError::InvalidArgument(format!(
4584 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
4585 )));
4586 }
4587 for condition in &request.must {
4588 self.validate_condition(condition)?;
4589 }
4590 if request.must.iter().any(|condition| {
4591 matches!(
4592 condition,
4593 Condition::Ann { .. }
4594 | Condition::SparseMatch { .. }
4595 | Condition::MinHashSimilar { .. }
4596 )
4597 }) {
4598 return Err(MongrelError::InvalidArgument(
4599 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
4600 .into(),
4601 ));
4602 }
4603 let mut names = std::collections::HashSet::new();
4604 for named in &request.retrievers {
4605 if named.name.is_empty()
4606 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
4607 || !names.insert(named.name.as_str())
4608 {
4609 return Err(MongrelError::InvalidArgument(format!(
4610 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
4611 crate::query::MAX_RETRIEVER_NAME_BYTES
4612 )));
4613 }
4614 if !named.weight.is_finite()
4615 || named.weight < 0.0
4616 || named.weight > MAX_RETRIEVER_WEIGHT
4617 {
4618 return Err(MongrelError::InvalidArgument(format!(
4619 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4620 )));
4621 }
4622 self.validate_retriever(&named.retriever)?;
4623 }
4624 let projection = request
4625 .projection
4626 .clone()
4627 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
4628 if projection.len() > MAX_PROJECTION_COLUMNS {
4629 return Err(MongrelError::InvalidArgument(format!(
4630 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
4631 )));
4632 }
4633 for column_id in &projection {
4634 if !self
4635 .schema
4636 .columns
4637 .iter()
4638 .any(|column| column.id == *column_id)
4639 {
4640 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
4641 }
4642 }
4643 if let Some(crate::query::Rerank::ExactVector {
4644 embedding_column,
4645 query,
4646 candidate_limit,
4647 weight,
4648 ..
4649 }) = &request.rerank
4650 {
4651 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
4652 {
4653 return Err(MongrelError::InvalidArgument(format!(
4654 "rerank candidate_limit must be between search limit and {}",
4655 crate::query::MAX_RETRIEVER_K
4656 )));
4657 }
4658 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
4659 return Err(MongrelError::InvalidArgument(format!(
4660 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
4661 )));
4662 }
4663 let column = self
4664 .schema
4665 .columns
4666 .iter()
4667 .find(|column| column.id == *embedding_column)
4668 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
4669 let crate::schema::TypeId::Embedding { dim } = column.ty else {
4670 return Err(MongrelError::InvalidArgument(format!(
4671 "rerank column {embedding_column} is not an embedding"
4672 )));
4673 };
4674 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
4675 return Err(MongrelError::InvalidArgument(format!(
4676 "rerank query must contain {dim} finite values"
4677 )));
4678 }
4679 }
4680
4681 let hard_filter_started = std::time::Instant::now();
4682 let hard_filter = if request.must.is_empty() {
4683 None
4684 } else {
4685 let mut sets = Vec::with_capacity(request.must.len());
4686 for condition in &request.must {
4687 if let Some(context) = context {
4688 context.checkpoint()?;
4689 }
4690 sets.push(self.resolve_condition(condition, snapshot)?);
4691 }
4692 Some(RowIdSet::intersect_many(sets))
4693 };
4694 crate::trace::QueryTrace::record(|trace| {
4695 trace.hard_filter_nanos = trace
4696 .hard_filter_nanos
4697 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
4698 });
4699 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
4700 return Ok(Vec::new());
4701 }
4702
4703 let constant = match request.fusion {
4704 Fusion::ReciprocalRank { constant } => constant,
4705 };
4706 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
4707 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
4708 let mut fusion_nanos = 0u64;
4709 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
4710 std::collections::HashMap::new();
4711 for named in retrievers {
4712 if named.weight == 0.0 {
4713 continue;
4714 }
4715 if let Some(context) = context {
4716 context.checkpoint()?;
4717 }
4718 let hits = self.retrieve_filtered(
4719 &named.retriever,
4720 snapshot,
4721 hard_filter.as_ref(),
4722 authorized,
4723 candidate_authorization,
4724 context,
4725 )?;
4726 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
4727 let fusion_started = std::time::Instant::now();
4728 for hit in hits {
4729 if let Some(context) = context {
4730 context.consume(1)?;
4731 }
4732 let contribution = named.weight / (constant as f64 + hit.rank as f64);
4733 if !contribution.is_finite() {
4734 return Err(MongrelError::InvalidArgument(
4735 "retriever contribution must be finite".into(),
4736 ));
4737 }
4738 let max_fused_candidates = context.map_or(
4739 crate::query::MAX_FUSED_CANDIDATES,
4740 crate::query::AiExecutionContext::max_fused_candidates,
4741 );
4742 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
4743 return Err(MongrelError::WorkBudgetExceeded);
4744 }
4745 let entry = fused.entry(hit.row_id).or_default();
4746 entry.0 += contribution;
4747 if !entry.0.is_finite() {
4748 return Err(MongrelError::InvalidArgument(
4749 "fused score must be finite".into(),
4750 ));
4751 }
4752 entry.1.push(ComponentScore {
4753 retriever_name: retriever_name.clone(),
4754 rank: hit.rank,
4755 raw_score: hit.score,
4756 contribution,
4757 });
4758 }
4759 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
4760 }
4761 let union_size = fused.len();
4762 let mut ranked: Vec<_> = fused
4763 .into_iter()
4764 .map(|(row_id, (fused_score, components))| {
4765 (row_id, fused_score, components, None, fused_score)
4766 })
4767 .collect();
4768 let order = |(a_row, _, _, _, a_score): &(
4769 RowId,
4770 f64,
4771 Vec<ComponentScore>,
4772 Option<f32>,
4773 f64,
4774 ),
4775 (b_row, _, _, _, b_score): &(
4776 RowId,
4777 f64,
4778 Vec<ComponentScore>,
4779 Option<f32>,
4780 f64,
4781 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
4782 if let Some(crate::query::Rerank::ExactVector {
4783 embedding_column,
4784 query,
4785 metric,
4786 candidate_limit,
4787 weight,
4788 }) = &request.rerank
4789 {
4790 let fused_order = |(a_row, a_score, ..): &(
4791 RowId,
4792 f64,
4793 Vec<ComponentScore>,
4794 Option<f32>,
4795 f64,
4796 ),
4797 (b_row, b_score, ..): &(
4798 RowId,
4799 f64,
4800 Vec<ComponentScore>,
4801 Option<f32>,
4802 f64,
4803 )| {
4804 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
4805 };
4806 let selection_started = std::time::Instant::now();
4807 if let Some(context) = context {
4808 context.consume(ranked.len())?;
4809 }
4810 if ranked.len() > *candidate_limit {
4811 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
4812 ranked.truncate(*candidate_limit);
4813 }
4814 ranked.sort_by(fused_order);
4815 fusion_nanos =
4816 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4817 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4818 if let Some(context) = context {
4819 context.consume(row_ids.len())?;
4820 }
4821 let query_now =
4822 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4823 let gather_started = std::time::Instant::now();
4824 let vectors = self.values_for_rids_batch_at_with_context(
4825 &row_ids,
4826 *embedding_column,
4827 snapshot,
4828 query_now,
4829 context,
4830 )?;
4831 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
4832 let vector_work =
4833 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
4834 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
4835 if let Some(context) = context {
4836 context.consume(vector_work)?;
4837 }
4838 query
4839 .iter()
4840 .map(|value| f64::from(*value).powi(2))
4841 .sum::<f64>()
4842 .sqrt()
4843 } else {
4844 0.0
4845 };
4846 let score_started = std::time::Instant::now();
4847 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
4848 for (row_id, value) in vectors {
4849 let Value::Embedding(vector) = value else {
4850 continue;
4851 };
4852 let score = match metric {
4853 crate::query::VectorMetric::DotProduct => {
4854 if let Some(context) = context {
4855 context.consume(vector_work)?;
4856 }
4857 query
4858 .iter()
4859 .zip(&vector)
4860 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4861 .sum::<f64>()
4862 }
4863 crate::query::VectorMetric::Cosine => {
4864 if let Some(context) = context {
4865 context.consume(vector_work.saturating_mul(2))?;
4866 }
4867 let dot = query
4868 .iter()
4869 .zip(&vector)
4870 .map(|(left, right)| f64::from(*left) * f64::from(*right))
4871 .sum::<f64>();
4872 let norm = vector
4873 .iter()
4874 .map(|value| f64::from(*value).powi(2))
4875 .sum::<f64>()
4876 .sqrt();
4877 if query_norm == 0.0 || norm == 0.0 {
4878 0.0
4879 } else {
4880 dot / (query_norm * norm)
4881 }
4882 }
4883 crate::query::VectorMetric::Euclidean => {
4884 if let Some(context) = context {
4885 context.consume(vector_work)?;
4886 }
4887 query
4888 .iter()
4889 .zip(&vector)
4890 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
4891 .sum::<f64>()
4892 .sqrt()
4893 }
4894 };
4895 if !score.is_finite() {
4896 return Err(MongrelError::InvalidArgument(
4897 "exact rerank score must be finite".into(),
4898 ));
4899 }
4900 scores.insert(row_id, score as f32);
4901 }
4902 let mut reranked = Vec::with_capacity(ranked.len());
4903 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
4904 let Some(score) = scores.get(&row_id).copied() else {
4905 continue;
4906 };
4907 let ordering_score = match metric {
4908 crate::query::VectorMetric::Euclidean => -f64::from(score),
4909 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
4910 f64::from(score)
4911 }
4912 };
4913 let final_score = fused_score + *weight * ordering_score;
4914 if !final_score.is_finite() {
4915 return Err(MongrelError::InvalidArgument(
4916 "final rerank score must be finite".into(),
4917 ));
4918 }
4919 reranked.push((row_id, fused_score, components, Some(score), final_score));
4920 }
4921 ranked = reranked;
4922 ranked.sort_by(order);
4923 crate::trace::QueryTrace::record(|trace| {
4924 trace.exact_vector_gather_nanos =
4925 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
4926 trace.exact_vector_score_nanos = trace
4927 .exact_vector_score_nanos
4928 .saturating_add(score_started.elapsed().as_nanos() as u64);
4929 });
4930 }
4931 if let Some(after) = after {
4932 ranked.retain(|(row_id, _, _, _, final_score)| {
4933 final_score.total_cmp(&after.final_score).is_lt()
4934 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
4935 });
4936 }
4937 let projection_started = std::time::Instant::now();
4938 let sentinel = projection
4939 .first()
4940 .copied()
4941 .or_else(|| self.schema.columns.first().map(|column| column.id));
4942 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
4943 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
4944 let mut projection_rows = 0usize;
4945 let mut projection_cells = 0usize;
4946 while out.len() < request.limit && !ranked.is_empty() {
4947 if let Some(context) = context {
4948 context.checkpoint()?;
4949 context.consume(ranked.len())?;
4950 }
4951 let needed = request.limit - out.len();
4952 let window_size = ranked
4953 .len()
4954 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
4955 let selection_started = std::time::Instant::now();
4956 let mut remainder = if ranked.len() > window_size {
4957 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
4958 ranked.split_off(window_size)
4959 } else {
4960 Vec::new()
4961 };
4962 ranked.sort_by(order);
4963 fusion_nanos =
4964 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
4965 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
4966 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
4967 if let Some(context) = context {
4968 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
4969 }
4970 projection_rows = projection_rows.saturating_add(row_ids.len());
4971 projection_cells =
4972 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
4973 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
4974 std::collections::HashMap::new();
4975 if let Some(column_id) = sentinel {
4976 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4977 &row_ids, column_id, snapshot, query_now, context,
4978 )? {
4979 cells.entry(row_id).or_default().insert(column_id, value);
4980 }
4981 }
4982 for &column_id in &projection {
4983 if Some(column_id) == sentinel {
4984 continue;
4985 }
4986 for (row_id, value) in self.values_for_rids_batch_at_with_context(
4987 &row_ids, column_id, snapshot, query_now, context,
4988 )? {
4989 cells.entry(row_id).or_default().insert(column_id, value);
4990 }
4991 }
4992 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
4993 ranked.drain(..)
4994 {
4995 let Some(row_cells) = cells.remove(&row_id) else {
4996 continue;
4997 };
4998 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
4999 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
5000 out.push(SearchHit {
5001 row_id,
5002 cells: projection
5003 .iter()
5004 .filter_map(|column_id| {
5005 row_cells
5006 .get(column_id)
5007 .cloned()
5008 .map(|value| (*column_id, value))
5009 })
5010 .collect(),
5011 components,
5012 fused_score,
5013 exact_rerank_score,
5014 final_score,
5015 final_rank,
5016 });
5017 if out.len() == request.limit {
5018 break;
5019 }
5020 }
5021 ranked.append(&mut remainder);
5022 }
5023 crate::trace::QueryTrace::record(|trace| {
5024 trace.union_size = union_size;
5025 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
5026 trace.projection_nanos = trace
5027 .projection_nanos
5028 .saturating_add(projection_started.elapsed().as_nanos() as u64);
5029 trace.total_nanos = trace
5030 .total_nanos
5031 .saturating_add(total_started.elapsed().as_nanos() as u64);
5032 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
5033 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
5034 if let Some(context) = context {
5035 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
5036 }
5037 });
5038 Ok(out)
5039 }
5040
5041 pub fn set_similarity(
5044 &mut self,
5045 request: &crate::query::SetSimilarityRequest,
5046 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5047 self.set_similarity_with_allowed(request, None)
5048 }
5049
5050 pub fn set_similarity_at(
5051 &mut self,
5052 request: &crate::query::SetSimilarityRequest,
5053 snapshot: Snapshot,
5054 allowed: Option<&std::collections::HashSet<RowId>>,
5055 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5056 self.set_similarity_explained_at(request, snapshot, allowed)
5057 .map(|(hits, _)| hits)
5058 }
5059
5060 pub fn ann_rerank(
5062 &mut self,
5063 request: &crate::query::AnnRerankRequest,
5064 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5065 self.ann_rerank_with_allowed(request, None)
5066 }
5067
5068 pub fn ann_rerank_with_allowed(
5069 &mut self,
5070 request: &crate::query::AnnRerankRequest,
5071 allowed: Option<&std::collections::HashSet<RowId>>,
5072 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5073 self.ann_rerank_at(request, self.snapshot(), allowed)
5074 }
5075
5076 pub fn ann_rerank_at(
5077 &mut self,
5078 request: &crate::query::AnnRerankRequest,
5079 snapshot: Snapshot,
5080 allowed: Option<&std::collections::HashSet<RowId>>,
5081 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5082 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
5083 }
5084
5085 pub fn ann_rerank_at_with_context(
5086 &mut self,
5087 request: &crate::query::AnnRerankRequest,
5088 snapshot: Snapshot,
5089 allowed: Option<&std::collections::HashSet<RowId>>,
5090 context: Option<&crate::query::AiExecutionContext>,
5091 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5092 self.ensure_indexes_complete()?;
5093 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
5094 }
5095
5096 pub fn ann_rerank_at_with_candidate_authorization_and_context(
5097 &mut self,
5098 request: &crate::query::AnnRerankRequest,
5099 snapshot: Snapshot,
5100 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5101 context: Option<&crate::query::AiExecutionContext>,
5102 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5103 self.ensure_indexes_complete()?;
5104 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
5105 }
5106
5107 #[doc(hidden)]
5108 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
5109 &self,
5110 request: &crate::query::AnnRerankRequest,
5111 snapshot: Snapshot,
5112 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5113 context: Option<&crate::query::AiExecutionContext>,
5114 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5115 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
5116 }
5117
5118 fn ann_rerank_at_with_filters_and_context(
5119 &self,
5120 request: &crate::query::AnnRerankRequest,
5121 snapshot: Snapshot,
5122 allowed: Option<&std::collections::HashSet<RowId>>,
5123 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5124 context: Option<&crate::query::AiExecutionContext>,
5125 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5126 use crate::query::{
5127 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
5128 };
5129 if request.candidate_k == 0 || request.limit == 0 {
5130 return Err(MongrelError::InvalidArgument(
5131 "candidate_k and limit must be > 0".into(),
5132 ));
5133 }
5134 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
5135 return Err(MongrelError::InvalidArgument(format!(
5136 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
5137 )));
5138 }
5139 let retriever = Retriever::Ann {
5140 column_id: request.column_id,
5141 query: request.query.clone(),
5142 k: request.candidate_k,
5143 };
5144 self.require_select()?;
5145 self.validate_retriever(&retriever)?;
5146 let hits = self.retrieve_filtered(
5147 &retriever,
5148 snapshot,
5149 None,
5150 allowed,
5151 candidate_authorization,
5152 context,
5153 )?;
5154 let distances: std::collections::HashMap<_, _> = hits
5155 .iter()
5156 .filter_map(|hit| match hit.score {
5157 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
5158 _ => None,
5159 })
5160 .collect();
5161 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
5162 if let Some(context) = context {
5163 context.consume(row_ids.len())?;
5164 }
5165 let gather_started = std::time::Instant::now();
5166 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5167 let values = self.values_for_rids_batch_at_with_context(
5168 &row_ids,
5169 request.column_id,
5170 snapshot,
5171 query_now,
5172 context,
5173 )?;
5174 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5175 let score_started = std::time::Instant::now();
5176 let vector_work =
5177 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
5178 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
5179 if let Some(context) = context {
5180 context.consume(vector_work)?;
5181 }
5182 request
5183 .query
5184 .iter()
5185 .map(|value| f64::from(*value).powi(2))
5186 .sum::<f64>()
5187 .sqrt()
5188 } else {
5189 0.0
5190 };
5191 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
5192 for (row_id, value) in values {
5193 let Value::Embedding(vector) = value else {
5194 continue;
5195 };
5196 let exact_score = match request.metric {
5197 VectorMetric::DotProduct => {
5198 if let Some(context) = context {
5199 context.consume(vector_work)?;
5200 }
5201 request
5202 .query
5203 .iter()
5204 .zip(&vector)
5205 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5206 .sum::<f64>()
5207 }
5208 VectorMetric::Cosine => {
5209 if let Some(context) = context {
5210 context.consume(vector_work.saturating_mul(2))?;
5211 }
5212 let dot = request
5213 .query
5214 .iter()
5215 .zip(&vector)
5216 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5217 .sum::<f64>();
5218 let norm = vector
5219 .iter()
5220 .map(|value| f64::from(*value).powi(2))
5221 .sum::<f64>()
5222 .sqrt();
5223 if query_norm == 0.0 || norm == 0.0 {
5224 0.0
5225 } else {
5226 dot / (query_norm * norm)
5227 }
5228 }
5229 VectorMetric::Euclidean => {
5230 if let Some(context) = context {
5231 context.consume(vector_work)?;
5232 }
5233 request
5234 .query
5235 .iter()
5236 .zip(&vector)
5237 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
5238 .sum::<f64>()
5239 .sqrt()
5240 }
5241 };
5242 let exact_score = exact_score as f32;
5243 if !exact_score.is_finite() {
5244 return Err(MongrelError::InvalidArgument(
5245 "exact ANN score must be finite".into(),
5246 ));
5247 }
5248 reranked.push(AnnRerankHit {
5249 row_id,
5250 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
5251 exact_score,
5252 });
5253 }
5254 reranked.sort_by(|left, right| {
5255 let score = match request.metric {
5256 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
5257 VectorMetric::Cosine | VectorMetric::DotProduct => {
5258 right.exact_score.total_cmp(&left.exact_score)
5259 }
5260 };
5261 score.then_with(|| left.row_id.cmp(&right.row_id))
5262 });
5263 reranked.truncate(request.limit);
5264 crate::trace::QueryTrace::record(|trace| {
5265 trace.exact_vector_gather_nanos =
5266 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
5267 trace.exact_vector_score_nanos = trace
5268 .exact_vector_score_nanos
5269 .saturating_add(score_started.elapsed().as_nanos() as u64);
5270 });
5271 Ok(reranked)
5272 }
5273
5274 pub fn set_similarity_with_allowed(
5275 &mut self,
5276 request: &crate::query::SetSimilarityRequest,
5277 allowed: Option<&std::collections::HashSet<RowId>>,
5278 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5279 self.set_similarity_explained_at(request, self.snapshot(), allowed)
5280 .map(|(hits, _)| hits)
5281 }
5282
5283 pub fn set_similarity_explained(
5284 &mut self,
5285 request: &crate::query::SetSimilarityRequest,
5286 ) -> Result<(
5287 Vec<crate::query::SetSimilarityHit>,
5288 crate::query::SetSimilarityTrace,
5289 )> {
5290 self.set_similarity_explained_at(request, self.snapshot(), None)
5291 }
5292
5293 fn set_similarity_explained_at(
5294 &mut self,
5295 request: &crate::query::SetSimilarityRequest,
5296 snapshot: Snapshot,
5297 allowed: Option<&std::collections::HashSet<RowId>>,
5298 ) -> Result<(
5299 Vec<crate::query::SetSimilarityHit>,
5300 crate::query::SetSimilarityTrace,
5301 )> {
5302 self.ensure_indexes_complete()?;
5303 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
5304 }
5305
5306 pub fn set_similarity_at_with_context(
5307 &mut self,
5308 request: &crate::query::SetSimilarityRequest,
5309 snapshot: Snapshot,
5310 allowed: Option<&std::collections::HashSet<RowId>>,
5311 context: Option<&crate::query::AiExecutionContext>,
5312 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5313 self.ensure_indexes_complete()?;
5314 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
5315 .map(|(hits, _)| hits)
5316 }
5317
5318 pub fn set_similarity_at_with_candidate_authorization_and_context(
5319 &mut self,
5320 request: &crate::query::SetSimilarityRequest,
5321 snapshot: Snapshot,
5322 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5323 context: Option<&crate::query::AiExecutionContext>,
5324 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5325 self.ensure_indexes_complete()?;
5326 self.set_similarity_explained_at_with_context(
5327 request,
5328 snapshot,
5329 None,
5330 authorization,
5331 context,
5332 )
5333 .map(|(hits, _)| hits)
5334 }
5335
5336 #[doc(hidden)]
5337 pub fn set_similarity_at_with_candidate_authorization_on_generation(
5338 &self,
5339 request: &crate::query::SetSimilarityRequest,
5340 snapshot: Snapshot,
5341 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5342 context: Option<&crate::query::AiExecutionContext>,
5343 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5344 self.set_similarity_explained_at_with_context(
5345 request,
5346 snapshot,
5347 None,
5348 authorization,
5349 context,
5350 )
5351 .map(|(hits, _)| hits)
5352 }
5353
5354 fn set_similarity_explained_at_with_context(
5355 &self,
5356 request: &crate::query::SetSimilarityRequest,
5357 snapshot: Snapshot,
5358 allowed: Option<&std::collections::HashSet<RowId>>,
5359 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5360 context: Option<&crate::query::AiExecutionContext>,
5361 ) -> Result<(
5362 Vec<crate::query::SetSimilarityHit>,
5363 crate::query::SetSimilarityTrace,
5364 )> {
5365 use crate::query::{
5366 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
5367 MAX_SET_MEMBERS,
5368 };
5369 let mut trace = crate::query::SetSimilarityTrace::default();
5370 if request.members.is_empty() {
5371 return Ok((Vec::new(), trace));
5372 }
5373 if request.candidate_k == 0 || request.limit == 0 {
5374 return Err(MongrelError::InvalidArgument(
5375 "candidate_k and limit must be > 0".into(),
5376 ));
5377 }
5378 if request.candidate_k > MAX_RETRIEVER_K
5379 || request.limit > MAX_FINAL_LIMIT
5380 || request.members.len() > MAX_SET_MEMBERS
5381 {
5382 return Err(MongrelError::InvalidArgument(format!(
5383 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
5384 )));
5385 }
5386 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
5387 return Err(MongrelError::InvalidArgument(
5388 "min_jaccard must be finite and between 0 and 1".into(),
5389 ));
5390 }
5391 let started = std::time::Instant::now();
5392 let retriever = Retriever::MinHash {
5393 column_id: request.column_id,
5394 members: request.members.clone(),
5395 k: request.candidate_k,
5396 };
5397 self.require_select()?;
5398 self.validate_retriever(&retriever)?;
5399 let hits = self.retrieve_filtered(
5400 &retriever,
5401 snapshot,
5402 None,
5403 allowed,
5404 candidate_authorization,
5405 context,
5406 )?;
5407 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
5408 trace.candidate_count = hits.len();
5409 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
5410 if let Some(context) = context {
5411 context.consume(row_ids.len())?;
5412 }
5413 let started = std::time::Instant::now();
5414 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5415 let values = self.values_for_rids_batch_at_with_context(
5416 &row_ids,
5417 request.column_id,
5418 snapshot,
5419 query_now,
5420 context,
5421 )?;
5422 trace.gather_us = started.elapsed().as_micros() as u64;
5423 if let Some(context) = context {
5424 context.consume(request.members.len())?;
5425 }
5426 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
5427 let estimates: std::collections::HashMap<_, _> = hits
5428 .into_iter()
5429 .filter_map(|hit| match hit.score {
5430 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
5431 _ => None,
5432 })
5433 .collect();
5434 let started = std::time::Instant::now();
5435 let mut parsed = Vec::with_capacity(values.len());
5436 for (row_id, value) in values {
5437 let Value::Bytes(bytes) = value else {
5438 continue;
5439 };
5440 if let Some(context) = context {
5441 context.consume(crate::query::work_units(
5442 bytes.len(),
5443 crate::query::PARSE_WORK_QUANTUM,
5444 ))?;
5445 }
5446 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
5447 continue;
5448 };
5449 if let Some(context) = context {
5450 context.consume(members.len())?;
5451 }
5452 let stored = members
5453 .into_iter()
5454 .filter_map(|member| match member {
5455 serde_json::Value::String(value) => {
5456 Some(crate::query::SetMember::String(value))
5457 }
5458 serde_json::Value::Number(value) => {
5459 Some(crate::query::SetMember::Number(value))
5460 }
5461 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
5462 _ => None,
5463 })
5464 .collect::<std::collections::HashSet<_>>();
5465 parsed.push((row_id, stored));
5466 }
5467 trace.parse_us = started.elapsed().as_micros() as u64;
5468 trace.verified_count = parsed.len();
5469 let started = std::time::Instant::now();
5470 let mut exact = Vec::new();
5471 for (row_id, stored) in parsed {
5472 if let Some(context) = context {
5473 context.consume(query.len().saturating_add(stored.len()))?;
5474 }
5475 let union = query.union(&stored).count();
5476 let score = if union == 0 {
5477 1.0
5478 } else {
5479 query.intersection(&stored).count() as f32 / union as f32
5480 };
5481 if score >= request.min_jaccard {
5482 exact.push(SetSimilarityHit {
5483 row_id,
5484 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
5485 exact_jaccard: score,
5486 });
5487 }
5488 }
5489 exact.sort_by(|a, b| {
5490 b.exact_jaccard
5491 .total_cmp(&a.exact_jaccard)
5492 .then_with(|| a.row_id.cmp(&b.row_id))
5493 });
5494 exact.truncate(request.limit);
5495 trace.score_us = started.elapsed().as_micros() as u64;
5496 crate::trace::QueryTrace::record(|query_trace| {
5497 query_trace.exact_set_gather_nanos = query_trace
5498 .exact_set_gather_nanos
5499 .saturating_add(trace.gather_us.saturating_mul(1_000));
5500 query_trace.exact_set_parse_nanos = query_trace
5501 .exact_set_parse_nanos
5502 .saturating_add(trace.parse_us.saturating_mul(1_000));
5503 query_trace.exact_set_score_nanos = query_trace
5504 .exact_set_score_nanos
5505 .saturating_add(trace.score_us.saturating_mul(1_000));
5506 });
5507 Ok((exact, trace))
5508 }
5509
5510 fn values_for_rids_batch_at(
5512 &self,
5513 row_ids: &[u64],
5514 column_id: u16,
5515 snapshot: Snapshot,
5516 now: i64,
5517 ) -> Result<Vec<(RowId, Value)>> {
5518 if self.ttl.is_none()
5519 && self.memtable.is_empty()
5520 && self.mutable_run.is_empty()
5521 && self.run_refs.len() == 1
5522 {
5523 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5524 if row_ids.len().saturating_mul(24) < reader.row_count() {
5529 let mut values = Vec::with_capacity(row_ids.len());
5530 for &raw_row_id in row_ids {
5531 let row_id = RowId(raw_row_id);
5532 if let Some((_, false, Some(value))) =
5533 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5534 {
5535 values.push((row_id, value));
5536 }
5537 }
5538 return Ok(values);
5539 }
5540 let (positions, visible_row_ids) =
5541 reader.visible_positions_with_rids(snapshot.epoch)?;
5542 let requested: Vec<(RowId, usize)> = row_ids
5543 .iter()
5544 .filter_map(|raw| {
5545 visible_row_ids
5546 .binary_search(&(*raw as i64))
5547 .ok()
5548 .map(|index| (RowId(*raw), positions[index]))
5549 })
5550 .collect();
5551 let values = reader.gather_column(
5552 column_id,
5553 &requested
5554 .iter()
5555 .map(|(_, position)| *position)
5556 .collect::<Vec<_>>(),
5557 )?;
5558 return Ok(requested
5559 .into_iter()
5560 .zip(values)
5561 .map(|((row_id, _), value)| (row_id, value))
5562 .collect());
5563 }
5564 self.values_for_rids_at(row_ids, column_id, snapshot, now)
5565 }
5566
5567 fn values_for_rids_batch_at_with_context(
5568 &self,
5569 row_ids: &[u64],
5570 column_id: u16,
5571 snapshot: Snapshot,
5572 now: i64,
5573 context: Option<&crate::query::AiExecutionContext>,
5574 ) -> Result<Vec<(RowId, Value)>> {
5575 let Some(context) = context else {
5576 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
5577 };
5578 let mut values = Vec::with_capacity(row_ids.len());
5579 for chunk in row_ids.chunks(256) {
5580 context.checkpoint()?;
5581 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
5582 }
5583 Ok(values)
5584 }
5585
5586 fn values_for_rids_at(
5588 &self,
5589 row_ids: &[u64],
5590 column_id: u16,
5591 snapshot: Snapshot,
5592 now: i64,
5593 ) -> Result<Vec<(RowId, Value)>> {
5594 let mut readers: Vec<_> = self
5595 .run_refs
5596 .iter()
5597 .map(|run| self.open_reader(run.run_id))
5598 .collect::<Result<_>>()?;
5599 let mut out = Vec::with_capacity(row_ids.len());
5600 for &raw_row_id in row_ids {
5601 let row_id = RowId(raw_row_id);
5602 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5603 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5604 let overlay = match (mem, mutable) {
5605 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
5606 (a_epoch, a)
5607 } else {
5608 (b_epoch, b)
5609 }),
5610 (Some(value), None) | (None, Some(value)) => Some(value),
5611 (None, None) => None,
5612 };
5613 if let Some((_, row)) = overlay {
5614 if !row.deleted && !self.row_expired_at(&row, now) {
5615 if let Some(value) = row.columns.get(&column_id) {
5616 out.push((row_id, value.clone()));
5617 }
5618 }
5619 continue;
5620 }
5621
5622 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
5623 for (index, reader) in readers.iter_mut().enumerate() {
5624 if let Some((epoch, deleted, value)) =
5625 reader.get_version_column(row_id, snapshot.epoch, column_id)?
5626 {
5627 if best
5628 .as_ref()
5629 .map(|(best_epoch, ..)| epoch > *best_epoch)
5630 .unwrap_or(true)
5631 {
5632 best = Some((epoch, deleted, value, index));
5633 }
5634 }
5635 }
5636 let Some((_, false, Some(value), reader_index)) = best else {
5637 continue;
5638 };
5639 if let Some(ttl) = self.ttl {
5640 if ttl.column_id != column_id {
5641 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5642 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5643 {
5644 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5645 continue;
5646 }
5647 }
5648 } else if let Value::Int64(timestamp) = value {
5649 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5650 continue;
5651 }
5652 }
5653 }
5654 out.push((row_id, value));
5655 }
5656 Ok(out)
5657 }
5658
5659 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
5664 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now())
5665 }
5666
5667 pub fn rows_for_rids_with_context(
5668 &self,
5669 rids: &[u64],
5670 snapshot: Snapshot,
5671 context: &crate::query::AiExecutionContext,
5672 ) -> Result<Vec<Row>> {
5673 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
5674 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos())
5675 }
5676
5677 fn rows_for_rids_at_time(
5678 &self,
5679 rids: &[u64],
5680 snapshot: Snapshot,
5681 ttl_now: i64,
5682 ) -> Result<Vec<Row>> {
5683 use std::collections::HashMap;
5684 let mut rows = Vec::with_capacity(rids.len());
5685 let tier_size = self.memtable.len() + self.mutable_run.len();
5702 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
5703 if rids.len().saturating_mul(24) < tier_size {
5704 for &rid in rids {
5705 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
5706 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
5707 let newest = match (mem, mrun) {
5708 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
5709 (Some((_, mr)), None) => Some(mr),
5710 (None, Some((_, rr))) => Some(rr),
5711 (None, None) => None,
5712 };
5713 if let Some(row) = newest {
5714 overlay.insert(rid, row);
5715 }
5716 }
5717 } else {
5718 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
5719 overlay
5720 .entry(row.row_id.0)
5721 .and_modify(|e| {
5722 if row.committed_epoch > e.committed_epoch {
5723 *e = row.clone();
5724 }
5725 })
5726 .or_insert(row);
5727 };
5728 for row in self.memtable.visible_versions(snapshot.epoch) {
5729 fold_newest(row, &mut overlay);
5730 }
5731 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5732 fold_newest(row, &mut overlay);
5733 }
5734 }
5735 if self.run_refs.len() == 1 {
5736 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5737 if rids.len().saturating_mul(24) < reader.row_count() {
5745 for &rid in rids {
5746 if let Some(r) = overlay.get(&rid) {
5747 if !r.deleted {
5748 rows.push(r.clone());
5749 }
5750 continue;
5751 }
5752 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
5753 if !row.deleted {
5754 rows.push(row);
5755 }
5756 }
5757 }
5758 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5759 return Ok(rows);
5760 }
5761 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
5770 enum Src {
5773 Overlay,
5774 Run,
5775 }
5776 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
5777 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
5778 for rid in rids {
5779 if overlay.contains_key(rid) {
5780 plan.push(Src::Overlay);
5781 continue;
5782 }
5783 match vis_rids.binary_search(&(*rid as i64)) {
5784 Ok(i) => {
5785 plan.push(Src::Run);
5786 fetch.push(positions[i]);
5787 }
5788 Err(_) => { }
5789 }
5790 }
5791 let fetched = reader.materialize_batch(&fetch)?;
5792 let mut fetched_iter = fetched.into_iter();
5793 for (rid, src) in rids.iter().zip(plan) {
5794 match src {
5795 Src::Overlay => {
5796 if let Some(r) = overlay.get(rid) {
5797 if !r.deleted {
5798 rows.push(r.clone());
5799 }
5800 }
5801 }
5802 Src::Run => {
5803 if let Some(row) = fetched_iter.next() {
5804 if !row.deleted {
5805 rows.push(row);
5806 }
5807 }
5808 }
5809 }
5810 }
5811 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5812 return Ok(rows);
5813 }
5814 let mut readers: Vec<_> = self
5818 .run_refs
5819 .iter()
5820 .map(|rr| self.open_reader(rr.run_id))
5821 .collect::<Result<Vec<_>>>()?;
5822 for rid in rids {
5823 if let Some(r) = overlay.get(rid) {
5824 if !r.deleted {
5825 rows.push(r.clone());
5826 }
5827 continue;
5828 }
5829 let mut best: Option<(Epoch, Row)> = None;
5830 for reader in readers.iter_mut() {
5831 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
5832 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
5833 best = Some((epoch, row));
5834 }
5835 }
5836 }
5837 if let Some((_, r)) = best {
5838 if !r.deleted {
5839 rows.push(r);
5840 }
5841 }
5842 }
5843 rows.retain(|row| !self.row_expired_at(row, ttl_now));
5844 Ok(rows)
5845 }
5846
5847 pub fn indexes_complete(&self) -> bool {
5857 self.indexes_complete
5858 }
5859
5860 pub fn index_build_policy(&self) -> IndexBuildPolicy {
5862 self.index_build_policy
5863 }
5864
5865 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
5869 self.index_build_policy = policy;
5870 }
5871
5872 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
5877 if !self.indexes_complete {
5881 return None;
5882 }
5883 let b = self.bitmap.get(&column_id)?;
5884 let result: Vec<Vec<u8>> = b
5885 .keys()
5886 .into_iter()
5887 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
5888 .collect();
5889 Some(result)
5890 }
5891
5892 pub fn fk_join_row_ids(
5893 &self,
5894 fk_column_id: u16,
5895 pk_values: &[Vec<u8>],
5896 fk_conditions: &[crate::query::Condition],
5897 snapshot: Snapshot,
5898 ) -> Result<Vec<u64>> {
5899 let Some(b) = self.bitmap.get(&fk_column_id) else {
5900 return Ok(Vec::new());
5901 };
5902 let mut join_set = {
5903 let mut acc = roaring::RoaringBitmap::new();
5904 for v in pk_values {
5905 acc |= b.get(v);
5906 }
5907 RowIdSet::from_roaring(acc)
5908 };
5909 if !fk_conditions.is_empty() {
5910 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5911 sets.push(join_set);
5912 for c in fk_conditions {
5913 sets.push(self.resolve_condition(c, snapshot)?);
5914 }
5915 join_set = RowIdSet::intersect_many(sets);
5916 }
5917 Ok(join_set.into_sorted_vec())
5918 }
5919
5920 pub fn fk_join_count(
5926 &self,
5927 fk_column_id: u16,
5928 pk_values: &[Vec<u8>],
5929 fk_conditions: &[crate::query::Condition],
5930 snapshot: Snapshot,
5931 ) -> Result<u64> {
5932 let Some(b) = self.bitmap.get(&fk_column_id) else {
5933 return Ok(0);
5934 };
5935 let mut acc = roaring::RoaringBitmap::new();
5936 for v in pk_values {
5937 acc |= b.get(v);
5938 }
5939 if fk_conditions.is_empty() {
5940 return Ok(acc.len());
5941 }
5942 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
5943 sets.push(RowIdSet::from_roaring(acc));
5944 for c in fk_conditions {
5945 sets.push(self.resolve_condition(c, snapshot)?);
5946 }
5947 Ok(RowIdSet::intersect_many(sets).len() as u64)
5948 }
5949
5950 fn resolve_condition(
5955 &self,
5956 c: &crate::query::Condition,
5957 snapshot: Snapshot,
5958 ) -> Result<RowIdSet> {
5959 self.resolve_condition_with_allowed(c, snapshot, None)
5960 }
5961
5962 fn resolve_condition_with_allowed(
5963 &self,
5964 c: &crate::query::Condition,
5965 snapshot: Snapshot,
5966 allowed: Option<&std::collections::HashSet<RowId>>,
5967 ) -> Result<RowIdSet> {
5968 use crate::query::Condition;
5969 self.validate_condition(c)?;
5970 Ok(match c {
5971 Condition::Pk(key) => {
5972 let lookup = self
5973 .schema
5974 .primary_key()
5975 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5976 .unwrap_or_else(|| key.clone());
5977 self.hot
5978 .get(&lookup)
5979 .map(|r| RowIdSet::one(r.0))
5980 .unwrap_or_else(RowIdSet::empty)
5981 }
5982 Condition::BitmapEq { column_id, value } => {
5983 let lookup = self.index_lookup_key_bytes(*column_id, value);
5984 self.bitmap
5985 .get(column_id)
5986 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5987 .unwrap_or_else(RowIdSet::empty)
5988 }
5989 Condition::BitmapIn { column_id, values } => {
5990 let bm = self.bitmap.get(column_id);
5991 let mut acc = roaring::RoaringBitmap::new();
5992 if let Some(b) = bm {
5993 for v in values {
5994 let lookup = self.index_lookup_key_bytes(*column_id, v);
5995 acc |= b.get(&lookup);
5996 }
5997 }
5998 RowIdSet::from_roaring(acc)
5999 }
6000 Condition::BytesPrefix { column_id, prefix } => {
6001 if let Some(b) = self.bitmap.get(column_id) {
6006 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
6007 let mut acc = roaring::RoaringBitmap::new();
6008 for key in b.keys() {
6009 if key.starts_with(&lookup_prefix) {
6010 acc |= b.get(&key);
6011 }
6012 }
6013 RowIdSet::from_roaring(acc)
6014 } else {
6015 RowIdSet::empty()
6016 }
6017 }
6018 Condition::FmContains { column_id, pattern } => self
6019 .fm
6020 .get(column_id)
6021 .map(|f| {
6022 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
6023 })
6024 .unwrap_or_else(RowIdSet::empty),
6025 Condition::FmContainsAll {
6026 column_id,
6027 patterns,
6028 } => {
6029 if let Some(f) = self.fm.get(column_id) {
6032 let sets: Vec<RowIdSet> = patterns
6033 .iter()
6034 .map(|pat| {
6035 RowIdSet::from_unsorted(
6036 f.locate(pat).into_iter().map(|r| r.0).collect(),
6037 )
6038 })
6039 .collect();
6040 RowIdSet::intersect_many(sets)
6041 } else {
6042 RowIdSet::empty()
6043 }
6044 }
6045 Condition::Ann {
6046 column_id,
6047 query,
6048 k,
6049 } => RowIdSet::from_unsorted(
6050 self.retrieve_filtered(
6051 &crate::query::Retriever::Ann {
6052 column_id: *column_id,
6053 query: query.clone(),
6054 k: *k,
6055 },
6056 snapshot,
6057 None,
6058 allowed,
6059 None,
6060 None,
6061 )?
6062 .into_iter()
6063 .map(|hit| hit.row_id.0)
6064 .collect(),
6065 ),
6066 Condition::SparseMatch {
6067 column_id,
6068 query,
6069 k,
6070 } => RowIdSet::from_unsorted(
6071 self.retrieve_filtered(
6072 &crate::query::Retriever::Sparse {
6073 column_id: *column_id,
6074 query: query.clone(),
6075 k: *k,
6076 },
6077 snapshot,
6078 None,
6079 allowed,
6080 None,
6081 None,
6082 )?
6083 .into_iter()
6084 .map(|hit| hit.row_id.0)
6085 .collect(),
6086 ),
6087 Condition::MinHashSimilar {
6088 column_id,
6089 query,
6090 k,
6091 } => match self.minhash.get(column_id) {
6092 Some(index) => {
6093 let candidates = index.candidate_row_ids(query);
6094 let eligible =
6095 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
6096 RowIdSet::from_unsorted(
6097 index
6098 .search_filtered(query, *k, |row_id| {
6099 eligible.contains(&row_id)
6100 && allowed.map_or(true, |allowed| allowed.contains(&row_id))
6101 })
6102 .into_iter()
6103 .map(|(row_id, _)| row_id.0)
6104 .collect(),
6105 )
6106 }
6107 None => RowIdSet::empty(),
6108 },
6109 Condition::Range { column_id, lo, hi } => {
6110 let mut set = if let Some(li) = self.learned_range.get(column_id) {
6119 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
6120 } else if self.run_refs.len() == 1 {
6121 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6122 r.range_row_id_set_i64(*column_id, *lo, *hi)?
6123 } else {
6124 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
6125 };
6126 set.remove_many(self.overlay_rid_set(snapshot));
6127 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
6128 set
6129 }
6130 Condition::RangeF64 {
6131 column_id,
6132 lo,
6133 lo_inclusive,
6134 hi,
6135 hi_inclusive,
6136 } => {
6137 let mut set = if let Some(li) = self.learned_range.get(column_id) {
6140 RowIdSet::from_unsorted(
6141 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
6142 .into_iter()
6143 .collect(),
6144 )
6145 } else if self.run_refs.len() == 1 {
6146 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6147 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
6148 } else {
6149 return self.range_scan_f64(
6150 *column_id,
6151 *lo,
6152 *lo_inclusive,
6153 *hi,
6154 *hi_inclusive,
6155 snapshot,
6156 );
6157 };
6158 set.remove_many(self.overlay_rid_set(snapshot));
6159 self.range_scan_overlay_f64(
6160 &mut set,
6161 *column_id,
6162 *lo,
6163 *lo_inclusive,
6164 *hi,
6165 *hi_inclusive,
6166 snapshot,
6167 );
6168 set
6169 }
6170 Condition::IsNull { column_id } => {
6171 let mut set = if self.run_refs.len() == 1 {
6172 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6173 r.null_row_id_set(*column_id, true)?
6174 } else {
6175 return self.null_scan(*column_id, true, snapshot);
6176 };
6177 set.remove_many(self.overlay_rid_set(snapshot));
6178 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
6179 set
6180 }
6181 Condition::IsNotNull { column_id } => {
6182 let mut set = if self.run_refs.len() == 1 {
6183 let mut r = self.open_reader(self.run_refs[0].run_id)?;
6184 r.null_row_id_set(*column_id, false)?
6185 } else {
6186 return self.null_scan(*column_id, false, snapshot);
6187 };
6188 set.remove_many(self.overlay_rid_set(snapshot));
6189 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
6190 set
6191 }
6192 })
6193 }
6194
6195 fn range_scan_i64(
6203 &self,
6204 column_id: u16,
6205 lo: i64,
6206 hi: i64,
6207 snapshot: Snapshot,
6208 ) -> Result<RowIdSet> {
6209 let mut row_ids = Vec::new();
6210 let overlay_rids = self.overlay_rid_set(snapshot);
6211 for rr in &self.run_refs {
6212 let mut reader = self.open_reader(rr.run_id)?;
6213 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
6214 for rid in matched {
6215 if !overlay_rids.contains(&rid) {
6216 row_ids.push(rid);
6217 }
6218 }
6219 }
6220 let mut s = RowIdSet::from_unsorted(row_ids);
6221 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
6222 Ok(s)
6223 }
6224
6225 fn range_scan_f64(
6228 &self,
6229 column_id: u16,
6230 lo: f64,
6231 lo_inclusive: bool,
6232 hi: f64,
6233 hi_inclusive: bool,
6234 snapshot: Snapshot,
6235 ) -> Result<RowIdSet> {
6236 let mut row_ids = Vec::new();
6237 let overlay_rids = self.overlay_rid_set(snapshot);
6238 for rr in &self.run_refs {
6239 let mut reader = self.open_reader(rr.run_id)?;
6240 let matched = reader.range_row_ids_visible_f64(
6241 column_id,
6242 lo,
6243 lo_inclusive,
6244 hi,
6245 hi_inclusive,
6246 snapshot.epoch,
6247 )?;
6248 for rid in matched {
6249 if !overlay_rids.contains(&rid) {
6250 row_ids.push(rid);
6251 }
6252 }
6253 }
6254 let mut s = RowIdSet::from_unsorted(row_ids);
6255 self.range_scan_overlay_f64(
6256 &mut s,
6257 column_id,
6258 lo,
6259 lo_inclusive,
6260 hi,
6261 hi_inclusive,
6262 snapshot,
6263 );
6264 Ok(s)
6265 }
6266
6267 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
6269 let mut s = HashSet::new();
6270 for row in self.memtable.visible_versions(snapshot.epoch) {
6271 s.insert(row.row_id.0);
6272 }
6273 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6274 s.insert(row.row_id.0);
6275 }
6276 s
6277 }
6278
6279 fn range_scan_overlay_i64(
6280 &self,
6281 s: &mut RowIdSet,
6282 column_id: u16,
6283 lo: i64,
6284 hi: i64,
6285 snapshot: Snapshot,
6286 ) {
6287 let mut newest: HashMap<u64, &Row> = HashMap::new();
6292 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6293 let memtable = self.memtable.visible_versions(snapshot.epoch);
6294 for r in &mutable {
6295 newest.entry(r.row_id.0).or_insert(r);
6296 }
6297 for r in &memtable {
6298 newest.insert(r.row_id.0, r);
6299 }
6300 for row in newest.values() {
6301 if !row.deleted {
6302 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
6303 if *v >= lo && *v <= hi {
6304 s.insert(row.row_id.0);
6305 }
6306 }
6307 }
6308 }
6309 }
6310
6311 #[allow(clippy::too_many_arguments)]
6312 fn range_scan_overlay_f64(
6313 &self,
6314 s: &mut RowIdSet,
6315 column_id: u16,
6316 lo: f64,
6317 lo_inclusive: bool,
6318 hi: f64,
6319 hi_inclusive: bool,
6320 snapshot: Snapshot,
6321 ) {
6322 let mut newest: HashMap<u64, &Row> = HashMap::new();
6325 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6326 let memtable = self.memtable.visible_versions(snapshot.epoch);
6327 for r in &mutable {
6328 newest.entry(r.row_id.0).or_insert(r);
6329 }
6330 for r in &memtable {
6331 newest.insert(r.row_id.0, r);
6332 }
6333 for row in newest.values() {
6334 if !row.deleted {
6335 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
6336 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
6337 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
6338 if ok_lo && ok_hi {
6339 s.insert(row.row_id.0);
6340 }
6341 }
6342 }
6343 }
6344 }
6345
6346 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
6349 let mut row_ids = Vec::new();
6350 let overlay_rids = self.overlay_rid_set(snapshot);
6351 for rr in &self.run_refs {
6352 let mut reader = self.open_reader(rr.run_id)?;
6353 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
6354 for rid in matched {
6355 if !overlay_rids.contains(&rid) {
6356 row_ids.push(rid);
6357 }
6358 }
6359 }
6360 let mut s = RowIdSet::from_unsorted(row_ids);
6361 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
6362 Ok(s)
6363 }
6364
6365 fn null_scan_overlay(
6369 &self,
6370 s: &mut RowIdSet,
6371 column_id: u16,
6372 want_nulls: bool,
6373 snapshot: Snapshot,
6374 ) {
6375 let mut newest: HashMap<u64, &Row> = HashMap::new();
6376 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
6377 let memtable = self.memtable.visible_versions(snapshot.epoch);
6378 for r in &mutable {
6379 newest.entry(r.row_id.0).or_insert(r);
6380 }
6381 for r in &memtable {
6382 newest.insert(r.row_id.0, r);
6383 }
6384 for row in newest.values() {
6385 if row.deleted {
6386 continue;
6387 }
6388 let is_null = !row.columns.contains_key(&column_id)
6389 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
6390 if is_null == want_nulls {
6391 s.insert(row.row_id.0);
6392 }
6393 }
6394 }
6395
6396 pub fn snapshot(&self) -> Snapshot {
6397 Snapshot::at(self.epoch.visible())
6398 }
6399
6400 pub fn data_generation(&self) -> u64 {
6402 self.data_generation
6403 }
6404
6405 pub(crate) fn bump_data_generation(&mut self) {
6406 self.data_generation = self.data_generation.wrapping_add(1);
6407 }
6408
6409 pub(crate) fn table_id(&self) -> u64 {
6410 self.table_id
6411 }
6412
6413 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
6414 self.ensure_indexes_complete()?;
6415 self.memtable.seal();
6416 self.mutable_run.seal();
6417 self.hot.seal();
6418 for index in self.bitmap.values_mut() {
6419 index.seal();
6420 }
6421 for index in self.ann.values_mut() {
6422 index.seal();
6423 }
6424 for index in self.fm.values_mut() {
6425 index.seal();
6426 }
6427 for index in self.sparse.values_mut() {
6428 index.seal();
6429 }
6430 for index in self.minhash.values_mut() {
6431 index.seal();
6432 }
6433 self.pk_by_row.seal();
6434 let mut generation = self.clone();
6435 generation.read_only = true;
6436 generation.wal = WalSink::ReadOnly;
6437 generation.pending_delete_rids.clear();
6438 generation.pending_put_cols.clear();
6439 generation.pending_rows.clear();
6440 generation.pending_rows_auto_inc.clear();
6441 generation.pending_dels.clear();
6442 generation.pending_truncate = None;
6443 generation.agg_cache = Arc::new(HashMap::new());
6444 Ok(generation)
6445 }
6446
6447 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
6448 (std::mem::size_of::<Self>() as u64)
6449 .saturating_add(self.memtable.approx_bytes())
6450 .saturating_add(self.mutable_run.approx_bytes())
6451 .saturating_add(self.live_count.saturating_mul(64))
6452 }
6453
6454 pub fn pin_snapshot(&mut self) -> Snapshot {
6457 let e = self.epoch.visible();
6458 *self.pinned.entry(e).or_insert(0) += 1;
6459 Snapshot::at(e)
6460 }
6461
6462 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
6464 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
6465 *count -= 1;
6466 if *count == 0 {
6467 self.pinned.remove(&snap.epoch);
6468 }
6469 }
6470 }
6471
6472 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
6482 let local = self.pinned.keys().next().copied();
6483 let global = self.snapshots.min_pinned();
6484 let history = self.snapshots.history_floor(self.current_epoch());
6485 [local, global, history].into_iter().flatten().min()
6486 }
6487
6488 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
6492 self.ensure_writable()?;
6493 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
6494 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
6495 }
6496
6497 pub fn clear_ttl(&mut self) -> Result<()> {
6498 self.ensure_writable()?;
6499 self.apply_ttl_policy_at(None, self.current_epoch())
6500 }
6501
6502 pub fn ttl(&self) -> Option<TtlPolicy> {
6503 self.ttl
6504 }
6505
6506 pub(crate) fn prepare_ttl_policy(
6507 &self,
6508 column_name: &str,
6509 duration_nanos: u64,
6510 ) -> Result<TtlPolicy> {
6511 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
6512 return Err(MongrelError::InvalidArgument(
6513 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
6514 ));
6515 }
6516 let column = self
6517 .schema
6518 .columns
6519 .iter()
6520 .find(|column| column.name == column_name)
6521 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
6522 if column.ty != TypeId::TimestampNanos {
6523 return Err(MongrelError::Schema(format!(
6524 "TTL column {column_name} must be TimestampNanos, is {:?}",
6525 column.ty
6526 )));
6527 }
6528 Ok(TtlPolicy {
6529 column_id: column.id,
6530 duration_nanos,
6531 })
6532 }
6533
6534 pub(crate) fn apply_ttl_policy_at(
6535 &mut self,
6536 policy: Option<TtlPolicy>,
6537 epoch: Epoch,
6538 ) -> Result<()> {
6539 if let Some(policy) = policy {
6540 let column = self
6541 .schema
6542 .columns
6543 .iter()
6544 .find(|column| column.id == policy.column_id)
6545 .ok_or_else(|| {
6546 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
6547 })?;
6548 if column.ty != TypeId::TimestampNanos
6549 || policy.duration_nanos == 0
6550 || policy.duration_nanos > i64::MAX as u64
6551 {
6552 return Err(MongrelError::Schema("invalid TTL policy".into()));
6553 }
6554 }
6555 self.ttl = policy;
6556 self.agg_cache = Arc::new(HashMap::new());
6557 self.clear_result_cache();
6558 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6559 self.persist_manifest(epoch)
6560 }
6561
6562 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
6563 let Some(policy) = self.ttl else {
6564 return false;
6565 };
6566 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
6567 return false;
6568 };
6569 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
6570 }
6571
6572 pub fn current_epoch(&self) -> Epoch {
6573 self.epoch.visible()
6574 }
6575
6576 pub fn memtable_len(&self) -> usize {
6577 self.memtable.len()
6578 }
6579
6580 pub fn count(&self) -> u64 {
6583 if self.ttl.is_none()
6584 && self.pending_put_cols.is_empty()
6585 && self.pending_delete_rids.is_empty()
6586 && self.pending_rows.is_empty()
6587 && self.pending_dels.is_empty()
6588 && self.pending_truncate.is_none()
6589 {
6590 self.live_count
6591 } else {
6592 self.visible_rows(self.snapshot())
6593 .map(|rows| rows.len() as u64)
6594 .unwrap_or(self.live_count)
6595 }
6596 }
6597
6598 pub fn count_conditions(
6602 &mut self,
6603 conditions: &[crate::query::Condition],
6604 snapshot: Snapshot,
6605 ) -> Result<Option<u64>> {
6606 use crate::query::Condition;
6607 if self.ttl.is_some() {
6608 if conditions.is_empty() {
6609 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
6610 }
6611 let mut sets = Vec::with_capacity(conditions.len());
6612 for condition in conditions {
6613 sets.push(self.resolve_condition(condition, snapshot)?);
6614 }
6615 let survivors = RowIdSet::intersect_many(sets);
6616 let rows = self.visible_rows(snapshot)?;
6617 return Ok(Some(
6618 rows.into_iter()
6619 .filter(|row| survivors.contains(row.row_id.0))
6620 .count() as u64,
6621 ));
6622 }
6623 if conditions.is_empty() {
6624 return Ok(Some(self.count()));
6625 }
6626 let served = |c: &Condition| {
6627 matches!(
6628 c,
6629 Condition::Pk(_)
6630 | Condition::BitmapEq { .. }
6631 | Condition::BitmapIn { .. }
6632 | Condition::BytesPrefix { .. }
6633 | Condition::FmContains { .. }
6634 | Condition::FmContainsAll { .. }
6635 | Condition::Ann { .. }
6636 | Condition::Range { .. }
6637 | Condition::RangeF64 { .. }
6638 | Condition::SparseMatch { .. }
6639 | Condition::MinHashSimilar { .. }
6640 | Condition::IsNull { .. }
6641 | Condition::IsNotNull { .. }
6642 )
6643 };
6644 if !conditions.iter().all(served) {
6645 return Ok(None);
6646 }
6647 self.ensure_indexes_complete()?;
6648 if !self.pending_put_cols.is_empty()
6649 || !self.pending_delete_rids.is_empty()
6650 || !self.pending_rows.is_empty()
6651 || !self.pending_dels.is_empty()
6652 || self.pending_truncate.is_some()
6653 {
6654 let mut sets = Vec::with_capacity(conditions.len());
6655 for condition in conditions {
6656 sets.push(self.resolve_condition(condition, snapshot)?);
6657 }
6658 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6659 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
6660 }
6661 let mut sets = Vec::with_capacity(conditions.len());
6662 for condition in conditions {
6663 sets.push(self.resolve_condition(condition, snapshot)?);
6664 }
6665 let mut rids = RowIdSet::intersect_many(sets);
6666 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
6676 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
6677 }
6678 let count = rids.len() as u64;
6679 crate::trace::QueryTrace::record(|t| {
6680 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
6681 t.survivor_count = Some(count as usize);
6682 t.conditions_pushed = conditions.len();
6683 });
6684 Ok(Some(count))
6685 }
6686
6687 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
6692 let mut out = Vec::new();
6693 for row in self.memtable.visible_versions(snapshot.epoch) {
6694 if row.deleted {
6695 out.push(row.row_id.0);
6696 }
6697 }
6698 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6699 if row.deleted {
6700 out.push(row.row_id.0);
6701 }
6702 }
6703 out
6704 }
6705
6706 pub fn bulk_load_columns(
6715 &mut self,
6716 user_columns: Vec<(u16, columnar::NativeColumn)>,
6717 ) -> Result<Epoch> {
6718 self.bulk_load_columns_with(user_columns, 3, false, true)
6719 }
6720
6721 pub fn bulk_load_fast(
6728 &mut self,
6729 user_columns: Vec<(u16, columnar::NativeColumn)>,
6730 ) -> Result<Epoch> {
6731 self.bulk_load_columns_with(user_columns, -1, true, false)
6732 }
6733
6734 fn bulk_load_columns_with(
6735 &mut self,
6736 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
6737 zstd_level: i32,
6738 force_plain: bool,
6739 lz4: bool,
6740 ) -> Result<Epoch> {
6741 let epoch = self.commit()?;
6742 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
6743 if n == 0 {
6744 return Ok(epoch);
6745 }
6746 let live_before = self.live_count;
6747 self.spill_mutable_run(epoch)?;
6749 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6750 && self.indexes_complete
6751 && self.run_refs.is_empty()
6752 && self.memtable.is_empty()
6753 && self.mutable_run.is_empty();
6754 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6757 self.validate_columns_not_null(&user_columns, n)?;
6758 let winner_idx = self
6759 .bulk_pk_winner_indices(&user_columns, n)
6760 .filter(|idx| idx.len() != n);
6761 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6762 match winner_idx.as_deref() {
6763 Some(idx) => {
6764 let compacted = user_columns
6765 .iter()
6766 .map(|(id, c)| (*id, c.gather(idx)))
6767 .collect();
6768 (compacted, idx.len())
6769 }
6770 None => (user_columns, n),
6771 };
6772 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6773 let first = self.allocator.alloc_range(write_n as u64).0;
6774 for rid in first..first + write_n as u64 {
6775 self.reservoir.offer(rid);
6776 }
6777 let run_id = self.next_run_id;
6778 self.next_run_id += 1;
6779 let path = self.run_path(run_id);
6780 let mut writer =
6781 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
6782 if force_plain {
6783 writer = writer.with_plain();
6784 } else if lz4 {
6785 writer = writer.with_lz4();
6788 } else {
6789 writer = writer.with_zstd_level(zstd_level);
6790 }
6791 if let Some(kek) = &self.kek {
6792 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6793 }
6794 let header = writer.write_native(&path, &write_columns, write_n, first)?;
6795 self.run_refs.push(RunRef {
6796 run_id: run_id as u128,
6797 level: 0,
6798 epoch_created: epoch.0,
6799 row_count: header.row_count,
6800 });
6801 self.live_count = self.live_count.saturating_add(write_n as u64);
6802 if eager_index_build {
6803 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6804 self.index_columns_bulk(&write_columns, &row_ids);
6805 self.indexes_complete = true;
6806 self.build_learned_ranges()?;
6807 } else {
6808 self.indexes_complete = false;
6812 }
6813 self.mark_flushed(epoch)?;
6814 self.persist_manifest(epoch)?;
6815 if eager_index_build {
6816 self.checkpoint_indexes(epoch);
6817 }
6818 self.clear_result_cache();
6819 self.data_generation = self.data_generation.wrapping_add(1);
6820 Ok(epoch)
6821 }
6822
6823 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
6841 let n = row_ids.len();
6842 if n == 0 {
6843 return;
6844 }
6845 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
6846 columns.iter().map(|(id, c)| (*id, c)).collect();
6847 let ty_of: std::collections::HashMap<u16, TypeId> = self
6848 .schema
6849 .columns
6850 .iter()
6851 .map(|c| (c.id, c.ty.clone()))
6852 .collect();
6853 let pk_id = self.schema.primary_key().map(|c| c.id);
6854
6855 for (i, &rid) in row_ids.iter().enumerate() {
6856 let row_id = RowId(rid);
6857 if let Some(pid) = pk_id {
6858 if let Some(col) = by_id.get(&pid) {
6859 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
6860 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
6861 self.insert_hot_pk(key, row_id);
6862 }
6863 }
6864 }
6865 for idef in &self.schema.indexes {
6866 let Some(col) = by_id.get(&idef.column_id) else {
6867 continue;
6868 };
6869 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
6870 match idef.kind {
6871 IndexKind::Bitmap => {
6872 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
6873 if let Some(key) =
6874 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
6875 {
6876 b.insert(key, row_id);
6877 }
6878 }
6879 }
6880 IndexKind::FmIndex => {
6881 if let Some(f) = self.fm.get_mut(&idef.column_id) {
6882 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6883 f.insert(bytes.to_vec(), row_id);
6884 }
6885 }
6886 }
6887 IndexKind::Sparse => {
6888 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
6889 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6890 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
6891 s.insert(&terms, row_id);
6892 }
6893 }
6894 }
6895 }
6896 IndexKind::MinHash => {
6897 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
6898 if let Some(bytes) = columnar::native_bytes_at(col, i) {
6899 let tokens = crate::index::token_hashes_from_bytes(bytes);
6900 mh.insert(&tokens, row_id);
6901 }
6902 }
6903 }
6904 _ => {}
6905 }
6906 }
6907 }
6908 }
6909
6910 pub fn visible_columns_native(
6915 &self,
6916 snapshot: Snapshot,
6917 projection: Option<&[u16]>,
6918 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
6919 self.visible_columns_native_inner(snapshot, projection, None)
6920 }
6921
6922 pub fn visible_columns_native_with_control(
6923 &self,
6924 snapshot: Snapshot,
6925 projection: Option<&[u16]>,
6926 control: &crate::ExecutionControl,
6927 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
6928 self.visible_columns_native_inner(snapshot, projection, Some(control))
6929 }
6930
6931 fn visible_columns_native_inner(
6932 &self,
6933 snapshot: Snapshot,
6934 projection: Option<&[u16]>,
6935 control: Option<&crate::ExecutionControl>,
6936 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
6937 execution_checkpoint(control, 0)?;
6938 let wanted: Vec<u16> = match projection {
6939 Some(p) => p.to_vec(),
6940 None => self.schema.columns.iter().map(|c| c.id).collect(),
6941 };
6942 if self.ttl.is_none()
6943 && self.memtable.is_empty()
6944 && self.mutable_run.is_empty()
6945 && self.run_refs.len() == 1
6946 {
6947 let rr = self.run_refs[0].clone();
6948 let mut reader = self.open_reader(rr.run_id)?;
6949 let idxs = reader.visible_indices_native(snapshot.epoch)?;
6950 execution_checkpoint(control, 0)?;
6951 let all_visible = idxs.len() == reader.row_count();
6952 if reader.has_mmap() && control.is_none() {
6958 use rayon::prelude::*;
6959 let valid: Vec<u16> = wanted
6962 .iter()
6963 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
6964 .copied()
6965 .collect();
6966 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
6968 .par_iter()
6969 .filter_map(|cid| {
6970 reader
6971 .column_native_shared(*cid)
6972 .ok()
6973 .map(|col| (*cid, col))
6974 })
6975 .collect();
6976 let cols = decoded
6977 .into_iter()
6978 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
6979 .collect();
6980 return Ok(cols);
6981 }
6982 let mut cols = Vec::with_capacity(wanted.len());
6983 for (index, cid) in wanted.iter().enumerate() {
6984 execution_checkpoint(control, index)?;
6985 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
6986 Some(c) => c,
6987 None => continue,
6988 };
6989 let col = reader.column_native(cdef.id)?;
6990 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
6991 }
6992 return Ok(cols);
6993 }
6994 let vcols = self.visible_columns(snapshot)?;
6995 execution_checkpoint(control, 0)?;
6996 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
6997 let out: Vec<(u16, columnar::NativeColumn)> = vcols
6998 .into_iter()
6999 .filter(|(id, _)| want_set.contains(id))
7000 .map(|(id, vals)| {
7001 let ty = self
7002 .schema
7003 .columns
7004 .iter()
7005 .find(|c| c.id == id)
7006 .map(|c| c.ty.clone())
7007 .unwrap_or(TypeId::Bytes);
7008 (id, columnar::values_to_native(ty, &vals))
7009 })
7010 .collect();
7011 Ok(out)
7012 }
7013
7014 pub fn run_count(&self) -> usize {
7015 self.run_refs.len()
7016 }
7017
7018 pub fn memtable_is_empty(&self) -> bool {
7020 self.memtable.is_empty()
7021 }
7022
7023 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
7027 self.page_cache.stats()
7028 }
7029
7030 pub fn reset_page_cache_stats(&self) {
7032 self.page_cache.reset_stats();
7033 }
7034
7035 pub fn run_ids(&self) -> Vec<u128> {
7038 self.run_refs.iter().map(|r| r.run_id).collect()
7039 }
7040
7041 pub fn single_run_is_clean(&self) -> bool {
7045 if self.ttl.is_some() || self.run_refs.len() != 1 {
7046 return false;
7047 }
7048 self.open_reader(self.run_refs[0].run_id)
7049 .map(|r| r.is_clean())
7050 .unwrap_or(false)
7051 }
7052
7053 fn resolve_footprint(
7060 &self,
7061 conditions: &[crate::query::Condition],
7062 snapshot: Snapshot,
7063 ) -> roaring::RoaringBitmap {
7064 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7065 return roaring::RoaringBitmap::new();
7066 }
7067 if self.run_refs.is_empty() {
7068 return roaring::RoaringBitmap::new();
7069 }
7070 if self.run_refs.len() == 1 {
7072 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
7073 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
7074 return rids.to_roaring_lossy();
7075 }
7076 }
7077 }
7078 roaring::RoaringBitmap::new()
7079 }
7080
7081 pub fn query_columns_native_cached(
7092 &mut self,
7093 conditions: &[crate::query::Condition],
7094 projection: Option<&[u16]>,
7095 snapshot: Snapshot,
7096 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7097 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
7098 }
7099
7100 pub fn query_columns_native_cached_with_control(
7101 &mut self,
7102 conditions: &[crate::query::Condition],
7103 projection: Option<&[u16]>,
7104 snapshot: Snapshot,
7105 control: &crate::ExecutionControl,
7106 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7107 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
7108 }
7109
7110 fn query_columns_native_cached_inner(
7111 &mut self,
7112 conditions: &[crate::query::Condition],
7113 projection: Option<&[u16]>,
7114 snapshot: Snapshot,
7115 control: Option<&crate::ExecutionControl>,
7116 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7117 execution_checkpoint(control, 0)?;
7118 if self.ttl.is_some() {
7121 return self.query_columns_native_inner(conditions, projection, snapshot, control);
7122 }
7123 if conditions.is_empty() {
7124 return self.query_columns_native_inner(conditions, projection, snapshot, control);
7125 }
7126 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
7130 if let Some(hit) = self.result_cache.lock().get_columns(key) {
7131 crate::trace::QueryTrace::record(|t| {
7132 t.result_cache_hit = true;
7133 t.scan_mode = crate::trace::ScanMode::NativePushdown;
7134 });
7135 return Ok(Some((*hit).clone()));
7136 }
7137 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
7138 execution_checkpoint(control, 0)?;
7139 if let Some(cols) = &res {
7140 let footprint = self.resolve_footprint(conditions, snapshot);
7141 let condition_cols = crate::query::condition_columns(conditions);
7142 execution_checkpoint(control, 0)?;
7143 self.result_cache.lock().insert(
7144 key,
7145 CachedEntry {
7146 data: CachedData::Columns(Arc::new(cols.clone())),
7147 footprint,
7148 condition_cols,
7149 },
7150 );
7151 }
7152 Ok(res)
7153 }
7154
7155 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
7160 if self.ttl.is_some() {
7161 return self.query(q);
7162 }
7163 if q.conditions.is_empty() {
7164 return self.query(q);
7165 }
7166 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
7167 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
7168 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
7169 if let Some(hit) = self.result_cache.lock().get_rows(key) {
7170 crate::trace::QueryTrace::record(|t| {
7171 t.result_cache_hit = true;
7172 t.scan_mode = crate::trace::ScanMode::Materialized;
7173 });
7174 return Ok((*hit).clone());
7175 }
7176 let rows = self.query(q)?;
7177 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
7178 let condition_cols = crate::query::condition_columns(&q.conditions);
7179 self.result_cache.lock().insert(
7180 key,
7181 CachedEntry {
7182 data: CachedData::Rows(Arc::new(rows.clone())),
7183 footprint,
7184 condition_cols,
7185 },
7186 );
7187 Ok(rows)
7188 }
7189
7190 #[allow(clippy::type_complexity)]
7205 pub fn query_columns_native_traced(
7206 &mut self,
7207 conditions: &[crate::query::Condition],
7208 projection: Option<&[u16]>,
7209 snapshot: Snapshot,
7210 ) -> Result<(
7211 Option<Vec<(u16, columnar::NativeColumn)>>,
7212 crate::trace::QueryTrace,
7213 )> {
7214 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7215 self.query_columns_native(conditions, projection, snapshot)
7216 });
7217 Ok((result?, trace))
7218 }
7219
7220 #[allow(clippy::type_complexity)]
7223 pub fn query_columns_native_cached_traced(
7224 &mut self,
7225 conditions: &[crate::query::Condition],
7226 projection: Option<&[u16]>,
7227 snapshot: Snapshot,
7228 ) -> Result<(
7229 Option<Vec<(u16, columnar::NativeColumn)>>,
7230 crate::trace::QueryTrace,
7231 )> {
7232 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7233 self.query_columns_native_cached(conditions, projection, snapshot)
7234 });
7235 Ok((result?, trace))
7236 }
7237
7238 pub fn native_page_cursor_traced(
7240 &self,
7241 snapshot: Snapshot,
7242 projection: Vec<(u16, TypeId)>,
7243 conditions: &[crate::query::Condition],
7244 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
7245 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7246 self.native_page_cursor(snapshot, projection, conditions)
7247 });
7248 Ok((result?, trace))
7249 }
7250
7251 pub fn native_multi_run_cursor_traced(
7253 &self,
7254 snapshot: Snapshot,
7255 projection: Vec<(u16, TypeId)>,
7256 conditions: &[crate::query::Condition],
7257 ) -> Result<(
7258 Option<crate::cursor::MultiRunCursor>,
7259 crate::trace::QueryTrace,
7260 )> {
7261 let (result, trace) = crate::trace::QueryTrace::capture(|| {
7262 self.native_multi_run_cursor(snapshot, projection, conditions)
7263 });
7264 Ok((result?, trace))
7265 }
7266
7267 pub fn count_conditions_traced(
7269 &mut self,
7270 conditions: &[crate::query::Condition],
7271 snapshot: Snapshot,
7272 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
7273 let (result, trace) =
7274 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
7275 Ok((result?, trace))
7276 }
7277
7278 pub fn query_traced(
7280 &mut self,
7281 q: &crate::query::Query,
7282 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
7283 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
7284 Ok((result?, trace))
7285 }
7286
7287 pub fn query_columns_native(
7292 &mut self,
7293 conditions: &[crate::query::Condition],
7294 projection: Option<&[u16]>,
7295 snapshot: Snapshot,
7296 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7297 self.query_columns_native_inner(conditions, projection, snapshot, None)
7298 }
7299
7300 pub fn query_columns_native_with_control(
7301 &mut self,
7302 conditions: &[crate::query::Condition],
7303 projection: Option<&[u16]>,
7304 snapshot: Snapshot,
7305 control: &crate::ExecutionControl,
7306 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7307 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
7308 }
7309
7310 fn query_columns_native_inner(
7311 &mut self,
7312 conditions: &[crate::query::Condition],
7313 projection: Option<&[u16]>,
7314 snapshot: Snapshot,
7315 control: Option<&crate::ExecutionControl>,
7316 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
7317 use crate::query::Condition;
7318 execution_checkpoint(control, 0)?;
7319 if self.ttl.is_some() {
7322 return Ok(None);
7323 }
7324 if conditions.is_empty() {
7325 return Ok(None);
7326 }
7327 self.ensure_indexes_complete()?;
7328
7329 let served = |c: &Condition| {
7334 matches!(
7335 c,
7336 Condition::Pk(_)
7337 | Condition::BitmapEq { .. }
7338 | Condition::BitmapIn { .. }
7339 | Condition::BytesPrefix { .. }
7340 | Condition::FmContains { .. }
7341 | Condition::FmContainsAll { .. }
7342 | Condition::Ann { .. }
7343 | Condition::Range { .. }
7344 | Condition::RangeF64 { .. }
7345 | Condition::SparseMatch { .. }
7346 | Condition::MinHashSimilar { .. }
7347 | Condition::IsNull { .. }
7348 | Condition::IsNotNull { .. }
7349 )
7350 };
7351 if !conditions.iter().all(served) {
7352 return Ok(None);
7353 }
7354 let fast_path =
7355 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
7356 crate::trace::QueryTrace::record(|t| {
7357 t.run_count = self.run_refs.len();
7358 t.memtable_rows = self.memtable.len();
7359 t.mutable_run_rows = self.mutable_run.len();
7360 t.conditions_pushed = conditions.len();
7361 t.learned_range_used = conditions.iter().any(|c| match c {
7362 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
7363 self.learned_range.contains_key(column_id)
7364 }
7365 _ => false,
7366 });
7367 });
7368 let col_ids: Vec<u16> = projection
7370 .map(|p| p.to_vec())
7371 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
7372 let proj_pairs: Vec<(u16, TypeId)> = col_ids
7373 .iter()
7374 .map(|&cid| {
7375 let ty = self
7376 .schema
7377 .columns
7378 .iter()
7379 .find(|c| c.id == cid)
7380 .map(|c| c.ty.clone())
7381 .unwrap_or(TypeId::Bytes);
7382 (cid, ty)
7383 })
7384 .collect();
7385
7386 if fast_path {
7392 let needs_column = conditions.iter().any(|c| match c {
7395 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
7396 Condition::RangeF64 { column_id, .. } => {
7397 !self.learned_range.contains_key(column_id)
7398 }
7399 _ => false,
7400 });
7401 let mut reader_opt: Option<RunReader> = if needs_column {
7402 Some(self.open_reader(self.run_refs[0].run_id)?)
7403 } else {
7404 None
7405 };
7406 let mut sets: Vec<RowIdSet> = Vec::new();
7407 for (index, c) in conditions.iter().enumerate() {
7408 execution_checkpoint(control, index)?;
7409 let s = match c {
7410 Condition::Range { column_id, lo, hi }
7411 if !self.learned_range.contains_key(column_id) =>
7412 {
7413 if reader_opt.is_none() {
7414 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7415 }
7416 reader_opt
7417 .as_mut()
7418 .expect("reader opened for range")
7419 .range_row_id_set_i64(*column_id, *lo, *hi)?
7420 }
7421 Condition::RangeF64 {
7422 column_id,
7423 lo,
7424 lo_inclusive,
7425 hi,
7426 hi_inclusive,
7427 } if !self.learned_range.contains_key(column_id) => {
7428 if reader_opt.is_none() {
7429 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
7430 }
7431 reader_opt
7432 .as_mut()
7433 .expect("reader opened for range")
7434 .range_row_id_set_f64(
7435 *column_id,
7436 *lo,
7437 *lo_inclusive,
7438 *hi,
7439 *hi_inclusive,
7440 )?
7441 }
7442 _ => self.resolve_condition(c, snapshot)?,
7443 };
7444 sets.push(s);
7445 }
7446 let candidates = RowIdSet::intersect_many(sets);
7447 crate::trace::QueryTrace::record(|t| {
7448 t.survivor_count = Some(candidates.len());
7449 });
7450 if candidates.is_empty() {
7451 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
7452 .iter()
7453 .map(|&id| {
7454 (
7455 id,
7456 columnar::null_native(
7457 proj_pairs
7458 .iter()
7459 .find(|(c, _)| c == &id)
7460 .map(|(_, t)| t.clone())
7461 .unwrap_or(TypeId::Bytes),
7462 0,
7463 ),
7464 )
7465 })
7466 .collect();
7467 return Ok(Some(cols));
7468 }
7469 let mut reader = match reader_opt.take() {
7470 Some(r) => r,
7471 None => self.open_reader(self.run_refs[0].run_id)?,
7472 };
7473 let candidate_ids = candidates.into_sorted_vec();
7474 let (positions, fast_rid) = if let Some(positions) =
7475 reader.positions_for_row_ids_fast(&candidate_ids)
7476 {
7477 (positions, true)
7478 } else {
7479 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
7480 match col {
7481 columnar::NativeColumn::Int64 { data, .. } => {
7482 let mut p = Vec::with_capacity(candidate_ids.len());
7483 for (index, rid) in candidate_ids.iter().enumerate() {
7484 execution_checkpoint(control, index)?;
7485 if let Ok(position) = data.binary_search(&(*rid as i64)) {
7486 p.push(position);
7487 }
7488 }
7489 p.sort_unstable();
7490 (p, false)
7491 }
7492 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
7493 }
7494 };
7495 crate::trace::QueryTrace::record(|t| {
7496 t.scan_mode = crate::trace::ScanMode::NativePushdown;
7497 t.fast_row_id_map = fast_rid;
7498 });
7499 let mut cols = Vec::with_capacity(col_ids.len());
7500 for (index, cid) in col_ids.iter().enumerate() {
7501 execution_checkpoint(control, index)?;
7502 let col = reader.column_native(*cid)?;
7503 cols.push((*cid, col.gather(&positions)));
7504 }
7505 return Ok(Some(cols));
7506 }
7507
7508 if !self.run_refs.is_empty() {
7521 use crate::cursor::{
7522 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
7523 };
7524 let remaining: usize;
7525 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
7526 let c = self
7527 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
7528 .expect("single-run cursor should build when run_refs.len() == 1");
7529 remaining = c.remaining_rows();
7530 Box::new(c)
7531 } else {
7532 let c = self
7533 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
7534 .expect("multi-run cursor should build when run_refs.len() >= 1");
7535 remaining = c.remaining_rows();
7536 Box::new(c)
7537 };
7538 crate::trace::QueryTrace::record(|t| {
7539 if t.survivor_count.is_none() {
7540 t.survivor_count = Some(remaining);
7541 }
7542 });
7543 let cols = match control {
7544 Some(control) => {
7545 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
7546 }
7547 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
7548 };
7549 return Ok(Some(cols));
7550 }
7551
7552 crate::trace::QueryTrace::record(|t| {
7557 t.scan_mode = crate::trace::ScanMode::Materialized;
7558 t.row_materialized = true;
7559 });
7560 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7561 for (index, c) in conditions.iter().enumerate() {
7562 execution_checkpoint(control, index)?;
7563 sets.push(self.resolve_condition(c, snapshot)?);
7564 }
7565 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7566 let rows = self.rows_for_rids(&rids, snapshot)?;
7567 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
7568 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
7569 execution_checkpoint(control, index)?;
7570 let vals: Vec<Value> = rows
7571 .iter()
7572 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
7573 .collect();
7574 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
7575 }
7576 Ok(Some(cols))
7577 }
7578
7579 pub fn native_page_cursor(
7594 &self,
7595 snapshot: Snapshot,
7596 projection: Vec<(u16, TypeId)>,
7597 conditions: &[crate::query::Condition],
7598 ) -> Result<Option<NativePageCursor>> {
7599 use crate::cursor::build_page_plans;
7600 if self.ttl.is_some() {
7601 return Ok(None);
7602 }
7603 if !conditions.is_empty() && !self.indexes_complete {
7606 return Ok(None);
7607 }
7608 if self.run_refs.len() != 1 {
7609 return Ok(None);
7610 }
7611 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7612 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7613
7614 let overlay_rids: HashSet<u64> = {
7617 let mut s = HashSet::new();
7618 for row in self.memtable.visible_versions(snapshot.epoch) {
7619 s.insert(row.row_id.0);
7620 }
7621 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7622 s.insert(row.row_id.0);
7623 }
7624 s
7625 };
7626
7627 let survivors = if conditions.is_empty() {
7631 None
7632 } else {
7633 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
7634 };
7635
7636 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
7643 survivors.clone()
7644 } else if let Some(s) = &survivors {
7645 let mut run_set = s.clone();
7646 run_set.remove_many(overlay_rids.iter().copied());
7647 Some(run_set)
7648 } else {
7649 Some(RowIdSet::from_unsorted(
7650 rids.iter()
7651 .map(|&r| r as u64)
7652 .filter(|r| !overlay_rids.contains(r))
7653 .collect(),
7654 ))
7655 };
7656
7657 let overlay_rows = if overlay_rids.is_empty() {
7658 Vec::new()
7659 } else {
7660 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7661 self.overlay_visible_rows(snapshot, bound)
7662 };
7663
7664 let plans = if positions.is_empty() {
7666 Vec::new()
7667 } else {
7668 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
7669 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
7670 };
7671
7672 let overlay = if overlay_rows.is_empty() {
7674 None
7675 } else {
7676 let filtered =
7677 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7678 if filtered.is_empty() {
7679 None
7680 } else {
7681 Some(self.materialize_overlay(&filtered, &projection))
7682 }
7683 };
7684
7685 let overlay_row_count = overlay
7686 .as_ref()
7687 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7688 .unwrap_or(0);
7689 crate::trace::QueryTrace::record(|t| {
7690 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
7691 t.run_count = self.run_refs.len();
7692 t.memtable_rows = self.memtable.len();
7693 t.mutable_run_rows = self.mutable_run.len();
7694 t.overlay_rows = overlay_row_count;
7695 t.conditions_pushed = conditions.len();
7696 t.pages_decoded = plans
7697 .iter()
7698 .map(|p| p.positions.len())
7699 .sum::<usize>()
7700 .min(1);
7701 });
7702
7703 Ok(Some(NativePageCursor::new_with_overlay(
7704 reader, projection, plans, overlay,
7705 )))
7706 }
7707 #[allow(clippy::type_complexity)]
7717 pub fn native_multi_run_cursor(
7718 &self,
7719 snapshot: Snapshot,
7720 projection: Vec<(u16, TypeId)>,
7721 conditions: &[crate::query::Condition],
7722 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
7723 use crate::cursor::{MultiRunCursor, RunStream};
7724 use crate::sorted_run::SYS_ROW_ID;
7725 use std::collections::{BinaryHeap, HashMap, HashSet};
7726 if self.ttl.is_some() {
7727 return Ok(None);
7728 }
7729 if !conditions.is_empty() && !self.indexes_complete {
7732 return Ok(None);
7733 }
7734 if self.run_refs.is_empty() {
7735 return Ok(None);
7736 }
7737
7738 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
7740 Vec::with_capacity(self.run_refs.len());
7741 for rr in &self.run_refs {
7742 let mut reader = self.open_reader(rr.run_id)?;
7743 let (rids, eps, del) = reader.system_columns_native()?;
7744 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
7745 run_meta.push((reader, rids, eps, del, page_rows));
7746 }
7747
7748 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
7752 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
7753 for i in 0..rids.len() {
7754 let rid = rids[i] as u64;
7755 let e = eps[i] as u64;
7756 if e > snapshot.epoch.0 {
7757 continue;
7758 }
7759 let is_del = del[i] != 0;
7760 best.entry(rid)
7761 .and_modify(|cur| {
7762 if e > cur.0 {
7763 *cur = (e, run_idx, i, is_del);
7764 }
7765 })
7766 .or_insert((e, run_idx, i, is_del));
7767 }
7768 }
7769
7770 let overlay_rids: HashSet<u64> = {
7772 let mut s = HashSet::new();
7773 for row in self.memtable.visible_versions(snapshot.epoch) {
7774 s.insert(row.row_id.0);
7775 }
7776 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7777 s.insert(row.row_id.0);
7778 }
7779 s
7780 };
7781
7782 let survivors: Option<RowIdSet> = if conditions.is_empty() {
7784 None
7785 } else {
7786 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7787 for c in conditions {
7788 sets.push(self.resolve_condition(c, snapshot)?);
7789 }
7790 Some(RowIdSet::intersect_many(sets))
7791 };
7792
7793 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
7797 for (rid, (_, run_idx, pos, deleted)) in &best {
7798 if *deleted {
7799 continue;
7800 }
7801 if overlay_rids.contains(rid) {
7802 continue;
7803 }
7804 if let Some(s) = &survivors {
7805 if !s.contains(*rid) {
7806 continue;
7807 }
7808 }
7809 per_run[*run_idx].push((*rid, *pos));
7810 }
7811 for v in per_run.iter_mut() {
7812 v.sort_unstable_by_key(|&(rid, _)| rid);
7813 }
7814
7815 let mut streams = Vec::with_capacity(run_meta.len());
7817 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
7818 let mut total = 0usize;
7819 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
7820 let mut starts = Vec::with_capacity(page_rows.len());
7821 let mut acc = 0usize;
7822 for &r in &page_rows {
7823 starts.push(acc);
7824 acc += r;
7825 }
7826 let mut survivors_vec: Vec<(u64, usize, usize)> =
7827 Vec::with_capacity(per_run[run_idx].len());
7828 for &(rid, pos) in &per_run[run_idx] {
7829 let page_seq = match starts.partition_point(|&s| s <= pos) {
7830 0 => continue,
7831 p => p - 1,
7832 };
7833 let within = pos - starts[page_seq];
7834 survivors_vec.push((rid, page_seq, within));
7835 }
7836 total += survivors_vec.len();
7837 if let Some(&(rid, _, _)) = survivors_vec.first() {
7838 heap.push(std::cmp::Reverse((rid, run_idx)));
7839 }
7840 streams.push(RunStream::new(reader, survivors_vec, page_rows));
7841 }
7842
7843 let overlay_rows = if overlay_rids.is_empty() {
7845 Vec::new()
7846 } else {
7847 let bound = Self::overlay_materialization_bound(conditions, &survivors);
7848 self.overlay_visible_rows(snapshot, bound)
7849 };
7850 let overlay = if overlay_rows.is_empty() {
7851 None
7852 } else {
7853 let filtered =
7854 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
7855 if filtered.is_empty() {
7856 None
7857 } else {
7858 Some(self.materialize_overlay(&filtered, &projection))
7859 }
7860 };
7861
7862 let overlay_row_count = overlay
7863 .as_ref()
7864 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
7865 .unwrap_or(0);
7866 crate::trace::QueryTrace::record(|t| {
7867 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
7868 t.run_count = self.run_refs.len();
7869 t.memtable_rows = self.memtable.len();
7870 t.mutable_run_rows = self.mutable_run.len();
7871 t.overlay_rows = overlay_row_count;
7872 t.conditions_pushed = conditions.len();
7873 t.survivor_count = Some(total);
7874 });
7875
7876 Ok(Some(MultiRunCursor::new(
7877 streams, projection, heap, total, overlay,
7878 )))
7879 }
7880
7881 fn overlay_materialization_bound<'a>(
7893 conditions: &[crate::query::Condition],
7894 survivors: &'a Option<RowIdSet>,
7895 ) -> Option<&'a RowIdSet> {
7896 use crate::query::Condition;
7897 let has_range = conditions
7898 .iter()
7899 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7900 if has_range {
7901 None
7902 } else {
7903 survivors.as_ref()
7904 }
7905 }
7906
7907 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
7919 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
7920 let mut fold = |row: Row| {
7921 if let Some(b) = bound {
7922 if !b.contains(row.row_id.0) {
7923 return;
7924 }
7925 }
7926 best.entry(row.row_id.0)
7927 .and_modify(|(be, br)| {
7928 if row.committed_epoch > *be {
7929 *be = row.committed_epoch;
7930 *br = row.clone();
7931 }
7932 })
7933 .or_insert_with(|| (row.committed_epoch, row));
7934 };
7935 for row in self.memtable.visible_versions(snapshot.epoch) {
7936 fold(row);
7937 }
7938 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7939 fold(row);
7940 }
7941 let mut out: Vec<Row> = best
7942 .into_values()
7943 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
7944 .collect();
7945 out.sort_by_key(|r| r.row_id);
7946 out
7947 }
7948
7949 fn filter_overlay_rows(
7957 &self,
7958 rows: Vec<Row>,
7959 conditions: &[crate::query::Condition],
7960 survivors: Option<&RowIdSet>,
7961 snapshot: Snapshot,
7962 ) -> Result<Vec<Row>> {
7963 if conditions.is_empty() {
7964 return Ok(rows);
7965 }
7966 use crate::query::Condition;
7967 let all_index_served = !conditions
7971 .iter()
7972 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
7973 if all_index_served {
7974 return Ok(rows
7975 .into_iter()
7976 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
7977 .collect());
7978 }
7979 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
7982 for c in conditions {
7983 let s = match c {
7984 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
7985 _ => self.resolve_condition(c, snapshot)?,
7986 };
7987 per_cond_sets.push(s);
7988 }
7989 Ok(rows
7990 .into_iter()
7991 .filter(|row| {
7992 conditions.iter().enumerate().all(|(i, c)| match c {
7993 Condition::Range { column_id, lo, hi } => {
7994 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
7995 }
7996 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
7997 match row.columns.get(column_id) {
7998 Some(Value::Float64(v)) => {
7999 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
8000 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
8001 lo_ok && hi_ok
8002 }
8003 _ => false,
8004 }
8005 }
8006 _ => per_cond_sets[i].contains(row.row_id.0),
8007 })
8008 })
8009 .collect())
8010 }
8011
8012 fn materialize_overlay(
8015 &self,
8016 rows: &[Row],
8017 projection: &[(u16, TypeId)],
8018 ) -> Vec<columnar::NativeColumn> {
8019 if projection.is_empty() {
8020 return vec![columnar::null_native(TypeId::Int64, rows.len())];
8021 }
8022 let mut cols = Vec::with_capacity(projection.len());
8023 for (cid, ty) in projection {
8024 let vals: Vec<Value> = rows
8025 .iter()
8026 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8027 .collect();
8028 cols.push(columnar::values_to_native(ty.clone(), &vals));
8029 }
8030 cols
8031 }
8032
8033 fn resolve_survivor_rids(
8038 &self,
8039 conditions: &[crate::query::Condition],
8040 reader: &mut RunReader,
8041 snapshot: Snapshot,
8042 ) -> Result<RowIdSet> {
8043 use crate::query::Condition;
8044 let mut sets: Vec<RowIdSet> = Vec::new();
8045 for c in conditions {
8046 self.validate_condition(c)?;
8047 let s: RowIdSet = match c {
8048 Condition::Pk(key) => {
8049 let lookup = self
8050 .schema
8051 .primary_key()
8052 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
8053 .unwrap_or_else(|| key.clone());
8054 self.hot
8055 .get(&lookup)
8056 .map(|r| RowIdSet::one(r.0))
8057 .unwrap_or_else(RowIdSet::empty)
8058 }
8059 Condition::BitmapEq { column_id, value } => {
8060 let lookup = self.index_lookup_key_bytes(*column_id, value);
8061 self.bitmap
8062 .get(column_id)
8063 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
8064 .unwrap_or_else(RowIdSet::empty)
8065 }
8066 Condition::BitmapIn { column_id, values } => {
8067 let bm = self.bitmap.get(column_id);
8068 let mut acc = roaring::RoaringBitmap::new();
8069 if let Some(b) = bm {
8070 for v in values {
8071 let lookup = self.index_lookup_key_bytes(*column_id, v);
8072 acc |= b.get(&lookup);
8073 }
8074 }
8075 RowIdSet::from_roaring(acc)
8076 }
8077 Condition::BytesPrefix { column_id, prefix } => {
8078 if let Some(b) = self.bitmap.get(column_id) {
8079 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
8080 let mut acc = roaring::RoaringBitmap::new();
8081 for key in b.keys() {
8082 if key.starts_with(&lookup_prefix) {
8083 acc |= b.get(&key);
8084 }
8085 }
8086 RowIdSet::from_roaring(acc)
8087 } else {
8088 RowIdSet::empty()
8089 }
8090 }
8091 Condition::FmContains { column_id, pattern } => self
8092 .fm
8093 .get(column_id)
8094 .map(|f| {
8095 RowIdSet::from_unsorted(
8096 f.locate(pattern).into_iter().map(|r| r.0).collect(),
8097 )
8098 })
8099 .unwrap_or_else(RowIdSet::empty),
8100 Condition::FmContainsAll {
8101 column_id,
8102 patterns,
8103 } => {
8104 if let Some(f) = self.fm.get(column_id) {
8105 let sets: Vec<RowIdSet> = patterns
8106 .iter()
8107 .map(|pat| {
8108 RowIdSet::from_unsorted(
8109 f.locate(pat).into_iter().map(|r| r.0).collect(),
8110 )
8111 })
8112 .collect();
8113 RowIdSet::intersect_many(sets)
8114 } else {
8115 RowIdSet::empty()
8116 }
8117 }
8118 Condition::Ann {
8119 column_id,
8120 query,
8121 k,
8122 } => RowIdSet::from_unsorted(
8123 self.retrieve_filtered(
8124 &crate::query::Retriever::Ann {
8125 column_id: *column_id,
8126 query: query.clone(),
8127 k: *k,
8128 },
8129 snapshot,
8130 None,
8131 None,
8132 None,
8133 None,
8134 )?
8135 .into_iter()
8136 .map(|hit| hit.row_id.0)
8137 .collect(),
8138 ),
8139 Condition::SparseMatch {
8140 column_id,
8141 query,
8142 k,
8143 } => RowIdSet::from_unsorted(
8144 self.retrieve_filtered(
8145 &crate::query::Retriever::Sparse {
8146 column_id: *column_id,
8147 query: query.clone(),
8148 k: *k,
8149 },
8150 snapshot,
8151 None,
8152 None,
8153 None,
8154 None,
8155 )?
8156 .into_iter()
8157 .map(|hit| hit.row_id.0)
8158 .collect(),
8159 ),
8160 Condition::MinHashSimilar {
8161 column_id,
8162 query,
8163 k,
8164 } => match self.minhash.get(column_id) {
8165 Some(index) => {
8166 let candidates = index.candidate_row_ids(query);
8167 let eligible =
8168 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
8169 RowIdSet::from_unsorted(
8170 index
8171 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
8172 .into_iter()
8173 .map(|(row_id, _)| row_id.0)
8174 .collect(),
8175 )
8176 }
8177 None => RowIdSet::empty(),
8178 },
8179 Condition::Range { column_id, lo, hi } => {
8180 if let Some(li) = self.learned_range.get(column_id) {
8181 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
8182 } else {
8183 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
8184 }
8185 }
8186 Condition::RangeF64 {
8187 column_id,
8188 lo,
8189 lo_inclusive,
8190 hi,
8191 hi_inclusive,
8192 } => {
8193 if let Some(li) = self.learned_range.get(column_id) {
8194 RowIdSet::from_unsorted(
8195 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
8196 .into_iter()
8197 .collect(),
8198 )
8199 } else {
8200 reader.range_row_id_set_f64(
8201 *column_id,
8202 *lo,
8203 *lo_inclusive,
8204 *hi,
8205 *hi_inclusive,
8206 )?
8207 }
8208 }
8209 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
8210 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
8211 };
8212 sets.push(s);
8213 }
8214 Ok(RowIdSet::intersect_many(sets))
8215 }
8216
8217 pub fn scan_cursor(
8238 &self,
8239 snapshot: Snapshot,
8240 projection: Vec<(u16, TypeId)>,
8241 conditions: &[crate::query::Condition],
8242 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
8243 if self.ttl.is_some() {
8244 return Ok(None);
8245 }
8246 if !conditions.is_empty() && !self.indexes_complete {
8252 return Ok(None);
8253 }
8254 if self.run_refs.len() == 1 {
8255 Ok(self
8256 .native_page_cursor(snapshot, projection, conditions)?
8257 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
8258 } else {
8259 Ok(self
8260 .native_multi_run_cursor(snapshot, projection, conditions)?
8261 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
8262 }
8263 }
8264
8265 pub fn aggregate_native(
8279 &self,
8280 snapshot: Snapshot,
8281 column: Option<u16>,
8282 conditions: &[crate::query::Condition],
8283 agg: NativeAgg,
8284 ) -> Result<Option<NativeAggResult>> {
8285 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
8286 }
8287
8288 pub fn aggregate_native_with_control(
8289 &self,
8290 snapshot: Snapshot,
8291 column: Option<u16>,
8292 conditions: &[crate::query::Condition],
8293 agg: NativeAgg,
8294 control: &crate::ExecutionControl,
8295 ) -> Result<Option<NativeAggResult>> {
8296 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
8297 }
8298
8299 fn aggregate_native_inner(
8300 &self,
8301 snapshot: Snapshot,
8302 column: Option<u16>,
8303 conditions: &[crate::query::Condition],
8304 agg: NativeAgg,
8305 control: Option<&crate::ExecutionControl>,
8306 ) -> Result<Option<NativeAggResult>> {
8307 execution_checkpoint(control, 0)?;
8308 if self.ttl.is_some() {
8309 return Ok(None);
8310 }
8311 if self.run_refs.len() == 1 && conditions.is_empty() {
8313 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
8314 return Ok(Some(res));
8315 }
8316 }
8317 if matches!(agg, NativeAgg::Count) && column.is_none() {
8319 return Ok(self
8320 .scan_cursor(snapshot, Vec::new(), conditions)?
8321 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
8322 }
8323 let cid = match column {
8326 Some(c) => c,
8327 None => return Ok(None),
8328 };
8329 let ty = self.column_type(cid);
8330 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
8331 else {
8332 return Ok(None);
8333 };
8334 execution_checkpoint(control, 0)?;
8335 match ty {
8336 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8337 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
8338 Ok(Some(pack_int(agg, count, sum, mn, mx)))
8339 }
8340 TypeId::Float64 => {
8341 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
8342 Ok(Some(pack_float(agg, count, sum, mn, mx)))
8343 }
8344 _ => Ok(None),
8345 }
8346 }
8347
8348 fn aggregate_from_stats(
8356 &self,
8357 snapshot: Snapshot,
8358 column: Option<u16>,
8359 agg: NativeAgg,
8360 ) -> Result<Option<NativeAggResult>> {
8361 let cid = match (agg, column) {
8362 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
8363 _ => return Ok(None), };
8365 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
8366 return Ok(None);
8367 };
8368 let Some(cs) = stats.get(&cid) else {
8369 return Ok(None);
8370 };
8371 match agg {
8372 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
8374 self.live_count.saturating_sub(cs.null_count),
8375 ))),
8376 NativeAgg::Min | NativeAgg::Max => {
8377 let bound = if agg == NativeAgg::Min {
8378 &cs.min
8379 } else {
8380 &cs.max
8381 };
8382 match bound {
8383 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
8384 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
8385 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
8390 None => Ok(None),
8391 }
8392 }
8393 _ => Ok(None),
8394 }
8395 }
8396
8397 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
8406 if self.ttl.is_some() {
8407 return Ok(None);
8408 }
8409 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
8410 return Ok(None);
8411 }
8412 self.ensure_indexes_complete()?;
8415 let reader = self.open_reader(self.run_refs[0].run_id)?;
8416 if self.live_count != reader.row_count() as u64 {
8417 return Ok(None);
8418 }
8419 let Some(bm) = self.bitmap.get(&column_id) else {
8420 return Ok(None); };
8422 let mut distinct = bm.value_count() as u64;
8423 if !bm.get(&Value::Null.encode_key()).is_empty() {
8426 distinct = distinct.saturating_sub(1);
8427 }
8428 Ok(Some(distinct))
8429 }
8430
8431 pub fn aggregate_incremental(
8443 &mut self,
8444 cache_key: u64,
8445 conditions: &[crate::query::Condition],
8446 column: Option<u16>,
8447 agg: NativeAgg,
8448 ) -> Result<IncrementalAggResult> {
8449 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
8450 }
8451
8452 pub fn aggregate_incremental_with_control(
8453 &mut self,
8454 cache_key: u64,
8455 conditions: &[crate::query::Condition],
8456 column: Option<u16>,
8457 agg: NativeAgg,
8458 control: &crate::ExecutionControl,
8459 ) -> Result<IncrementalAggResult> {
8460 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
8461 }
8462
8463 fn aggregate_incremental_inner(
8464 &mut self,
8465 cache_key: u64,
8466 conditions: &[crate::query::Condition],
8467 column: Option<u16>,
8468 agg: NativeAgg,
8469 control: Option<&crate::ExecutionControl>,
8470 ) -> Result<IncrementalAggResult> {
8471 execution_checkpoint(control, 0)?;
8472 let snap = self.snapshot();
8473 let cur_wm = self.allocator.current().0;
8474 let cur_epoch = snap.epoch.0;
8475 let incremental_ok = self.ttl.is_none()
8482 && !self.had_deletes
8483 && self.memtable.is_empty()
8484 && self.mutable_run.is_empty();
8485
8486 if incremental_ok {
8489 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
8490 if cached.epoch == cur_epoch {
8491 return Ok(IncrementalAggResult {
8492 state: cached.state,
8493 incremental: true,
8494 delta_rows: 0,
8495 });
8496 }
8497 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
8498 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
8499 let mut delta_rids = Vec::with_capacity(delta_len);
8500 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
8501 execution_checkpoint(control, index)?;
8502 delta_rids.push(row_id);
8503 }
8504 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
8505 execution_checkpoint(control, 0)?;
8506 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8507 let delta_state = agg_state_from_rows(
8508 &delta_rows,
8509 conditions,
8510 &index_sets,
8511 column,
8512 agg,
8513 &self.schema,
8514 control,
8515 )?;
8516 let merged = cached.state.merge(delta_state);
8517 let delta_n = delta_rids.len() as u64;
8518 Arc::make_mut(&mut self.agg_cache).insert(
8519 cache_key,
8520 CachedAgg {
8521 state: merged.clone(),
8522 watermark: cur_wm,
8523 epoch: cur_epoch,
8524 },
8525 );
8526 return Ok(IncrementalAggResult {
8527 state: merged,
8528 incremental: true,
8529 delta_rows: delta_n,
8530 });
8531 }
8532 }
8533 }
8534
8535 let cursor_ok =
8540 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8541 let state = if cursor_ok && agg != NativeAgg::Avg {
8542 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
8543 Some(result) => {
8544 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
8545 }
8546 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
8547 }
8548 } else {
8549 self.agg_state_full_scan(conditions, column, agg, snap, control)?
8550 };
8551 if incremental_ok {
8553 Arc::make_mut(&mut self.agg_cache).insert(
8554 cache_key,
8555 CachedAgg {
8556 state: state.clone(),
8557 watermark: cur_wm,
8558 epoch: cur_epoch,
8559 },
8560 );
8561 }
8562 Ok(IncrementalAggResult {
8563 state,
8564 incremental: false,
8565 delta_rows: 0,
8566 })
8567 }
8568
8569 fn agg_state_full_scan(
8572 &self,
8573 conditions: &[crate::query::Condition],
8574 column: Option<u16>,
8575 agg: NativeAgg,
8576 snap: Snapshot,
8577 control: Option<&crate::ExecutionControl>,
8578 ) -> Result<AggState> {
8579 execution_checkpoint(control, 0)?;
8580 let rows = self.visible_rows(snap)?;
8581 execution_checkpoint(control, 0)?;
8582 let index_sets = self.resolve_index_conditions(conditions, snap)?;
8583 agg_state_from_rows(
8584 &rows,
8585 conditions,
8586 &index_sets,
8587 column,
8588 agg,
8589 &self.schema,
8590 control,
8591 )
8592 }
8593
8594 fn resolve_index_conditions(
8597 &self,
8598 conditions: &[crate::query::Condition],
8599 snapshot: Snapshot,
8600 ) -> Result<Vec<RowIdSet>> {
8601 use crate::query::Condition;
8602 let mut sets = Vec::new();
8603 for c in conditions {
8604 if matches!(
8605 c,
8606 Condition::Ann { .. }
8607 | Condition::SparseMatch { .. }
8608 | Condition::MinHashSimilar { .. }
8609 ) {
8610 sets.push(self.resolve_condition(c, snapshot)?);
8611 }
8612 }
8613 Ok(sets)
8614 }
8615
8616 fn column_type(&self, cid: u16) -> TypeId {
8617 self.schema
8618 .columns
8619 .iter()
8620 .find(|c| c.id == cid)
8621 .map(|c| c.ty.clone())
8622 .unwrap_or(TypeId::Bytes)
8623 }
8624
8625 pub fn approx_aggregate(
8634 &mut self,
8635 conditions: &[crate::query::Condition],
8636 column: Option<u16>,
8637 agg: ApproxAgg,
8638 z: f64,
8639 ) -> Result<Option<ApproxResult>> {
8640 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
8641 }
8642
8643 pub fn approx_aggregate_with_candidate_authorization(
8646 &mut self,
8647 conditions: &[crate::query::Condition],
8648 column: Option<u16>,
8649 agg: ApproxAgg,
8650 z: f64,
8651 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8652 ) -> Result<Option<ApproxResult>> {
8653 use crate::query::Condition;
8654 self.ensure_reservoir_complete()?;
8655 let snapshot = self.snapshot();
8656 let n_pop = self.count();
8657 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
8658 if sample_rids.is_empty() {
8659 return Ok(None);
8660 }
8661 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
8663 let s = live_sample.len();
8664 if s == 0 {
8665 return Ok(None);
8666 }
8667 let authorized = authorization
8668 .map(|authorization| {
8669 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
8670 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
8671 })
8672 .transpose()?;
8673
8674 let mut index_sets: Vec<RowIdSet> = Vec::new();
8677 for c in conditions {
8678 if matches!(
8679 c,
8680 Condition::Ann { .. }
8681 | Condition::SparseMatch { .. }
8682 | Condition::MinHashSimilar { .. }
8683 ) {
8684 index_sets.push(self.resolve_condition(c, snapshot)?);
8685 }
8686 }
8687
8688 let cid = match (agg, column) {
8690 (ApproxAgg::Count, _) => None,
8691 (_, Some(c)) => Some(c),
8692 _ => return Ok(None),
8693 };
8694 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
8695 for r in &live_sample {
8696 if authorized
8697 .as_ref()
8698 .is_some_and(|authorized| !authorized.contains(&r.row_id))
8699 {
8700 continue;
8701 }
8702 if !conditions
8704 .iter()
8705 .all(|c| condition_matches_row(c, r, &self.schema))
8706 {
8707 continue;
8708 }
8709 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
8711 continue;
8712 }
8713 if let Some(cid) = cid {
8714 let mut cells = r
8715 .columns
8716 .get(&cid)
8717 .cloned()
8718 .map(|value| vec![(cid, value)])
8719 .unwrap_or_default();
8720 if let Some(authorization) = authorization {
8721 authorization.security.apply_masks_to_cells(
8722 authorization.table,
8723 &mut cells,
8724 authorization.principal,
8725 );
8726 }
8727 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
8728 passing_vals.push(v);
8729 } } else {
8731 passing_vals.push(0.0); }
8733 }
8734 let m = passing_vals.len();
8735
8736 let (point, half) = match agg {
8737 ApproxAgg::Count => {
8738 let p = m as f64 / s as f64;
8740 let point = n_pop as f64 * p;
8741 let var = if s > 1 {
8742 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
8743 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8744 } else {
8745 0.0
8746 };
8747 (point, z * var.sqrt())
8748 }
8749 ApproxAgg::Sum => {
8750 let y: Vec<f64> = live_sample
8752 .iter()
8753 .map(|r| {
8754 let passes_row = authorized
8755 .as_ref()
8756 .map_or(true, |authorized| authorized.contains(&r.row_id))
8757 && conditions
8758 .iter()
8759 .all(|c| condition_matches_row(c, r, &self.schema))
8760 && index_sets.iter().all(|set| set.contains(r.row_id.0));
8761 if passes_row {
8762 cid.and_then(|cid| {
8763 let mut cells = r
8764 .columns
8765 .get(&cid)
8766 .cloned()
8767 .map(|value| vec![(cid, value)])
8768 .unwrap_or_default();
8769 if let Some(authorization) = authorization {
8770 authorization.security.apply_masks_to_cells(
8771 authorization.table,
8772 &mut cells,
8773 authorization.principal,
8774 );
8775 }
8776 as_f64(cells.first().map(|(_, value)| value))
8777 })
8778 .unwrap_or(0.0)
8779 } else {
8780 0.0
8781 }
8782 })
8783 .collect();
8784 let mean_y = y.iter().sum::<f64>() / s as f64;
8785 let point = n_pop as f64 * mean_y;
8786 let var = if s > 1 {
8787 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
8788 let var_y = ss / (s - 1) as f64;
8789 n_pop as f64 * n_pop as f64 * var_y / s as f64
8790 * (1.0 - s as f64 / n_pop as f64).max(0.0)
8791 } else {
8792 0.0
8793 };
8794 (point, z * var.sqrt())
8795 }
8796 ApproxAgg::Avg => {
8797 if m == 0 {
8798 return Ok(Some(ApproxResult {
8799 point: 0.0,
8800 ci_low: 0.0,
8801 ci_high: 0.0,
8802 n_population: n_pop,
8803 n_sample_live: s,
8804 n_passing: 0,
8805 }));
8806 }
8807 let mean = passing_vals.iter().sum::<f64>() / m as f64;
8808 let half = if m > 1 {
8809 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
8810 let sd = (ss / (m - 1) as f64).sqrt();
8811 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
8812 z * sd / (m as f64).sqrt() * fpc.sqrt()
8813 } else {
8814 0.0
8815 };
8816 (mean, half)
8817 }
8818 };
8819
8820 Ok(Some(ApproxResult {
8821 point,
8822 ci_low: point - half,
8823 ci_high: point + half,
8824 n_population: n_pop,
8825 n_sample_live: s,
8826 n_passing: m,
8827 }))
8828 }
8829
8830 pub fn exact_column_stats(
8838 &self,
8839 _snapshot: Snapshot,
8840 projection: &[u16],
8841 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
8842 if self.ttl.is_some()
8843 || !(self.memtable.is_empty()
8844 && self.mutable_run.is_empty()
8845 && self.run_refs.len() == 1)
8846 {
8847 return Ok(None);
8848 }
8849 let reader = self.open_reader(self.run_refs[0].run_id)?;
8850 if self.live_count != reader.row_count() as u64 {
8851 return Ok(None);
8852 }
8853 let mut out = HashMap::new();
8854 for &cid in projection {
8855 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
8856 Some(c) => c,
8857 None => continue,
8858 };
8859 let Some(stats) = reader.column_page_stats(cid) else {
8861 out.insert(
8862 cid,
8863 ColumnStat {
8864 min: None,
8865 max: None,
8866 null_count: self.live_count,
8867 },
8868 );
8869 continue;
8870 };
8871 let stat = match cdef.ty {
8872 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
8873 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
8874 min: mn.map(Value::Int64),
8875 max: mx.map(Value::Int64),
8876 null_count: n,
8877 })
8878 }
8879 TypeId::Float64 => {
8880 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
8881 min: mn.map(Value::Float64),
8882 max: mx.map(Value::Float64),
8883 null_count: n,
8884 })
8885 }
8886 _ => None,
8887 };
8888 if let Some(s) = stat {
8889 out.insert(cid, s);
8890 }
8891 }
8892 Ok(Some(out))
8893 }
8894
8895 pub fn dir(&self) -> &Path {
8896 &self.dir
8897 }
8898
8899 pub fn schema(&self) -> &Schema {
8900 &self.schema
8901 }
8902
8903 pub(crate) fn set_catalog_name(&mut self, name: String) {
8904 self.name = name;
8905 }
8906
8907 pub(crate) fn prepare_alter_column(
8908 &mut self,
8909 column_name: &str,
8910 change: &AlterColumn,
8911 ) -> Result<ColumnDef> {
8912 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
8913 return Err(MongrelError::InvalidArgument(
8914 "ALTER COLUMN requires committing staged writes first".into(),
8915 ));
8916 }
8917 let old = self
8918 .schema
8919 .columns
8920 .iter()
8921 .find(|c| c.name == column_name)
8922 .cloned()
8923 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
8924 let mut next = old.clone();
8925
8926 if let Some(name) = &change.name {
8927 let trimmed = name.trim();
8928 if trimmed.is_empty() {
8929 return Err(MongrelError::InvalidArgument(
8930 "ALTER COLUMN name must not be empty".into(),
8931 ));
8932 }
8933 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
8934 return Err(MongrelError::Schema(format!(
8935 "column {trimmed} already exists"
8936 )));
8937 }
8938 next.name = trimmed.to_string();
8939 }
8940
8941 if let Some(ty) = &change.ty {
8942 next.ty = ty.clone();
8943 }
8944 if let Some(flags) = change.flags {
8945 validate_alter_column_flags(old.flags, flags)?;
8946 next.flags = flags;
8947 }
8948
8949 if let Some(default_change) = &change.default_value {
8950 next.default_value = default_change.clone();
8951 }
8952
8953 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
8954 if old.flags.contains(ColumnFlags::NULLABLE)
8955 && !next.flags.contains(ColumnFlags::NULLABLE)
8956 && self.column_has_nulls(old.id)?
8957 {
8958 return Err(MongrelError::InvalidArgument(format!(
8959 "column '{}' contains NULL values",
8960 old.name
8961 )));
8962 }
8963 Ok(next)
8964 }
8965
8966 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
8967 let idx = self
8968 .schema
8969 .columns
8970 .iter()
8971 .position(|c| c.id == column.id)
8972 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
8973 if self.schema.columns[idx] == column {
8974 return Ok(());
8975 }
8976 self.schema.columns[idx] = column;
8977 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
8978 self.schema.validate_auto_increment()?;
8979 self.schema.validate_defaults()?;
8980 self.auto_inc = resolve_auto_inc(&self.schema);
8981 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
8982 write_schema(&self.dir, &self.schema)?;
8983 self.clear_result_cache();
8984 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8985 self.persist_manifest(self.current_epoch())?;
8986 Ok(())
8987 }
8988
8989 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
8990 self.ensure_writable()?;
8991 let column = self.prepare_alter_column(column_name, &change)?;
8992 self.apply_altered_column(column.clone())?;
8993 Ok(column)
8994 }
8995
8996 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
8997 if self.live_count == 0 {
8998 return Ok(false);
8999 }
9000 let snap = self.snapshot();
9001 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
9002 Ok(columns
9003 .first()
9004 .map(|(_, col)| col.null_count(col.len()) != 0)
9005 .unwrap_or(true))
9006 }
9007
9008 fn has_stored_versions(&self) -> bool {
9009 !self.memtable.is_empty()
9010 || !self.mutable_run.is_empty()
9011 || self.run_refs.iter().any(|r| r.row_count > 0)
9012 || !self.retiring.is_empty()
9013 }
9014
9015 pub fn add_column(
9020 &mut self,
9021 name: &str,
9022 ty: TypeId,
9023 flags: ColumnFlags,
9024 default_value: Option<crate::schema::DefaultExpr>,
9025 ) -> Result<u16> {
9026 self.add_column_with_id(name, ty, flags, default_value, None)
9027 }
9028
9029 pub fn add_column_with_id(
9030 &mut self,
9031 name: &str,
9032 ty: TypeId,
9033 flags: ColumnFlags,
9034 default_value: Option<crate::schema::DefaultExpr>,
9035 requested_id: Option<u16>,
9036 ) -> Result<u16> {
9037 self.ensure_writable()?;
9038 if self.schema.columns.iter().any(|c| c.name == name) {
9039 return Err(MongrelError::Schema(format!(
9040 "column {name} already exists"
9041 )));
9042 }
9043 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
9044 if self.schema.columns.iter().any(|c| c.id == id) {
9045 return Err(MongrelError::Schema(format!(
9046 "column id {id} already exists"
9047 )));
9048 }
9049 id
9050 } else {
9051 self.schema
9052 .columns
9053 .iter()
9054 .map(|c| c.id)
9055 .max()
9056 .unwrap_or(0)
9057 .checked_add(1)
9058 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
9059 };
9060 self.schema.columns.push(ColumnDef {
9061 id,
9062 name: name.to_string(),
9063 ty,
9064 flags,
9065 default_value,
9066 });
9067 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
9068 self.schema.validate_auto_increment()?;
9069 self.schema.validate_defaults()?;
9070 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
9071 self.auto_inc = resolve_auto_inc(&self.schema);
9072 }
9073 write_schema(&self.dir, &self.schema)?;
9074 self.clear_result_cache();
9075 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
9077 self.persist_manifest(self.current_epoch())?;
9078 Ok(id)
9079 }
9080
9081 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
9090 self.ensure_writable()?;
9091 let cid = self
9092 .schema
9093 .columns
9094 .iter()
9095 .find(|c| c.name == column_name)
9096 .map(|c| c.id)
9097 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
9098 let ty = self
9099 .schema
9100 .columns
9101 .iter()
9102 .find(|c| c.id == cid)
9103 .map(|c| c.ty.clone())
9104 .unwrap_or(TypeId::Int64);
9105 if !matches!(
9106 ty,
9107 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
9108 ) {
9109 return Err(MongrelError::Schema(format!(
9110 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
9111 )));
9112 }
9113 if self
9114 .schema
9115 .indexes
9116 .iter()
9117 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
9118 {
9119 return Ok(()); }
9121 self.schema.indexes.push(IndexDef {
9122 name: format!("{}_learned_range", column_name),
9123 column_id: cid,
9124 kind: IndexKind::LearnedRange,
9125 predicate: None,
9126 options: Default::default(),
9127 });
9128 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
9129 write_schema(&self.dir, &self.schema)?;
9130 self.build_learned_ranges()?;
9131 Ok(())
9132 }
9133
9134 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
9137 self.sync_byte_threshold = threshold;
9138 if let WalSink::Private(w) = &mut self.wal {
9139 w.set_sync_byte_threshold(threshold);
9140 }
9141 }
9142
9143 pub fn page_cache_flush(&self) {
9147 self.page_cache.flush_to_disk();
9148 }
9149
9150 pub fn page_cache_len(&self) -> usize {
9152 self.page_cache.len()
9153 }
9154
9155 pub fn decoded_cache_len(&self) -> usize {
9158 self.decoded_cache.len()
9159 }
9160
9161 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
9164 self.memtable.drain_sorted()
9165 }
9166
9167 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
9168 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
9169 }
9170
9171 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
9172 self.run_refs.iter().map(|run| run.run_id)
9173 }
9174
9175 pub(crate) fn table_dir(&self) -> &Path {
9176 &self.dir
9177 }
9178
9179 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
9180 &self.schema
9181 }
9182
9183 pub(crate) fn alloc_run_id(&mut self) -> u64 {
9184 let id = self.next_run_id;
9185 self.next_run_id += 1;
9186 id
9187 }
9188
9189 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
9190 self.run_refs.push(run_ref);
9191 }
9192
9193 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
9203 self.retiring.push(crate::manifest::RetiredRun {
9204 run_id,
9205 retire_epoch,
9206 });
9207 }
9208
9209 pub(crate) fn reap_retiring(
9213 &mut self,
9214 min_active: Epoch,
9215 backup_pinned: &std::collections::HashSet<u128>,
9216 ) -> Result<usize> {
9217 if self.retiring.is_empty() {
9218 return Ok(0);
9219 }
9220 let mut reaped = 0;
9221 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
9222 for r in std::mem::take(&mut self.retiring) {
9228 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
9229 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
9230 reaped += 1;
9231 } else {
9232 kept.push(r);
9233 }
9234 }
9235 self.retiring = kept;
9236 if reaped > 0 {
9237 self.persist_manifest(self.current_epoch())?;
9238 }
9239 Ok(reaped)
9240 }
9241
9242 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
9243 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
9244 return false;
9245 }
9246 self.live_count = self.live_count.saturating_add(run_ref.row_count);
9247 self.run_refs.push(run_ref);
9248 self.indexes_complete = false;
9249 true
9250 }
9251
9252 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
9253 self.kek.as_ref()
9254 }
9255
9256 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
9257 let mut reader = RunReader::open_with_cache(
9258 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
9259 self.schema.clone(),
9260 self.kek.clone(),
9261 Some(self.page_cache.clone()),
9262 Some(self.decoded_cache.clone()),
9263 self.table_id,
9264 Some(&self.verified_runs),
9265 )?;
9266 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
9270 reader.set_uniform_epoch(Epoch(rr.epoch_created));
9271 }
9272 Ok(reader)
9273 }
9274
9275 pub(crate) fn run_refs(&self) -> &[RunRef] {
9276 &self.run_refs
9277 }
9278
9279 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
9280 self.retiring.iter().map(|run| run.run_id)
9281 }
9282
9283 pub(crate) fn runs_dir(&self) -> PathBuf {
9284 self.dir.join(RUNS_DIR)
9285 }
9286
9287 pub(crate) fn wal_dir(&self) -> PathBuf {
9288 self.dir.join(WAL_DIR)
9289 }
9290
9291 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
9292 self.run_refs = refs;
9293 }
9294
9295 pub(crate) fn next_run_id(&self) -> u64 {
9296 self.next_run_id
9297 }
9298
9299 pub(crate) fn compaction_zstd_level(&self) -> i32 {
9300 self.compaction_zstd_level
9301 }
9302
9303 pub(crate) fn bump_next_run_id(&mut self) {
9304 self.next_run_id += 1;
9305 }
9306
9307 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
9308 self.kek.clone()
9309 }
9310
9311 #[cfg(feature = "encryption")]
9315 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
9316 self.kek.as_ref().map(|k| k.derive_idx_key())
9317 }
9318
9319 #[cfg(not(feature = "encryption"))]
9320 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
9321 None
9322 }
9323
9324 #[cfg(feature = "encryption")]
9328 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
9329 self.kek.as_ref().map(|k| *k.derive_meta_key())
9330 }
9331
9332 #[cfg(not(feature = "encryption"))]
9333 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
9334 None
9335 }
9336
9337 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
9340 self.column_keys
9341 .iter()
9342 .map(|(&id, &(_, scheme))| (id, scheme))
9343 .collect()
9344 }
9345
9346 #[cfg(feature = "encryption")]
9351 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
9352 self.tokenize_value_enc(column_id, v)
9353 }
9354
9355 #[cfg(feature = "encryption")]
9356 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
9357 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
9358 let (key, scheme) = self.column_keys.get(&column_id)?;
9359 let token: Vec<u8> = match (*scheme, v) {
9360 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
9361 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
9362 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
9363 _ => hmac_token(key, &v.encode_key()).to_vec(),
9364 };
9365 Some(Value::Bytes(token))
9366 }
9367
9368 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
9370 self.index_lookup_key_bytes(column_id, &v.encode_key())
9371 }
9372
9373 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
9376 #[cfg(feature = "encryption")]
9377 {
9378 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
9379 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
9380 if *scheme == SCHEME_HMAC_EQ {
9381 return hmac_token(key, encoded).to_vec();
9382 }
9383 }
9384 }
9385 let _ = column_id;
9386 encoded.to_vec()
9387 }
9388}
9389
9390fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
9391 let columnar::NativeColumn::Int64 { data, validity } = col else {
9392 return false;
9393 };
9394 if data.len() < n || !columnar::all_non_null(validity, n) {
9395 return false;
9396 }
9397 data.iter()
9398 .take(n)
9399 .zip(data.iter().skip(1))
9400 .all(|(a, b)| a < b)
9401}
9402
9403#[derive(Debug, Clone)]
9407pub struct ColumnStat {
9408 pub min: Option<Value>,
9409 pub max: Option<Value>,
9410 pub null_count: u64,
9411}
9412
9413#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9415pub enum NativeAgg {
9416 Count,
9417 Sum,
9418 Min,
9419 Max,
9420 Avg,
9421}
9422
9423#[derive(Debug, Clone, PartialEq)]
9425pub enum NativeAggResult {
9426 Count(u64),
9427 Int(i64),
9428 Float(f64),
9429 Null,
9431}
9432
9433#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9435pub enum ApproxAgg {
9436 Count,
9437 Sum,
9438 Avg,
9439}
9440
9441#[derive(Debug, Clone)]
9445pub struct ApproxResult {
9446 pub point: f64,
9448 pub ci_low: f64,
9450 pub ci_high: f64,
9452 pub n_population: u64,
9454 pub n_sample_live: usize,
9456 pub n_passing: usize,
9458}
9459
9460#[derive(Debug, Clone, PartialEq)]
9465pub enum AggState {
9466 Count(u64),
9468 SumI {
9470 sum: i128,
9471 count: u64,
9472 },
9473 SumF {
9475 sum: f64,
9476 count: u64,
9477 },
9478 AvgI {
9480 sum: i128,
9481 count: u64,
9482 },
9483 AvgF {
9485 sum: f64,
9486 count: u64,
9487 },
9488 MinI(i64),
9490 MaxI(i64),
9491 MinF(f64),
9493 MaxF(f64),
9494 Empty,
9496}
9497
9498impl AggState {
9499 pub fn merge(self, other: AggState) -> AggState {
9501 use AggState::*;
9502 match (self, other) {
9503 (Empty, x) | (x, Empty) => x,
9504 (Count(a), Count(b)) => Count(a + b),
9505 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
9506 sum: sa + sb,
9507 count: ca + cb,
9508 },
9509 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
9510 sum: sa + sb,
9511 count: ca + cb,
9512 },
9513 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
9514 sum: sa + sb,
9515 count: ca + cb,
9516 },
9517 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
9518 sum: sa + sb,
9519 count: ca + cb,
9520 },
9521 (MinI(a), MinI(b)) => MinI(a.min(b)),
9522 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
9523 (MinF(a), MinF(b)) => MinF(a.min(b)),
9524 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
9525 _ => Empty, }
9527 }
9528
9529 pub fn point(&self) -> Option<f64> {
9531 match self {
9532 AggState::Count(n) => Some(*n as f64),
9533 AggState::SumI { sum, .. } => Some(*sum as f64),
9534 AggState::SumF { sum, .. } => Some(*sum),
9535 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
9536 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
9537 AggState::MinI(n) => Some(*n as f64),
9538 AggState::MaxI(n) => Some(*n as f64),
9539 AggState::MinF(n) => Some(*n),
9540 AggState::MaxF(n) => Some(*n),
9541 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
9542 }
9543 }
9544
9545 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
9549 let is_float = matches!(ty, Some(TypeId::Float64));
9550 match (agg, result) {
9551 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
9552 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
9553 sum: x as i128,
9554 count: 1, },
9556 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
9557 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
9558 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
9559 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
9560 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
9561 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
9562 (NativeAgg::Count, _) => AggState::Empty,
9563 (_, NativeAggResult::Null) => AggState::Empty,
9564 _ => {
9565 let _ = is_float;
9566 AggState::Empty
9567 }
9568 }
9569 }
9570}
9571
9572#[derive(Debug, Clone)]
9575pub struct CachedAgg {
9576 pub state: AggState,
9577 pub watermark: u64,
9578 pub epoch: u64,
9579}
9580
9581#[derive(Debug, Clone)]
9583pub struct IncrementalAggResult {
9584 pub state: AggState,
9586 pub incremental: bool,
9589 pub delta_rows: u64,
9591}
9592
9593fn agg_state_from_rows(
9597 rows: &[Row],
9598 conditions: &[crate::query::Condition],
9599 index_sets: &[RowIdSet],
9600 column: Option<u16>,
9601 agg: NativeAgg,
9602 schema: &Schema,
9603 control: Option<&crate::ExecutionControl>,
9604) -> Result<AggState> {
9605 let mut count: u64 = 0;
9606 let mut sum_i: i128 = 0;
9607 let mut sum_f: f64 = 0.0;
9608 let mut mn_i: i64 = i64::MAX;
9609 let mut mx_i: i64 = i64::MIN;
9610 let mut mn_f: f64 = f64::INFINITY;
9611 let mut mx_f: f64 = f64::NEG_INFINITY;
9612 let mut saw_int = false;
9613 let mut saw_float = false;
9614 for (index, r) in rows.iter().enumerate() {
9615 execution_checkpoint(control, index)?;
9616 if !conditions
9617 .iter()
9618 .all(|c| condition_matches_row(c, r, schema))
9619 {
9620 continue;
9621 }
9622 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
9623 continue;
9624 }
9625 match agg {
9626 NativeAgg::Count => match column {
9627 None => count += 1,
9629 Some(cid) => match r.columns.get(&cid) {
9632 None | Some(Value::Null) => {}
9633 Some(_) => count += 1,
9634 },
9635 },
9636 _ => match column.and_then(|cid| r.columns.get(&cid)) {
9637 Some(Value::Int64(n)) => {
9638 count += 1;
9639 sum_i += *n as i128;
9640 mn_i = mn_i.min(*n);
9641 mx_i = mx_i.max(*n);
9642 saw_int = true;
9643 }
9644 Some(Value::Float64(f)) => {
9645 count += 1;
9646 sum_f += f;
9647 mn_f = mn_f.min(*f);
9648 mx_f = mx_f.max(*f);
9649 saw_float = true;
9650 }
9651 _ => {}
9652 },
9653 }
9654 }
9655 Ok(match agg {
9656 NativeAgg::Count => {
9657 if count == 0 {
9658 AggState::Empty
9659 } else {
9660 AggState::Count(count)
9661 }
9662 }
9663 NativeAgg::Sum => {
9664 if count == 0 {
9665 AggState::Empty
9666 } else if saw_int {
9667 AggState::SumI { sum: sum_i, count }
9668 } else {
9669 AggState::SumF { sum: sum_f, count }
9670 }
9671 }
9672 NativeAgg::Avg => {
9673 if count == 0 {
9674 AggState::Empty
9675 } else if saw_int {
9676 AggState::AvgI { sum: sum_i, count }
9677 } else {
9678 AggState::AvgF { sum: sum_f, count }
9679 }
9680 }
9681 NativeAgg::Min => {
9682 if !saw_int && !saw_float {
9683 AggState::Empty
9684 } else if saw_int {
9685 AggState::MinI(mn_i)
9686 } else {
9687 AggState::MinF(mn_f)
9688 }
9689 }
9690 NativeAgg::Max => {
9691 if !saw_int && !saw_float {
9692 AggState::Empty
9693 } else if saw_int {
9694 AggState::MaxI(mx_i)
9695 } else {
9696 AggState::MaxF(mx_f)
9697 }
9698 }
9699 })
9700}
9701
9702fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
9706 use crate::query::Condition;
9707 match c {
9708 Condition::Pk(key) => match schema.primary_key() {
9709 Some(pk) => row
9710 .columns
9711 .get(&pk.id)
9712 .map(|v| v.encode_key() == *key)
9713 .unwrap_or(false),
9714 None => false,
9715 },
9716 Condition::BitmapEq { column_id, value } => row
9717 .columns
9718 .get(column_id)
9719 .map(|v| v.encode_key() == *value)
9720 .unwrap_or(false),
9721 Condition::BitmapIn { column_id, values } => {
9722 let key = row.columns.get(column_id).map(|v| v.encode_key());
9723 match key {
9724 Some(k) => values.contains(&k),
9725 None => false,
9726 }
9727 }
9728 Condition::BytesPrefix { column_id, prefix } => row
9729 .columns
9730 .get(column_id)
9731 .map(|v| v.encode_key().starts_with(prefix))
9732 .unwrap_or(false),
9733 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
9734 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
9735 _ => false,
9736 },
9737 Condition::RangeF64 {
9738 column_id,
9739 lo,
9740 lo_inclusive,
9741 hi,
9742 hi_inclusive,
9743 } => match row.columns.get(column_id) {
9744 Some(Value::Float64(n)) => {
9745 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
9746 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
9747 lo_ok && hi_ok
9748 }
9749 _ => false,
9750 },
9751 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
9752 Some(Value::Bytes(b)) => {
9753 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
9754 }
9755 _ => false,
9756 },
9757 Condition::FmContainsAll {
9758 column_id,
9759 patterns,
9760 } => match row.columns.get(column_id) {
9761 Some(Value::Bytes(b)) => patterns
9762 .iter()
9763 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
9764 _ => false,
9765 },
9766 Condition::Ann { .. }
9767 | Condition::SparseMatch { .. }
9768 | Condition::MinHashSimilar { .. } => true,
9769 Condition::IsNull { column_id } => {
9770 matches!(row.columns.get(column_id), Some(Value::Null) | None)
9771 }
9772 Condition::IsNotNull { column_id } => {
9773 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
9774 }
9775 }
9776}
9777
9778fn as_f64(v: Option<&Value>) -> Option<f64> {
9780 match v {
9781 Some(Value::Int64(n)) => Some(*n as f64),
9782 Some(Value::Float64(f)) => Some(*f),
9783 _ => None,
9784 }
9785}
9786
9787fn accumulate_int(
9791 cursor: &mut dyn crate::cursor::Cursor,
9792 control: Option<&crate::ExecutionControl>,
9793) -> Result<(u64, i128, i64, i64)> {
9794 let mut count: u64 = 0;
9795 let mut sum: i128 = 0;
9796 let mut mn: i64 = i64::MAX;
9797 let mut mx: i64 = i64::MIN;
9798 while let Some(cols) = cursor.next_batch()? {
9799 execution_checkpoint(control, 0)?;
9800 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
9801 if crate::columnar::all_non_null(validity, data.len()) {
9802 count += data.len() as u64;
9804 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
9805 execution_checkpoint(control, chunk_index * 1024)?;
9806 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
9807 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
9808 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
9809 }
9810 } else {
9811 for (i, &v) in data.iter().enumerate() {
9812 execution_checkpoint(control, i)?;
9813 if crate::columnar::validity_bit(validity, i) {
9814 count += 1;
9815 sum += v as i128;
9816 mn = mn.min(v);
9817 mx = mx.max(v);
9818 }
9819 }
9820 }
9821 }
9822 }
9823 Ok((count, sum, mn, mx))
9824}
9825
9826fn accumulate_float(
9828 cursor: &mut dyn crate::cursor::Cursor,
9829 control: Option<&crate::ExecutionControl>,
9830) -> Result<(u64, f64, f64, f64)> {
9831 let mut count: u64 = 0;
9832 let mut sum: f64 = 0.0;
9833 let mut mn: f64 = f64::INFINITY;
9834 let mut mx: f64 = f64::NEG_INFINITY;
9835 while let Some(cols) = cursor.next_batch()? {
9836 execution_checkpoint(control, 0)?;
9837 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
9838 if crate::columnar::all_non_null(validity, data.len()) {
9839 count += data.len() as u64;
9840 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
9841 execution_checkpoint(control, chunk_index * 1024)?;
9842 sum += chunk.iter().sum::<f64>();
9843 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
9844 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
9845 }
9846 } else {
9847 for (i, &v) in data.iter().enumerate() {
9848 execution_checkpoint(control, i)?;
9849 if crate::columnar::validity_bit(validity, i) {
9850 count += 1;
9851 sum += v;
9852 mn = mn.min(v);
9853 mx = mx.max(v);
9854 }
9855 }
9856 }
9857 }
9858 }
9859 Ok((count, sum, mn, mx))
9860}
9861
9862#[inline]
9863fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
9864 if index % 256 == 0 {
9865 control
9866 .map(crate::ExecutionControl::checkpoint)
9867 .transpose()?;
9868 }
9869 Ok(())
9870}
9871
9872fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
9873 if count == 0 && !matches!(agg, NativeAgg::Count) {
9874 return NativeAggResult::Null;
9875 }
9876 match agg {
9877 NativeAgg::Count => NativeAggResult::Count(count),
9878 NativeAgg::Sum => match sum.try_into() {
9881 Ok(v) => NativeAggResult::Int(v),
9882 Err(_) => NativeAggResult::Null,
9883 },
9884 NativeAgg::Min => NativeAggResult::Int(mn),
9885 NativeAgg::Max => NativeAggResult::Int(mx),
9886 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
9887 }
9888}
9889
9890fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
9891 if count == 0 && !matches!(agg, NativeAgg::Count) {
9892 return NativeAggResult::Null;
9893 }
9894 match agg {
9895 NativeAgg::Count => NativeAggResult::Count(count),
9896 NativeAgg::Sum => NativeAggResult::Float(sum),
9897 NativeAgg::Min => NativeAggResult::Float(mn),
9898 NativeAgg::Max => NativeAggResult::Float(mx),
9899 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
9900 }
9901}
9902
9903fn agg_int(
9906 stats: &[crate::page::PageStat],
9907 decode: fn(Option<&[u8]>) -> Option<i64>,
9908) -> Option<(Option<i64>, Option<i64>, u64)> {
9909 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
9910 let mut any = false;
9911 for s in stats {
9912 if let Some(v) = decode(s.min.as_deref()) {
9913 mn = mn.min(v);
9914 any = true;
9915 }
9916 if let Some(v) = decode(s.max.as_deref()) {
9917 mx = mx.max(v);
9918 any = true;
9919 }
9920 nulls += s.null_count;
9921 }
9922 any.then_some((Some(mn), Some(mx), nulls))
9923}
9924
9925fn agg_float(
9927 stats: &[crate::page::PageStat],
9928 decode: fn(Option<&[u8]>) -> Option<f64>,
9929) -> Option<(Option<f64>, Option<f64>, u64)> {
9930 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
9931 let mut any = false;
9932 for s in stats {
9933 if let Some(v) = decode(s.min.as_deref()) {
9934 mn = mn.min(v);
9935 any = true;
9936 }
9937 if let Some(v) = decode(s.max.as_deref()) {
9938 mx = mx.max(v);
9939 any = true;
9940 }
9941 nulls += s.null_count;
9942 }
9943 any.then_some((Some(mn), Some(mx), nulls))
9944}
9945
9946type SecondaryIndexes = (
9948 HashMap<u16, BitmapIndex>,
9949 HashMap<u16, AnnIndex>,
9950 HashMap<u16, FmIndex>,
9951 HashMap<u16, SparseIndex>,
9952 HashMap<u16, MinHashIndex>,
9953);
9954
9955fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
9956 let mut bitmap = HashMap::new();
9957 let mut ann = HashMap::new();
9958 let mut fm = HashMap::new();
9959 let mut sparse = HashMap::new();
9960 let mut minhash = HashMap::new();
9961 for idef in &schema.indexes {
9962 match idef.kind {
9963 IndexKind::Bitmap => {
9964 bitmap.insert(idef.column_id, BitmapIndex::new());
9965 }
9966 IndexKind::Ann => {
9967 let dim = schema
9968 .columns
9969 .iter()
9970 .find(|c| c.id == idef.column_id)
9971 .and_then(|c| match c.ty {
9972 TypeId::Embedding { dim } => Some(dim as usize),
9973 _ => None,
9974 })
9975 .unwrap_or(0);
9976 let options = idef.options.ann.clone().unwrap_or_default();
9977 ann.insert(
9978 idef.column_id,
9979 AnnIndex::with_options(
9980 dim,
9981 options.m,
9982 options.ef_construction,
9983 options.ef_search,
9984 ),
9985 );
9986 }
9987 IndexKind::FmIndex => {
9988 fm.insert(idef.column_id, FmIndex::new());
9989 }
9990 IndexKind::Sparse => {
9991 sparse.insert(idef.column_id, SparseIndex::new());
9992 }
9993 IndexKind::MinHash => {
9994 let options = idef.options.minhash.clone().unwrap_or_default();
9995 minhash.insert(
9996 idef.column_id,
9997 MinHashIndex::with_options(options.permutations, options.bands),
9998 );
9999 }
10000 _ => {}
10001 }
10002 }
10003 (bitmap, ann, fm, sparse, minhash)
10004}
10005
10006const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
10007 | ColumnFlags::AUTO_INCREMENT
10008 | ColumnFlags::ENCRYPTED
10009 | ColumnFlags::ENCRYPTED_INDEXABLE
10010 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
10011
10012fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
10013 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
10014 return Err(MongrelError::Schema(
10015 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
10016 ));
10017 }
10018 Ok(())
10019}
10020
10021fn validate_alter_column_type(
10022 schema: &Schema,
10023 old: &ColumnDef,
10024 next: &ColumnDef,
10025 has_stored_versions: bool,
10026) -> Result<()> {
10027 if old.ty == next.ty {
10028 return Ok(());
10029 }
10030 if schema.indexes.iter().any(|i| i.column_id == old.id) {
10031 return Err(MongrelError::Schema(format!(
10032 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
10033 old.name
10034 )));
10035 }
10036 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
10037 return Ok(());
10038 }
10039 Err(MongrelError::Schema(format!(
10040 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
10041 old.ty, next.ty
10042 )))
10043}
10044
10045fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
10046 matches!(
10047 (old, new),
10048 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
10049 )
10050}
10051
10052fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
10058 let mut prev: Option<i64> = None;
10059 for r in rows {
10060 match r.columns.get(&pk_id) {
10061 Some(Value::Int64(v)) => {
10062 if prev.is_some_and(|p| p >= *v) {
10063 return false;
10064 }
10065 prev = Some(*v);
10066 }
10067 _ => return false,
10068 }
10069 }
10070 true
10071}
10072
10073#[allow(clippy::too_many_arguments)]
10074fn index_into(
10075 schema: &Schema,
10076 row: &Row,
10077 hot: &mut HotIndex,
10078 bitmap: &mut HashMap<u16, BitmapIndex>,
10079 ann: &mut HashMap<u16, AnnIndex>,
10080 fm: &mut HashMap<u16, FmIndex>,
10081 sparse: &mut HashMap<u16, SparseIndex>,
10082 minhash: &mut HashMap<u16, MinHashIndex>,
10083) {
10084 for idef in &schema.indexes {
10085 let Some(val) = row.columns.get(&idef.column_id) else {
10086 continue;
10087 };
10088 match idef.kind {
10089 IndexKind::Bitmap => {
10090 if let Some(b) = bitmap.get_mut(&idef.column_id) {
10091 b.insert(val.encode_key(), row.row_id);
10092 }
10093 }
10094 IndexKind::Ann => {
10095 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
10096 a.insert_validated(v, row.row_id);
10097 }
10098 }
10099 IndexKind::FmIndex => {
10100 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
10101 f.insert(b.clone(), row.row_id);
10102 }
10103 }
10104 IndexKind::Sparse => {
10105 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
10106 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
10109 s.insert(&terms, row.row_id);
10110 }
10111 }
10112 }
10113 IndexKind::MinHash => {
10114 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
10115 let tokens = crate::index::token_hashes_from_bytes(b);
10118 mh.insert(&tokens, row.row_id);
10119 }
10120 }
10121 _ => {}
10122 }
10123 }
10124 if let Some(pk_col) = schema.primary_key() {
10125 if let Some(pk_val) = row.columns.get(&pk_col.id) {
10126 hot.insert(pk_val.encode_key(), row.row_id);
10127 }
10128 }
10129}
10130
10131#[allow(clippy::too_many_arguments)]
10134fn index_into_single(
10135 idef: &IndexDef,
10136 _schema: &Schema,
10137 row: &Row,
10138 _hot: &mut HotIndex,
10139 bitmap: &mut HashMap<u16, BitmapIndex>,
10140 ann: &mut HashMap<u16, AnnIndex>,
10141 fm: &mut HashMap<u16, FmIndex>,
10142 sparse: &mut HashMap<u16, SparseIndex>,
10143 minhash: &mut HashMap<u16, MinHashIndex>,
10144) {
10145 let Some(val) = row.columns.get(&idef.column_id) else {
10146 return;
10147 };
10148 match idef.kind {
10149 IndexKind::Bitmap => {
10150 if let Some(b) = bitmap.get_mut(&idef.column_id) {
10151 b.insert(val.encode_key(), row.row_id);
10152 }
10153 }
10154 IndexKind::Ann => {
10155 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
10156 a.insert_validated(v, row.row_id);
10157 }
10158 }
10159 IndexKind::FmIndex => {
10160 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
10161 f.insert(b.clone(), row.row_id);
10162 }
10163 }
10164 IndexKind::Sparse => {
10165 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
10166 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
10167 s.insert(&terms, row.row_id);
10168 }
10169 }
10170 }
10171 IndexKind::MinHash => {
10172 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
10173 let tokens = crate::index::token_hashes_from_bytes(b);
10174 mh.insert(&tokens, row.row_id);
10175 }
10176 }
10177 _ => {}
10178 }
10179}
10180
10181fn eval_partial_predicate(
10187 pred: &str,
10188 columns_map: &HashMap<u16, &Value>,
10189 name_to_id: &HashMap<&str, u16>,
10190) -> bool {
10191 let lower = pred.trim().to_ascii_lowercase();
10192 if let Some(rest) = lower.strip_suffix(" is not null") {
10194 let col_name = rest.trim();
10195 if let Some(col_id) = name_to_id.get(col_name) {
10196 return columns_map
10197 .get(col_id)
10198 .is_some_and(|v| !matches!(v, Value::Null));
10199 }
10200 }
10201 if let Some(rest) = lower.strip_suffix(" is null") {
10203 let col_name = rest.trim();
10204 if let Some(col_id) = name_to_id.get(col_name) {
10205 return columns_map
10206 .get(col_id)
10207 .map_or(true, |v| matches!(v, Value::Null));
10208 }
10209 }
10210 true
10213}
10214
10215#[allow(dead_code)]
10221fn bulk_index_key(
10222 column_keys: &HashMap<u16, ([u8; 32], u8)>,
10223 column_id: u16,
10224 ty: TypeId,
10225 col: &columnar::NativeColumn,
10226 i: usize,
10227) -> Option<Vec<u8>> {
10228 let encoded = columnar::encode_key_native(ty, col, i)?;
10229 #[cfg(feature = "encryption")]
10230 {
10231 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10232 if let Some((key, scheme)) = column_keys.get(&column_id) {
10233 return Some(match (*scheme, col) {
10234 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
10235 (_, columnar::NativeColumn::Int64 { data, .. }) => {
10236 ope_token_i64(key, data[i]).to_vec()
10237 }
10238 (_, columnar::NativeColumn::Float64 { data, .. }) => {
10239 ope_token_f64(key, data[i]).to_vec()
10240 }
10241 _ => hmac_token(key, &encoded).to_vec(),
10242 });
10243 }
10244 }
10245 #[cfg(not(feature = "encryption"))]
10246 {
10247 let _ = (column_id, column_keys, col);
10248 }
10249 Some(encoded)
10250}
10251
10252pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
10253 let json = serde_json::to_string_pretty(schema)
10254 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
10255 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
10256 Ok(())
10257}
10258
10259fn read_schema(dir: &Path) -> Result<Schema> {
10260 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
10261 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
10262}
10263
10264fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
10265 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
10266}
10267
10268fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
10269 let n = list_wal_numbers(wal_dir)?;
10270 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
10271}
10272
10273fn next_wal_number(wal_dir: &Path) -> Result<u32> {
10274 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
10275}
10276
10277fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
10278 let _ = std::fs::create_dir_all(wal_dir);
10279 let mut max_n = None;
10280 for entry in std::fs::read_dir(wal_dir)? {
10281 let entry = entry?;
10282 let fname = entry.file_name();
10283 let Some(s) = fname.to_str() else {
10284 continue;
10285 };
10286 let Some(stripped) = s.strip_prefix("seg-") else {
10287 continue;
10288 };
10289 let Some(stripped) = stripped.strip_suffix(".wal") else {
10290 continue;
10291 };
10292 if let Ok(n) = stripped.parse::<u32>() {
10293 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
10294 }
10295 }
10296 Ok(max_n)
10297}