1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, EpochGuard, MaintenanceReceipt, 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, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::cmp::Reverse;
30use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
31use std::path::{Path, PathBuf};
32use std::sync::atomic::AtomicBool;
33use std::sync::Arc;
34use zeroize::Zeroizing;
35
36pub const WAL_DIR: &str = "_wal";
37pub const RUNS_DIR: &str = "_runs";
38pub const CACHE_DIR: &str = "_cache";
39pub const META_DIR: &str = "_meta";
40pub const RCACHE_DIR: &str = "_rcache";
41pub const KEYS_FILENAME: &str = "keys";
42pub const SCHEMA_FILENAME: &str = "schema.json";
43
44fn derive_next_run_id(
45 dir: &Path,
46 runs_root: Option<&crate::durable_file::DurableRoot>,
47 active: &[RunRef],
48 retiring: &[crate::manifest::RetiredRun],
49) -> Result<u64> {
50 let mut maximum = 0_u64;
51 for run_id in active
52 .iter()
53 .map(|run| run.run_id)
54 .chain(retiring.iter().map(|run| run.run_id))
55 {
56 let run_id = u64::try_from(run_id)
57 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
58 maximum = maximum.max(run_id);
59 }
60 let names = match runs_root {
61 Some(root) => root.list_regular_files(".")?,
62 None => std::fs::read_dir(dir.join(RUNS_DIR))?
63 .map(|entry| entry.map(|entry| entry.file_name()))
64 .collect::<std::io::Result<Vec<_>>>()?,
65 };
66 for name in names {
67 let Some(name) = name.to_str() else {
68 continue;
69 };
70 let Some(digits) = name
71 .strip_prefix("r-")
72 .and_then(|name| name.strip_suffix(".sr"))
73 else {
74 continue;
75 };
76 let Ok(run_id) = digits.parse::<u64>() else {
77 continue;
78 };
79 if name == format!("r-{run_id}.sr") {
80 maximum = maximum.max(run_id);
81 }
82 }
83 maximum
84 .checked_add(1)
85 .map(|next| next.max(1))
86 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
87}
88
89enum ControlledVisibleCandidate {
90 Memory(Row),
91 Run(RunVisibleVersion),
92}
93
94impl ControlledVisibleCandidate {
95 fn row_id(&self) -> RowId {
96 match self {
97 Self::Memory(row) => row.row_id,
98 Self::Run(version) => version.row_id,
99 }
100 }
101
102 fn committed_epoch(&self) -> Epoch {
103 match self {
104 Self::Memory(row) => row.committed_epoch,
105 Self::Run(version) => version.committed_epoch,
106 }
107 }
108
109 fn deleted(&self) -> bool {
110 match self {
111 Self::Memory(row) => row.deleted,
112 Self::Run(version) => version.deleted,
113 }
114 }
115}
116
117enum ControlledVisibleCursor {
118 Memory(std::vec::IntoIter<Row>),
119 Run(Box<RunVisibleVersionCursor>),
120 #[cfg(test)]
121 Synthetic {
122 next: u64,
123 end: u64,
124 },
125}
126
127struct ControlledVisibleSource {
128 cursor: ControlledVisibleCursor,
129 current: Option<ControlledVisibleCandidate>,
130}
131
132impl ControlledVisibleSource {
133 fn memory(rows: Vec<Row>) -> Self {
134 Self {
135 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
136 current: None,
137 }
138 }
139
140 fn run(cursor: RunVisibleVersionCursor) -> Self {
141 Self {
142 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
143 current: None,
144 }
145 }
146
147 #[cfg(test)]
148 fn synthetic(end: u64) -> Self {
149 Self {
150 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
151 current: None,
152 }
153 }
154
155 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
156 self.current = match &mut self.cursor {
157 ControlledVisibleCursor::Memory(rows) => {
158 rows.next().map(ControlledVisibleCandidate::Memory)
159 }
160 ControlledVisibleCursor::Run(cursor) => cursor
161 .next_visible_version(control)?
162 .map(ControlledVisibleCandidate::Run),
163 #[cfg(test)]
164 ControlledVisibleCursor::Synthetic { next, end } => {
165 if *next > *end {
166 None
167 } else {
168 let row = Row::new(RowId(*next), Epoch(1));
169 *next += 1;
170 Some(ControlledVisibleCandidate::Memory(row))
171 }
172 }
173 };
174 Ok(())
175 }
176
177 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate> {
178 let current = self.current.take().ok_or_else(|| {
179 MongrelError::Other("controlled visible source was not primed".into())
180 })?;
181 self.advance(control)?;
182 Ok(current)
183 }
184
185 fn materialize(
186 &mut self,
187 candidate: ControlledVisibleCandidate,
188 control: &crate::ExecutionControl,
189 ) -> Result<Row> {
190 match candidate {
191 ControlledVisibleCandidate::Memory(row) => Ok(row),
192 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
193 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
194 _ => Err(MongrelError::Other(
195 "run candidate escaped its controlled cursor".into(),
196 )),
197 },
198 }
199 }
200}
201
202fn merge_controlled_visible_sources(
203 sources: &mut [ControlledVisibleSource],
204 control: &crate::ExecutionControl,
205 mut expired: impl FnMut(&Row) -> bool,
206 mut visit: impl FnMut(Row) -> Result<()>,
207) -> Result<()> {
208 let mut heap = BinaryHeap::new();
209 for (source_index, source) in sources.iter_mut().enumerate() {
210 source.advance(control)?;
211 if let Some(candidate) = &source.current {
212 heap.push(Reverse((candidate.row_id(), source_index)));
213 }
214 }
215 let mut merged = 0_usize;
216 while let Some(Reverse((row_id, source_index))) = heap.pop() {
217 if merged.is_multiple_of(256) {
218 control.checkpoint()?;
219 }
220 merged += 1;
221 let mut best_source = source_index;
222 let mut best = sources[source_index].pop(control)?;
223 if let Some(next) = &sources[source_index].current {
224 heap.push(Reverse((next.row_id(), source_index)));
225 }
226 while heap
227 .peek()
228 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
229 {
230 let Some(Reverse((_, source_index))) = heap.pop() else {
231 break;
232 };
233 let candidate = sources[source_index].pop(control)?;
234 if candidate.committed_epoch() > best.committed_epoch() {
235 best = candidate;
236 best_source = source_index;
237 }
238 if let Some(next) = &sources[source_index].current {
239 heap.push(Reverse((next.row_id(), source_index)));
240 }
241 }
242 if best.deleted() {
243 continue;
244 }
245 let row = sources[best_source].materialize(best, control)?;
246 if !expired(&row) {
247 visit(row)?;
248 }
249 }
250 control.checkpoint()
251}
252
253#[cfg(test)]
254mod controlled_visible_cursor_tests {
255 use super::*;
256
257 #[test]
258 fn streams_more_than_one_million_rows_without_a_source_cap() {
259 let control = crate::ExecutionControl::new(None);
260 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
261 let mut count = 0_u64;
262 let mut last = 0_u64;
263 merge_controlled_visible_sources(
264 &mut sources,
265 &control,
266 |_| false,
267 |row| {
268 count += 1;
269 assert!(row.row_id.0 > last);
270 last = row.row_id.0;
271 Ok(())
272 },
273 )
274 .unwrap();
275 assert_eq!(count, 1_000_001);
276 assert_eq!(last, 1_000_001);
277 }
278
279 #[test]
280 fn merge_orders_rows_and_honors_newest_tombstones() {
281 let control = crate::ExecutionControl::new(None);
282 let older = vec![
283 Row::new(RowId(1), Epoch(1)),
284 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
285 Row::new(RowId(4), Epoch(1)),
286 ];
287 let mut deleted = Row::new(RowId(1), Epoch(2));
288 deleted.deleted = true;
289 let newer = vec![
290 deleted,
291 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
292 Row::new(RowId(3), Epoch(2)),
293 ];
294 let mut sources = vec![
295 ControlledVisibleSource::memory(older),
296 ControlledVisibleSource::memory(newer),
297 ];
298 let mut rows = Vec::new();
299 merge_controlled_visible_sources(
300 &mut sources,
301 &control,
302 |_| false,
303 |row| {
304 rows.push(row);
305 Ok(())
306 },
307 )
308 .unwrap();
309 assert_eq!(
310 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
311 vec![2, 3, 4]
312 );
313 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
314 }
315}
316
317fn iso_now_bytes() -> Vec<u8> {
320 let secs = std::time::SystemTime::now()
321 .duration_since(std::time::UNIX_EPOCH)
322 .map(|d| d.as_secs() as i64)
323 .unwrap_or(0);
324 let days = secs.div_euclid(86_400);
325 let rem = secs.rem_euclid(86_400);
326 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
327 let (year, month, day) = civil_from_days(days);
328 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
329}
330
331pub(crate) fn unix_nanos_now() -> i64 {
332 std::time::SystemTime::now()
333 .duration_since(std::time::UNIX_EPOCH)
334 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
335 .unwrap_or(0)
336}
337
338fn ann_candidate_cap(
339 index_len: usize,
340 context: Option<&crate::query::AiExecutionContext>,
341) -> usize {
342 index_len
343 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
344 .min(context.map_or(
345 crate::query::MAX_RAW_INDEX_CANDIDATES,
346 crate::query::AiExecutionContext::max_fused_candidates,
347 ))
348}
349
350#[cfg(test)]
351mod ann_candidate_cap_tests {
352 use super::*;
353
354 #[test]
355 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
356 assert_eq!(
357 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
358 crate::query::MAX_RAW_INDEX_CANDIDATES,
359 );
360 let context = crate::query::AiExecutionContext::with_limits(
361 std::time::Duration::from_secs(1),
362 usize::MAX,
363 17,
364 );
365 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
366 }
367}
368
369fn civil_from_days(z: i64) -> (i64, u32, u32) {
370 let z = z + 719_468;
371 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
372 let doe = z - era * 146_097;
373 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
374 let y = yoe + era * 400;
375 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
376 let mp = (5 * doy + 2) / 153;
377 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
378 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
379 (if m <= 2 { y + 1 } else { y }, m, d)
380}
381
382const 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;
389
390#[derive(Clone, Copy, Debug)]
405struct AutoIncState {
406 column_id: u16,
407 next: i64,
408 seeded: bool,
409}
410
411pub(crate) struct RecoveryMetadataPlan {
412 live_count: u64,
413 auto_inc: Option<AutoIncState>,
414 changed: bool,
415}
416
417type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
418
419fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
422 schema.auto_increment_column().map(|c| AutoIncState {
423 column_id: c.id,
424 next: 0,
425 seeded: false,
426 })
427}
428
429#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
442pub enum IndexBuildPolicy {
443 #[default]
445 Deferred,
446 Eager,
448}
449
450#[derive(Clone)]
451struct ReversePkSegment {
452 values: HashMap<RowId, Vec<u8>>,
453 removed: HashSet<RowId>,
454}
455
456#[derive(Clone)]
457struct ReversePkMap {
458 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
459 active: ReversePkSegment,
460}
461
462impl ReversePkMap {
463 fn new() -> Self {
464 Self {
465 frozen: Arc::new(Vec::new()),
466 active: ReversePkSegment {
467 values: HashMap::new(),
468 removed: HashSet::new(),
469 },
470 }
471 }
472
473 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
474 let mut map = Self::new();
475 map.active.values.extend(entries);
476 map
477 }
478
479 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
480 self.active.removed.remove(&row_id);
481 self.active.values.insert(row_id, key);
482 }
483
484 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
485 if let Some(key) = self.active.values.get(row_id) {
486 return Some(key);
487 }
488 if self.active.removed.contains(row_id) {
489 return None;
490 }
491 for segment in self.frozen.iter().rev() {
492 if let Some(key) = segment.values.get(row_id) {
493 return Some(key);
494 }
495 if segment.removed.contains(row_id) {
496 return None;
497 }
498 }
499 None
500 }
501
502 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
503 let previous = self.get(row_id).cloned();
504 self.active.values.remove(row_id);
505 self.active.removed.insert(*row_id);
506 previous
507 }
508
509 fn clear(&mut self) {
510 *self = Self::new();
511 }
512
513 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
514 let mut entries = HashMap::new();
515 for segment in self
516 .frozen
517 .iter()
518 .map(Arc::as_ref)
519 .chain(std::iter::once(&self.active))
520 {
521 for row_id in &segment.removed {
522 entries.remove(row_id);
523 }
524 entries.extend(
525 segment
526 .values
527 .iter()
528 .map(|(row_id, key)| (*row_id, key.clone())),
529 );
530 }
531 entries
532 }
533
534 fn seal(&mut self) {
535 if self.active.values.is_empty() && self.active.removed.is_empty() {
536 return;
537 }
538 let active = std::mem::replace(
539 &mut self.active,
540 ReversePkSegment {
541 values: HashMap::new(),
542 removed: HashSet::new(),
543 },
544 );
545 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
546 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
547 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
548 values: self.entries(),
549 removed: HashSet::new(),
550 })]);
551 }
552 }
553}
554
555#[derive(Clone)]
557pub struct Table {
558 dir: PathBuf,
559 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
560 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
561 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
562 table_id: u64,
563 name: String,
567 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
572 read_only: bool,
575 durable_commit_failed: bool,
579 wal: WalSink,
580 memtable: Memtable,
581 mutable_run: MutableRun,
586 mutable_run_spill_bytes: u64,
588 compaction_zstd_level: i32,
591 allocator: RowIdAllocator,
592 epoch: Arc<EpochAuthority>,
593 data_generation: u64,
596 schema: Schema,
597 hot: HotIndex,
598 kek: Option<Arc<Kek>>,
601 column_keys: HashMap<u16, ([u8; 32], u8)>,
605 run_refs: Vec<RunRef>,
606 retiring: Vec<crate::manifest::RetiredRun>,
609 next_run_id: u64,
610 sync_byte_threshold: u64,
611 current_txn_id: u64,
616 pending_private_mutations: bool,
620 bitmap: HashMap<u16, BitmapIndex>,
621 ann: HashMap<u16, AnnIndex>,
622 fm: HashMap<u16, FmIndex>,
623 sparse: HashMap<u16, SparseIndex>,
624 minhash: HashMap<u16, MinHashIndex>,
625 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
628 pk_by_row: ReversePkMap,
630 pinned: BTreeMap<Epoch, usize>,
633 pub(crate) live_count: u64,
636 reservoir: crate::reservoir::Reservoir,
639 reservoir_complete: bool,
647 had_deletes: bool,
651 agg_cache: Arc<HashMap<u64, CachedAgg>>,
655 global_idx_epoch: u64,
659 indexes_complete: bool,
664 index_build_policy: IndexBuildPolicy,
666 pk_by_row_complete: bool,
673 flushed_epoch: u64,
676 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
679 snapshots: Arc<crate::retention::SnapshotRegistry>,
682 commit_lock: Arc<parking_lot::Mutex<()>>,
684 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
687 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
697 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
706 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
708 pending_delete_rids: roaring::RoaringBitmap,
711 pending_put_cols: std::collections::HashSet<u16>,
714 pending_rows: Vec<Row>,
720 pending_rows_auto_inc: Vec<bool>,
721 pending_dels: Vec<RowId>,
724 pending_truncate: Option<Epoch>,
728 auto_inc: Option<AutoIncState>,
731 ttl: Option<TtlPolicy>,
734}
735
736const _: () = {
743 const fn assert_sync<T: ?Sized + Sync>() {}
744 assert_sync::<Table>();
745};
746
747enum CachedData {
753 Rows(Arc<Vec<Row>>),
754 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
755}
756
757impl CachedData {
758 fn approx_bytes(&self) -> u64 {
759 match self {
760 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
761 CachedData::Columns(c) => c
762 .iter()
763 .map(|(_, c)| c.approx_bytes())
764 .sum::<u64>()
765 .saturating_add(c.len() as u64 * 16),
766 }
767 }
768}
769
770struct CachedEntry {
774 data: CachedData,
775 footprint: roaring::RoaringBitmap,
776 condition_cols: Vec<u16>,
777}
778
779struct ResultCache {
790 entries: std::collections::HashMap<u64, CachedEntry>,
791 order: std::collections::VecDeque<u64>,
792 bytes: u64,
793 max_bytes: u64,
794 dir: Option<std::path::PathBuf>,
795 #[allow(dead_code)]
796 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
797}
798
799#[derive(serde::Serialize, serde::Deserialize)]
801struct SerializedEntry {
802 condition_cols: Vec<u16>,
803 footprint_bits: Vec<u32>,
804 data: SerializedData,
805}
806
807#[derive(serde::Serialize, serde::Deserialize)]
808enum SerializedData {
809 Rows(Vec<Row>),
810 Columns(Vec<(u16, columnar::NativeColumn)>),
811}
812
813impl SerializedEntry {
814 fn from_entry(entry: &CachedEntry) -> Self {
815 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
816 let data = match &entry.data {
817 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
818 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
819 };
820 Self {
821 condition_cols: entry.condition_cols.clone(),
822 footprint_bits,
823 data,
824 }
825 }
826
827 fn into_entry(self) -> Option<CachedEntry> {
828 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
829 let data = match self.data {
830 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
831 SerializedData::Columns(c) => {
832 if !c.iter().all(|(_, col)| col.validate()) {
835 return None;
836 }
837 CachedData::Columns(Arc::new(c))
838 }
839 };
840 Some(CachedEntry {
841 data,
842 footprint,
843 condition_cols: self.condition_cols,
844 })
845 }
846}
847
848impl ResultCache {
849 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
850
851 fn new() -> Self {
852 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
853 }
854
855 fn with_max_bytes(max_bytes: u64) -> Self {
856 Self {
857 entries: std::collections::HashMap::new(),
858 order: std::collections::VecDeque::new(),
859 bytes: 0,
860 max_bytes,
861 dir: None,
862 cache_dek: None,
863 }
864 }
865
866 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
867 let _ = std::fs::create_dir_all(&dir);
868 self.dir = Some(dir);
869 self
870 }
871
872 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
873 self.cache_dek = dek;
874 self
875 }
876
877 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
878 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
879 }
880
881 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
885 let Some(path) = self.disk_path(key) else {
886 return;
887 };
888 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
889 Ok(s) => s,
890 Err(_) => return,
891 };
892 let on_disk = if let Some(dek) = &self.cache_dek {
894 match self.encrypt_cache(&serialized, dek) {
895 Some(b) => b,
896 None => return,
897 }
898 } else {
899 serialized
900 };
901 let tmp = path.with_extension("tmp");
902 use std::io::Write;
903 let write = || -> std::io::Result<()> {
904 let mut f = std::fs::File::create(&tmp)?;
905 f.write_all(&on_disk)?;
906 f.flush()?;
907 Ok(())
908 };
909 if write().is_err() {
910 let _ = std::fs::remove_file(&tmp);
911 return;
912 }
913 let _ = std::fs::rename(&tmp, &path);
914 }
915
916 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
918 let path = self.disk_path(key)?;
919 let bytes = std::fs::read(&path).ok()?;
920 let plaintext = if let Some(dek) = &self.cache_dek {
921 self.decrypt_cache(&bytes, dek)?
922 } else {
923 bytes
924 };
925 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
926 serialized.into_entry()
927 }
928
929 fn remove_from_disk(&self, key: u64) {
931 if let Some(path) = self.disk_path(key) {
932 let _ = std::fs::remove_file(&path);
933 }
934 }
935
936 #[cfg(feature = "encryption")]
938 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
939 use crate::encryption::Cipher;
940 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
941 let mut nonce = [0u8; 12];
942 crate::encryption::fill_random(&mut nonce).ok()?;
943 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
944 let mut out = Vec::with_capacity(12 + ct.len());
945 out.extend_from_slice(&nonce);
946 out.extend_from_slice(&ct);
947 Some(out)
948 }
949
950 #[cfg(not(feature = "encryption"))]
951 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
952 None
953 }
954
955 #[cfg(feature = "encryption")]
957 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
958 use crate::encryption::Cipher;
959 if bytes.len() < 28 {
960 return None;
961 }
962 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
963 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
964 let ct = &bytes[12..];
965 cipher.decrypt_page(&nonce, ct).ok()
966 }
967
968 #[cfg(not(feature = "encryption"))]
969 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
970 None
971 }
972
973 fn load_persistent(&mut self) {
976 let Some(dir) = self.dir.as_ref().cloned() else {
977 return;
978 };
979 let entries = match std::fs::read_dir(&dir) {
980 Ok(e) => e,
981 Err(_) => return,
982 };
983 for entry in entries.flatten() {
984 let path = entry.path();
985 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
987 let _ = std::fs::remove_file(&path);
988 continue;
989 }
990 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
991 continue;
992 }
993 let stem = match path.file_stem().and_then(|s| s.to_str()) {
994 Some(s) => s,
995 None => continue,
996 };
997 let key = match u64::from_str_radix(stem, 16) {
998 Ok(k) => k,
999 Err(_) => continue,
1000 };
1001 let bytes = match std::fs::read(&path) {
1002 Ok(b) => b,
1003 Err(_) => continue,
1004 };
1005 let plaintext = if let Some(dek) = &self.cache_dek {
1007 match self.decrypt_cache(&bytes, dek) {
1008 Some(p) => p,
1009 None => {
1010 let _ = std::fs::remove_file(&path);
1011 continue;
1012 }
1013 }
1014 } else {
1015 bytes
1016 };
1017 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1018 Ok(serialized) => {
1019 if let Some(entry) = serialized.into_entry() {
1020 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1021 self.entries.insert(key, entry);
1022 self.order.push_back(key);
1023 } else {
1024 let _ = std::fs::remove_file(&path);
1025 }
1026 }
1027 Err(_) => {
1028 let _ = std::fs::remove_file(&path);
1029 }
1030 }
1031 }
1032 self.evict();
1033 }
1034
1035 fn set_max_bytes(&mut self, max_bytes: u64) {
1036 self.max_bytes = max_bytes;
1037 self.evict();
1038 }
1039
1040 fn touch(&mut self, key: u64) {
1042 self.order.retain(|k| *k != key);
1043 self.order.push_back(key);
1044 }
1045
1046 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1047 let res = self.entries.get(&key).and_then(|e| match &e.data {
1048 CachedData::Rows(r) => Some(r.clone()),
1049 CachedData::Columns(_) => None,
1050 });
1051 if res.is_some() {
1052 self.touch(key);
1053 return res;
1054 }
1055 if let Some(entry) = self.load_from_disk(key) {
1057 let res = match &entry.data {
1058 CachedData::Rows(r) => Some(r.clone()),
1059 CachedData::Columns(_) => None,
1060 };
1061 if res.is_some() {
1062 let approx = entry.data.approx_bytes();
1063 self.bytes = self.bytes.saturating_add(approx);
1064 self.entries.insert(key, entry);
1065 self.order.push_back(key);
1066 self.evict();
1067 return res;
1068 }
1069 }
1070 None
1071 }
1072
1073 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1074 let res = self.entries.get(&key).and_then(|e| match &e.data {
1075 CachedData::Columns(c) => Some(c.clone()),
1076 CachedData::Rows(_) => None,
1077 });
1078 if res.is_some() {
1079 self.touch(key);
1080 return res;
1081 }
1082 if let Some(entry) = self.load_from_disk(key) {
1084 let res = match &entry.data {
1085 CachedData::Columns(c) => Some(c.clone()),
1086 CachedData::Rows(_) => None,
1087 };
1088 if res.is_some() {
1089 let approx = entry.data.approx_bytes();
1090 self.bytes = self.bytes.saturating_add(approx);
1091 self.entries.insert(key, entry);
1092 self.order.push_back(key);
1093 self.evict();
1094 return res;
1095 }
1096 }
1097 None
1098 }
1099
1100 fn insert(&mut self, key: u64, entry: CachedEntry) {
1101 let approx = entry.data.approx_bytes();
1102 if self.entries.remove(&key).is_some() {
1103 self.order.retain(|k| *k != key);
1104 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1105 }
1106 self.store_to_disk(key, &entry);
1108 self.bytes = self.bytes.saturating_add(approx);
1109 self.entries.insert(key, entry);
1110 self.order.push_back(key);
1111 self.evict();
1112 }
1113
1114 fn invalidate(
1123 &mut self,
1124 delete_rids: &roaring::RoaringBitmap,
1125 put_cols: &std::collections::HashSet<u16>,
1126 ) {
1127 if self.entries.is_empty() {
1128 return;
1129 }
1130 let has_deletes = !delete_rids.is_empty();
1131 let to_remove: std::collections::HashSet<u64> = self
1132 .entries
1133 .iter()
1134 .filter(|(_, e)| {
1135 let delete_hit = if e.footprint.is_empty() {
1136 has_deletes
1137 } else {
1138 e.footprint.intersection_len(delete_rids) > 0
1139 };
1140 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1141 delete_hit || col_hit
1142 })
1143 .map(|(&k, _)| k)
1144 .collect();
1145 for key in &to_remove {
1146 if let Some(e) = self.entries.remove(key) {
1147 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1148 }
1149 self.remove_from_disk(*key);
1150 }
1151 if !to_remove.is_empty() {
1152 self.order.retain(|k| !to_remove.contains(k));
1153 }
1154 }
1155
1156 fn clear(&mut self) {
1157 if let Some(dir) = &self.dir {
1159 if let Ok(entries) = std::fs::read_dir(dir) {
1160 for entry in entries.flatten() {
1161 let path = entry.path();
1162 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1163 let _ = std::fs::remove_file(&path);
1164 }
1165 }
1166 }
1167 }
1168 self.entries.clear();
1169 self.order.clear();
1170 self.bytes = 0;
1171 }
1172
1173 fn evict(&mut self) {
1174 while self.bytes > self.max_bytes {
1175 let Some(k) = self.order.pop_front() else {
1176 break;
1177 };
1178 if let Some(e) = self.entries.remove(&k) {
1179 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1180 self.remove_from_disk(k);
1184 }
1185 }
1186 }
1187}
1188
1189type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1196
1197fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1198 let _ = kek;
1199 #[cfg(feature = "encryption")]
1200 {
1201 if let Some(k) = kek {
1202 return (
1203 Some(k.derive_table_wal_key(_table_id)),
1204 Some(k.derive_cache_key()),
1205 );
1206 }
1207 }
1208 (None, None)
1209}
1210
1211#[cfg(feature = "encryption")]
1212fn read_table_encryption_salt_root(
1213 root: &crate::durable_file::DurableRoot,
1214) -> Result<[u8; crate::encryption::SALT_LEN]> {
1215 use std::io::Read;
1216
1217 let mut file = root
1218 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1219 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1220 let length = file.metadata()?.len();
1221 if length != crate::encryption::SALT_LEN as u64 {
1222 return Err(MongrelError::InvalidArgument(format!(
1223 "salt file is {length} bytes, expected {}",
1224 crate::encryption::SALT_LEN
1225 )));
1226 }
1227 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1228 file.read_exact(&mut salt)?;
1229 Ok(salt)
1230}
1231
1232#[cfg(feature = "encryption")]
1234fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1235 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1236}
1237
1238#[cfg(not(feature = "encryption"))]
1239fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1240 Box::new(crate::encryption::PlaintextCipher)
1241}
1242
1243fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1244 let Some(kek) = kek else {
1245 return HashMap::new();
1246 };
1247 #[cfg(feature = "encryption")]
1248 {
1249 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1250 schema
1251 .columns
1252 .iter()
1253 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1254 .map(|c| {
1255 let scheme = if schema
1256 .indexes
1257 .iter()
1258 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1259 {
1260 SCHEME_OPE_RANGE
1261 } else {
1262 SCHEME_HMAC_EQ
1263 };
1264 let key: [u8; 32] = *kek.derive_column_key(c.id);
1265 (c.id, (key, scheme))
1266 })
1267 .collect()
1268 }
1269 #[cfg(not(feature = "encryption"))]
1270 {
1271 let _ = (kek, schema);
1272 HashMap::new()
1273 }
1274}
1275
1276pub(crate) struct SharedCtx {
1281 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1282 pub epoch: Arc<EpochAuthority>,
1283 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1284 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1285 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1286 pub kek: Option<Arc<Kek>>,
1287 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1293 pub shared: Option<SharedWalCtx>,
1297 pub table_name: Option<String>,
1300 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1303 pub read_only: bool,
1305}
1306
1307#[derive(Clone)]
1313pub(crate) struct SharedWalCtx {
1314 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1315 pub group: Arc<GroupCommit>,
1316 pub poisoned: Arc<AtomicBool>,
1317 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1318 pub change_wake: tokio::sync::broadcast::Sender<()>,
1319}
1320
1321enum WalSink {
1324 Private(Wal),
1325 Shared(SharedWalCtx),
1326 ReadOnly,
1327}
1328
1329impl Clone for WalSink {
1330 fn clone(&self) -> Self {
1331 match self {
1332 Self::Shared(shared) => Self::Shared(shared.clone()),
1333 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1334 }
1335 }
1336}
1337
1338impl SharedCtx {
1339 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1343 let n_shards = if cache_dir.is_some() {
1347 1
1348 } else {
1349 crate::cache::CACHE_SHARDS
1350 };
1351 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1352 let page_cache = if let Some(d) = cache_dir {
1353 Arc::new(crate::cache::Sharded::new(1, || {
1354 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1355 }))
1356 } else {
1357 Arc::new(crate::cache::Sharded::new(n_shards, || {
1358 crate::cache::PageCache::new(per_shard)
1359 }))
1360 };
1361 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1362 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1363 crate::cache::CACHE_SHARDS,
1364 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1365 ));
1366 Self {
1367 root_guard: None,
1368 epoch: Arc::new(EpochAuthority::new(0)),
1369 page_cache,
1370 decoded_cache,
1371 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1372 kek,
1373 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1374 shared: None,
1375 table_name: None,
1376 auth: None,
1377 read_only: false,
1378 }
1379 }
1380}
1381
1382fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1386 use crate::query::Condition;
1387 match c {
1388 Condition::Pk(_)
1390 | Condition::BitmapEq { .. }
1391 | Condition::BitmapIn { .. }
1392 | Condition::BytesPrefix { .. }
1393 | Condition::IsNull { .. }
1394 | Condition::IsNotNull { .. } => 0,
1395 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1397 1
1398 }
1399 Condition::FmContains { .. }
1401 | Condition::FmContainsAll { .. }
1402 | Condition::Ann { .. }
1403 | Condition::SparseMatch { .. } => 2,
1404 }
1405}
1406
1407impl Table {
1408 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1409 let dir = dir.as_ref().to_path_buf();
1410 crate::durable_file::create_directory_all(&dir)?;
1411 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1412 let pinned = root.io_path()?;
1413 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1414 ctx.root_guard = Some(root);
1415 Self::create_in(&pinned, schema, table_id, ctx)
1416 }
1417
1418 #[cfg(feature = "encryption")]
1429 pub fn create_encrypted(
1430 dir: impl AsRef<Path>,
1431 schema: Schema,
1432 table_id: u64,
1433 passphrase: &str,
1434 ) -> Result<Self> {
1435 let dir = dir.as_ref().to_path_buf();
1436 crate::durable_file::create_directory_all(&dir)?;
1437 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1438 root.create_directory_all(META_DIR)?;
1439 let salt = crate::encryption::random_salt()?;
1440 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1441 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1442 let pinned = root.io_path()?;
1443 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1444 ctx.root_guard = Some(root);
1445 Self::create_in(&pinned, schema, table_id, ctx)
1446 }
1447
1448 #[cfg(feature = "encryption")]
1453 pub fn create_with_key(
1454 dir: impl AsRef<Path>,
1455 schema: Schema,
1456 table_id: u64,
1457 key: &[u8],
1458 ) -> Result<Self> {
1459 let dir = dir.as_ref().to_path_buf();
1460 crate::durable_file::create_directory_all(&dir)?;
1461 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1462 root.create_directory_all(META_DIR)?;
1463 let salt = crate::encryption::random_salt()?;
1464 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1465 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1466 let pinned = root.io_path()?;
1467 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1468 ctx.root_guard = Some(root);
1469 Self::create_in(&pinned, schema, table_id, ctx)
1470 }
1471
1472 #[cfg(feature = "encryption")]
1474 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1475 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1476 let salt = read_table_encryption_salt_root(&root)?;
1477 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1478 let pinned = root.io_path()?;
1479 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1480 ctx.root_guard = Some(root);
1481 Self::open_in(&pinned, ctx)
1482 }
1483
1484 pub(crate) fn create_in(
1485 dir: impl AsRef<Path>,
1486 schema: Schema,
1487 table_id: u64,
1488 ctx: SharedCtx,
1489 ) -> Result<Self> {
1490 schema.validate_auto_increment()?;
1491 schema.validate_defaults()?;
1492 schema.validate_ai()?;
1493 for index in &schema.indexes {
1494 index.validate_options()?;
1495 }
1496 let dir = dir.as_ref().to_path_buf();
1497 let runs_root = match ctx.root_guard.as_ref() {
1498 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1499 None => {
1500 crate::durable_file::create_directory_all(&dir)?;
1501 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1502 None
1503 }
1504 };
1505 match ctx.root_guard.as_deref() {
1506 Some(root) => write_schema_durable(root, &schema)?,
1507 None => write_schema(&dir, &schema)?,
1508 }
1509 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1510 let (wal, current_txn_id) = match ctx.shared.clone() {
1513 Some(s) => (WalSink::Shared(s), 0),
1514 None => {
1515 let pinned_wal_root = match ctx.root_guard.as_deref() {
1516 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1517 None => None,
1518 };
1519 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1520 root.io_path()?
1521 } else {
1522 let wal_dir = dir.join(WAL_DIR);
1523 crate::durable_file::create_directory_all(&wal_dir)?;
1524 wal_dir
1525 };
1526 let mut w = if let Some(ref dk) = wal_dek {
1527 Wal::create_with_cipher(
1528 wal_dir.join("seg-000000.wal"),
1529 Epoch(0),
1530 Some(make_cipher(dk)),
1531 0,
1532 )?
1533 } else {
1534 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1535 };
1536 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1537 (WalSink::Private(w), 1)
1538 }
1539 };
1540 let mut manifest = Manifest::new(table_id, schema.schema_id);
1541 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1546 match ctx.root_guard.as_deref() {
1547 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1548 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1549 }
1550 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1551 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1552 let auto_inc = resolve_auto_inc(&schema);
1553 let rcache_dir = dir.join(RCACHE_DIR);
1554 Ok(Self {
1555 dir,
1556 _root_guard: ctx.root_guard,
1557 runs_root,
1558 idx_root: None,
1559 table_id,
1560 name: ctx.table_name.unwrap_or_default(),
1561 auth: ctx.auth,
1562 read_only: ctx.read_only,
1563 durable_commit_failed: false,
1564 wal,
1565 memtable: Memtable::new(),
1566 mutable_run: MutableRun::new(),
1567 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1568 compaction_zstd_level: 3,
1569 allocator: RowIdAllocator::new(0),
1570 epoch: ctx.epoch,
1571 data_generation: 0,
1572 schema,
1573 hot: HotIndex::new(),
1574 kek: ctx.kek,
1575 column_keys,
1576 run_refs: Vec::new(),
1577 retiring: Vec::new(),
1578 next_run_id: 1,
1579 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1580 current_txn_id,
1581 pending_private_mutations: false,
1582 bitmap,
1583 ann,
1584 fm,
1585 sparse,
1586 minhash,
1587 learned_range: Arc::new(HashMap::new()),
1588 pk_by_row: ReversePkMap::new(),
1589 pinned: BTreeMap::new(),
1590 live_count: 0,
1591 reservoir: crate::reservoir::Reservoir::default(),
1592 reservoir_complete: true,
1593 had_deletes: false,
1594 agg_cache: Arc::new(HashMap::new()),
1595 global_idx_epoch: 0,
1596 indexes_complete: true,
1597 index_build_policy: IndexBuildPolicy::default(),
1598 pk_by_row_complete: false,
1599 flushed_epoch: 0,
1600 page_cache: ctx.page_cache,
1601 decoded_cache: ctx.decoded_cache,
1602 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1603 snapshots: ctx.snapshots,
1604 commit_lock: ctx.commit_lock,
1605 result_cache: Arc::new(parking_lot::Mutex::new(
1606 ResultCache::new()
1607 .with_dir(rcache_dir)
1608 .with_cache_dek(cache_dek.clone()),
1609 )),
1610 pending_delete_rids: roaring::RoaringBitmap::new(),
1611 pending_put_cols: std::collections::HashSet::new(),
1612 pending_rows: Vec::new(),
1613 pending_rows_auto_inc: Vec::new(),
1614 pending_dels: Vec::new(),
1615 pending_truncate: None,
1616 wal_dek,
1617 auto_inc,
1618 ttl: None,
1619 })
1620 }
1621
1622 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1626 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1627 let pinned = root.io_path()?;
1628 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1629 ctx.root_guard = Some(root);
1630 Self::open_in(&pinned, ctx)
1631 }
1632
1633 #[cfg(feature = "encryption")]
1636 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1637 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1638 let salt = read_table_encryption_salt_root(&root)?;
1639 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1640 let pinned = root.io_path()?;
1641 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1642 ctx.root_guard = Some(root);
1643 let t = Self::open_in(&pinned, ctx)?;
1644 Ok(t)
1645 }
1646
1647 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1648 let dir = dir.as_ref().to_path_buf();
1649 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1650 let mut manifest = match ctx.root_guard.as_ref() {
1651 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1652 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1653 };
1654 let schema: Schema = match ctx.root_guard.as_ref() {
1655 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1656 None => read_schema(&dir)?,
1657 };
1658 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1664 let runs_root = match ctx.root_guard.as_ref() {
1665 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1666 None => None,
1667 };
1668 let idx_root = match ctx.root_guard.as_ref() {
1669 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1670 Ok(root) => Some(Arc::new(root)),
1671 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1672 Err(error) => return Err(error.into()),
1673 },
1674 None => None,
1675 };
1676 schema.validate_auto_increment()?;
1677 schema.validate_defaults()?;
1678 schema.validate_ai()?;
1679 for index in &schema.indexes {
1680 index.validate_options()?;
1681 }
1682 let replay_epoch = Epoch(manifest.current_epoch);
1683 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1684 let private_replayed = if ctx.shared.is_none() {
1685 match latest_wal_segment(&dir.join(WAL_DIR))? {
1686 Some(path) => {
1687 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1688 crate::wal::replay_with_cipher(path, cipher)?
1689 }
1690 None => Vec::new(),
1691 }
1692 } else {
1693 Vec::new()
1694 };
1695 if ctx.shared.is_none() {
1696 preflight_standalone_open(
1697 &dir,
1698 runs_root.as_deref(),
1699 idx_root.as_deref(),
1700 &manifest,
1701 &schema,
1702 &private_replayed,
1703 ctx.kek.clone(),
1704 )?;
1705 }
1706 let next_run_id = derive_next_run_id(
1707 &dir,
1708 runs_root.as_deref(),
1709 &manifest.runs,
1710 &manifest.retiring,
1711 )?;
1712 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1716 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1717 None => {
1718 let replayed = private_replayed;
1719 let wal_dir = dir.join(WAL_DIR);
1725 crate::durable_file::create_directory_all(&wal_dir)?;
1726 let segment = next_wal_segment(&wal_dir)?;
1727 let segment_no = wal_segment_number(&segment).unwrap_or(0);
1728 let temporary = wal_dir.join(format!(
1729 ".recovery-{}-{}-{segment_no:06}.tmp",
1730 std::process::id(),
1731 std::time::SystemTime::now()
1732 .duration_since(std::time::UNIX_EPOCH)
1733 .unwrap_or_default()
1734 .as_nanos()
1735 ));
1736 let mut w = Wal::create_with_cipher(
1737 &temporary,
1738 replay_epoch,
1739 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1740 segment_no,
1741 )?;
1742 for record in &replayed {
1743 w.append_txn(record.txn_id, record.op.clone())?;
1744 }
1745 let mut w = w.publish_as(segment)?;
1746 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1747 let next_txn_id = replayed
1748 .iter()
1749 .map(|record| record.txn_id)
1750 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
1751 .max()
1752 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
1753 .unwrap_or(1);
1754 (WalSink::Private(w), replayed, next_txn_id)
1755 }
1756 };
1757
1758 let mut memtable = Memtable::new();
1759 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1760 let persisted_epoch = manifest.current_epoch;
1761 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1768 s.next = manifest.auto_inc_next;
1769 s.seeded = manifest.auto_inc_next > 0;
1770 s
1771 });
1772
1773 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1780 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1781 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
1782 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1783 std::collections::BTreeMap::new();
1784 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1785 let mut recovered_epoch = manifest.current_epoch;
1786 let mut recovered_manifest_dirty = schema_manifest_repair;
1787 let mut saw_delete = false;
1788 for record in replayed {
1789 let txn_id = record.txn_id;
1790 match record.op {
1791 Op::Put { rows, .. } => {
1792 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1793 for row in &rows {
1794 allocator.advance_to(row.row_id)?;
1795 if let Some(ai) = auto_inc.as_mut() {
1796 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1797 let next = n.checked_add(1).ok_or_else(|| {
1798 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
1799 })?;
1800 if next > ai.next {
1801 ai.next = next;
1802 }
1803 }
1804 }
1805 }
1806 staged_puts.entry(txn_id).or_default().extend(rows);
1807 }
1808 Op::Delete { row_ids, .. } => {
1809 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1810 }
1811 Op::TxnCommit { epoch, .. } => {
1812 let commit_epoch = Epoch(epoch);
1813 recovered_epoch = recovered_epoch.max(epoch);
1814 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
1815 memtable = Memtable::new();
1816 replayed_puts.clear();
1817 replayed_deletes.clear();
1818 manifest.runs.clear();
1819 manifest.retiring.clear();
1820 manifest.live_count = 0;
1821 manifest.global_idx_epoch = 0;
1822 manifest.current_epoch = manifest.current_epoch.max(epoch);
1823 recovered_manifest_dirty = true;
1824 saw_delete = true;
1825 }
1826 if let Some(puts) = staged_puts.remove(&txn_id) {
1827 if commit_epoch.0 > manifest.flushed_epoch {
1828 for row in &puts {
1829 memtable.upsert(row.clone());
1830 }
1831 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1832 }
1833 }
1834 if let Some(dels) = staged_deletes.remove(&txn_id) {
1835 saw_delete = true;
1836 if commit_epoch.0 > manifest.flushed_epoch {
1837 for rid in dels {
1838 memtable.tombstone(rid, commit_epoch);
1839 replayed_deletes.push((rid, commit_epoch));
1840 }
1841 }
1842 }
1843 }
1844 Op::TxnAbort => {
1845 staged_puts.remove(&txn_id);
1846 staged_deletes.remove(&txn_id);
1847 staged_truncates.remove(&txn_id);
1848 }
1849 Op::TruncateTable { .. } => {
1850 staged_puts.remove(&txn_id);
1851 staged_deletes.remove(&txn_id);
1852 staged_truncates.insert(txn_id);
1853 }
1854 Op::ExternalTableState { .. }
1855 | Op::Flush { .. }
1856 | Op::Ddl(_)
1857 | Op::BeforeImage { .. }
1858 | Op::CommitTimestamp { .. }
1859 | Op::SpilledRows { .. } => {}
1860 }
1861 }
1862
1863 let rcache_dir = dir.join(RCACHE_DIR);
1864 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1865 let mut db = Self {
1866 dir,
1867 _root_guard: ctx.root_guard,
1868 runs_root,
1869 idx_root,
1870 table_id: manifest.table_id,
1871 name: ctx.table_name.unwrap_or_default(),
1872 auth: ctx.auth,
1873 read_only: ctx.read_only,
1874 durable_commit_failed: false,
1875 wal,
1876 memtable,
1877 mutable_run: MutableRun::new(),
1878 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1879 compaction_zstd_level: 3,
1880 allocator,
1881 epoch: ctx.epoch,
1882 data_generation: persisted_epoch,
1883 schema,
1884 hot: HotIndex::new(),
1885 kek: ctx.kek,
1886 column_keys,
1887 run_refs: manifest.runs.clone(),
1888 retiring: manifest.retiring.clone(),
1889 next_run_id,
1890 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1891 current_txn_id,
1892 pending_private_mutations: false,
1893 bitmap: HashMap::new(),
1894 ann: HashMap::new(),
1895 fm: HashMap::new(),
1896 sparse: HashMap::new(),
1897 minhash: HashMap::new(),
1898 learned_range: Arc::new(HashMap::new()),
1899 pk_by_row: ReversePkMap::new(),
1900 pinned: BTreeMap::new(),
1901 live_count: manifest.live_count,
1902 reservoir: crate::reservoir::Reservoir::default(),
1903 reservoir_complete: false,
1904 had_deletes: saw_delete
1905 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1906 != manifest.live_count,
1907 agg_cache: Arc::new(HashMap::new()),
1908 global_idx_epoch: manifest.global_idx_epoch,
1909 indexes_complete: true,
1910 index_build_policy: IndexBuildPolicy::default(),
1911 pk_by_row_complete: false,
1912 flushed_epoch: manifest.flushed_epoch,
1913 page_cache: ctx.page_cache,
1914 decoded_cache: ctx.decoded_cache,
1915 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1916 snapshots: ctx.snapshots,
1917 commit_lock: ctx.commit_lock,
1918 result_cache: Arc::new(parking_lot::Mutex::new(
1919 ResultCache::new()
1920 .with_dir(rcache_dir)
1921 .with_cache_dek(cache_dek.clone()),
1922 )),
1923 pending_delete_rids: roaring::RoaringBitmap::new(),
1924 pending_put_cols: std::collections::HashSet::new(),
1925 pending_rows: Vec::new(),
1926 pending_rows_auto_inc: Vec::new(),
1927 pending_dels: Vec::new(),
1928 pending_truncate: None,
1929 wal_dek,
1930 auto_inc,
1931 ttl: manifest.ttl,
1932 };
1933
1934 db.epoch.advance_recovered(Epoch(recovered_epoch));
1937
1938 let checkpoint = match db.idx_root.as_deref() {
1943 Some(root) => {
1944 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
1945 }
1946 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
1947 };
1948 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1949 c.epoch_built == manifest.global_idx_epoch
1950 && manifest.global_idx_epoch > 0
1951 && manifest
1952 .runs
1953 .iter()
1954 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1955 });
1956 if let Some(loaded) = checkpoint {
1957 if checkpoint_valid {
1958 db.hot = loaded.hot;
1959 db.bitmap = loaded.bitmap;
1960 db.ann = loaded.ann;
1961 db.fm = loaded.fm;
1962 db.sparse = loaded.sparse;
1963 db.minhash = loaded.minhash;
1964 db.learned_range = Arc::new(loaded.learned_range);
1965 }
1968 }
1969 if !checkpoint_valid {
1970 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1971 db.bitmap = bitmap;
1972 db.ann = ann;
1973 db.fm = fm;
1974 db.sparse = sparse;
1975 db.minhash = minhash;
1976 db.rebuild_indexes_from_runs()?;
1977 db.build_learned_ranges()?;
1978 }
1979
1980 for (epoch, group) in replayed_puts {
1985 let (losers, winner_pks) = db.partition_pk_winners(&group);
1986 for (key, &row_id) in &winner_pks {
1987 if let Some(old_rid) = db.hot.get(key) {
1988 if old_rid != row_id {
1989 db.tombstone_row(old_rid, epoch, false);
1990 }
1991 }
1992 }
1993 for &loser_rid in &losers {
1994 db.tombstone_row(loser_rid, epoch, false);
1995 }
1996 for (key, row_id) in winner_pks {
1997 db.insert_hot_pk(key, row_id);
1998 }
1999 if db.schema.primary_key().is_none() {
2000 for r in &group {
2001 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2002 }
2003 }
2004 for r in &group {
2005 if !losers.contains(&r.row_id) {
2006 db.index_row(r);
2007 }
2008 }
2009 }
2010 for (rid, epoch) in &replayed_deletes {
2014 db.remove_hot_for_row(*rid, *epoch);
2015 }
2016
2017 if recovered_manifest_dirty {
2018 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2019 db.live_count = rows.len() as u64;
2020 db.persist_manifest(Epoch(recovered_epoch))?;
2021 }
2022
2023 db.result_cache.lock().load_persistent();
2030 Ok(db)
2031 }
2032
2033 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2039 if self.reservoir_complete {
2040 return Ok(());
2041 }
2042 self.rebuild_reservoir()?;
2043 self.reservoir_complete = true;
2044 Ok(())
2045 }
2046
2047 fn rebuild_reservoir(&mut self) -> Result<()> {
2050 let snap = self.snapshot();
2051 let rows = self.visible_rows(snap)?;
2052 self.reservoir.reset();
2053 for r in rows {
2054 self.reservoir.offer(r.row_id.0);
2055 }
2056 Ok(())
2057 }
2058
2059 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2060 self.rebuild_indexes_from_runs_inner(None)
2061 }
2062
2063 fn rebuild_indexes_from_runs_inner(
2064 &mut self,
2065 control: Option<&crate::ExecutionControl>,
2066 ) -> Result<()> {
2067 self.hot = HotIndex::new();
2068 self.pk_by_row.clear();
2069 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2070 self.bitmap = bitmap;
2071 self.ann = ann;
2072 self.fm = fm;
2073 self.sparse = sparse;
2074 self.minhash = minhash;
2075 let snapshot = Epoch(u64::MAX);
2076 let ttl_now = unix_nanos_now();
2077 let mut scanned = 0_usize;
2078 for rr in self.run_refs.clone() {
2079 if let Some(control) = control {
2080 control.checkpoint()?;
2081 }
2082 let mut reader = self.open_reader(rr.run_id)?;
2083 for row in reader.visible_rows(snapshot)? {
2084 if scanned.is_multiple_of(256) {
2085 if let Some(control) = control {
2086 control.checkpoint()?;
2087 }
2088 }
2089 scanned += 1;
2090 if self.row_expired_at(&row, ttl_now) {
2091 continue;
2092 }
2093 let tok_row = self.tokenized_for_indexes(&row);
2094 index_into(
2095 &self.schema,
2096 &tok_row,
2097 &mut self.hot,
2098 &mut self.bitmap,
2099 &mut self.ann,
2100 &mut self.fm,
2101 &mut self.sparse,
2102 &mut self.minhash,
2103 );
2104 }
2105 }
2106 for row in self.mutable_run.visible_versions(snapshot) {
2107 if scanned.is_multiple_of(256) {
2108 if let Some(control) = control {
2109 control.checkpoint()?;
2110 }
2111 }
2112 scanned += 1;
2113 if row.deleted {
2114 self.remove_hot_for_row(row.row_id, snapshot);
2115 } else if !self.row_expired_at(&row, ttl_now) {
2116 self.index_row(&row);
2117 }
2118 }
2119 for row in self.memtable.visible_versions(snapshot) {
2120 if scanned.is_multiple_of(256) {
2121 if let Some(control) = control {
2122 control.checkpoint()?;
2123 }
2124 }
2125 scanned += 1;
2126 if row.deleted {
2127 self.remove_hot_for_row(row.row_id, snapshot);
2128 } else if !self.row_expired_at(&row, ttl_now) {
2129 self.index_row(&row);
2130 }
2131 }
2132 self.refresh_pk_by_row_from_hot();
2133 Ok(())
2134 }
2135
2136 fn refresh_pk_by_row_from_hot(&mut self) {
2137 self.pk_by_row_complete = true;
2138 if self.schema.primary_key().is_none() {
2139 self.pk_by_row.clear();
2140 return;
2141 }
2142 self.pk_by_row = ReversePkMap::from_entries(
2148 self.hot
2149 .entries()
2150 .into_iter()
2151 .map(|(key, row_id)| (row_id, key)),
2152 );
2153 }
2154
2155 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2156 if self.schema.primary_key().is_some() {
2157 self.pk_by_row.insert(row_id, key.clone());
2158 }
2159 self.hot.insert(key, row_id);
2160 }
2161
2162 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2166 self.build_learned_ranges_inner(None)
2167 }
2168
2169 fn build_learned_ranges_inner(
2170 &mut self,
2171 control: Option<&crate::ExecutionControl>,
2172 ) -> Result<()> {
2173 self.learned_range = Arc::new(HashMap::new());
2174 if self.run_refs.len() != 1 {
2175 return Ok(());
2176 }
2177 let cols: Vec<(u16, usize)> = self
2178 .schema
2179 .indexes
2180 .iter()
2181 .filter(|i| i.kind == IndexKind::LearnedRange)
2182 .map(|i| {
2183 (
2184 i.column_id,
2185 i.options
2186 .learned_range
2187 .as_ref()
2188 .map(|options| options.epsilon)
2189 .unwrap_or(16),
2190 )
2191 })
2192 .collect();
2193 if cols.is_empty() {
2194 return Ok(());
2195 }
2196 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2197 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2198 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2199 _ => return Ok(()),
2200 };
2201 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2202 if column_index % 256 == 0 {
2203 if let Some(control) = control {
2204 control.checkpoint()?;
2205 }
2206 }
2207 let ty = self
2208 .schema
2209 .columns
2210 .iter()
2211 .find(|c| c.id == cid)
2212 .map(|c| c.ty.clone())
2213 .unwrap_or(TypeId::Int64);
2214 match ty {
2215 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2216 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2217 let pairs: Vec<(i64, u64)> = data
2218 .iter()
2219 .zip(row_ids.iter())
2220 .map(|(v, r)| (*v, *r))
2221 .collect();
2222 Arc::make_mut(&mut self.learned_range).insert(
2223 cid,
2224 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2225 );
2226 }
2227 }
2228 TypeId::Float64 => {
2229 if let columnar::NativeColumn::Float64 { data, .. } =
2230 reader.column_native(cid)?
2231 {
2232 let pairs: Vec<(f64, u64)> = data
2233 .iter()
2234 .zip(row_ids.iter())
2235 .map(|(v, r)| (*v, *r))
2236 .collect();
2237 Arc::make_mut(&mut self.learned_range).insert(
2238 cid,
2239 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2240 );
2241 }
2242 }
2243 _ => {}
2244 }
2245 }
2246 Ok(())
2247 }
2248
2249 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2256 if self.indexes_complete {
2257 crate::trace::QueryTrace::record(|t| {
2258 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2259 });
2260 return Ok(());
2261 }
2262 crate::trace::QueryTrace::record(|t| {
2263 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2264 });
2265 self.rebuild_indexes_from_runs()?;
2266 self.build_learned_ranges()?;
2267 self.indexes_complete = true;
2268 let epoch = self.current_epoch();
2269 self.checkpoint_indexes(epoch);
2270 Ok(())
2271 }
2272
2273 #[doc(hidden)]
2276 pub fn ensure_indexes_complete_controlled<F>(
2277 &mut self,
2278 control: &crate::ExecutionControl,
2279 before_publish: F,
2280 ) -> Result<bool>
2281 where
2282 F: FnOnce() -> bool,
2283 {
2284 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2285 .map(|(changed, _)| changed)
2286 }
2287
2288 #[doc(hidden)]
2291 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2292 &mut self,
2293 control: &crate::ExecutionControl,
2294 before_publish: F,
2295 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2296 where
2297 F: FnOnce() -> bool,
2298 {
2299 if self.indexes_complete {
2300 crate::trace::QueryTrace::record(|trace| {
2301 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2302 });
2303 return Ok((false, None));
2304 }
2305 crate::trace::QueryTrace::record(|trace| {
2306 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2307 });
2308 control.checkpoint()?;
2309 let maintenance_epoch = self.current_epoch();
2310 self.rebuild_indexes_from_runs_inner(Some(control))?;
2311 self.build_learned_ranges_inner(Some(control))?;
2312 control.checkpoint()?;
2313 if !before_publish() {
2314 return Err(MongrelError::Cancelled);
2315 }
2316 self.indexes_complete = true;
2317 self.checkpoint_indexes(maintenance_epoch);
2318 Ok((
2319 true,
2320 Some(MaintenanceReceipt {
2321 epoch: maintenance_epoch,
2322 }),
2323 ))
2324 }
2325
2326 fn pending_epoch(&self) -> Epoch {
2327 Epoch(self.epoch.visible().0 + 1)
2328 }
2329
2330 fn is_shared(&self) -> bool {
2333 matches!(self.wal, WalSink::Shared(_))
2334 }
2335
2336 fn ensure_txn_id(&mut self) -> Result<u64> {
2340 if self.current_txn_id == 0 {
2341 let id = match &self.wal {
2342 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2343 WalSink::Private(_) => {
2344 return Err(MongrelError::Full(
2345 "standalone transaction id namespace exhausted".into(),
2346 ))
2347 }
2348 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2349 };
2350 self.current_txn_id = id;
2351 }
2352 Ok(self.current_txn_id)
2353 }
2354
2355 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2358 self.ensure_writable()?;
2359 let txn_id = self.ensure_txn_id()?;
2360 let table_id = self.table_id;
2361 match &mut self.wal {
2362 WalSink::Private(w) => {
2363 w.append_txn(txn_id, op)?;
2364 self.pending_private_mutations = true;
2365 }
2366 WalSink::Shared(s) => {
2367 s.wal.lock().append(txn_id, table_id, op)?;
2368 }
2369 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2370 }
2371 Ok(())
2372 }
2373
2374 fn ensure_writable(&self) -> Result<()> {
2375 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2376 return Err(MongrelError::ReadOnlyReplica);
2377 }
2378 if self.durable_commit_failed {
2379 return Err(MongrelError::Other(
2380 "table poisoned by post-commit failure; reopen required".into(),
2381 ));
2382 }
2383 Ok(())
2384 }
2385
2386 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2397 match &self.auth {
2398 Some(checker) => checker.check(&self.name, perm),
2399 None => Ok(()),
2400 }
2401 }
2402 pub fn require_select(&self) -> Result<()> {
2407 self.require(crate::auth_state::RequiredPermission::Select)
2408 }
2409 fn require_insert(&self) -> Result<()> {
2410 self.require(crate::auth_state::RequiredPermission::Insert)
2411 }
2412 #[allow(dead_code)]
2416 fn require_update(&self) -> Result<()> {
2417 self.require(crate::auth_state::RequiredPermission::Update)
2418 }
2419 fn require_delete(&self) -> Result<()> {
2420 self.require(crate::auth_state::RequiredPermission::Delete)
2421 }
2422
2423 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2426 self.require_insert()?;
2427 Ok(self.put_returning(columns)?.0)
2428 }
2429
2430 pub fn put_returning(
2435 &mut self,
2436 mut columns: Vec<(u16, Value)>,
2437 ) -> Result<(RowId, Option<i64>)> {
2438 self.require_insert()?;
2439 let assigned = self.fill_auto_inc(&mut columns)?;
2440 self.apply_defaults(&mut columns)?;
2441 self.schema.validate_values(&columns)?;
2442 let row_id = if self.schema.clustered {
2447 self.derive_clustered_row_id(&columns)?
2448 } else {
2449 self.allocator.alloc()?
2450 };
2451 let epoch = self.pending_epoch();
2452 let mut row = Row::new(row_id, epoch);
2453 for (col_id, val) in columns {
2454 row.columns.insert(col_id, val);
2455 }
2456 self.commit_rows(vec![row], assigned.is_some())?;
2457 Ok((row_id, assigned))
2458 }
2459
2460 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2463 self.require_insert()?;
2464 Ok(self
2465 .put_batch_returning(batch)?
2466 .into_iter()
2467 .map(|(r, _)| r)
2468 .collect())
2469 }
2470
2471 pub fn put_batch_returning(
2474 &mut self,
2475 batch: Vec<Vec<(u16, Value)>>,
2476 ) -> Result<Vec<(RowId, Option<i64>)>> {
2477 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2478 for mut cols in batch {
2479 let assigned = self.fill_auto_inc(&mut cols)?;
2480 self.apply_defaults(&mut cols)?;
2481 filled.push((cols, assigned));
2482 }
2483 for (cols, _) in &filled {
2484 self.schema.validate_values(cols)?;
2485 }
2486 let epoch = self.pending_epoch();
2487 let mut rows = Vec::with_capacity(filled.len());
2488 let mut ids = Vec::with_capacity(filled.len());
2489 let first_row_id = if self.schema.clustered {
2490 None
2491 } else {
2492 let count = u64::try_from(filled.len())
2493 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2494 Some(self.allocator.alloc_range(count)?.0)
2495 };
2496 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2497 let row_id = match first_row_id {
2498 Some(first) => RowId(first + row_index as u64),
2499 None => self.derive_clustered_row_id(&cols)?,
2500 };
2501 let mut row = Row::new(row_id, epoch);
2502 for (c, v) in cols {
2503 row.columns.insert(c, v);
2504 }
2505 ids.push((row_id, assigned));
2506 rows.push(row);
2507 }
2508 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2509 self.commit_rows(rows, all_auto_generated)?;
2510 Ok(ids)
2511 }
2512
2513 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2519 self.ensure_writable()?;
2520 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2521 return Ok(None);
2522 };
2523 let pos = columns.iter().position(|(c, _)| *c == cid);
2524 let assigned = match pos {
2525 Some(i) => match &columns[i].1 {
2526 Value::Null => {
2527 let next = self.alloc_auto_inc_value()?;
2528 columns[i].1 = Value::Int64(next);
2529 Some(next)
2530 }
2531 Value::Int64(n) => {
2532 self.advance_auto_inc_past(*n)?;
2533 None
2534 }
2535 other => {
2536 return Err(MongrelError::InvalidArgument(format!(
2537 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2538 other
2539 )))
2540 }
2541 },
2542 None => {
2543 let next = self.alloc_auto_inc_value()?;
2544 columns.push((cid, Value::Int64(next)));
2545 Some(next)
2546 }
2547 };
2548 Ok(assigned)
2549 }
2550
2551 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2557 for col in &self.schema.columns {
2558 let Some(expr) = &col.default_value else {
2559 continue;
2560 };
2561 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2563 continue;
2564 }
2565 let pos = columns.iter().position(|(c, _)| *c == col.id);
2566 let needs_default = match pos {
2567 None => true,
2568 Some(i) => matches!(columns[i].1, Value::Null),
2569 };
2570 if !needs_default {
2571 continue;
2572 }
2573 let v = match expr {
2574 crate::schema::DefaultExpr::Static(v) => v.clone(),
2575 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2576 crate::schema::DefaultExpr::Uuid => {
2577 let mut buf = [0u8; 16];
2578 getrandom::getrandom(&mut buf)
2579 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2580 Value::Uuid(buf)
2581 }
2582 };
2583 match pos {
2584 None => columns.push((col.id, v)),
2585 Some(i) => columns[i].1 = v,
2586 }
2587 }
2588 Ok(())
2589 }
2590
2591 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2593 self.ensure_auto_inc_seeded()?;
2594 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2596 let v = ai.next;
2597 ai.next = ai
2598 .next
2599 .checked_add(1)
2600 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2601 Ok(v)
2602 }
2603
2604 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2607 self.ensure_auto_inc_seeded()?;
2608 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2609 let floor = used
2610 .checked_add(1)
2611 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2612 .max(1);
2613 if ai.next < floor {
2614 ai.next = floor;
2615 }
2616 Ok(())
2617 }
2618
2619 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2624 let needs_seed = match self.auto_inc {
2625 Some(ai) => !ai.seeded,
2626 None => return Ok(()),
2627 };
2628 if !needs_seed {
2629 return Ok(());
2630 }
2631 if self.seed_empty_auto_inc() {
2632 return Ok(());
2633 }
2634 let cid = self
2635 .auto_inc
2636 .as_ref()
2637 .expect("auto-inc column present")
2638 .column_id;
2639 let max = self.scan_max_int64(cid)?;
2640 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2641 let floor = max
2642 .checked_add(1)
2643 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2644 .max(1);
2645 if ai.next < floor {
2646 ai.next = floor;
2647 }
2648 ai.seeded = true;
2649 Ok(())
2650 }
2651
2652 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2653 if n == 0 || self.auto_inc.is_none() {
2654 return Ok(None);
2655 }
2656 self.ensure_auto_inc_seeded()?;
2657 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2658 let start = ai.next;
2659 let count = i64::try_from(n)
2660 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2661 ai.next = ai
2662 .next
2663 .checked_add(count)
2664 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2665 Ok(Some(start))
2666 }
2667
2668 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2672 let mut max: i64 = 0;
2673 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2674 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2675 if *n > max {
2676 max = *n;
2677 }
2678 }
2679 }
2680 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2681 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2682 if *n > max {
2683 max = *n;
2684 }
2685 }
2686 }
2687 for rr in self.run_refs.clone() {
2688 let reader = self.open_reader(rr.run_id)?;
2689 if let Some(stats) = reader.column_page_stats(column_id) {
2690 for s in stats {
2691 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2692 if n > max {
2693 max = n;
2694 }
2695 }
2696 }
2697 } else if reader.has_column(column_id) {
2698 if let columnar::NativeColumn::Int64 { data, validity } =
2699 reader.column_native_shared(column_id)?
2700 {
2701 for (i, n) in data.iter().enumerate() {
2702 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2703 {
2704 max = *n;
2705 }
2706 }
2707 }
2708 }
2709 }
2710 Ok(max)
2711 }
2712
2713 fn seed_empty_auto_inc(&mut self) -> bool {
2714 let Some(ai) = self.auto_inc.as_mut() else {
2715 return false;
2716 };
2717 if ai.seeded || self.live_count != 0 {
2718 return false;
2719 }
2720 if ai.next < 1 {
2721 ai.next = 1;
2722 }
2723 ai.seeded = true;
2724 true
2725 }
2726
2727 fn advance_auto_inc_from_native_columns(
2728 &mut self,
2729 columns: &[(u16, columnar::NativeColumn)],
2730 n: usize,
2731 live_before: u64,
2732 ) -> Result<()> {
2733 let Some(ai) = self.auto_inc.as_mut() else {
2734 return Ok(());
2735 };
2736 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2737 return Ok(());
2738 };
2739 let columnar::NativeColumn::Int64 { data, validity } = col else {
2740 return Err(MongrelError::InvalidArgument(format!(
2741 "AUTO_INCREMENT column {} must be Int64",
2742 ai.column_id
2743 )));
2744 };
2745 let max = if native_int64_strictly_increasing(col, n) {
2746 data.get(n.saturating_sub(1)).copied()
2747 } else {
2748 data.iter()
2749 .take(n)
2750 .enumerate()
2751 .filter_map(|(i, v)| {
2752 if validity.is_empty() || columnar::validity_bit(validity, i) {
2753 Some(*v)
2754 } else {
2755 None
2756 }
2757 })
2758 .max()
2759 };
2760 if let Some(max) = max {
2761 let floor = max
2762 .checked_add(1)
2763 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2764 .max(1);
2765 if ai.next < floor {
2766 ai.next = floor;
2767 }
2768 if ai.seeded || live_before == 0 {
2769 ai.seeded = true;
2770 }
2771 }
2772 Ok(())
2773 }
2774
2775 fn fill_auto_inc_native_columns(
2776 &mut self,
2777 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2778 n: usize,
2779 ) -> Result<()> {
2780 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2781 return Ok(());
2782 };
2783 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2784 if let Some(start) = self.alloc_auto_inc_range(n)? {
2785 columns.push((
2786 cid,
2787 columnar::NativeColumn::Int64 {
2788 data: (start..start.saturating_add(n as i64)).collect(),
2789 validity: vec![0xFF; n.div_ceil(8)],
2790 },
2791 ));
2792 }
2793 return Ok(());
2794 };
2795
2796 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2797 return Err(MongrelError::InvalidArgument(format!(
2798 "AUTO_INCREMENT column {cid} must be Int64"
2799 )));
2800 };
2801 if data.len() < n {
2802 return Err(MongrelError::InvalidArgument(format!(
2803 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2804 data.len()
2805 )));
2806 }
2807 if columnar::all_non_null(validity, n) {
2808 return Ok(());
2809 }
2810 if validity.iter().all(|b| *b == 0) {
2811 if let Some(start) = self.alloc_auto_inc_range(n)? {
2812 for (i, slot) in data.iter_mut().take(n).enumerate() {
2813 *slot = start.saturating_add(i as i64);
2814 }
2815 *validity = vec![0xFF; n.div_ceil(8)];
2816 }
2817 return Ok(());
2818 }
2819
2820 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2821 for (i, slot) in data.iter_mut().enumerate().take(n) {
2822 if columnar::validity_bit(validity, i) {
2823 self.advance_auto_inc_past(*slot)?;
2824 } else {
2825 *slot = self.alloc_auto_inc_value()?;
2826 }
2827 }
2828 *validity = new_validity;
2829 Ok(())
2830 }
2831
2832 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2846 self.ensure_writable()?;
2847 if self.auto_inc.is_none() {
2848 return Ok(None);
2849 }
2850 Ok(Some(self.alloc_auto_inc_value()?))
2851 }
2852
2853 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2859 let payload = bincode::serialize(&rows)?;
2860 self.wal_append_data(Op::Put {
2861 table_id: self.table_id,
2862 rows: payload,
2863 })?;
2864 if self.is_shared() {
2865 self.pending_rows_auto_inc
2866 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
2867 self.pending_rows.extend(rows);
2868 } else {
2869 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2870 }
2871 Ok(())
2872 }
2873
2874 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
2877 self.ensure_indexes_complete()
2878 }
2879
2880 pub(crate) fn prepare_durable_publish_controlled(
2881 &mut self,
2882 control: &crate::ExecutionControl,
2883 ) -> Result<()> {
2884 self.ensure_indexes_complete_controlled(control, || true)?;
2885 Ok(())
2886 }
2887
2888 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
2889 self.apply_put_rows_inner_prepared(rows, true);
2890 }
2891
2892 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2893 if check_existing_pk {
2894 self.ensure_indexes_complete()?;
2895 }
2896 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
2897 Ok(())
2898 }
2899
2900 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
2904 if rows.len() == 1 {
2908 let row = rows.into_iter().next().expect("len checked");
2909 self.apply_put_row_single(row, check_existing_pk);
2910 return;
2911 }
2912 let pk_id = self.schema.primary_key().map(|c| c.id);
2929 let probe = match pk_id {
2930 Some(pid) => {
2931 check_existing_pk
2932 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2933 }
2934 None => false,
2935 };
2936 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2939 for r in rows {
2940 for &cid in r.columns.keys() {
2941 self.pending_put_cols.insert(cid);
2942 }
2943 match pk_id {
2944 Some(pid) if probe || maintain_pk_by_row => {
2945 if let Some(pk_val) = r.columns.get(&pid) {
2946 let key = self.index_lookup_key(pid, pk_val);
2947 if probe {
2948 if let Some(old_rid) = self.hot.get(&key) {
2949 if old_rid != r.row_id {
2950 self.tombstone_row(old_rid, r.committed_epoch, true);
2951 }
2952 }
2953 }
2954 if maintain_pk_by_row {
2955 self.pk_by_row.insert(r.row_id, key);
2956 }
2957 }
2958 }
2959 Some(_) => {}
2960 None => {
2961 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2962 }
2963 }
2964 self.index_row(&r);
2965 self.reservoir.offer(r.row_id.0);
2966 self.memtable.upsert(r);
2967 self.live_count = self.live_count.saturating_add(1);
2970 }
2971 self.data_generation = self.data_generation.wrapping_add(1);
2972 }
2973
2974 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
2978 for &cid in row.columns.keys() {
2979 self.pending_put_cols.insert(cid);
2980 }
2981 let epoch = row.committed_epoch;
2982 if let Some(pk_col) = self.schema.primary_key() {
2983 let pk_id = pk_col.id;
2984 if let Some(pk_val) = row.columns.get(&pk_id) {
2985 let maintain_pk_by_row = self.pk_by_row_complete;
2989 if check_existing_pk || maintain_pk_by_row {
2990 let key = self.index_lookup_key(pk_id, pk_val);
2991 if check_existing_pk {
2992 if let Some(old_rid) = self.hot.get(&key) {
2993 if old_rid != row.row_id {
2994 self.tombstone_row(old_rid, epoch, true);
2995 }
2996 }
2997 }
2998 if maintain_pk_by_row {
2999 self.pk_by_row.insert(row.row_id, key);
3000 }
3001 }
3002 }
3003 } else {
3004 self.hot
3005 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3006 }
3007 self.index_row(&row);
3008 self.reservoir.offer(row.row_id.0);
3009 self.memtable.upsert(row);
3010 self.live_count = self.live_count.saturating_add(1);
3011 self.data_generation = self.data_generation.wrapping_add(1);
3012 }
3013
3014 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3017 self.allocator.alloc()
3018 }
3019
3020 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3026 let pk = self.schema.primary_key().ok_or_else(|| {
3027 MongrelError::Schema("clustered table requires a single-column primary key".into())
3028 })?;
3029 let pk_val = columns
3030 .iter()
3031 .find(|(id, _)| *id == pk.id)
3032 .map(|(_, v)| v)
3033 .ok_or_else(|| {
3034 MongrelError::Schema(format!(
3035 "clustered table missing primary key column {} ({})",
3036 pk.id, pk.name
3037 ))
3038 })?;
3039 let key_bytes = pk_val.encode_key();
3040 let mut hash: u64 = 0xcbf29ce484222325;
3042 for &b in &key_bytes {
3043 hash ^= b as u64;
3044 hash = hash.wrapping_mul(0x100000001b3);
3045 }
3046 Ok(RowId(hash.max(1)))
3049 }
3050
3051 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3059 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3060 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3061 }
3062 let n = rows.len();
3063 for r in rows {
3064 for &cid in r.columns.keys() {
3065 self.pending_put_cols.insert(cid);
3066 }
3067 }
3068 let (losers, winner_pks) = self.partition_pk_winners(rows);
3069 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3070 for (key, &row_id) in &winner_pks {
3072 if let Some(old_rid) = self.hot.get(key) {
3073 if old_rid != row_id {
3074 self.tombstone_row(old_rid, epoch, true);
3075 }
3076 }
3077 }
3078 for &loser_rid in &losers {
3081 self.tombstone_row(loser_rid, epoch, false);
3082 }
3083 for (key, row_id) in winner_pks {
3085 self.insert_hot_pk(key, row_id);
3086 }
3087 if self.schema.primary_key().is_none() {
3088 for r in rows {
3089 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3090 }
3091 }
3092 for r in rows {
3093 self.allocator.advance_to(r.row_id)?;
3094 if !losers.contains(&r.row_id) {
3095 self.index_row(r);
3096 }
3097 }
3098 for r in rows {
3099 if !losers.contains(&r.row_id) {
3100 self.reservoir.offer(r.row_id.0);
3101 }
3102 }
3103 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3104 self.data_generation = self.data_generation.wrapping_add(1);
3105 Ok(())
3106 }
3107
3108 pub(crate) fn recover_apply(
3113 &mut self,
3114 rows: Vec<Row>,
3115 deletes: Vec<(RowId, Epoch)>,
3116 ) -> Result<()> {
3117 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3121 std::collections::BTreeMap::new();
3122 for row in rows {
3123 if row.row_id.0 >= u64::MAX - 1 {
3124 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3125 }
3126 self.allocator.advance_to(row.row_id)?;
3127 if let Some(ai) = self.auto_inc.as_mut() {
3132 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3133 let next = n.checked_add(1).ok_or_else(|| {
3134 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3135 })?;
3136 if next > ai.next {
3137 ai.next = next;
3138 }
3139 }
3140 }
3141 by_epoch.entry(row.committed_epoch).or_default().push(row);
3142 }
3143 for (epoch, group) in by_epoch {
3144 let (losers, winner_pks) = self.partition_pk_winners(&group);
3145 for (key, &row_id) in &winner_pks {
3147 if let Some(old_rid) = self.hot.get(key) {
3148 if old_rid != row_id {
3149 self.tombstone_row(old_rid, epoch, false);
3150 }
3151 }
3152 }
3153 for (key, row_id) in winner_pks {
3154 self.insert_hot_pk(key, row_id);
3155 }
3156 if self.schema.primary_key().is_none() {
3157 for r in &group {
3158 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3159 }
3160 }
3161 for r in &group {
3162 if !losers.contains(&r.row_id) {
3163 self.memtable.upsert(r.clone());
3164 self.index_row(r);
3165 }
3166 }
3167 }
3168 for (rid, epoch) in deletes {
3169 self.memtable.tombstone(rid, epoch);
3170 self.remove_hot_for_row(rid, epoch);
3171 }
3172 self.reservoir_complete = false;
3175 Ok(())
3176 }
3177
3178 pub(crate) fn flushed_epoch(&self) -> u64 {
3180 self.flushed_epoch
3181 }
3182
3183 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3184 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3185 }
3186
3187 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3189 self.schema.validate_values(cells)
3190 }
3191
3192 fn validate_columns_not_null(
3196 &self,
3197 columns: &[(u16, columnar::NativeColumn)],
3198 n: usize,
3199 ) -> Result<()> {
3200 let by_id: HashMap<u16, &columnar::NativeColumn> =
3201 columns.iter().map(|(id, c)| (*id, c)).collect();
3202 for col in &self.schema.columns {
3203 if !col.flags.contains(ColumnFlags::NULLABLE) {
3204 match by_id.get(&col.id) {
3205 None => {
3206 return Err(MongrelError::InvalidArgument(format!(
3207 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3208 col.name, col.id
3209 )));
3210 }
3211 Some(c) => {
3212 if c.null_count(n) != 0 {
3213 return Err(MongrelError::InvalidArgument(format!(
3214 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3215 col.name, col.id
3216 )));
3217 }
3218 }
3219 }
3220 }
3221 if let TypeId::Enum { variants } = &col.ty {
3222 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3223 if by_id.contains_key(&col.id) {
3224 return Err(MongrelError::InvalidArgument(format!(
3225 "column '{}' ({}) enum requires a bytes column",
3226 col.name, col.id
3227 )));
3228 }
3229 continue;
3230 };
3231 for index in 0..n {
3232 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3233 continue;
3234 };
3235 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3236 return Err(MongrelError::InvalidArgument(format!(
3237 "column '{}' ({}) enum value {:?} is not one of {:?}",
3238 col.name,
3239 col.id,
3240 String::from_utf8_lossy(value),
3241 variants
3242 )));
3243 }
3244 }
3245 }
3246 }
3247 Ok(())
3248 }
3249
3250 fn bulk_pk_winner_indices(
3255 &self,
3256 columns: &[(u16, columnar::NativeColumn)],
3257 n: usize,
3258 ) -> Option<Vec<usize>> {
3259 let pk_col = self.schema.primary_key()?;
3260 let pk_id = pk_col.id;
3261 let pk_ty = pk_col.ty.clone();
3262 let by_id: HashMap<u16, &columnar::NativeColumn> =
3263 columns.iter().map(|(id, c)| (*id, c)).collect();
3264 let pk_native = by_id.get(&pk_id)?;
3265 if native_int64_strictly_increasing(pk_native, n) {
3266 return None;
3267 }
3268 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3270 let mut null_pk_rows: Vec<usize> = Vec::new();
3271 for i in 0..n {
3272 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3273 Some(key) => {
3274 last.insert(key, i);
3275 }
3276 None => null_pk_rows.push(i),
3277 }
3278 }
3279 let mut winners: HashSet<usize> = last.values().copied().collect();
3280 for i in null_pk_rows {
3281 winners.insert(i);
3282 }
3283 Some((0..n).filter(|i| winners.contains(i)).collect())
3284 }
3285
3286 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3288 self.require_delete()?;
3289 let epoch = self.pending_epoch();
3290 self.wal_append_data(Op::Delete {
3291 table_id: self.table_id,
3292 row_ids: vec![row_id],
3293 })?;
3294 if self.is_shared() {
3295 self.pending_dels.push(row_id);
3296 } else {
3297 self.apply_delete(row_id, epoch);
3298 }
3299 Ok(())
3300 }
3301
3302 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3303 let pre = self.get(row_id, self.snapshot());
3304 self.delete(row_id)?;
3305 Ok(pre.map(|row| {
3306 let mut columns: Vec<_> = row.columns.into_iter().collect();
3307 columns.sort_by_key(|(id, _)| *id);
3308 OwnedRow { columns }
3309 }))
3310 }
3311
3312 pub fn truncate(&mut self) -> Result<()> {
3314 self.require_delete()?;
3315 let epoch = self.pending_epoch();
3316 self.wal_append_data(Op::TruncateTable {
3317 table_id: self.table_id,
3318 })?;
3319 self.pending_rows.clear();
3320 self.pending_rows_auto_inc.clear();
3321 self.pending_dels.clear();
3322 self.pending_truncate = Some(epoch);
3323 Ok(())
3324 }
3325
3326 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3328 self.run_refs.clear();
3333 self.retiring.clear();
3334 self.memtable = Memtable::new();
3335 self.mutable_run = MutableRun::new();
3336 self.hot = HotIndex::new();
3337 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3338 self.bitmap = bitmap;
3339 self.ann = ann;
3340 self.fm = fm;
3341 self.sparse = sparse;
3342 self.minhash = minhash;
3343 self.learned_range = Arc::new(HashMap::new());
3344 self.pk_by_row.clear();
3345 self.pk_by_row_complete = false;
3346 self.live_count = 0;
3347 self.reservoir = crate::reservoir::Reservoir::default();
3348 self.reservoir_complete = true;
3349 self.had_deletes = true;
3350 self.agg_cache = Arc::new(HashMap::new());
3351 self.global_idx_epoch = 0;
3352 self.indexes_complete = true;
3353 self.pending_delete_rids.clear();
3354 self.pending_put_cols.clear();
3355 self.pending_rows.clear();
3356 self.pending_rows_auto_inc.clear();
3357 self.pending_dels.clear();
3358 self.clear_result_cache();
3359 self.invalidate_index_checkpoint();
3360 self.data_generation = self.data_generation.wrapping_add(1);
3361 }
3362
3363 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3366 self.remove_hot_for_row(row_id, epoch);
3367 self.tombstone_row(row_id, epoch, true);
3368 self.data_generation = self.data_generation.wrapping_add(1);
3369 }
3370
3371 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3375 let tombstone = Row {
3376 row_id,
3377 committed_epoch: epoch,
3378 columns: std::collections::HashMap::new(),
3379 deleted: true,
3380 };
3381 self.memtable.upsert(tombstone);
3382 self.pk_by_row.remove(&row_id);
3383 if adjust_live_count {
3384 self.live_count = self.live_count.saturating_sub(1);
3385 }
3386 self.pending_delete_rids.insert(row_id.0 as u32);
3388 self.had_deletes = true;
3391 self.agg_cache = Arc::new(HashMap::new());
3392 }
3393
3394 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3398 let Some(pk_col) = self.schema.primary_key() else {
3399 return;
3400 };
3401 if self.pk_by_row_complete {
3404 if let Some(key) = self.pk_by_row.remove(&row_id) {
3405 if self.hot.get(&key) == Some(row_id) {
3406 self.hot.remove(&key);
3407 }
3408 }
3409 return;
3410 }
3411 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3430 if self.indexes_complete {
3431 let pk_val = self
3432 .memtable
3433 .get_version(row_id, lookup_epoch)
3434 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3435 .or_else(|| {
3436 self.mutable_run
3437 .get_version(row_id, lookup_epoch)
3438 .filter(|(_, r)| !r.deleted)
3439 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3440 })
3441 .or_else(|| {
3442 self.run_refs.iter().find_map(|rr| {
3443 let mut reader = self.open_reader(rr.run_id).ok()?;
3444 let (_, deleted, val) = reader
3445 .get_version_column(row_id, lookup_epoch, pk_col.id)
3446 .ok()??;
3447 if deleted {
3448 return None;
3449 }
3450 val
3451 })
3452 });
3453 if let Some(pk_val) = pk_val {
3454 let key = self.index_lookup_key(pk_col.id, &pk_val);
3455 if self.hot.get(&key) == Some(row_id) {
3456 self.hot.remove(&key);
3457 }
3458 return;
3459 }
3460 }
3461 self.refresh_pk_by_row_from_hot();
3466 if let Some(key) = self.pk_by_row.remove(&row_id) {
3467 if self.hot.get(&key) == Some(row_id) {
3468 self.hot.remove(&key);
3469 }
3470 }
3471 }
3472
3473 fn partition_pk_winners(
3478 &self,
3479 rows: &[Row],
3480 ) -> (
3481 std::collections::HashSet<RowId>,
3482 std::collections::HashMap<Vec<u8>, RowId>,
3483 ) {
3484 let mut losers = std::collections::HashSet::new();
3485 let Some(pk_col) = self.schema.primary_key() else {
3486 return (losers, std::collections::HashMap::new());
3487 };
3488 let pk_id = pk_col.id;
3489 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3490 std::collections::HashMap::new();
3491 for r in rows {
3492 let Some(pk_val) = r.columns.get(&pk_id) else {
3493 continue;
3494 };
3495 let key = self.index_lookup_key(pk_id, pk_val);
3496 if let Some(&old_rid) = winners.get(&key) {
3497 losers.insert(old_rid);
3498 }
3499 winners.insert(key, r.row_id);
3500 }
3501 (losers, winners)
3502 }
3503
3504 fn index_row(&mut self, row: &Row) {
3505 if row.deleted {
3506 return;
3507 }
3508 let any_predicate = self
3516 .schema
3517 .indexes
3518 .iter()
3519 .any(|idx| idx.predicate.is_some());
3520 if any_predicate {
3521 let columns_map: HashMap<u16, &Value> =
3522 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3523 let name_to_id: HashMap<&str, u16> = self
3524 .schema
3525 .columns
3526 .iter()
3527 .map(|c| (c.name.as_str(), c.id))
3528 .collect();
3529 for idx in &self.schema.indexes {
3530 if let Some(pred) = &idx.predicate {
3531 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3532 continue; }
3534 }
3535 index_into_single(
3537 idx,
3538 &self.schema,
3539 row,
3540 &mut self.hot,
3541 &mut self.bitmap,
3542 &mut self.ann,
3543 &mut self.fm,
3544 &mut self.sparse,
3545 &mut self.minhash,
3546 );
3547 }
3548 return;
3549 }
3550 if self.column_keys.is_empty() {
3554 index_into(
3555 &self.schema,
3556 row,
3557 &mut self.hot,
3558 &mut self.bitmap,
3559 &mut self.ann,
3560 &mut self.fm,
3561 &mut self.sparse,
3562 &mut self.minhash,
3563 );
3564 return;
3565 }
3566 let effective_row = self.tokenized_for_indexes(row);
3567 index_into(
3568 &self.schema,
3569 &effective_row,
3570 &mut self.hot,
3571 &mut self.bitmap,
3572 &mut self.ann,
3573 &mut self.fm,
3574 &mut self.sparse,
3575 &mut self.minhash,
3576 );
3577 }
3578
3579 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3585 if self.column_keys.is_empty() {
3586 return row.clone();
3587 }
3588 #[cfg(feature = "encryption")]
3589 {
3590 use crate::encryption::SCHEME_HMAC_EQ;
3591 let mut tok = row.clone();
3592 for (&cid, &(_, scheme)) in &self.column_keys {
3593 if scheme != SCHEME_HMAC_EQ {
3594 continue;
3595 }
3596 if let Some(v) = tok.columns.get(&cid).cloned() {
3597 if let Some(t) = self.tokenize_value(cid, &v) {
3598 tok.columns.insert(cid, t);
3599 }
3600 }
3601 }
3602 tok
3603 }
3604 #[cfg(not(feature = "encryption"))]
3605 {
3606 row.clone()
3607 }
3608 }
3609
3610 pub fn commit(&mut self) -> Result<Epoch> {
3615 self.commit_inner(None)
3616 }
3617
3618 #[doc(hidden)]
3621 pub fn commit_controlled<F>(
3622 &mut self,
3623 control: &crate::ExecutionControl,
3624 mut before_commit: F,
3625 ) -> Result<Epoch>
3626 where
3627 F: FnMut() -> Result<()>,
3628 {
3629 self.commit_inner(Some((control, &mut before_commit)))
3630 }
3631
3632 fn commit_inner(
3633 &mut self,
3634 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3635 ) -> Result<Epoch> {
3636 self.ensure_writable()?;
3637 if !self.has_pending_mutations() {
3638 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3639 return Err(MongrelError::Full(
3640 "standalone transaction id namespace exhausted".into(),
3641 ));
3642 }
3643 return Ok(self.epoch.visible());
3644 }
3645 self.commit_new_epoch_inner(controlled)
3646 }
3647
3648 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3652 self.commit_new_epoch_inner(None)
3653 }
3654
3655 fn commit_new_epoch_inner(
3656 &mut self,
3657 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3658 ) -> Result<Epoch> {
3659 self.ensure_writable()?;
3660 if self.is_shared() {
3661 self.commit_shared(controlled)
3662 } else {
3663 self.commit_private(controlled)
3664 }
3665 }
3666
3667 fn commit_private(
3669 &mut self,
3670 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3671 ) -> Result<Epoch> {
3672 let commit_lock = Arc::clone(&self.commit_lock);
3676 let _g = commit_lock.lock();
3677 let txn_id = self.ensure_txn_id()?;
3680 if let Some((control, before_commit)) = controlled {
3681 control.checkpoint()?;
3682 before_commit()?;
3683 }
3684 let new_epoch = self.epoch.bump_assigned();
3685 let epoch_authority = Arc::clone(&self.epoch);
3686 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3687 let wal_result = match &mut self.wal {
3691 WalSink::Private(w) => w
3692 .append_txn(
3693 txn_id,
3694 Op::TxnCommit {
3695 epoch: new_epoch.0,
3696 added_runs: Vec::new(),
3697 },
3698 )
3699 .and_then(|_| w.sync()),
3700 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3701 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
3702 };
3703 if let Err(error) = wal_result {
3704 self.durable_commit_failed = true;
3705 return Err(MongrelError::CommitOutcomeUnknown {
3706 epoch: new_epoch.0,
3707 message: error.to_string(),
3708 });
3709 }
3710 if let Some(epoch) = self.pending_truncate.take() {
3713 self.apply_truncate(epoch);
3714 }
3715 self.invalidate_pending_cache();
3716 let publish_result = self.persist_manifest(new_epoch);
3717 self.epoch.publish_in_order(new_epoch);
3721 epoch_guard.disarm();
3722 if let Err(error) = publish_result {
3723 self.durable_commit_failed = true;
3724 return Err(MongrelError::DurableCommit {
3725 epoch: new_epoch.0,
3726 message: error.to_string(),
3727 });
3728 }
3729 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
3730 self.pending_private_mutations = false;
3731 self.data_generation = self.data_generation.wrapping_add(1);
3732 Ok(new_epoch)
3733 }
3734
3735 fn commit_shared(
3741 &mut self,
3742 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3743 ) -> Result<Epoch> {
3744 use std::sync::atomic::Ordering;
3745 let s = match &self.wal {
3746 WalSink::Shared(s) => s.clone(),
3747 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3748 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3749 };
3750 if s.poisoned.load(Ordering::Relaxed) {
3751 return Err(MongrelError::Other(
3752 "database poisoned by fsync error".into(),
3753 ));
3754 }
3755 let commit_lock = Arc::clone(&self.commit_lock);
3762 let _g = commit_lock.lock();
3763 if !self.pending_rows.is_empty() {
3764 match controlled.as_ref() {
3765 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
3766 None => self.prepare_durable_publish()?,
3767 }
3768 }
3769 let txn_id = self.ensure_txn_id()?;
3772 let mut wal = s.wal.lock();
3773 if let Some((control, before_commit)) = controlled {
3774 control.checkpoint()?;
3775 before_commit()?;
3776 }
3777 let new_epoch = self.epoch.bump_assigned();
3778 let epoch_authority = Arc::clone(&self.epoch);
3779 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3780 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
3781 Ok(commit_seq) => commit_seq,
3782 Err(error) => {
3783 s.poisoned.store(true, Ordering::Relaxed);
3784 return Err(MongrelError::CommitOutcomeUnknown {
3785 epoch: new_epoch.0,
3786 message: error.to_string(),
3787 });
3788 }
3789 };
3790 drop(wal);
3791 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
3792 s.poisoned.store(true, Ordering::Relaxed);
3793 return Err(MongrelError::CommitOutcomeUnknown {
3794 epoch: new_epoch.0,
3795 message: error.to_string(),
3796 });
3797 }
3798
3799 if self.pending_truncate.take().is_some() {
3802 self.apply_truncate(new_epoch);
3803 }
3804 let mut rows = std::mem::take(&mut self.pending_rows);
3805 if !rows.is_empty() {
3806 for r in &mut rows {
3807 r.committed_epoch = new_epoch;
3808 }
3809 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3810 let all_auto_generated =
3811 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3812 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
3813 } else {
3814 self.pending_rows_auto_inc.clear();
3815 }
3816 let dels = std::mem::take(&mut self.pending_dels);
3817 for rid in dels {
3818 self.apply_delete(rid, new_epoch);
3819 }
3820
3821 self.invalidate_pending_cache();
3822 let publish_result = self.persist_manifest(new_epoch);
3823 self.epoch.publish_in_order(new_epoch);
3824 epoch_guard.disarm();
3825 let _ = s.change_wake.send(());
3826 if let Err(error) = publish_result {
3827 self.durable_commit_failed = true;
3828 s.poisoned.store(true, Ordering::Relaxed);
3829 return Err(MongrelError::DurableCommit {
3830 epoch: new_epoch.0,
3831 message: error.to_string(),
3832 });
3833 }
3834 self.current_txn_id = 0;
3836 self.data_generation = self.data_generation.wrapping_add(1);
3837 Ok(new_epoch)
3838 }
3839
3840 pub fn flush(&mut self) -> Result<Epoch> {
3848 self.flush_with_outcome().map(|(epoch, _)| epoch)
3849 }
3850
3851 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
3853 self.flush_with_outcome_inner(None)
3854 }
3855
3856 #[doc(hidden)]
3859 pub fn flush_with_outcome_controlled<F>(
3860 &mut self,
3861 control: &crate::ExecutionControl,
3862 mut before_commit: F,
3863 ) -> Result<(Epoch, bool)>
3864 where
3865 F: FnMut() -> Result<()>,
3866 {
3867 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
3868 }
3869
3870 fn flush_with_outcome_inner(
3871 &mut self,
3872 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3873 ) -> Result<(Epoch, bool)> {
3874 match controlled.as_ref() {
3875 Some((control, _)) => {
3876 self.ensure_indexes_complete_controlled(control, || true)?;
3877 }
3878 None => self.ensure_indexes_complete()?,
3879 }
3880 let committed = self.has_pending_mutations();
3881 let epoch = self.commit_inner(controlled)?;
3882 let finish: Result<(Epoch, bool)> = (|| {
3883 let rows = self.memtable.drain_sorted();
3884 if !rows.is_empty() {
3885 self.mutable_run.insert_many(rows);
3886 }
3887 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3888 self.spill_mutable_run(epoch)?;
3889 self.mark_flushed(epoch)?;
3893 self.persist_manifest(epoch)?;
3894 self.build_learned_ranges()?;
3895 self.checkpoint_indexes(epoch);
3898 }
3899 Ok((epoch, committed))
3902 })();
3903 match finish {
3904 Err(error) if committed => Err(MongrelError::DurableCommit {
3905 epoch: epoch.0,
3906 message: error.to_string(),
3907 }),
3908 result => result,
3909 }
3910 }
3911
3912 fn has_pending_mutations(&self) -> bool {
3913 self.pending_private_mutations
3914 || !self.pending_rows.is_empty()
3915 || !self.pending_dels.is_empty()
3916 || self.pending_truncate.is_some()
3917 }
3918
3919 pub fn has_pending_writes(&self) -> bool {
3920 self.has_pending_mutations()
3921 }
3922
3923 pub fn force_flush(&mut self) -> Result<Epoch> {
3932 let saved = self.mutable_run_spill_bytes;
3933 self.mutable_run_spill_bytes = 1;
3934 let result = self.flush();
3935 self.mutable_run_spill_bytes = saved;
3936 result
3937 }
3938
3939 pub fn close(&mut self) -> Result<()> {
3946 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3947 self.force_flush()?;
3948 }
3949 Ok(())
3950 }
3951
3952 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3959 let op = Op::Flush {
3960 table_id: self.table_id,
3961 flushed_epoch: epoch.0,
3962 };
3963 match &mut self.wal {
3964 WalSink::Private(w) => {
3965 w.append_system(op)?;
3966 w.sync()?;
3967 }
3968 WalSink::Shared(s) => {
3969 s.wal.lock().append_system(op)?;
3974 }
3975 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3976 }
3977 self.flushed_epoch = epoch.0;
3978 if matches!(self.wal, WalSink::Private(_)) {
3979 self.rotate_wal(epoch)?;
3980 }
3981 Ok(())
3982 }
3983
3984 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3988 if self.mutable_run.is_empty() {
3989 return Ok(());
3990 }
3991 let run_id = self.alloc_run_id()?;
3992 let rows = self.mutable_run.drain_sorted();
3993 if rows.is_empty() {
3994 return Ok(());
3995 }
3996 let path = self.run_path(run_id);
3997 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3998 if let Some(kek) = &self.kek {
3999 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4000 }
4001 let header = match self.create_run_file(run_id)? {
4002 Some(file) => writer.write_file(file, &rows)?,
4003 None => writer.write(&path, &rows)?,
4004 };
4005 self.run_refs.push(RunRef {
4006 run_id: run_id as u128,
4007 level: 0,
4008 epoch_created: epoch.0,
4009 row_count: header.row_count,
4010 });
4011 Ok(())
4012 }
4013
4014 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4018 self.mutable_run_spill_bytes = bytes.max(1);
4019 }
4020
4021 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4025 self.compaction_zstd_level = level;
4026 }
4027
4028 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4032 self.result_cache.lock().set_max_bytes(max_bytes);
4033 }
4034
4035 pub(crate) fn clear_result_cache(&mut self) {
4039 self.result_cache.lock().clear();
4040 }
4041
4042 pub fn mutable_run_len(&self) -> usize {
4044 self.mutable_run.len()
4045 }
4046
4047 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4050 self.mutable_run.drain_sorted()
4051 }
4052
4053 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4055 let mut snapshot = self.mutable_run.clone();
4056 snapshot.drain_sorted()
4057 }
4058
4059 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4064 self.ensure_writable()?;
4065 let n = batch.len();
4066 if n == 0 {
4067 return Ok(self.current_epoch());
4068 }
4069 for row in &batch {
4070 self.schema.validate_values(row)?;
4071 }
4072 let epoch = self.commit_new_epoch()?;
4073 let live_before = self.live_count;
4074 self.spill_mutable_run(epoch)?;
4078 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4079 && self.indexes_complete
4080 && self.run_refs.is_empty()
4081 && self.memtable.is_empty()
4082 && self.mutable_run.is_empty();
4083 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4089 use rayon::prelude::*;
4090 self.schema
4091 .columns
4092 .par_iter()
4093 .map(|cdef| {
4094 (
4095 cdef.id,
4096 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4097 )
4098 })
4099 .collect::<Vec<_>>()
4100 };
4101 drop(batch);
4102 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4107 self.validate_columns_not_null(&user_columns, n)?;
4108 let winner_idx = self
4109 .bulk_pk_winner_indices(&user_columns, n)
4110 .filter(|idx| idx.len() != n);
4111 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4112 match winner_idx.as_deref() {
4113 Some(idx) => {
4114 let compacted = user_columns
4115 .iter()
4116 .map(|(id, c)| (*id, c.gather(idx)))
4117 .collect();
4118 (compacted, idx.len())
4119 }
4120 None => (std::mem::take(&mut user_columns), n),
4121 };
4122 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4123 let first = self.allocator.alloc_range(write_n as u64)?.0;
4124 for rid in first..first + write_n as u64 {
4125 self.reservoir.offer(rid);
4126 }
4127 let run_id = self.alloc_run_id()?;
4128 let path = self.run_path(run_id);
4129 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4130 .clean(true)
4131 .with_lz4()
4132 .with_native_endian();
4133 if let Some(kek) = &self.kek {
4134 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4135 }
4136 let header = match self.create_run_file(run_id)? {
4137 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4138 None => writer.write_native(&path, &write_columns, write_n, first)?,
4139 };
4140 self.run_refs.push(RunRef {
4141 run_id: run_id as u128,
4142 level: 0,
4143 epoch_created: epoch.0,
4144 row_count: header.row_count,
4145 });
4146 self.live_count = self.live_count.saturating_add(write_n as u64);
4147 if eager_index_build {
4148 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4149 self.index_columns_bulk(&write_columns, &row_ids);
4150 self.indexes_complete = true;
4151 self.build_learned_ranges()?;
4152 } else {
4153 self.indexes_complete = false;
4154 }
4155 self.mark_flushed(epoch)?;
4156 self.persist_manifest(epoch)?;
4157 if eager_index_build {
4158 self.checkpoint_indexes(epoch);
4159 }
4160 self.clear_result_cache();
4161 Ok(epoch)
4162 }
4163
4164 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4167 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4168 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4169 let segment_no = segment
4172 .file_stem()
4173 .and_then(|s| s.to_str())
4174 .and_then(|s| s.strip_prefix("seg-"))
4175 .and_then(|s| s.parse::<u64>().ok())
4176 .unwrap_or(0);
4177 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4178 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4179 wal.sync()?;
4180 self.wal = WalSink::Private(wal);
4181 Ok(())
4182 }
4183
4184 pub(crate) fn invalidate_pending_cache(&mut self) {
4189 self.result_cache
4190 .lock()
4191 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4192 self.pending_delete_rids.clear();
4193 self.pending_put_cols.clear();
4194 }
4195
4196 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4197 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4198 m.current_epoch = epoch.0;
4199 m.next_row_id = self.allocator.current().0;
4200 m.runs = self.run_refs.clone();
4201 m.live_count = self.live_count;
4202 m.global_idx_epoch = self.global_idx_epoch;
4203 m.flushed_epoch = self.flushed_epoch;
4204 m.retiring = self.retiring.clone();
4205 m.auto_inc_next = match self.auto_inc {
4209 Some(ai) if ai.seeded => ai.next,
4210 _ => 0,
4211 };
4212 m.ttl = self.ttl;
4213 let meta_dek = self.manifest_meta_dek();
4214 match self._root_guard.as_deref() {
4215 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4216 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4217 }
4218 Ok(())
4219 }
4220
4221 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4222 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4226 let live_count = u64::try_from(rows.len())
4227 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4228 let auto_inc = match self.auto_inc {
4229 Some(mut state) => {
4230 let maximum = self.scan_max_int64(state.column_id)?;
4231 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4232 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4233 })?;
4234 state.next = state.next.max(after_maximum).max(1);
4235 state.seeded = true;
4236 Some(state)
4237 }
4238 None => None,
4239 };
4240 Ok(RecoveryMetadataPlan {
4241 live_count,
4242 auto_inc,
4243 changed: live_count != self.live_count
4244 || auto_inc.is_some_and(|planned| {
4245 self.auto_inc.is_none_or(|current| {
4246 current.next != planned.next || current.seeded != planned.seeded
4247 })
4248 }),
4249 })
4250 }
4251
4252 pub(crate) fn apply_recovered_metadata(
4253 &mut self,
4254 plan: RecoveryMetadataPlan,
4255 epoch: Epoch,
4256 ) -> Result<()> {
4257 if !plan.changed {
4258 return Ok(());
4259 }
4260 self.live_count = plan.live_count;
4261 self.auto_inc = plan.auto_inc;
4262 self.persist_manifest(epoch)
4263 }
4264
4265 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4271 if !self.indexes_complete {
4274 return;
4275 }
4276 if self.idx_root.is_none() {
4277 if let Some(root) = self._root_guard.as_ref() {
4278 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4279 return;
4280 };
4281 self.idx_root = Some(Arc::new(idx_root));
4282 }
4283 }
4284 let snap = global_idx::IndexSnapshot {
4285 hot: &self.hot,
4286 bitmap: &self.bitmap,
4287 ann: &self.ann,
4288 fm: &self.fm,
4289 sparse: &self.sparse,
4290 minhash: &self.minhash,
4291 learned_range: &self.learned_range,
4292 };
4293 let idx_dek = self.idx_dek();
4295 let written = match self.idx_root.as_deref() {
4296 Some(root) => global_idx::write_atomic_root(
4297 root,
4298 self.table_id,
4299 epoch.0,
4300 snap,
4301 idx_dek.as_deref(),
4302 ),
4303 None => global_idx::write_atomic(
4304 &self.dir,
4305 self.table_id,
4306 epoch.0,
4307 snap,
4308 idx_dek.as_deref(),
4309 ),
4310 };
4311 if written.is_ok() {
4312 self.global_idx_epoch = epoch.0;
4313 let _ = self.persist_manifest(epoch);
4314 }
4315 }
4316
4317 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4320 self.global_idx_epoch = 0;
4321 if let Some(root) = self.idx_root.as_deref() {
4322 let _ = root.remove_file(global_idx::IDX_FILENAME);
4323 } else {
4324 global_idx::remove(&self.dir);
4325 }
4326 let _ = self.persist_manifest(self.epoch.visible());
4327 }
4328
4329 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4334 self.indexes_complete = false;
4335 self.global_idx_epoch = 0;
4336 if let Some(root) = self.idx_root.as_deref() {
4337 let _ = root.remove_file(global_idx::IDX_FILENAME);
4338 } else {
4339 global_idx::remove(&self.dir);
4340 }
4341 }
4342
4343 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4344 self.indexes_complete = true;
4345 }
4346
4347 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4353 self.durable_commit_failed = true;
4354 if let WalSink::Shared(shared) = &self.wal {
4355 shared
4356 .poisoned
4357 .store(true, std::sync::atomic::Ordering::Relaxed);
4358 }
4359 }
4360
4361 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4365 self.durable_commit_failed = true;
4366 }
4367
4368 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4371 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4372 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4373 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4374 best = Some((epoch, row));
4375 }
4376 }
4377 for rr in &self.run_refs {
4378 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4379 continue;
4380 };
4381 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4382 continue;
4383 };
4384 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4385 best = Some((epoch, row));
4386 }
4387 }
4388 let now_nanos = unix_nanos_now();
4389 match best {
4390 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4391 Some((_, r)) => Some(r),
4392 None => None,
4393 }
4394 }
4395
4396 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4400 self.visible_rows_at_time(snapshot, unix_nanos_now())
4401 }
4402
4403 #[doc(hidden)]
4406 pub fn visible_rows_controlled(
4407 &self,
4408 snapshot: Snapshot,
4409 control: &crate::ExecutionControl,
4410 ) -> Result<Vec<Row>> {
4411 let mut out = Vec::new();
4412 self.for_each_visible_row_controlled(snapshot, control, |row| {
4413 out.push(row);
4414 Ok(())
4415 })?;
4416 Ok(out)
4417 }
4418
4419 #[doc(hidden)]
4422 pub fn for_each_visible_row_controlled<F>(
4423 &self,
4424 snapshot: Snapshot,
4425 control: &crate::ExecutionControl,
4426 visit: F,
4427 ) -> Result<()>
4428 where
4429 F: FnMut(Row) -> Result<()>,
4430 {
4431 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4432 control.checkpoint()?;
4433 let memtable = self.memtable.visible_versions(snapshot.epoch);
4434 if !memtable.is_empty() {
4435 sources.push(ControlledVisibleSource::memory(memtable));
4436 }
4437 control.checkpoint()?;
4438 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4439 if !mutable.is_empty() {
4440 sources.push(ControlledVisibleSource::memory(mutable));
4441 }
4442 for run in &self.run_refs {
4443 control.checkpoint()?;
4444 let reader = self.open_reader(run.run_id)?;
4445 sources.push(ControlledVisibleSource::run(
4446 reader.into_visible_version_cursor(snapshot.epoch)?,
4447 ));
4448 }
4449 let now_nanos = unix_nanos_now();
4450 merge_controlled_visible_sources(
4451 &mut sources,
4452 control,
4453 |row| self.row_expired_at(row, now_nanos),
4454 visit,
4455 )
4456 }
4457
4458 #[doc(hidden)]
4459 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4460 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4461 let mut fold = |row: Row| {
4462 best.entry(row.row_id.0)
4463 .and_modify(|e| {
4464 if row.committed_epoch > e.0 {
4465 *e = (row.committed_epoch, row.clone());
4466 }
4467 })
4468 .or_insert_with(|| (row.committed_epoch, row));
4469 };
4470 for row in self.memtable.visible_versions(snapshot.epoch) {
4471 fold(row);
4472 }
4473 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4474 fold(row);
4475 }
4476 for rr in &self.run_refs {
4477 let mut reader = self.open_reader(rr.run_id)?;
4478 for row in reader.visible_versions(snapshot.epoch)? {
4479 fold(row);
4480 }
4481 }
4482 let mut out: Vec<Row> = best
4483 .into_values()
4484 .filter_map(|(_, r)| {
4485 if r.deleted || self.row_expired_at(&r, now_nanos) {
4486 None
4487 } else {
4488 Some(r)
4489 }
4490 })
4491 .collect();
4492 out.sort_by_key(|r| r.row_id);
4493 Ok(out)
4494 }
4495
4496 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4503 if self.ttl.is_none()
4504 && self.memtable.is_empty()
4505 && self.mutable_run.is_empty()
4506 && self.run_refs.len() == 1
4507 {
4508 let rr = self.run_refs[0].clone();
4509 let mut reader = self.open_reader(rr.run_id)?;
4510 let idxs = reader.visible_indices(snapshot.epoch)?;
4511 let mut cols = Vec::with_capacity(self.schema.columns.len());
4512 for cdef in &self.schema.columns {
4513 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4514 }
4515 return Ok(cols);
4516 }
4517 let rows = self.visible_rows(snapshot)?;
4519 let mut cols: Vec<(u16, Vec<Value>)> = self
4520 .schema
4521 .columns
4522 .iter()
4523 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4524 .collect();
4525 for r in &rows {
4526 for (cid, vec) in cols.iter_mut() {
4527 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4528 }
4529 }
4530 Ok(cols)
4531 }
4532
4533 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4535 let row_id = self.hot.get(key)?;
4536 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4537 Some(row_id)
4538 } else {
4539 None
4540 }
4541 }
4542
4543 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4548 self.query_at_with_allowed(q, self.snapshot(), None)
4549 }
4550
4551 pub fn query_controlled(
4554 &mut self,
4555 q: &crate::query::Query,
4556 control: &crate::ExecutionControl,
4557 ) -> Result<Vec<Row>> {
4558 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4559 }
4560
4561 pub fn query_at_with_allowed(
4564 &mut self,
4565 q: &crate::query::Query,
4566 snapshot: Snapshot,
4567 allowed: Option<&std::collections::HashSet<RowId>>,
4568 ) -> Result<Vec<Row>> {
4569 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4570 }
4571
4572 #[doc(hidden)]
4573 pub fn query_at_with_allowed_controlled(
4574 &mut self,
4575 q: &crate::query::Query,
4576 snapshot: Snapshot,
4577 allowed: Option<&std::collections::HashSet<RowId>>,
4578 control: &crate::ExecutionControl,
4579 ) -> Result<Vec<Row>> {
4580 self.require_select()?;
4581 self.ensure_indexes_complete_controlled(control, || true)?;
4582 self.validate_native_query(q)?;
4583 self.query_conditions_at(
4584 &q.conditions,
4585 snapshot,
4586 allowed,
4587 q.limit,
4588 q.offset,
4589 None,
4590 unix_nanos_now(),
4591 Some(control),
4592 )
4593 }
4594
4595 #[doc(hidden)]
4596 pub fn query_at_with_allowed_after(
4597 &mut self,
4598 q: &crate::query::Query,
4599 snapshot: Snapshot,
4600 allowed: Option<&std::collections::HashSet<RowId>>,
4601 after_row_id: Option<RowId>,
4602 ) -> Result<Vec<Row>> {
4603 self.query_at_with_allowed_after_at_time(
4604 q,
4605 snapshot,
4606 allowed,
4607 after_row_id,
4608 unix_nanos_now(),
4609 )
4610 }
4611
4612 #[doc(hidden)]
4613 pub fn query_at_with_allowed_after_at_time(
4614 &mut self,
4615 q: &crate::query::Query,
4616 snapshot: Snapshot,
4617 allowed: Option<&std::collections::HashSet<RowId>>,
4618 after_row_id: Option<RowId>,
4619 query_time_nanos: i64,
4620 ) -> Result<Vec<Row>> {
4621 self.require_select()?;
4622 self.ensure_indexes_complete()?;
4623 self.validate_native_query(q)?;
4624 self.query_conditions_at(
4625 &q.conditions,
4626 snapshot,
4627 allowed,
4628 q.limit,
4629 q.offset,
4630 after_row_id,
4631 query_time_nanos,
4632 None,
4633 )
4634 }
4635
4636 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4637 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4638 return Err(MongrelError::InvalidArgument(format!(
4639 "query exceeds {} conditions",
4640 crate::query::MAX_HARD_CONDITIONS
4641 )));
4642 }
4643 if let Some(limit) = q.limit {
4644 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4645 return Err(MongrelError::InvalidArgument(format!(
4646 "query limit must be between 1 and {}",
4647 crate::query::MAX_FINAL_LIMIT
4648 )));
4649 }
4650 }
4651 if q.offset > crate::query::MAX_QUERY_OFFSET {
4652 return Err(MongrelError::InvalidArgument(format!(
4653 "query offset exceeds {}",
4654 crate::query::MAX_QUERY_OFFSET
4655 )));
4656 }
4657 Ok(())
4658 }
4659
4660 #[doc(hidden)]
4663 pub fn query_all_at(
4664 &mut self,
4665 conditions: &[crate::query::Condition],
4666 snapshot: Snapshot,
4667 ) -> Result<Vec<Row>> {
4668 self.require_select()?;
4669 self.ensure_indexes_complete()?;
4670 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4671 return Err(MongrelError::InvalidArgument(format!(
4672 "query exceeds {} conditions",
4673 crate::query::MAX_HARD_CONDITIONS
4674 )));
4675 }
4676 self.query_conditions_at(
4677 conditions,
4678 snapshot,
4679 None,
4680 None,
4681 0,
4682 None,
4683 unix_nanos_now(),
4684 None,
4685 )
4686 }
4687
4688 #[allow(clippy::too_many_arguments)]
4689 fn query_conditions_at(
4690 &self,
4691 conditions: &[crate::query::Condition],
4692 snapshot: Snapshot,
4693 allowed: Option<&std::collections::HashSet<RowId>>,
4694 limit: Option<usize>,
4695 offset: usize,
4696 after_row_id: Option<RowId>,
4697 query_time_nanos: i64,
4698 control: Option<&crate::ExecutionControl>,
4699 ) -> Result<Vec<Row>> {
4700 control
4701 .map(crate::ExecutionControl::checkpoint)
4702 .transpose()?;
4703 crate::trace::QueryTrace::record(|t| {
4704 t.run_count = self.run_refs.len();
4705 t.memtable_rows = self.memtable.len();
4706 t.mutable_run_rows = self.mutable_run.len();
4707 });
4708 if conditions.is_empty() {
4712 crate::trace::QueryTrace::record(|t| {
4713 t.scan_mode = crate::trace::ScanMode::Materialized;
4714 t.row_materialized = true;
4715 });
4716 let mut rows = match control {
4717 Some(control) => self.visible_rows_controlled(snapshot, control)?,
4718 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
4719 };
4720 if let Some(allowed) = allowed {
4721 let mut filtered = Vec::with_capacity(rows.len());
4722 for (index, row) in rows.into_iter().enumerate() {
4723 if index & 255 == 0 {
4724 control
4725 .map(crate::ExecutionControl::checkpoint)
4726 .transpose()?;
4727 }
4728 if allowed.contains(&row.row_id) {
4729 filtered.push(row);
4730 }
4731 }
4732 rows = filtered;
4733 }
4734 if let Some(after_row_id) = after_row_id {
4735 rows.retain(|row| row.row_id > after_row_id);
4736 }
4737 rows.drain(..offset.min(rows.len()));
4738 if let Some(limit) = limit {
4739 rows.truncate(limit);
4740 }
4741 return Ok(rows);
4742 }
4743 crate::trace::QueryTrace::record(|t| {
4744 t.conditions_pushed = conditions.len();
4745 t.scan_mode = crate::trace::ScanMode::Materialized;
4746 t.row_materialized = true;
4747 });
4748 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
4755 ordered.sort_by_key(|c| condition_cost_rank(c));
4756 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
4757 for c in &ordered {
4758 control
4759 .map(crate::ExecutionControl::checkpoint)
4760 .transpose()?;
4761 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
4762 let empty = s.is_empty();
4763 sets.push(s);
4764 if empty {
4765 break;
4766 }
4767 }
4768 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4769 if let Some(allowed) = allowed {
4770 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
4771 }
4772 if let Some(after_row_id) = after_row_id {
4773 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
4774 rids.drain(..first);
4775 }
4776 rids.drain(..offset.min(rids.len()));
4777 if let Some(limit) = limit {
4778 rids.truncate(limit);
4779 }
4780 control
4781 .map(crate::ExecutionControl::checkpoint)
4782 .transpose()?;
4783 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
4784 }
4785
4786 pub fn retrieve(
4788 &mut self,
4789 retriever: &crate::query::Retriever,
4790 ) -> Result<Vec<crate::query::RetrieverHit>> {
4791 self.retrieve_with_allowed(retriever, None)
4792 }
4793
4794 pub fn retrieve_at(
4795 &mut self,
4796 retriever: &crate::query::Retriever,
4797 snapshot: Snapshot,
4798 allowed: Option<&std::collections::HashSet<RowId>>,
4799 ) -> Result<Vec<crate::query::RetrieverHit>> {
4800 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
4801 }
4802
4803 pub fn retrieve_with_allowed(
4806 &mut self,
4807 retriever: &crate::query::Retriever,
4808 allowed: Option<&std::collections::HashSet<RowId>>,
4809 ) -> Result<Vec<crate::query::RetrieverHit>> {
4810 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
4811 }
4812
4813 pub fn retrieve_at_with_allowed(
4814 &mut self,
4815 retriever: &crate::query::Retriever,
4816 snapshot: Snapshot,
4817 allowed: Option<&std::collections::HashSet<RowId>>,
4818 ) -> Result<Vec<crate::query::RetrieverHit>> {
4819 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
4820 }
4821
4822 pub fn retrieve_at_with_allowed_and_context(
4823 &mut self,
4824 retriever: &crate::query::Retriever,
4825 snapshot: Snapshot,
4826 allowed: Option<&std::collections::HashSet<RowId>>,
4827 context: Option<&crate::query::AiExecutionContext>,
4828 ) -> Result<Vec<crate::query::RetrieverHit>> {
4829 self.require_select()?;
4830 self.ensure_indexes_complete()?;
4831 self.validate_retriever(retriever)?;
4832 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
4833 }
4834
4835 pub fn retrieve_at_with_candidate_authorization_and_context(
4836 &mut self,
4837 retriever: &crate::query::Retriever,
4838 snapshot: Snapshot,
4839 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4840 context: Option<&crate::query::AiExecutionContext>,
4841 ) -> Result<Vec<crate::query::RetrieverHit>> {
4842 self.require_select()?;
4843 self.ensure_indexes_complete()?;
4844 self.retrieve_at_with_candidate_authorization_on_generation(
4845 retriever,
4846 snapshot,
4847 authorization,
4848 context,
4849 )
4850 }
4851
4852 #[doc(hidden)]
4853 pub fn retrieve_at_with_candidate_authorization_on_generation(
4854 &self,
4855 retriever: &crate::query::Retriever,
4856 snapshot: Snapshot,
4857 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4858 context: Option<&crate::query::AiExecutionContext>,
4859 ) -> Result<Vec<crate::query::RetrieverHit>> {
4860 self.require_select()?;
4861 self.validate_retriever(retriever)?;
4862 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
4863 }
4864
4865 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
4866 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
4867 let (column_id, k) = match retriever {
4868 Retriever::Ann {
4869 column_id,
4870 query,
4871 k,
4872 } => {
4873 let index = self.ann.get(column_id).ok_or_else(|| {
4874 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
4875 })?;
4876 if query.len() != index.dim() {
4877 return Err(MongrelError::InvalidArgument(format!(
4878 "ANN query dimension must be {}, got {}",
4879 index.dim(),
4880 query.len()
4881 )));
4882 }
4883 if query.iter().any(|value| !value.is_finite()) {
4884 return Err(MongrelError::InvalidArgument(
4885 "ANN query values must be finite".into(),
4886 ));
4887 }
4888 (*column_id, *k)
4889 }
4890 Retriever::Sparse {
4891 column_id,
4892 query,
4893 k,
4894 } => {
4895 if !self.sparse.contains_key(column_id) {
4896 return Err(MongrelError::InvalidArgument(format!(
4897 "column {column_id} has no Sparse index"
4898 )));
4899 }
4900 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
4901 return Err(MongrelError::InvalidArgument(
4902 "Sparse query must be non-empty with finite weights".into(),
4903 ));
4904 }
4905 if query.len() > MAX_SPARSE_TERMS {
4906 return Err(MongrelError::InvalidArgument(format!(
4907 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
4908 )));
4909 }
4910 (*column_id, *k)
4911 }
4912 Retriever::MinHash {
4913 column_id,
4914 members,
4915 k,
4916 } => {
4917 if !self.minhash.contains_key(column_id) {
4918 return Err(MongrelError::InvalidArgument(format!(
4919 "column {column_id} has no MinHash index"
4920 )));
4921 }
4922 if members.is_empty() {
4923 return Err(MongrelError::InvalidArgument(
4924 "MinHash members must not be empty".into(),
4925 ));
4926 }
4927 if members.len() > MAX_SET_MEMBERS {
4928 return Err(MongrelError::InvalidArgument(format!(
4929 "MinHash query exceeds {MAX_SET_MEMBERS} members"
4930 )));
4931 }
4932 let mut total_bytes = 0usize;
4933 for member in members {
4934 let bytes = member.encoded_len();
4935 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
4936 return Err(MongrelError::InvalidArgument(format!(
4937 "MinHash member exceeds {} bytes",
4938 crate::query::MAX_SET_MEMBER_BYTES
4939 )));
4940 }
4941 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
4942 MongrelError::InvalidArgument("MinHash input size overflow".into())
4943 })?;
4944 }
4945 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
4946 return Err(MongrelError::InvalidArgument(format!(
4947 "MinHash input exceeds {} bytes",
4948 crate::query::MAX_SET_INPUT_BYTES
4949 )));
4950 }
4951 (*column_id, *k)
4952 }
4953 };
4954 if k == 0 {
4955 return Err(MongrelError::InvalidArgument(
4956 "retriever k must be > 0".into(),
4957 ));
4958 }
4959 if k > MAX_RETRIEVER_K {
4960 return Err(MongrelError::InvalidArgument(format!(
4961 "retriever k exceeds {MAX_RETRIEVER_K}"
4962 )));
4963 }
4964 debug_assert!(self
4965 .schema
4966 .columns
4967 .iter()
4968 .any(|column| column.id == column_id));
4969 Ok(())
4970 }
4971
4972 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
4973 use crate::query::Condition;
4974 match condition {
4975 Condition::Ann {
4976 column_id,
4977 query,
4978 k,
4979 } => self.validate_retriever(&crate::query::Retriever::Ann {
4980 column_id: *column_id,
4981 query: query.clone(),
4982 k: *k,
4983 }),
4984 Condition::SparseMatch {
4985 column_id,
4986 query,
4987 k,
4988 } => self.validate_retriever(&crate::query::Retriever::Sparse {
4989 column_id: *column_id,
4990 query: query.clone(),
4991 k: *k,
4992 }),
4993 Condition::MinHashSimilar {
4994 column_id,
4995 query,
4996 k,
4997 } => {
4998 if !self.minhash.contains_key(column_id) {
4999 return Err(MongrelError::InvalidArgument(format!(
5000 "column {column_id} has no MinHash index"
5001 )));
5002 }
5003 if query.is_empty() || *k == 0 {
5004 return Err(MongrelError::InvalidArgument(
5005 "MinHash query must be non-empty and k must be > 0".into(),
5006 ));
5007 }
5008 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5009 {
5010 return Err(MongrelError::InvalidArgument(format!(
5011 "MinHash query must have <= {} members and k <= {}",
5012 crate::query::MAX_SET_MEMBERS,
5013 crate::query::MAX_RETRIEVER_K
5014 )));
5015 }
5016 Ok(())
5017 }
5018 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5019 Err(MongrelError::InvalidArgument(format!(
5020 "bitmap IN exceeds {} values",
5021 crate::query::MAX_SET_MEMBERS
5022 )))
5023 }
5024 Condition::FmContainsAll { patterns, .. }
5025 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5026 {
5027 Err(MongrelError::InvalidArgument(format!(
5028 "FM query exceeds {} patterns",
5029 crate::query::MAX_HARD_CONDITIONS
5030 )))
5031 }
5032 _ => Ok(()),
5033 }
5034 }
5035
5036 fn retrieve_filtered(
5037 &self,
5038 retriever: &crate::query::Retriever,
5039 snapshot: Snapshot,
5040 hard_filter: Option<&RowIdSet>,
5041 allowed: Option<&std::collections::HashSet<RowId>>,
5042 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5043 context: Option<&crate::query::AiExecutionContext>,
5044 ) -> Result<Vec<crate::query::RetrieverHit>> {
5045 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5046 let started = std::time::Instant::now();
5047 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5048 Retriever::Ann {
5049 column_id,
5050 query,
5051 k,
5052 } => {
5053 let Some(index) = self.ann.get(column_id) else {
5054 return Ok(Vec::new());
5055 };
5056 let cap = ann_candidate_cap(index.len(), context);
5057 if cap == 0 {
5058 return Ok(Vec::new());
5059 }
5060 let mut breadth = (*k).max(1).min(cap);
5061 let mut eligibility = std::collections::HashMap::new();
5062 let mut filtered = loop {
5063 let mut seen = std::collections::HashSet::new();
5064 if let Some(context) = context {
5065 context.checkpoint()?;
5066 }
5067 let raw = index.search_with_context(query, breadth, context)?;
5068 let unchecked: Vec<_> = raw
5069 .iter()
5070 .map(|(row_id, _)| *row_id)
5071 .filter(|row_id| !eligibility.contains_key(row_id))
5072 .filter(|row_id| {
5073 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5074 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5075 })
5076 .collect();
5077 let eligible = self.eligible_and_authorized_candidate_ids(
5078 &unchecked,
5079 *column_id,
5080 snapshot,
5081 candidate_authorization,
5082 context,
5083 )?;
5084 for row_id in unchecked {
5085 eligibility.insert(row_id, eligible.contains(&row_id));
5086 }
5087 let filtered: Vec<_> = raw
5088 .into_iter()
5089 .filter(|(row_id, _)| {
5090 seen.insert(*row_id)
5091 && eligibility.get(row_id).copied().unwrap_or(false)
5092 })
5093 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
5094 .collect();
5095 if filtered.len() >= *k || breadth >= cap {
5096 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5097 crate::trace::QueryTrace::record(|trace| {
5098 trace.ann_candidate_cap_hit = true;
5099 });
5100 }
5101 break filtered;
5102 }
5103 breadth = breadth.saturating_mul(2).min(cap);
5104 };
5105 filtered.truncate(*k);
5106 filtered
5107 }
5108 Retriever::Sparse {
5109 column_id,
5110 query,
5111 k,
5112 } => self
5113 .sparse
5114 .get(column_id)
5115 .map(|index| -> Result<Vec<_>> {
5116 let mut breadth = (*k).max(1);
5117 let mut eligibility = std::collections::HashMap::new();
5118 loop {
5119 if let Some(context) = context {
5120 context.checkpoint()?;
5121 }
5122 let raw = index.search_with_context(query, breadth, context)?;
5123 let unchecked: Vec<_> = raw
5124 .iter()
5125 .map(|(row_id, _)| *row_id)
5126 .filter(|row_id| !eligibility.contains_key(row_id))
5127 .filter(|row_id| {
5128 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5129 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5130 })
5131 .collect();
5132 let eligible = self.eligible_and_authorized_candidate_ids(
5133 &unchecked,
5134 *column_id,
5135 snapshot,
5136 candidate_authorization,
5137 context,
5138 )?;
5139 for row_id in unchecked {
5140 eligibility.insert(row_id, eligible.contains(&row_id));
5141 }
5142 let filtered: Vec<_> = raw
5143 .iter()
5144 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5145 .take(*k)
5146 .map(|(row_id, score)| {
5147 (*row_id, RetrieverScore::SparseDotProduct(*score))
5148 })
5149 .collect();
5150 if filtered.len() >= *k || raw.len() < breadth {
5151 break Ok(filtered);
5152 }
5153 let next = breadth.saturating_mul(2);
5154 if next == breadth {
5155 break Ok(filtered);
5156 }
5157 breadth = next;
5158 }
5159 })
5160 .transpose()?
5161 .unwrap_or_default(),
5162 Retriever::MinHash {
5163 column_id,
5164 members,
5165 k,
5166 } => self
5167 .minhash
5168 .get(column_id)
5169 .map(|index| -> Result<Vec<_>> {
5170 let mut hashes = Vec::with_capacity(members.len());
5171 for member in members {
5172 if let Some(context) = context {
5173 context.consume(crate::query::work_units(
5174 member.encoded_len(),
5175 crate::query::PARSE_WORK_QUANTUM,
5176 ))?;
5177 }
5178 hashes.push(member.hash_v1());
5179 }
5180 let mut breadth = (*k).max(1);
5181 let mut eligibility = std::collections::HashMap::new();
5182 loop {
5183 if let Some(context) = context {
5184 context.checkpoint()?;
5185 }
5186 let raw = index.search_with_context(&hashes, breadth, context)?;
5187 let unchecked: Vec<_> = raw
5188 .iter()
5189 .map(|(row_id, _)| *row_id)
5190 .filter(|row_id| !eligibility.contains_key(row_id))
5191 .filter(|row_id| {
5192 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5193 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5194 })
5195 .collect();
5196 let eligible = self.eligible_and_authorized_candidate_ids(
5197 &unchecked,
5198 *column_id,
5199 snapshot,
5200 candidate_authorization,
5201 context,
5202 )?;
5203 for row_id in unchecked {
5204 eligibility.insert(row_id, eligible.contains(&row_id));
5205 }
5206 let filtered: Vec<_> = raw
5207 .iter()
5208 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5209 .take(*k)
5210 .map(|(row_id, score)| {
5211 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5212 })
5213 .collect();
5214 if filtered.len() >= *k || raw.len() < breadth {
5215 break Ok(filtered);
5216 }
5217 let next = breadth.saturating_mul(2);
5218 if next == breadth {
5219 break Ok(filtered);
5220 }
5221 breadth = next;
5222 }
5223 })
5224 .transpose()?
5225 .unwrap_or_default(),
5226 };
5227 let elapsed = started.elapsed().as_nanos() as u64;
5228 crate::trace::QueryTrace::record(|trace| {
5229 match retriever {
5230 Retriever::Ann { .. } => {
5231 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5232 }
5233 Retriever::Sparse { .. } => {
5234 trace.sparse_candidate_nanos =
5235 trace.sparse_candidate_nanos.saturating_add(elapsed)
5236 }
5237 Retriever::MinHash { .. } => {
5238 trace.minhash_candidate_nanos =
5239 trace.minhash_candidate_nanos.saturating_add(elapsed)
5240 }
5241 }
5242 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5243 });
5244 Ok(scored
5245 .into_iter()
5246 .enumerate()
5247 .map(|(rank, (row_id, score))| RetrieverHit {
5248 row_id,
5249 rank: rank + 1,
5250 score,
5251 })
5252 .collect())
5253 }
5254
5255 fn eligible_candidate_ids(
5256 &self,
5257 candidates: &[RowId],
5258 _column_id: u16,
5259 snapshot: Snapshot,
5260 context: Option<&crate::query::AiExecutionContext>,
5261 ) -> Result<std::collections::HashSet<RowId>> {
5262 if !self.had_deletes
5263 && self.ttl.is_none()
5264 && self.pending_put_cols.is_empty()
5265 && snapshot.epoch == self.snapshot().epoch
5266 {
5267 return Ok(candidates.iter().copied().collect());
5268 }
5269 let mut readers: Vec<_> = self
5270 .run_refs
5271 .iter()
5272 .map(|run| self.open_reader(run.run_id))
5273 .collect::<Result<_>>()?;
5274 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5275 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5276 for &row_id in candidates {
5277 if let Some(context) = context {
5278 context.consume(1)?;
5279 }
5280 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5281 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5282 let overlay = match (mem, mutable) {
5283 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5284 (Some(value), None) | (None, Some(value)) => Some(value),
5285 (None, None) => None,
5286 };
5287 if let Some((_, row)) = overlay {
5288 if !row.deleted && !self.row_expired_at(&row, now) {
5289 eligible.insert(row_id);
5290 }
5291 continue;
5292 }
5293 let mut best: Option<(Epoch, bool, usize)> = None;
5294 for (index, reader) in readers.iter_mut().enumerate() {
5295 if let Some((epoch, deleted)) =
5296 reader.get_version_visibility(row_id, snapshot.epoch)?
5297 {
5298 if best
5299 .as_ref()
5300 .map(|(best_epoch, ..)| epoch > *best_epoch)
5301 .unwrap_or(true)
5302 {
5303 best = Some((epoch, deleted, index));
5304 }
5305 }
5306 }
5307 let Some((_, false, reader_index)) = best else {
5308 continue;
5309 };
5310 if let Some(ttl) = self.ttl {
5311 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5312 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5313 {
5314 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5315 continue;
5316 }
5317 }
5318 }
5319 eligible.insert(row_id);
5320 }
5321 Ok(eligible)
5322 }
5323
5324 fn eligible_and_authorized_candidate_ids(
5325 &self,
5326 candidates: &[RowId],
5327 column_id: u16,
5328 snapshot: Snapshot,
5329 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5330 context: Option<&crate::query::AiExecutionContext>,
5331 ) -> Result<std::collections::HashSet<RowId>> {
5332 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5333 let Some(authorization) = authorization else {
5334 return Ok(eligible);
5335 };
5336 let candidates: Vec<_> = eligible.into_iter().collect();
5337 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5338 }
5339
5340 fn policy_allowed_candidate_ids(
5341 &self,
5342 candidates: &[RowId],
5343 snapshot: Snapshot,
5344 authorization: &crate::security::CandidateAuthorization<'_>,
5345 context: Option<&crate::query::AiExecutionContext>,
5346 ) -> Result<std::collections::HashSet<RowId>> {
5347 let started = std::time::Instant::now();
5348 if candidates.is_empty()
5349 || authorization.principal.is_admin
5350 || !authorization.security.rls_enabled(authorization.table)
5351 {
5352 return Ok(candidates.iter().copied().collect());
5353 }
5354 if let Some(context) = context {
5355 context.checkpoint()?;
5356 }
5357 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5358 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5359 .iter()
5360 .map(|row_id| {
5361 (
5362 *row_id,
5363 Row {
5364 row_id: *row_id,
5365 committed_epoch: snapshot.epoch,
5366 columns: std::collections::HashMap::new(),
5367 deleted: false,
5368 },
5369 )
5370 })
5371 .collect();
5372 let columns = authorization
5373 .security
5374 .select_policy_columns(authorization.table, authorization.principal);
5375 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5376 let mut decoded = 0usize;
5377 for column_id in &columns {
5378 if let Some(context) = context {
5379 context.checkpoint()?;
5380 }
5381 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5382 &row_ids, *column_id, snapshot, query_now, context,
5383 )? {
5384 if let Some(row) = rows.get_mut(&row_id) {
5385 row.columns.insert(*column_id, value);
5386 decoded = decoded.saturating_add(1);
5387 }
5388 }
5389 }
5390 if let Some(context) = context {
5391 context.consume(candidates.len().saturating_add(decoded))?;
5392 }
5393 let allowed = rows
5394 .into_values()
5395 .filter_map(|row| {
5396 authorization
5397 .security
5398 .row_allowed(
5399 authorization.table,
5400 crate::security::PolicyCommand::Select,
5401 &row,
5402 authorization.principal,
5403 false,
5404 )
5405 .then_some(row.row_id)
5406 })
5407 .collect();
5408 crate::trace::QueryTrace::record(|trace| {
5409 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5410 trace.rls_policy_columns_decoded =
5411 trace.rls_policy_columns_decoded.saturating_add(decoded);
5412 trace.authorization_nanos = trace
5413 .authorization_nanos
5414 .saturating_add(started.elapsed().as_nanos() as u64);
5415 });
5416 Ok(allowed)
5417 }
5418
5419 pub fn search(
5421 &mut self,
5422 request: &crate::query::SearchRequest,
5423 ) -> Result<Vec<crate::query::SearchHit>> {
5424 self.search_with_allowed(request, None)
5425 }
5426
5427 pub fn search_at(
5428 &mut self,
5429 request: &crate::query::SearchRequest,
5430 snapshot: Snapshot,
5431 authorized: Option<&std::collections::HashSet<RowId>>,
5432 ) -> Result<Vec<crate::query::SearchHit>> {
5433 self.search_at_with_allowed(request, snapshot, authorized)
5434 }
5435
5436 pub fn search_with_allowed(
5437 &mut self,
5438 request: &crate::query::SearchRequest,
5439 authorized: Option<&std::collections::HashSet<RowId>>,
5440 ) -> Result<Vec<crate::query::SearchHit>> {
5441 self.search_at_with_allowed(request, self.snapshot(), authorized)
5442 }
5443
5444 pub fn search_at_with_allowed(
5445 &mut self,
5446 request: &crate::query::SearchRequest,
5447 snapshot: Snapshot,
5448 authorized: Option<&std::collections::HashSet<RowId>>,
5449 ) -> Result<Vec<crate::query::SearchHit>> {
5450 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5451 }
5452
5453 pub fn search_at_with_allowed_and_context(
5454 &mut self,
5455 request: &crate::query::SearchRequest,
5456 snapshot: Snapshot,
5457 authorized: Option<&std::collections::HashSet<RowId>>,
5458 context: Option<&crate::query::AiExecutionContext>,
5459 ) -> Result<Vec<crate::query::SearchHit>> {
5460 self.ensure_indexes_complete()?;
5461 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5462 }
5463
5464 pub fn search_at_with_candidate_authorization_and_context(
5465 &mut self,
5466 request: &crate::query::SearchRequest,
5467 snapshot: Snapshot,
5468 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5469 context: Option<&crate::query::AiExecutionContext>,
5470 ) -> Result<Vec<crate::query::SearchHit>> {
5471 self.ensure_indexes_complete()?;
5472 self.search_at_with_filters_and_context(
5473 request,
5474 snapshot,
5475 None,
5476 authorization,
5477 context,
5478 None,
5479 )
5480 }
5481
5482 #[doc(hidden)]
5483 pub fn search_at_with_candidate_authorization_on_generation(
5484 &self,
5485 request: &crate::query::SearchRequest,
5486 snapshot: Snapshot,
5487 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5488 context: Option<&crate::query::AiExecutionContext>,
5489 ) -> Result<Vec<crate::query::SearchHit>> {
5490 self.search_at_with_filters_and_context(
5491 request,
5492 snapshot,
5493 None,
5494 authorization,
5495 context,
5496 None,
5497 )
5498 }
5499
5500 #[doc(hidden)]
5501 pub fn search_at_with_candidate_authorization_on_generation_after(
5502 &self,
5503 request: &crate::query::SearchRequest,
5504 snapshot: Snapshot,
5505 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5506 context: Option<&crate::query::AiExecutionContext>,
5507 after: Option<crate::query::SearchAfter>,
5508 ) -> Result<Vec<crate::query::SearchHit>> {
5509 self.search_at_with_filters_and_context(
5510 request,
5511 snapshot,
5512 None,
5513 authorization,
5514 context,
5515 after,
5516 )
5517 }
5518
5519 fn search_at_with_filters_and_context(
5520 &self,
5521 request: &crate::query::SearchRequest,
5522 snapshot: Snapshot,
5523 authorized: Option<&std::collections::HashSet<RowId>>,
5524 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5525 context: Option<&crate::query::AiExecutionContext>,
5526 after: Option<crate::query::SearchAfter>,
5527 ) -> Result<Vec<crate::query::SearchHit>> {
5528 use crate::query::{
5529 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5530 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5531 };
5532 let total_started = std::time::Instant::now();
5533 let rank_offset = after.map_or(0, |after| after.returned_count);
5534 self.require_select()?;
5535 if request.limit == 0 {
5536 return Err(MongrelError::InvalidArgument(
5537 "search limit must be > 0".into(),
5538 ));
5539 }
5540 if request.limit > MAX_FINAL_LIMIT {
5541 return Err(MongrelError::InvalidArgument(format!(
5542 "search limit exceeds {MAX_FINAL_LIMIT}"
5543 )));
5544 }
5545 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5546 return Err(MongrelError::InvalidArgument(
5547 "search-after score must be finite".into(),
5548 ));
5549 }
5550 if request.retrievers.is_empty() {
5551 return Err(MongrelError::InvalidArgument(
5552 "search requires at least one retriever".into(),
5553 ));
5554 }
5555 if request.retrievers.len() > MAX_RETRIEVERS {
5556 return Err(MongrelError::InvalidArgument(format!(
5557 "search exceeds {MAX_RETRIEVERS} retrievers"
5558 )));
5559 }
5560 if request.must.len() > MAX_HARD_CONDITIONS {
5561 return Err(MongrelError::InvalidArgument(format!(
5562 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5563 )));
5564 }
5565 for condition in &request.must {
5566 self.validate_condition(condition)?;
5567 }
5568 if request.must.iter().any(|condition| {
5569 matches!(
5570 condition,
5571 Condition::Ann { .. }
5572 | Condition::SparseMatch { .. }
5573 | Condition::MinHashSimilar { .. }
5574 )
5575 }) {
5576 return Err(MongrelError::InvalidArgument(
5577 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5578 .into(),
5579 ));
5580 }
5581 let mut names = std::collections::HashSet::new();
5582 for named in &request.retrievers {
5583 if named.name.is_empty()
5584 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5585 || !names.insert(named.name.as_str())
5586 {
5587 return Err(MongrelError::InvalidArgument(format!(
5588 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5589 crate::query::MAX_RETRIEVER_NAME_BYTES
5590 )));
5591 }
5592 if !named.weight.is_finite()
5593 || named.weight < 0.0
5594 || named.weight > MAX_RETRIEVER_WEIGHT
5595 {
5596 return Err(MongrelError::InvalidArgument(format!(
5597 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5598 )));
5599 }
5600 self.validate_retriever(&named.retriever)?;
5601 }
5602 let projection = request
5603 .projection
5604 .clone()
5605 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5606 if projection.len() > MAX_PROJECTION_COLUMNS {
5607 return Err(MongrelError::InvalidArgument(format!(
5608 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5609 )));
5610 }
5611 for column_id in &projection {
5612 if !self
5613 .schema
5614 .columns
5615 .iter()
5616 .any(|column| column.id == *column_id)
5617 {
5618 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5619 }
5620 }
5621 if let Some(crate::query::Rerank::ExactVector {
5622 embedding_column,
5623 query,
5624 candidate_limit,
5625 weight,
5626 ..
5627 }) = &request.rerank
5628 {
5629 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5630 {
5631 return Err(MongrelError::InvalidArgument(format!(
5632 "rerank candidate_limit must be between search limit and {}",
5633 crate::query::MAX_RETRIEVER_K
5634 )));
5635 }
5636 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5637 return Err(MongrelError::InvalidArgument(format!(
5638 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5639 )));
5640 }
5641 let column = self
5642 .schema
5643 .columns
5644 .iter()
5645 .find(|column| column.id == *embedding_column)
5646 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5647 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5648 return Err(MongrelError::InvalidArgument(format!(
5649 "rerank column {embedding_column} is not an embedding"
5650 )));
5651 };
5652 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5653 return Err(MongrelError::InvalidArgument(format!(
5654 "rerank query must contain {dim} finite values"
5655 )));
5656 }
5657 }
5658
5659 let hard_filter_started = std::time::Instant::now();
5660 let hard_filter = if request.must.is_empty() {
5661 None
5662 } else {
5663 let mut sets = Vec::with_capacity(request.must.len());
5664 for condition in &request.must {
5665 if let Some(context) = context {
5666 context.checkpoint()?;
5667 }
5668 sets.push(self.resolve_condition(condition, snapshot)?);
5669 }
5670 Some(RowIdSet::intersect_many(sets))
5671 };
5672 crate::trace::QueryTrace::record(|trace| {
5673 trace.hard_filter_nanos = trace
5674 .hard_filter_nanos
5675 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
5676 });
5677 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
5678 return Ok(Vec::new());
5679 }
5680
5681 let constant = match request.fusion {
5682 Fusion::ReciprocalRank { constant } => constant,
5683 };
5684 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
5685 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
5686 let mut fusion_nanos = 0u64;
5687 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
5688 std::collections::HashMap::new();
5689 for named in retrievers {
5690 if named.weight == 0.0 {
5691 continue;
5692 }
5693 if let Some(context) = context {
5694 context.checkpoint()?;
5695 }
5696 let hits = self.retrieve_filtered(
5697 &named.retriever,
5698 snapshot,
5699 hard_filter.as_ref(),
5700 authorized,
5701 candidate_authorization,
5702 context,
5703 )?;
5704 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
5705 let fusion_started = std::time::Instant::now();
5706 for hit in hits {
5707 if let Some(context) = context {
5708 context.consume(1)?;
5709 }
5710 let contribution = named.weight / (constant as f64 + hit.rank as f64);
5711 if !contribution.is_finite() {
5712 return Err(MongrelError::InvalidArgument(
5713 "retriever contribution must be finite".into(),
5714 ));
5715 }
5716 let max_fused_candidates = context.map_or(
5717 crate::query::MAX_FUSED_CANDIDATES,
5718 crate::query::AiExecutionContext::max_fused_candidates,
5719 );
5720 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
5721 return Err(MongrelError::WorkBudgetExceeded);
5722 }
5723 let entry = fused.entry(hit.row_id).or_default();
5724 entry.0 += contribution;
5725 if !entry.0.is_finite() {
5726 return Err(MongrelError::InvalidArgument(
5727 "fused score must be finite".into(),
5728 ));
5729 }
5730 entry.1.push(ComponentScore {
5731 retriever_name: retriever_name.clone(),
5732 rank: hit.rank,
5733 raw_score: hit.score,
5734 contribution,
5735 });
5736 }
5737 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
5738 }
5739 let union_size = fused.len();
5740 let mut ranked: Vec<_> = fused
5741 .into_iter()
5742 .map(|(row_id, (fused_score, components))| {
5743 (row_id, fused_score, components, None, fused_score)
5744 })
5745 .collect();
5746 let order = |(a_row, _, _, _, a_score): &(
5747 RowId,
5748 f64,
5749 Vec<ComponentScore>,
5750 Option<f32>,
5751 f64,
5752 ),
5753 (b_row, _, _, _, b_score): &(
5754 RowId,
5755 f64,
5756 Vec<ComponentScore>,
5757 Option<f32>,
5758 f64,
5759 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
5760 if let Some(crate::query::Rerank::ExactVector {
5761 embedding_column,
5762 query,
5763 metric,
5764 candidate_limit,
5765 weight,
5766 }) = &request.rerank
5767 {
5768 let fused_order = |(a_row, a_score, ..): &(
5769 RowId,
5770 f64,
5771 Vec<ComponentScore>,
5772 Option<f32>,
5773 f64,
5774 ),
5775 (b_row, b_score, ..): &(
5776 RowId,
5777 f64,
5778 Vec<ComponentScore>,
5779 Option<f32>,
5780 f64,
5781 )| {
5782 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
5783 };
5784 let selection_started = std::time::Instant::now();
5785 if let Some(context) = context {
5786 context.consume(ranked.len())?;
5787 }
5788 if ranked.len() > *candidate_limit {
5789 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
5790 ranked.truncate(*candidate_limit);
5791 }
5792 ranked.sort_by(fused_order);
5793 fusion_nanos =
5794 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5795 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5796 if let Some(context) = context {
5797 context.consume(row_ids.len())?;
5798 }
5799 let query_now =
5800 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5801 let gather_started = std::time::Instant::now();
5802 let vectors = self.values_for_rids_batch_at_with_context(
5803 &row_ids,
5804 *embedding_column,
5805 snapshot,
5806 query_now,
5807 context,
5808 )?;
5809 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5810 let vector_work =
5811 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
5812 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
5813 if let Some(context) = context {
5814 context.consume(vector_work)?;
5815 }
5816 query
5817 .iter()
5818 .map(|value| f64::from(*value).powi(2))
5819 .sum::<f64>()
5820 .sqrt()
5821 } else {
5822 0.0
5823 };
5824 let score_started = std::time::Instant::now();
5825 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
5826 for (row_id, value) in vectors {
5827 let Value::Embedding(vector) = value else {
5828 continue;
5829 };
5830 let score = match metric {
5831 crate::query::VectorMetric::DotProduct => {
5832 if let Some(context) = context {
5833 context.consume(vector_work)?;
5834 }
5835 query
5836 .iter()
5837 .zip(&vector)
5838 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5839 .sum::<f64>()
5840 }
5841 crate::query::VectorMetric::Cosine => {
5842 if let Some(context) = context {
5843 context.consume(vector_work.saturating_mul(2))?;
5844 }
5845 let dot = query
5846 .iter()
5847 .zip(&vector)
5848 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5849 .sum::<f64>();
5850 let norm = vector
5851 .iter()
5852 .map(|value| f64::from(*value).powi(2))
5853 .sum::<f64>()
5854 .sqrt();
5855 if query_norm == 0.0 || norm == 0.0 {
5856 0.0
5857 } else {
5858 dot / (query_norm * norm)
5859 }
5860 }
5861 crate::query::VectorMetric::Euclidean => {
5862 if let Some(context) = context {
5863 context.consume(vector_work)?;
5864 }
5865 query
5866 .iter()
5867 .zip(&vector)
5868 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
5869 .sum::<f64>()
5870 .sqrt()
5871 }
5872 };
5873 if !score.is_finite() {
5874 return Err(MongrelError::InvalidArgument(
5875 "exact rerank score must be finite".into(),
5876 ));
5877 }
5878 scores.insert(row_id, score as f32);
5879 }
5880 let mut reranked = Vec::with_capacity(ranked.len());
5881 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
5882 let Some(score) = scores.get(&row_id).copied() else {
5883 continue;
5884 };
5885 let ordering_score = match metric {
5886 crate::query::VectorMetric::Euclidean => -f64::from(score),
5887 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
5888 f64::from(score)
5889 }
5890 };
5891 let final_score = fused_score + *weight * ordering_score;
5892 if !final_score.is_finite() {
5893 return Err(MongrelError::InvalidArgument(
5894 "final rerank score must be finite".into(),
5895 ));
5896 }
5897 reranked.push((row_id, fused_score, components, Some(score), final_score));
5898 }
5899 ranked = reranked;
5900 ranked.sort_by(order);
5901 crate::trace::QueryTrace::record(|trace| {
5902 trace.exact_vector_gather_nanos =
5903 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
5904 trace.exact_vector_score_nanos = trace
5905 .exact_vector_score_nanos
5906 .saturating_add(score_started.elapsed().as_nanos() as u64);
5907 });
5908 }
5909 if let Some(after) = after {
5910 ranked.retain(|(row_id, _, _, _, final_score)| {
5911 final_score.total_cmp(&after.final_score).is_lt()
5912 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
5913 });
5914 }
5915 let projection_started = std::time::Instant::now();
5916 let sentinel = projection
5917 .first()
5918 .copied()
5919 .or_else(|| self.schema.columns.first().map(|column| column.id));
5920 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5921 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
5922 let mut projection_rows = 0usize;
5923 let mut projection_cells = 0usize;
5924 while out.len() < request.limit && !ranked.is_empty() {
5925 if let Some(context) = context {
5926 context.checkpoint()?;
5927 context.consume(ranked.len())?;
5928 }
5929 let needed = request.limit - out.len();
5930 let window_size = ranked
5931 .len()
5932 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
5933 let selection_started = std::time::Instant::now();
5934 let mut remainder = if ranked.len() > window_size {
5935 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
5936 ranked.split_off(window_size)
5937 } else {
5938 Vec::new()
5939 };
5940 ranked.sort_by(order);
5941 fusion_nanos =
5942 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5943 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5944 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
5945 if let Some(context) = context {
5946 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
5947 }
5948 projection_rows = projection_rows.saturating_add(row_ids.len());
5949 projection_cells =
5950 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
5951 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
5952 std::collections::HashMap::new();
5953 if let Some(column_id) = sentinel {
5954 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5955 &row_ids, column_id, snapshot, query_now, context,
5956 )? {
5957 cells.entry(row_id).or_default().insert(column_id, value);
5958 }
5959 }
5960 for &column_id in &projection {
5961 if Some(column_id) == sentinel {
5962 continue;
5963 }
5964 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5965 &row_ids, column_id, snapshot, query_now, context,
5966 )? {
5967 cells.entry(row_id).or_default().insert(column_id, value);
5968 }
5969 }
5970 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
5971 ranked.drain(..)
5972 {
5973 let Some(row_cells) = cells.remove(&row_id) else {
5974 continue;
5975 };
5976 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
5977 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
5978 out.push(SearchHit {
5979 row_id,
5980 cells: projection
5981 .iter()
5982 .filter_map(|column_id| {
5983 row_cells
5984 .get(column_id)
5985 .cloned()
5986 .map(|value| (*column_id, value))
5987 })
5988 .collect(),
5989 components,
5990 fused_score,
5991 exact_rerank_score,
5992 final_score,
5993 final_rank,
5994 });
5995 if out.len() == request.limit {
5996 break;
5997 }
5998 }
5999 ranked.append(&mut remainder);
6000 }
6001 crate::trace::QueryTrace::record(|trace| {
6002 trace.union_size = union_size;
6003 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6004 trace.projection_nanos = trace
6005 .projection_nanos
6006 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6007 trace.total_nanos = trace
6008 .total_nanos
6009 .saturating_add(total_started.elapsed().as_nanos() as u64);
6010 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6011 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6012 if let Some(context) = context {
6013 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6014 }
6015 });
6016 Ok(out)
6017 }
6018
6019 pub fn set_similarity(
6022 &mut self,
6023 request: &crate::query::SetSimilarityRequest,
6024 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6025 self.set_similarity_with_allowed(request, None)
6026 }
6027
6028 pub fn set_similarity_at(
6029 &mut self,
6030 request: &crate::query::SetSimilarityRequest,
6031 snapshot: Snapshot,
6032 allowed: Option<&std::collections::HashSet<RowId>>,
6033 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6034 self.set_similarity_explained_at(request, snapshot, allowed)
6035 .map(|(hits, _)| hits)
6036 }
6037
6038 pub fn ann_rerank(
6040 &mut self,
6041 request: &crate::query::AnnRerankRequest,
6042 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6043 self.ann_rerank_with_allowed(request, None)
6044 }
6045
6046 pub fn ann_rerank_with_allowed(
6047 &mut self,
6048 request: &crate::query::AnnRerankRequest,
6049 allowed: Option<&std::collections::HashSet<RowId>>,
6050 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6051 self.ann_rerank_at(request, self.snapshot(), allowed)
6052 }
6053
6054 pub fn ann_rerank_at(
6055 &mut self,
6056 request: &crate::query::AnnRerankRequest,
6057 snapshot: Snapshot,
6058 allowed: Option<&std::collections::HashSet<RowId>>,
6059 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6060 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6061 }
6062
6063 pub fn ann_rerank_at_with_context(
6064 &mut self,
6065 request: &crate::query::AnnRerankRequest,
6066 snapshot: Snapshot,
6067 allowed: Option<&std::collections::HashSet<RowId>>,
6068 context: Option<&crate::query::AiExecutionContext>,
6069 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6070 self.ensure_indexes_complete()?;
6071 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6072 }
6073
6074 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6075 &mut self,
6076 request: &crate::query::AnnRerankRequest,
6077 snapshot: Snapshot,
6078 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6079 context: Option<&crate::query::AiExecutionContext>,
6080 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6081 self.ensure_indexes_complete()?;
6082 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6083 }
6084
6085 #[doc(hidden)]
6086 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6087 &self,
6088 request: &crate::query::AnnRerankRequest,
6089 snapshot: Snapshot,
6090 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6091 context: Option<&crate::query::AiExecutionContext>,
6092 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6093 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6094 }
6095
6096 fn ann_rerank_at_with_filters_and_context(
6097 &self,
6098 request: &crate::query::AnnRerankRequest,
6099 snapshot: Snapshot,
6100 allowed: Option<&std::collections::HashSet<RowId>>,
6101 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6102 context: Option<&crate::query::AiExecutionContext>,
6103 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6104 use crate::query::{
6105 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6106 };
6107 if request.candidate_k == 0 || request.limit == 0 {
6108 return Err(MongrelError::InvalidArgument(
6109 "candidate_k and limit must be > 0".into(),
6110 ));
6111 }
6112 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6113 return Err(MongrelError::InvalidArgument(format!(
6114 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6115 )));
6116 }
6117 let retriever = Retriever::Ann {
6118 column_id: request.column_id,
6119 query: request.query.clone(),
6120 k: request.candidate_k,
6121 };
6122 self.require_select()?;
6123 self.validate_retriever(&retriever)?;
6124 let hits = self.retrieve_filtered(
6125 &retriever,
6126 snapshot,
6127 None,
6128 allowed,
6129 candidate_authorization,
6130 context,
6131 )?;
6132 let distances: std::collections::HashMap<_, _> = hits
6133 .iter()
6134 .filter_map(|hit| match hit.score {
6135 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
6136 _ => None,
6137 })
6138 .collect();
6139 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6140 if let Some(context) = context {
6141 context.consume(row_ids.len())?;
6142 }
6143 let gather_started = std::time::Instant::now();
6144 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6145 let values = self.values_for_rids_batch_at_with_context(
6146 &row_ids,
6147 request.column_id,
6148 snapshot,
6149 query_now,
6150 context,
6151 )?;
6152 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6153 let score_started = std::time::Instant::now();
6154 let vector_work =
6155 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6156 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6157 if let Some(context) = context {
6158 context.consume(vector_work)?;
6159 }
6160 request
6161 .query
6162 .iter()
6163 .map(|value| f64::from(*value).powi(2))
6164 .sum::<f64>()
6165 .sqrt()
6166 } else {
6167 0.0
6168 };
6169 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6170 for (row_id, value) in values {
6171 let Value::Embedding(vector) = value else {
6172 continue;
6173 };
6174 let exact_score = match request.metric {
6175 VectorMetric::DotProduct => {
6176 if let Some(context) = context {
6177 context.consume(vector_work)?;
6178 }
6179 request
6180 .query
6181 .iter()
6182 .zip(&vector)
6183 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6184 .sum::<f64>()
6185 }
6186 VectorMetric::Cosine => {
6187 if let Some(context) = context {
6188 context.consume(vector_work.saturating_mul(2))?;
6189 }
6190 let dot = request
6191 .query
6192 .iter()
6193 .zip(&vector)
6194 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6195 .sum::<f64>();
6196 let norm = vector
6197 .iter()
6198 .map(|value| f64::from(*value).powi(2))
6199 .sum::<f64>()
6200 .sqrt();
6201 if query_norm == 0.0 || norm == 0.0 {
6202 0.0
6203 } else {
6204 dot / (query_norm * norm)
6205 }
6206 }
6207 VectorMetric::Euclidean => {
6208 if let Some(context) = context {
6209 context.consume(vector_work)?;
6210 }
6211 request
6212 .query
6213 .iter()
6214 .zip(&vector)
6215 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6216 .sum::<f64>()
6217 .sqrt()
6218 }
6219 };
6220 let exact_score = exact_score as f32;
6221 if !exact_score.is_finite() {
6222 return Err(MongrelError::InvalidArgument(
6223 "exact ANN score must be finite".into(),
6224 ));
6225 }
6226 reranked.push(AnnRerankHit {
6227 row_id,
6228 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
6229 exact_score,
6230 });
6231 }
6232 reranked.sort_by(|left, right| {
6233 let score = match request.metric {
6234 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6235 VectorMetric::Cosine | VectorMetric::DotProduct => {
6236 right.exact_score.total_cmp(&left.exact_score)
6237 }
6238 };
6239 score.then_with(|| left.row_id.cmp(&right.row_id))
6240 });
6241 reranked.truncate(request.limit);
6242 crate::trace::QueryTrace::record(|trace| {
6243 trace.exact_vector_gather_nanos =
6244 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6245 trace.exact_vector_score_nanos = trace
6246 .exact_vector_score_nanos
6247 .saturating_add(score_started.elapsed().as_nanos() as u64);
6248 });
6249 Ok(reranked)
6250 }
6251
6252 pub fn set_similarity_with_allowed(
6253 &mut self,
6254 request: &crate::query::SetSimilarityRequest,
6255 allowed: Option<&std::collections::HashSet<RowId>>,
6256 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6257 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6258 .map(|(hits, _)| hits)
6259 }
6260
6261 pub fn set_similarity_explained(
6262 &mut self,
6263 request: &crate::query::SetSimilarityRequest,
6264 ) -> Result<(
6265 Vec<crate::query::SetSimilarityHit>,
6266 crate::query::SetSimilarityTrace,
6267 )> {
6268 self.set_similarity_explained_at(request, self.snapshot(), None)
6269 }
6270
6271 fn set_similarity_explained_at(
6272 &mut self,
6273 request: &crate::query::SetSimilarityRequest,
6274 snapshot: Snapshot,
6275 allowed: Option<&std::collections::HashSet<RowId>>,
6276 ) -> Result<(
6277 Vec<crate::query::SetSimilarityHit>,
6278 crate::query::SetSimilarityTrace,
6279 )> {
6280 self.ensure_indexes_complete()?;
6281 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6282 }
6283
6284 pub fn set_similarity_at_with_context(
6285 &mut self,
6286 request: &crate::query::SetSimilarityRequest,
6287 snapshot: Snapshot,
6288 allowed: Option<&std::collections::HashSet<RowId>>,
6289 context: Option<&crate::query::AiExecutionContext>,
6290 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6291 self.ensure_indexes_complete()?;
6292 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6293 .map(|(hits, _)| hits)
6294 }
6295
6296 pub fn set_similarity_at_with_candidate_authorization_and_context(
6297 &mut self,
6298 request: &crate::query::SetSimilarityRequest,
6299 snapshot: Snapshot,
6300 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6301 context: Option<&crate::query::AiExecutionContext>,
6302 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6303 self.ensure_indexes_complete()?;
6304 self.set_similarity_explained_at_with_context(
6305 request,
6306 snapshot,
6307 None,
6308 authorization,
6309 context,
6310 )
6311 .map(|(hits, _)| hits)
6312 }
6313
6314 #[doc(hidden)]
6315 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6316 &self,
6317 request: &crate::query::SetSimilarityRequest,
6318 snapshot: Snapshot,
6319 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6320 context: Option<&crate::query::AiExecutionContext>,
6321 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6322 self.set_similarity_explained_at_with_context(
6323 request,
6324 snapshot,
6325 None,
6326 authorization,
6327 context,
6328 )
6329 .map(|(hits, _)| hits)
6330 }
6331
6332 fn set_similarity_explained_at_with_context(
6333 &self,
6334 request: &crate::query::SetSimilarityRequest,
6335 snapshot: Snapshot,
6336 allowed: Option<&std::collections::HashSet<RowId>>,
6337 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6338 context: Option<&crate::query::AiExecutionContext>,
6339 ) -> Result<(
6340 Vec<crate::query::SetSimilarityHit>,
6341 crate::query::SetSimilarityTrace,
6342 )> {
6343 use crate::query::{
6344 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6345 MAX_SET_MEMBERS,
6346 };
6347 let mut trace = crate::query::SetSimilarityTrace::default();
6348 if request.members.is_empty() {
6349 return Ok((Vec::new(), trace));
6350 }
6351 if request.candidate_k == 0 || request.limit == 0 {
6352 return Err(MongrelError::InvalidArgument(
6353 "candidate_k and limit must be > 0".into(),
6354 ));
6355 }
6356 if request.candidate_k > MAX_RETRIEVER_K
6357 || request.limit > MAX_FINAL_LIMIT
6358 || request.members.len() > MAX_SET_MEMBERS
6359 {
6360 return Err(MongrelError::InvalidArgument(format!(
6361 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6362 )));
6363 }
6364 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6365 return Err(MongrelError::InvalidArgument(
6366 "min_jaccard must be finite and between 0 and 1".into(),
6367 ));
6368 }
6369 let started = std::time::Instant::now();
6370 let retriever = Retriever::MinHash {
6371 column_id: request.column_id,
6372 members: request.members.clone(),
6373 k: request.candidate_k,
6374 };
6375 self.require_select()?;
6376 self.validate_retriever(&retriever)?;
6377 let hits = self.retrieve_filtered(
6378 &retriever,
6379 snapshot,
6380 None,
6381 allowed,
6382 candidate_authorization,
6383 context,
6384 )?;
6385 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6386 trace.candidate_count = hits.len();
6387 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6388 if let Some(context) = context {
6389 context.consume(row_ids.len())?;
6390 }
6391 let started = std::time::Instant::now();
6392 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6393 let values = self.values_for_rids_batch_at_with_context(
6394 &row_ids,
6395 request.column_id,
6396 snapshot,
6397 query_now,
6398 context,
6399 )?;
6400 trace.gather_us = started.elapsed().as_micros() as u64;
6401 if let Some(context) = context {
6402 context.consume(request.members.len())?;
6403 }
6404 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6405 let estimates: std::collections::HashMap<_, _> = hits
6406 .into_iter()
6407 .filter_map(|hit| match hit.score {
6408 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6409 _ => None,
6410 })
6411 .collect();
6412 let started = std::time::Instant::now();
6413 let mut parsed = Vec::with_capacity(values.len());
6414 for (row_id, value) in values {
6415 let Value::Bytes(bytes) = value else {
6416 continue;
6417 };
6418 if let Some(context) = context {
6419 context.consume(crate::query::work_units(
6420 bytes.len(),
6421 crate::query::PARSE_WORK_QUANTUM,
6422 ))?;
6423 }
6424 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6425 continue;
6426 };
6427 if let Some(context) = context {
6428 context.consume(members.len())?;
6429 }
6430 let stored = members
6431 .into_iter()
6432 .filter_map(|member| match member {
6433 serde_json::Value::String(value) => {
6434 Some(crate::query::SetMember::String(value))
6435 }
6436 serde_json::Value::Number(value) => {
6437 Some(crate::query::SetMember::Number(value))
6438 }
6439 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6440 _ => None,
6441 })
6442 .collect::<std::collections::HashSet<_>>();
6443 parsed.push((row_id, stored));
6444 }
6445 trace.parse_us = started.elapsed().as_micros() as u64;
6446 trace.verified_count = parsed.len();
6447 let started = std::time::Instant::now();
6448 let mut exact = Vec::new();
6449 for (row_id, stored) in parsed {
6450 if let Some(context) = context {
6451 context.consume(query.len().saturating_add(stored.len()))?;
6452 }
6453 let union = query.union(&stored).count();
6454 let score = if union == 0 {
6455 1.0
6456 } else {
6457 query.intersection(&stored).count() as f32 / union as f32
6458 };
6459 if score >= request.min_jaccard {
6460 exact.push(SetSimilarityHit {
6461 row_id,
6462 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6463 exact_jaccard: score,
6464 });
6465 }
6466 }
6467 exact.sort_by(|a, b| {
6468 b.exact_jaccard
6469 .total_cmp(&a.exact_jaccard)
6470 .then_with(|| a.row_id.cmp(&b.row_id))
6471 });
6472 exact.truncate(request.limit);
6473 trace.score_us = started.elapsed().as_micros() as u64;
6474 crate::trace::QueryTrace::record(|query_trace| {
6475 query_trace.exact_set_gather_nanos = query_trace
6476 .exact_set_gather_nanos
6477 .saturating_add(trace.gather_us.saturating_mul(1_000));
6478 query_trace.exact_set_parse_nanos = query_trace
6479 .exact_set_parse_nanos
6480 .saturating_add(trace.parse_us.saturating_mul(1_000));
6481 query_trace.exact_set_score_nanos = query_trace
6482 .exact_set_score_nanos
6483 .saturating_add(trace.score_us.saturating_mul(1_000));
6484 });
6485 Ok((exact, trace))
6486 }
6487
6488 fn values_for_rids_batch_at(
6490 &self,
6491 row_ids: &[u64],
6492 column_id: u16,
6493 snapshot: Snapshot,
6494 now: i64,
6495 ) -> Result<Vec<(RowId, Value)>> {
6496 if self.ttl.is_none()
6497 && self.memtable.is_empty()
6498 && self.mutable_run.is_empty()
6499 && self.run_refs.len() == 1
6500 {
6501 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6502 if row_ids.len().saturating_mul(24) < reader.row_count() {
6507 let mut values = Vec::with_capacity(row_ids.len());
6508 for &raw_row_id in row_ids {
6509 let row_id = RowId(raw_row_id);
6510 if let Some((_, false, Some(value))) =
6511 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6512 {
6513 values.push((row_id, value));
6514 }
6515 }
6516 return Ok(values);
6517 }
6518 let (positions, visible_row_ids) =
6519 reader.visible_positions_with_rids(snapshot.epoch)?;
6520 let requested: Vec<(RowId, usize)> = row_ids
6521 .iter()
6522 .filter_map(|raw| {
6523 visible_row_ids
6524 .binary_search(&(*raw as i64))
6525 .ok()
6526 .map(|index| (RowId(*raw), positions[index]))
6527 })
6528 .collect();
6529 let values = reader.gather_column(
6530 column_id,
6531 &requested
6532 .iter()
6533 .map(|(_, position)| *position)
6534 .collect::<Vec<_>>(),
6535 )?;
6536 return Ok(requested
6537 .into_iter()
6538 .zip(values)
6539 .map(|((row_id, _), value)| (row_id, value))
6540 .collect());
6541 }
6542 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6543 }
6544
6545 fn values_for_rids_batch_at_with_context(
6546 &self,
6547 row_ids: &[u64],
6548 column_id: u16,
6549 snapshot: Snapshot,
6550 now: i64,
6551 context: Option<&crate::query::AiExecutionContext>,
6552 ) -> Result<Vec<(RowId, Value)>> {
6553 let Some(context) = context else {
6554 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6555 };
6556 let mut values = Vec::with_capacity(row_ids.len());
6557 for chunk in row_ids.chunks(256) {
6558 context.checkpoint()?;
6559 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6560 }
6561 Ok(values)
6562 }
6563
6564 fn values_for_rids_at(
6566 &self,
6567 row_ids: &[u64],
6568 column_id: u16,
6569 snapshot: Snapshot,
6570 now: i64,
6571 ) -> Result<Vec<(RowId, Value)>> {
6572 let mut readers: Vec<_> = self
6573 .run_refs
6574 .iter()
6575 .map(|run| self.open_reader(run.run_id))
6576 .collect::<Result<_>>()?;
6577 let mut out = Vec::with_capacity(row_ids.len());
6578 for &raw_row_id in row_ids {
6579 let row_id = RowId(raw_row_id);
6580 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6581 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6582 let overlay = match (mem, mutable) {
6583 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6584 (a_epoch, a)
6585 } else {
6586 (b_epoch, b)
6587 }),
6588 (Some(value), None) | (None, Some(value)) => Some(value),
6589 (None, None) => None,
6590 };
6591 if let Some((_, row)) = overlay {
6592 if !row.deleted && !self.row_expired_at(&row, now) {
6593 if let Some(value) = row.columns.get(&column_id) {
6594 out.push((row_id, value.clone()));
6595 }
6596 }
6597 continue;
6598 }
6599
6600 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6601 for (index, reader) in readers.iter_mut().enumerate() {
6602 if let Some((epoch, deleted, value)) =
6603 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6604 {
6605 if best
6606 .as_ref()
6607 .map(|(best_epoch, ..)| epoch > *best_epoch)
6608 .unwrap_or(true)
6609 {
6610 best = Some((epoch, deleted, value, index));
6611 }
6612 }
6613 }
6614 let Some((_, false, Some(value), reader_index)) = best else {
6615 continue;
6616 };
6617 if let Some(ttl) = self.ttl {
6618 if ttl.column_id != column_id {
6619 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6620 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6621 {
6622 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6623 continue;
6624 }
6625 }
6626 } else if let Value::Int64(timestamp) = value {
6627 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6628 continue;
6629 }
6630 }
6631 }
6632 out.push((row_id, value));
6633 }
6634 Ok(out)
6635 }
6636
6637 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6642 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6643 }
6644
6645 pub fn rows_for_rids_with_context(
6646 &self,
6647 rids: &[u64],
6648 snapshot: Snapshot,
6649 context: &crate::query::AiExecutionContext,
6650 ) -> Result<Vec<Row>> {
6651 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6652 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
6653 }
6654
6655 fn rows_for_rids_at_time(
6656 &self,
6657 rids: &[u64],
6658 snapshot: Snapshot,
6659 ttl_now: i64,
6660 control: Option<&crate::ExecutionControl>,
6661 ) -> Result<Vec<Row>> {
6662 use std::collections::HashMap;
6663 let mut rows = Vec::with_capacity(rids.len());
6664 let tier_size = self.memtable.len() + self.mutable_run.len();
6681 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
6682 if rids.len().saturating_mul(24) < tier_size {
6683 for &rid in rids {
6684 if overlay.len() & 255 == 0 {
6685 control
6686 .map(crate::ExecutionControl::checkpoint)
6687 .transpose()?;
6688 }
6689 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
6690 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
6691 let newest = match (mem, mrun) {
6692 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
6693 (Some((_, mr)), None) => Some(mr),
6694 (None, Some((_, rr))) => Some(rr),
6695 (None, None) => None,
6696 };
6697 if let Some(row) = newest {
6698 overlay.insert(rid, row);
6699 }
6700 }
6701 } else {
6702 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
6703 overlay
6704 .entry(row.row_id.0)
6705 .and_modify(|e| {
6706 if row.committed_epoch > e.committed_epoch {
6707 *e = row.clone();
6708 }
6709 })
6710 .or_insert(row);
6711 };
6712 for (index, row) in self
6713 .memtable
6714 .visible_versions(snapshot.epoch)
6715 .into_iter()
6716 .enumerate()
6717 {
6718 if index & 255 == 0 {
6719 control
6720 .map(crate::ExecutionControl::checkpoint)
6721 .transpose()?;
6722 }
6723 fold_newest(row, &mut overlay);
6724 }
6725 for (index, row) in self
6726 .mutable_run
6727 .visible_versions(snapshot.epoch)
6728 .into_iter()
6729 .enumerate()
6730 {
6731 if index & 255 == 0 {
6732 control
6733 .map(crate::ExecutionControl::checkpoint)
6734 .transpose()?;
6735 }
6736 fold_newest(row, &mut overlay);
6737 }
6738 }
6739 if self.run_refs.len() == 1 {
6740 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6741 if rids.len().saturating_mul(24) < reader.row_count() {
6749 for (index, &rid) in rids.iter().enumerate() {
6750 if index & 255 == 0 {
6751 control
6752 .map(crate::ExecutionControl::checkpoint)
6753 .transpose()?;
6754 }
6755 if let Some(r) = overlay.get(&rid) {
6756 if !r.deleted {
6757 rows.push(r.clone());
6758 }
6759 continue;
6760 }
6761 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
6762 if !row.deleted {
6763 rows.push(row);
6764 }
6765 }
6766 }
6767 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6768 return Ok(rows);
6769 }
6770 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6779 enum Src {
6782 Overlay,
6783 Run,
6784 }
6785 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
6786 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
6787 for (index, rid) in rids.iter().enumerate() {
6788 if index & 255 == 0 {
6789 control
6790 .map(crate::ExecutionControl::checkpoint)
6791 .transpose()?;
6792 }
6793 if overlay.contains_key(rid) {
6794 plan.push(Src::Overlay);
6795 continue;
6796 }
6797 match vis_rids.binary_search(&(*rid as i64)) {
6798 Ok(i) => {
6799 plan.push(Src::Run);
6800 fetch.push(positions[i]);
6801 }
6802 Err(_) => { }
6803 }
6804 }
6805 let fetched = reader.materialize_batch(&fetch)?;
6806 let mut fetched_iter = fetched.into_iter();
6807 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
6808 if index & 255 == 0 {
6809 control
6810 .map(crate::ExecutionControl::checkpoint)
6811 .transpose()?;
6812 }
6813 match src {
6814 Src::Overlay => {
6815 if let Some(r) = overlay.get(rid) {
6816 if !r.deleted {
6817 rows.push(r.clone());
6818 }
6819 }
6820 }
6821 Src::Run => {
6822 if let Some(row) = fetched_iter.next() {
6823 if !row.deleted {
6824 rows.push(row);
6825 }
6826 }
6827 }
6828 }
6829 }
6830 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6831 return Ok(rows);
6832 }
6833 let mut readers: Vec<_> = self
6837 .run_refs
6838 .iter()
6839 .map(|rr| self.open_reader(rr.run_id))
6840 .collect::<Result<Vec<_>>>()?;
6841 for (index, rid) in rids.iter().enumerate() {
6842 if index & 255 == 0 {
6843 control
6844 .map(crate::ExecutionControl::checkpoint)
6845 .transpose()?;
6846 }
6847 if let Some(r) = overlay.get(rid) {
6848 if !r.deleted {
6849 rows.push(r.clone());
6850 }
6851 continue;
6852 }
6853 let mut best: Option<(Epoch, Row)> = None;
6854 for reader in readers.iter_mut() {
6855 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
6856 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
6857 best = Some((epoch, row));
6858 }
6859 }
6860 }
6861 if let Some((_, r)) = best {
6862 if !r.deleted {
6863 rows.push(r);
6864 }
6865 }
6866 }
6867 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6868 Ok(rows)
6869 }
6870
6871 pub fn indexes_complete(&self) -> bool {
6881 self.indexes_complete
6882 }
6883
6884 pub fn index_build_policy(&self) -> IndexBuildPolicy {
6886 self.index_build_policy
6887 }
6888
6889 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
6893 self.index_build_policy = policy;
6894 }
6895
6896 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
6901 if !self.indexes_complete {
6905 return None;
6906 }
6907 let b = self.bitmap.get(&column_id)?;
6908 let result: Vec<Vec<u8>> = b
6909 .keys()
6910 .into_iter()
6911 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
6912 .collect();
6913 Some(result)
6914 }
6915
6916 pub fn fk_join_row_ids(
6917 &self,
6918 fk_column_id: u16,
6919 pk_values: &[Vec<u8>],
6920 fk_conditions: &[crate::query::Condition],
6921 snapshot: Snapshot,
6922 ) -> Result<Vec<u64>> {
6923 let Some(b) = self.bitmap.get(&fk_column_id) else {
6924 return Ok(Vec::new());
6925 };
6926 let mut join_set = {
6927 let mut acc = roaring::RoaringBitmap::new();
6928 for v in pk_values {
6929 acc |= b.get(v);
6930 }
6931 RowIdSet::from_roaring(acc)
6932 };
6933 if !fk_conditions.is_empty() {
6934 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
6935 sets.push(join_set);
6936 for c in fk_conditions {
6937 sets.push(self.resolve_condition(c, snapshot)?);
6938 }
6939 join_set = RowIdSet::intersect_many(sets);
6940 }
6941 Ok(join_set.into_sorted_vec())
6942 }
6943
6944 pub fn fk_join_count(
6950 &self,
6951 fk_column_id: u16,
6952 pk_values: &[Vec<u8>],
6953 fk_conditions: &[crate::query::Condition],
6954 snapshot: Snapshot,
6955 ) -> Result<u64> {
6956 let Some(b) = self.bitmap.get(&fk_column_id) else {
6957 return Ok(0);
6958 };
6959 let mut acc = roaring::RoaringBitmap::new();
6960 for v in pk_values {
6961 acc |= b.get(v);
6962 }
6963 if fk_conditions.is_empty() {
6964 return Ok(acc.len());
6965 }
6966 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
6967 sets.push(RowIdSet::from_roaring(acc));
6968 for c in fk_conditions {
6969 sets.push(self.resolve_condition(c, snapshot)?);
6970 }
6971 Ok(RowIdSet::intersect_many(sets).len() as u64)
6972 }
6973
6974 fn resolve_condition(
6979 &self,
6980 c: &crate::query::Condition,
6981 snapshot: Snapshot,
6982 ) -> Result<RowIdSet> {
6983 self.resolve_condition_with_allowed(c, snapshot, None)
6984 }
6985
6986 fn resolve_condition_with_allowed(
6987 &self,
6988 c: &crate::query::Condition,
6989 snapshot: Snapshot,
6990 allowed: Option<&std::collections::HashSet<RowId>>,
6991 ) -> Result<RowIdSet> {
6992 use crate::query::Condition;
6993 self.validate_condition(c)?;
6994 Ok(match c {
6995 Condition::Pk(key) => {
6996 let lookup = self
6997 .schema
6998 .primary_key()
6999 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7000 .unwrap_or_else(|| key.clone());
7001 self.hot
7002 .get(&lookup)
7003 .map(|r| RowIdSet::one(r.0))
7004 .unwrap_or_else(RowIdSet::empty)
7005 }
7006 Condition::BitmapEq { column_id, value } => {
7007 let lookup = self.index_lookup_key_bytes(*column_id, value);
7008 self.bitmap
7009 .get(column_id)
7010 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7011 .unwrap_or_else(RowIdSet::empty)
7012 }
7013 Condition::BitmapIn { column_id, values } => {
7014 let bm = self.bitmap.get(column_id);
7015 let mut acc = roaring::RoaringBitmap::new();
7016 if let Some(b) = bm {
7017 for v in values {
7018 let lookup = self.index_lookup_key_bytes(*column_id, v);
7019 acc |= b.get(&lookup);
7020 }
7021 }
7022 RowIdSet::from_roaring(acc)
7023 }
7024 Condition::BytesPrefix { column_id, prefix } => {
7025 if let Some(b) = self.bitmap.get(column_id) {
7030 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7031 let mut acc = roaring::RoaringBitmap::new();
7032 for key in b.keys() {
7033 if key.starts_with(&lookup_prefix) {
7034 acc |= b.get(&key);
7035 }
7036 }
7037 RowIdSet::from_roaring(acc)
7038 } else {
7039 RowIdSet::empty()
7040 }
7041 }
7042 Condition::FmContains { column_id, pattern } => self
7043 .fm
7044 .get(column_id)
7045 .map(|f| {
7046 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7047 })
7048 .unwrap_or_else(RowIdSet::empty),
7049 Condition::FmContainsAll {
7050 column_id,
7051 patterns,
7052 } => {
7053 if let Some(f) = self.fm.get(column_id) {
7056 let sets: Vec<RowIdSet> = patterns
7057 .iter()
7058 .map(|pat| {
7059 RowIdSet::from_unsorted(
7060 f.locate(pat).into_iter().map(|r| r.0).collect(),
7061 )
7062 })
7063 .collect();
7064 RowIdSet::intersect_many(sets)
7065 } else {
7066 RowIdSet::empty()
7067 }
7068 }
7069 Condition::Ann {
7070 column_id,
7071 query,
7072 k,
7073 } => RowIdSet::from_unsorted(
7074 self.retrieve_filtered(
7075 &crate::query::Retriever::Ann {
7076 column_id: *column_id,
7077 query: query.clone(),
7078 k: *k,
7079 },
7080 snapshot,
7081 None,
7082 allowed,
7083 None,
7084 None,
7085 )?
7086 .into_iter()
7087 .map(|hit| hit.row_id.0)
7088 .collect(),
7089 ),
7090 Condition::SparseMatch {
7091 column_id,
7092 query,
7093 k,
7094 } => RowIdSet::from_unsorted(
7095 self.retrieve_filtered(
7096 &crate::query::Retriever::Sparse {
7097 column_id: *column_id,
7098 query: query.clone(),
7099 k: *k,
7100 },
7101 snapshot,
7102 None,
7103 allowed,
7104 None,
7105 None,
7106 )?
7107 .into_iter()
7108 .map(|hit| hit.row_id.0)
7109 .collect(),
7110 ),
7111 Condition::MinHashSimilar {
7112 column_id,
7113 query,
7114 k,
7115 } => match self.minhash.get(column_id) {
7116 Some(index) => {
7117 let candidates = index.candidate_row_ids(query);
7118 let eligible =
7119 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7120 RowIdSet::from_unsorted(
7121 index
7122 .search_filtered(query, *k, |row_id| {
7123 eligible.contains(&row_id)
7124 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7125 })
7126 .into_iter()
7127 .map(|(row_id, _)| row_id.0)
7128 .collect(),
7129 )
7130 }
7131 None => RowIdSet::empty(),
7132 },
7133 Condition::Range { column_id, lo, hi } => {
7134 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7143 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7144 } else if self.run_refs.len() == 1 {
7145 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7146 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7147 } else {
7148 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7149 };
7150 set.remove_many(self.overlay_rid_set(snapshot));
7151 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7152 set
7153 }
7154 Condition::RangeF64 {
7155 column_id,
7156 lo,
7157 lo_inclusive,
7158 hi,
7159 hi_inclusive,
7160 } => {
7161 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7164 RowIdSet::from_unsorted(
7165 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7166 .into_iter()
7167 .collect(),
7168 )
7169 } else if self.run_refs.len() == 1 {
7170 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7171 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7172 } else {
7173 return self.range_scan_f64(
7174 *column_id,
7175 *lo,
7176 *lo_inclusive,
7177 *hi,
7178 *hi_inclusive,
7179 snapshot,
7180 );
7181 };
7182 set.remove_many(self.overlay_rid_set(snapshot));
7183 self.range_scan_overlay_f64(
7184 &mut set,
7185 *column_id,
7186 *lo,
7187 *lo_inclusive,
7188 *hi,
7189 *hi_inclusive,
7190 snapshot,
7191 );
7192 set
7193 }
7194 Condition::IsNull { column_id } => {
7195 let mut set = if self.run_refs.len() == 1 {
7196 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7197 r.null_row_id_set(*column_id, true)?
7198 } else {
7199 return self.null_scan(*column_id, true, snapshot);
7200 };
7201 set.remove_many(self.overlay_rid_set(snapshot));
7202 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7203 set
7204 }
7205 Condition::IsNotNull { column_id } => {
7206 let mut set = if self.run_refs.len() == 1 {
7207 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7208 r.null_row_id_set(*column_id, false)?
7209 } else {
7210 return self.null_scan(*column_id, false, snapshot);
7211 };
7212 set.remove_many(self.overlay_rid_set(snapshot));
7213 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7214 set
7215 }
7216 })
7217 }
7218
7219 fn range_scan_i64(
7227 &self,
7228 column_id: u16,
7229 lo: i64,
7230 hi: i64,
7231 snapshot: Snapshot,
7232 ) -> Result<RowIdSet> {
7233 let mut row_ids = Vec::new();
7234 let overlay_rids = self.overlay_rid_set(snapshot);
7235 for rr in &self.run_refs {
7236 let mut reader = self.open_reader(rr.run_id)?;
7237 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7238 for rid in matched {
7239 if !overlay_rids.contains(&rid) {
7240 row_ids.push(rid);
7241 }
7242 }
7243 }
7244 let mut s = RowIdSet::from_unsorted(row_ids);
7245 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7246 Ok(s)
7247 }
7248
7249 fn range_scan_f64(
7252 &self,
7253 column_id: u16,
7254 lo: f64,
7255 lo_inclusive: bool,
7256 hi: f64,
7257 hi_inclusive: bool,
7258 snapshot: Snapshot,
7259 ) -> Result<RowIdSet> {
7260 let mut row_ids = Vec::new();
7261 let overlay_rids = self.overlay_rid_set(snapshot);
7262 for rr in &self.run_refs {
7263 let mut reader = self.open_reader(rr.run_id)?;
7264 let matched = reader.range_row_ids_visible_f64(
7265 column_id,
7266 lo,
7267 lo_inclusive,
7268 hi,
7269 hi_inclusive,
7270 snapshot.epoch,
7271 )?;
7272 for rid in matched {
7273 if !overlay_rids.contains(&rid) {
7274 row_ids.push(rid);
7275 }
7276 }
7277 }
7278 let mut s = RowIdSet::from_unsorted(row_ids);
7279 self.range_scan_overlay_f64(
7280 &mut s,
7281 column_id,
7282 lo,
7283 lo_inclusive,
7284 hi,
7285 hi_inclusive,
7286 snapshot,
7287 );
7288 Ok(s)
7289 }
7290
7291 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7293 let mut s = HashSet::new();
7294 for row in self.memtable.visible_versions(snapshot.epoch) {
7295 s.insert(row.row_id.0);
7296 }
7297 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7298 s.insert(row.row_id.0);
7299 }
7300 s
7301 }
7302
7303 fn range_scan_overlay_i64(
7304 &self,
7305 s: &mut RowIdSet,
7306 column_id: u16,
7307 lo: i64,
7308 hi: i64,
7309 snapshot: Snapshot,
7310 ) {
7311 let mut newest: HashMap<u64, &Row> = HashMap::new();
7316 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7317 let memtable = self.memtable.visible_versions(snapshot.epoch);
7318 for r in &mutable {
7319 newest.entry(r.row_id.0).or_insert(r);
7320 }
7321 for r in &memtable {
7322 newest.insert(r.row_id.0, r);
7323 }
7324 for row in newest.values() {
7325 if !row.deleted {
7326 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7327 if *v >= lo && *v <= hi {
7328 s.insert(row.row_id.0);
7329 }
7330 }
7331 }
7332 }
7333 }
7334
7335 #[allow(clippy::too_many_arguments)]
7336 fn range_scan_overlay_f64(
7337 &self,
7338 s: &mut RowIdSet,
7339 column_id: u16,
7340 lo: f64,
7341 lo_inclusive: bool,
7342 hi: f64,
7343 hi_inclusive: bool,
7344 snapshot: Snapshot,
7345 ) {
7346 let mut newest: HashMap<u64, &Row> = HashMap::new();
7349 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7350 let memtable = self.memtable.visible_versions(snapshot.epoch);
7351 for r in &mutable {
7352 newest.entry(r.row_id.0).or_insert(r);
7353 }
7354 for r in &memtable {
7355 newest.insert(r.row_id.0, r);
7356 }
7357 for row in newest.values() {
7358 if !row.deleted {
7359 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7360 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7361 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7362 if ok_lo && ok_hi {
7363 s.insert(row.row_id.0);
7364 }
7365 }
7366 }
7367 }
7368 }
7369
7370 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7373 let mut row_ids = Vec::new();
7374 let overlay_rids = self.overlay_rid_set(snapshot);
7375 for rr in &self.run_refs {
7376 let mut reader = self.open_reader(rr.run_id)?;
7377 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7378 for rid in matched {
7379 if !overlay_rids.contains(&rid) {
7380 row_ids.push(rid);
7381 }
7382 }
7383 }
7384 let mut s = RowIdSet::from_unsorted(row_ids);
7385 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7386 Ok(s)
7387 }
7388
7389 fn null_scan_overlay(
7393 &self,
7394 s: &mut RowIdSet,
7395 column_id: u16,
7396 want_nulls: bool,
7397 snapshot: Snapshot,
7398 ) {
7399 let mut newest: HashMap<u64, &Row> = HashMap::new();
7400 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7401 let memtable = self.memtable.visible_versions(snapshot.epoch);
7402 for r in &mutable {
7403 newest.entry(r.row_id.0).or_insert(r);
7404 }
7405 for r in &memtable {
7406 newest.insert(r.row_id.0, r);
7407 }
7408 for row in newest.values() {
7409 if row.deleted {
7410 continue;
7411 }
7412 let is_null = !row.columns.contains_key(&column_id)
7413 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7414 if is_null == want_nulls {
7415 s.insert(row.row_id.0);
7416 }
7417 }
7418 }
7419
7420 pub fn snapshot(&self) -> Snapshot {
7421 Snapshot::at(self.epoch.visible())
7422 }
7423
7424 pub fn data_generation(&self) -> u64 {
7426 self.data_generation
7427 }
7428
7429 pub(crate) fn bump_data_generation(&mut self) {
7430 self.data_generation = self.data_generation.wrapping_add(1);
7431 }
7432
7433 pub(crate) fn table_id(&self) -> u64 {
7434 self.table_id
7435 }
7436
7437 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7438 self.ensure_indexes_complete()?;
7439 self.memtable.seal();
7440 self.mutable_run.seal();
7441 self.hot.seal();
7442 for index in self.bitmap.values_mut() {
7443 index.seal();
7444 }
7445 for index in self.ann.values_mut() {
7446 index.seal();
7447 }
7448 for index in self.fm.values_mut() {
7449 index.seal();
7450 }
7451 for index in self.sparse.values_mut() {
7452 index.seal();
7453 }
7454 for index in self.minhash.values_mut() {
7455 index.seal();
7456 }
7457 self.pk_by_row.seal();
7458 let mut generation = self.clone();
7459 generation.read_only = true;
7460 generation.wal = WalSink::ReadOnly;
7461 generation.pending_delete_rids.clear();
7462 generation.pending_put_cols.clear();
7463 generation.pending_rows.clear();
7464 generation.pending_rows_auto_inc.clear();
7465 generation.pending_dels.clear();
7466 generation.pending_truncate = None;
7467 generation.agg_cache = Arc::new(HashMap::new());
7468 Ok(generation)
7469 }
7470
7471 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7472 (std::mem::size_of::<Self>() as u64)
7473 .saturating_add(self.memtable.approx_bytes())
7474 .saturating_add(self.mutable_run.approx_bytes())
7475 .saturating_add(self.live_count.saturating_mul(64))
7476 }
7477
7478 pub fn pin_snapshot(&mut self) -> Snapshot {
7481 let e = self.epoch.visible();
7482 *self.pinned.entry(e).or_insert(0) += 1;
7483 Snapshot::at(e)
7484 }
7485
7486 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7488 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7489 *count -= 1;
7490 if *count == 0 {
7491 self.pinned.remove(&snap.epoch);
7492 }
7493 }
7494 }
7495
7496 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7506 let local = self.pinned.keys().next().copied();
7507 let global = self.snapshots.min_pinned();
7508 let history = self.snapshots.history_floor(self.current_epoch());
7509 [local, global, history].into_iter().flatten().min()
7510 }
7511
7512 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7516 self.ensure_writable()?;
7517 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7518 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7519 }
7520
7521 pub fn clear_ttl(&mut self) -> Result<()> {
7522 self.ensure_writable()?;
7523 self.apply_ttl_policy_at(None, self.current_epoch())
7524 }
7525
7526 pub fn ttl(&self) -> Option<TtlPolicy> {
7527 self.ttl
7528 }
7529
7530 pub(crate) fn prepare_ttl_policy(
7531 &self,
7532 column_name: &str,
7533 duration_nanos: u64,
7534 ) -> Result<TtlPolicy> {
7535 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7536 return Err(MongrelError::InvalidArgument(
7537 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7538 ));
7539 }
7540 let column = self
7541 .schema
7542 .columns
7543 .iter()
7544 .find(|column| column.name == column_name)
7545 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7546 if column.ty != TypeId::TimestampNanos {
7547 return Err(MongrelError::Schema(format!(
7548 "TTL column {column_name} must be TimestampNanos, is {:?}",
7549 column.ty
7550 )));
7551 }
7552 Ok(TtlPolicy {
7553 column_id: column.id,
7554 duration_nanos,
7555 })
7556 }
7557
7558 pub(crate) fn apply_ttl_policy_at(
7559 &mut self,
7560 policy: Option<TtlPolicy>,
7561 epoch: Epoch,
7562 ) -> Result<()> {
7563 if let Some(policy) = policy {
7564 let column = self
7565 .schema
7566 .columns
7567 .iter()
7568 .find(|column| column.id == policy.column_id)
7569 .ok_or_else(|| {
7570 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
7571 })?;
7572 if column.ty != TypeId::TimestampNanos
7573 || policy.duration_nanos == 0
7574 || policy.duration_nanos > i64::MAX as u64
7575 {
7576 return Err(MongrelError::Schema("invalid TTL policy".into()));
7577 }
7578 }
7579 self.ttl = policy;
7580 self.agg_cache = Arc::new(HashMap::new());
7581 self.clear_result_cache();
7582 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7583 self.persist_manifest(epoch)
7584 }
7585
7586 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
7587 let Some(policy) = self.ttl else {
7588 return false;
7589 };
7590 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
7591 return false;
7592 };
7593 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
7594 }
7595
7596 pub fn current_epoch(&self) -> Epoch {
7597 self.epoch.visible()
7598 }
7599
7600 pub fn memtable_len(&self) -> usize {
7601 self.memtable.len()
7602 }
7603
7604 pub fn count(&self) -> u64 {
7607 if self.ttl.is_none()
7608 && self.pending_put_cols.is_empty()
7609 && self.pending_delete_rids.is_empty()
7610 && self.pending_rows.is_empty()
7611 && self.pending_dels.is_empty()
7612 && self.pending_truncate.is_none()
7613 {
7614 self.live_count
7615 } else {
7616 self.visible_rows(self.snapshot())
7617 .map(|rows| rows.len() as u64)
7618 .unwrap_or(self.live_count)
7619 }
7620 }
7621
7622 pub fn count_conditions(
7626 &mut self,
7627 conditions: &[crate::query::Condition],
7628 snapshot: Snapshot,
7629 ) -> Result<Option<u64>> {
7630 use crate::query::Condition;
7631 if self.ttl.is_some() {
7632 if conditions.is_empty() {
7633 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
7634 }
7635 let mut sets = Vec::with_capacity(conditions.len());
7636 for condition in conditions {
7637 sets.push(self.resolve_condition(condition, snapshot)?);
7638 }
7639 let survivors = RowIdSet::intersect_many(sets);
7640 let rows = self.visible_rows(snapshot)?;
7641 return Ok(Some(
7642 rows.into_iter()
7643 .filter(|row| survivors.contains(row.row_id.0))
7644 .count() as u64,
7645 ));
7646 }
7647 if conditions.is_empty() {
7648 return Ok(Some(self.count()));
7649 }
7650 let served = |c: &Condition| {
7651 matches!(
7652 c,
7653 Condition::Pk(_)
7654 | Condition::BitmapEq { .. }
7655 | Condition::BitmapIn { .. }
7656 | Condition::BytesPrefix { .. }
7657 | Condition::FmContains { .. }
7658 | Condition::FmContainsAll { .. }
7659 | Condition::Ann { .. }
7660 | Condition::Range { .. }
7661 | Condition::RangeF64 { .. }
7662 | Condition::SparseMatch { .. }
7663 | Condition::MinHashSimilar { .. }
7664 | Condition::IsNull { .. }
7665 | Condition::IsNotNull { .. }
7666 )
7667 };
7668 if !conditions.iter().all(served) {
7669 return Ok(None);
7670 }
7671 self.ensure_indexes_complete()?;
7672 if !self.pending_put_cols.is_empty()
7673 || !self.pending_delete_rids.is_empty()
7674 || !self.pending_rows.is_empty()
7675 || !self.pending_dels.is_empty()
7676 || self.pending_truncate.is_some()
7677 {
7678 let mut sets = Vec::with_capacity(conditions.len());
7679 for condition in conditions {
7680 sets.push(self.resolve_condition(condition, snapshot)?);
7681 }
7682 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7683 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
7684 }
7685 let mut sets = Vec::with_capacity(conditions.len());
7686 for condition in conditions {
7687 sets.push(self.resolve_condition(condition, snapshot)?);
7688 }
7689 let mut rids = RowIdSet::intersect_many(sets);
7690 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7700 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
7701 }
7702 let count = rids.len() as u64;
7703 crate::trace::QueryTrace::record(|t| {
7704 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
7705 t.survivor_count = Some(count as usize);
7706 t.conditions_pushed = conditions.len();
7707 });
7708 Ok(Some(count))
7709 }
7710
7711 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
7716 let mut out = Vec::new();
7717 for row in self.memtable.visible_versions(snapshot.epoch) {
7718 if row.deleted {
7719 out.push(row.row_id.0);
7720 }
7721 }
7722 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7723 if row.deleted {
7724 out.push(row.row_id.0);
7725 }
7726 }
7727 out
7728 }
7729
7730 pub fn bulk_load_columns(
7739 &mut self,
7740 user_columns: Vec<(u16, columnar::NativeColumn)>,
7741 ) -> Result<Epoch> {
7742 self.bulk_load_columns_with(user_columns, 3, false, true)
7743 }
7744
7745 pub fn bulk_load_fast(
7752 &mut self,
7753 user_columns: Vec<(u16, columnar::NativeColumn)>,
7754 ) -> Result<Epoch> {
7755 self.bulk_load_columns_with(user_columns, -1, true, false)
7756 }
7757
7758 fn bulk_load_columns_with(
7759 &mut self,
7760 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
7761 zstd_level: i32,
7762 force_plain: bool,
7763 lz4: bool,
7764 ) -> Result<Epoch> {
7765 self.ensure_writable()?;
7766 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
7767 if n == 0 {
7768 return Ok(self.current_epoch());
7769 }
7770 let epoch = self.commit_new_epoch()?;
7771 let live_before = self.live_count;
7772 self.spill_mutable_run(epoch)?;
7774 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
7775 && self.indexes_complete
7776 && self.run_refs.is_empty()
7777 && self.memtable.is_empty()
7778 && self.mutable_run.is_empty();
7779 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
7782 self.validate_columns_not_null(&user_columns, n)?;
7783 let winner_idx = self
7784 .bulk_pk_winner_indices(&user_columns, n)
7785 .filter(|idx| idx.len() != n);
7786 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
7787 match winner_idx.as_deref() {
7788 Some(idx) => {
7789 let compacted = user_columns
7790 .iter()
7791 .map(|(id, c)| (*id, c.gather(idx)))
7792 .collect();
7793 (compacted, idx.len())
7794 }
7795 None => (user_columns, n),
7796 };
7797 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
7798 let first = self.allocator.alloc_range(write_n as u64)?.0;
7799 for rid in first..first + write_n as u64 {
7800 self.reservoir.offer(rid);
7801 }
7802 let run_id = self.alloc_run_id()?;
7803 let path = self.run_path(run_id);
7804 let mut writer =
7805 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
7806 if force_plain {
7807 writer = writer.with_plain();
7808 } else if lz4 {
7809 writer = writer.with_lz4();
7812 } else {
7813 writer = writer.with_zstd_level(zstd_level);
7814 }
7815 if let Some(kek) = &self.kek {
7816 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
7817 }
7818 let header = match self.create_run_file(run_id)? {
7819 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
7820 None => writer.write_native(&path, &write_columns, write_n, first)?,
7821 };
7822 self.run_refs.push(RunRef {
7823 run_id: run_id as u128,
7824 level: 0,
7825 epoch_created: epoch.0,
7826 row_count: header.row_count,
7827 });
7828 self.live_count = self.live_count.saturating_add(write_n as u64);
7829 if eager_index_build {
7830 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
7831 self.index_columns_bulk(&write_columns, &row_ids);
7832 self.indexes_complete = true;
7833 self.build_learned_ranges()?;
7834 } else {
7835 self.indexes_complete = false;
7839 }
7840 self.mark_flushed(epoch)?;
7841 self.persist_manifest(epoch)?;
7842 if eager_index_build {
7843 self.checkpoint_indexes(epoch);
7844 }
7845 self.clear_result_cache();
7846 self.data_generation = self.data_generation.wrapping_add(1);
7847 Ok(epoch)
7848 }
7849
7850 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
7868 let n = row_ids.len();
7869 if n == 0 {
7870 return;
7871 }
7872 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
7873 columns.iter().map(|(id, c)| (*id, c)).collect();
7874 let ty_of: std::collections::HashMap<u16, TypeId> = self
7875 .schema
7876 .columns
7877 .iter()
7878 .map(|c| (c.id, c.ty.clone()))
7879 .collect();
7880 let pk_id = self.schema.primary_key().map(|c| c.id);
7881
7882 for (i, &rid) in row_ids.iter().enumerate() {
7883 let row_id = RowId(rid);
7884 if let Some(pid) = pk_id {
7885 if let Some(col) = by_id.get(&pid) {
7886 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
7887 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
7888 self.insert_hot_pk(key, row_id);
7889 }
7890 }
7891 }
7892 for idef in &self.schema.indexes {
7893 let Some(col) = by_id.get(&idef.column_id) else {
7894 continue;
7895 };
7896 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
7897 match idef.kind {
7898 IndexKind::Bitmap => {
7899 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
7900 if let Some(key) =
7901 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
7902 {
7903 b.insert(key, row_id);
7904 }
7905 }
7906 }
7907 IndexKind::FmIndex => {
7908 if let Some(f) = self.fm.get_mut(&idef.column_id) {
7909 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7910 f.insert(bytes.to_vec(), row_id);
7911 }
7912 }
7913 }
7914 IndexKind::Sparse => {
7915 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
7916 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7917 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
7918 s.insert(&terms, row_id);
7919 }
7920 }
7921 }
7922 }
7923 IndexKind::MinHash => {
7924 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
7925 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7926 let tokens = crate::index::token_hashes_from_bytes(bytes);
7927 mh.insert(&tokens, row_id);
7928 }
7929 }
7930 }
7931 _ => {}
7932 }
7933 }
7934 }
7935 }
7936
7937 pub fn visible_columns_native(
7942 &self,
7943 snapshot: Snapshot,
7944 projection: Option<&[u16]>,
7945 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7946 self.visible_columns_native_inner(snapshot, projection, None)
7947 }
7948
7949 pub fn visible_columns_native_with_control(
7950 &self,
7951 snapshot: Snapshot,
7952 projection: Option<&[u16]>,
7953 control: &crate::ExecutionControl,
7954 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7955 self.visible_columns_native_inner(snapshot, projection, Some(control))
7956 }
7957
7958 fn visible_columns_native_inner(
7959 &self,
7960 snapshot: Snapshot,
7961 projection: Option<&[u16]>,
7962 control: Option<&crate::ExecutionControl>,
7963 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7964 execution_checkpoint(control, 0)?;
7965 let wanted: Vec<u16> = match projection {
7966 Some(p) => p.to_vec(),
7967 None => self.schema.columns.iter().map(|c| c.id).collect(),
7968 };
7969 if self.ttl.is_none()
7970 && self.memtable.is_empty()
7971 && self.mutable_run.is_empty()
7972 && self.run_refs.len() == 1
7973 {
7974 let rr = self.run_refs[0].clone();
7975 let mut reader = self.open_reader(rr.run_id)?;
7976 let idxs = reader.visible_indices_native(snapshot.epoch)?;
7977 execution_checkpoint(control, 0)?;
7978 let all_visible = idxs.len() == reader.row_count();
7979 if reader.has_mmap() && control.is_none() {
7985 use rayon::prelude::*;
7986 let valid: Vec<u16> = wanted
7989 .iter()
7990 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
7991 .copied()
7992 .collect();
7993 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
7995 .par_iter()
7996 .filter_map(|cid| {
7997 reader
7998 .column_native_shared(*cid)
7999 .ok()
8000 .map(|col| (*cid, col))
8001 })
8002 .collect();
8003 let cols = decoded
8004 .into_iter()
8005 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8006 .collect();
8007 return Ok(cols);
8008 }
8009 let mut cols = Vec::with_capacity(wanted.len());
8010 for (index, cid) in wanted.iter().enumerate() {
8011 execution_checkpoint(control, index)?;
8012 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8013 Some(c) => c,
8014 None => continue,
8015 };
8016 let col = reader.column_native(cdef.id)?;
8017 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8018 }
8019 return Ok(cols);
8020 }
8021 let vcols = self.visible_columns(snapshot)?;
8022 execution_checkpoint(control, 0)?;
8023 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8024 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8025 .into_iter()
8026 .filter(|(id, _)| want_set.contains(id))
8027 .map(|(id, vals)| {
8028 let ty = self
8029 .schema
8030 .columns
8031 .iter()
8032 .find(|c| c.id == id)
8033 .map(|c| c.ty.clone())
8034 .unwrap_or(TypeId::Bytes);
8035 (id, columnar::values_to_native(ty, &vals))
8036 })
8037 .collect();
8038 Ok(out)
8039 }
8040
8041 pub fn run_count(&self) -> usize {
8042 self.run_refs.len()
8043 }
8044
8045 pub fn memtable_is_empty(&self) -> bool {
8047 self.memtable.is_empty()
8048 }
8049
8050 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8054 self.page_cache.stats()
8055 }
8056
8057 pub fn reset_page_cache_stats(&self) {
8059 self.page_cache.reset_stats();
8060 }
8061
8062 pub fn run_ids(&self) -> Vec<u128> {
8065 self.run_refs.iter().map(|r| r.run_id).collect()
8066 }
8067
8068 pub fn single_run_is_clean(&self) -> bool {
8072 if self.ttl.is_some() || self.run_refs.len() != 1 {
8073 return false;
8074 }
8075 self.open_reader(self.run_refs[0].run_id)
8076 .map(|r| r.is_clean())
8077 .unwrap_or(false)
8078 }
8079
8080 fn resolve_footprint(
8087 &self,
8088 conditions: &[crate::query::Condition],
8089 snapshot: Snapshot,
8090 ) -> roaring::RoaringBitmap {
8091 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8092 return roaring::RoaringBitmap::new();
8093 }
8094 if self.run_refs.is_empty() {
8095 return roaring::RoaringBitmap::new();
8096 }
8097 if self.run_refs.len() == 1 {
8099 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8100 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8101 return rids.to_roaring_lossy();
8102 }
8103 }
8104 }
8105 roaring::RoaringBitmap::new()
8106 }
8107
8108 pub fn query_columns_native_cached(
8119 &mut self,
8120 conditions: &[crate::query::Condition],
8121 projection: Option<&[u16]>,
8122 snapshot: Snapshot,
8123 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8124 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8125 }
8126
8127 pub fn query_columns_native_cached_with_control(
8128 &mut self,
8129 conditions: &[crate::query::Condition],
8130 projection: Option<&[u16]>,
8131 snapshot: Snapshot,
8132 control: &crate::ExecutionControl,
8133 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8134 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8135 }
8136
8137 fn query_columns_native_cached_inner(
8138 &mut self,
8139 conditions: &[crate::query::Condition],
8140 projection: Option<&[u16]>,
8141 snapshot: Snapshot,
8142 control: Option<&crate::ExecutionControl>,
8143 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8144 execution_checkpoint(control, 0)?;
8145 if self.ttl.is_some() {
8148 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8149 }
8150 if conditions.is_empty() {
8151 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8152 }
8153 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8157 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8158 crate::trace::QueryTrace::record(|t| {
8159 t.result_cache_hit = true;
8160 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8161 });
8162 return Ok(Some((*hit).clone()));
8163 }
8164 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8165 execution_checkpoint(control, 0)?;
8166 if let Some(cols) = &res {
8167 let footprint = self.resolve_footprint(conditions, snapshot);
8168 let condition_cols = crate::query::condition_columns(conditions);
8169 execution_checkpoint(control, 0)?;
8170 self.result_cache.lock().insert(
8171 key,
8172 CachedEntry {
8173 data: CachedData::Columns(Arc::new(cols.clone())),
8174 footprint,
8175 condition_cols,
8176 },
8177 );
8178 }
8179 Ok(res)
8180 }
8181
8182 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8187 if self.ttl.is_some() {
8188 return self.query(q);
8189 }
8190 if q.conditions.is_empty() {
8191 return self.query(q);
8192 }
8193 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8194 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8195 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8196 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8197 crate::trace::QueryTrace::record(|t| {
8198 t.result_cache_hit = true;
8199 t.scan_mode = crate::trace::ScanMode::Materialized;
8200 });
8201 return Ok((*hit).clone());
8202 }
8203 let rows = self.query(q)?;
8204 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8205 let condition_cols = crate::query::condition_columns(&q.conditions);
8206 self.result_cache.lock().insert(
8207 key,
8208 CachedEntry {
8209 data: CachedData::Rows(Arc::new(rows.clone())),
8210 footprint,
8211 condition_cols,
8212 },
8213 );
8214 Ok(rows)
8215 }
8216
8217 #[allow(clippy::type_complexity)]
8232 pub fn query_columns_native_traced(
8233 &mut self,
8234 conditions: &[crate::query::Condition],
8235 projection: Option<&[u16]>,
8236 snapshot: Snapshot,
8237 ) -> Result<(
8238 Option<Vec<(u16, columnar::NativeColumn)>>,
8239 crate::trace::QueryTrace,
8240 )> {
8241 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8242 self.query_columns_native(conditions, projection, snapshot)
8243 });
8244 Ok((result?, trace))
8245 }
8246
8247 #[allow(clippy::type_complexity)]
8250 pub fn query_columns_native_cached_traced(
8251 &mut self,
8252 conditions: &[crate::query::Condition],
8253 projection: Option<&[u16]>,
8254 snapshot: Snapshot,
8255 ) -> Result<(
8256 Option<Vec<(u16, columnar::NativeColumn)>>,
8257 crate::trace::QueryTrace,
8258 )> {
8259 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8260 self.query_columns_native_cached(conditions, projection, snapshot)
8261 });
8262 Ok((result?, trace))
8263 }
8264
8265 pub fn native_page_cursor_traced(
8267 &self,
8268 snapshot: Snapshot,
8269 projection: Vec<(u16, TypeId)>,
8270 conditions: &[crate::query::Condition],
8271 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8272 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8273 self.native_page_cursor(snapshot, projection, conditions)
8274 });
8275 Ok((result?, trace))
8276 }
8277
8278 pub fn native_multi_run_cursor_traced(
8280 &self,
8281 snapshot: Snapshot,
8282 projection: Vec<(u16, TypeId)>,
8283 conditions: &[crate::query::Condition],
8284 ) -> Result<(
8285 Option<crate::cursor::MultiRunCursor>,
8286 crate::trace::QueryTrace,
8287 )> {
8288 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8289 self.native_multi_run_cursor(snapshot, projection, conditions)
8290 });
8291 Ok((result?, trace))
8292 }
8293
8294 pub fn count_conditions_traced(
8296 &mut self,
8297 conditions: &[crate::query::Condition],
8298 snapshot: Snapshot,
8299 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8300 let (result, trace) =
8301 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8302 Ok((result?, trace))
8303 }
8304
8305 pub fn query_traced(
8307 &mut self,
8308 q: &crate::query::Query,
8309 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8310 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8311 Ok((result?, trace))
8312 }
8313
8314 pub fn query_columns_native(
8319 &mut self,
8320 conditions: &[crate::query::Condition],
8321 projection: Option<&[u16]>,
8322 snapshot: Snapshot,
8323 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8324 self.query_columns_native_inner(conditions, projection, snapshot, None)
8325 }
8326
8327 pub fn query_columns_native_with_control(
8328 &mut self,
8329 conditions: &[crate::query::Condition],
8330 projection: Option<&[u16]>,
8331 snapshot: Snapshot,
8332 control: &crate::ExecutionControl,
8333 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8334 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8335 }
8336
8337 fn query_columns_native_inner(
8338 &mut self,
8339 conditions: &[crate::query::Condition],
8340 projection: Option<&[u16]>,
8341 snapshot: Snapshot,
8342 control: Option<&crate::ExecutionControl>,
8343 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8344 use crate::query::Condition;
8345 execution_checkpoint(control, 0)?;
8346 if self.ttl.is_some() {
8349 return Ok(None);
8350 }
8351 if conditions.is_empty() {
8352 return Ok(None);
8353 }
8354 self.ensure_indexes_complete()?;
8355
8356 let served = |c: &Condition| {
8361 matches!(
8362 c,
8363 Condition::Pk(_)
8364 | Condition::BitmapEq { .. }
8365 | Condition::BitmapIn { .. }
8366 | Condition::BytesPrefix { .. }
8367 | Condition::FmContains { .. }
8368 | Condition::FmContainsAll { .. }
8369 | Condition::Ann { .. }
8370 | Condition::Range { .. }
8371 | Condition::RangeF64 { .. }
8372 | Condition::SparseMatch { .. }
8373 | Condition::MinHashSimilar { .. }
8374 | Condition::IsNull { .. }
8375 | Condition::IsNotNull { .. }
8376 )
8377 };
8378 if !conditions.iter().all(served) {
8379 return Ok(None);
8380 }
8381 let fast_path =
8382 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8383 crate::trace::QueryTrace::record(|t| {
8384 t.run_count = self.run_refs.len();
8385 t.memtable_rows = self.memtable.len();
8386 t.mutable_run_rows = self.mutable_run.len();
8387 t.conditions_pushed = conditions.len();
8388 t.learned_range_used = conditions.iter().any(|c| match c {
8389 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8390 self.learned_range.contains_key(column_id)
8391 }
8392 _ => false,
8393 });
8394 });
8395 let col_ids: Vec<u16> = projection
8397 .map(|p| p.to_vec())
8398 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8399 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8400 .iter()
8401 .map(|&cid| {
8402 let ty = self
8403 .schema
8404 .columns
8405 .iter()
8406 .find(|c| c.id == cid)
8407 .map(|c| c.ty.clone())
8408 .unwrap_or(TypeId::Bytes);
8409 (cid, ty)
8410 })
8411 .collect();
8412
8413 if fast_path {
8419 let needs_column = conditions.iter().any(|c| match c {
8422 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8423 Condition::RangeF64 { column_id, .. } => {
8424 !self.learned_range.contains_key(column_id)
8425 }
8426 _ => false,
8427 });
8428 let mut reader_opt: Option<RunReader> = if needs_column {
8429 Some(self.open_reader(self.run_refs[0].run_id)?)
8430 } else {
8431 None
8432 };
8433 let mut sets: Vec<RowIdSet> = Vec::new();
8434 for (index, c) in conditions.iter().enumerate() {
8435 execution_checkpoint(control, index)?;
8436 let s = match c {
8437 Condition::Range { column_id, lo, hi }
8438 if !self.learned_range.contains_key(column_id) =>
8439 {
8440 if reader_opt.is_none() {
8441 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8442 }
8443 reader_opt
8444 .as_mut()
8445 .expect("reader opened for range")
8446 .range_row_id_set_i64(*column_id, *lo, *hi)?
8447 }
8448 Condition::RangeF64 {
8449 column_id,
8450 lo,
8451 lo_inclusive,
8452 hi,
8453 hi_inclusive,
8454 } if !self.learned_range.contains_key(column_id) => {
8455 if reader_opt.is_none() {
8456 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8457 }
8458 reader_opt
8459 .as_mut()
8460 .expect("reader opened for range")
8461 .range_row_id_set_f64(
8462 *column_id,
8463 *lo,
8464 *lo_inclusive,
8465 *hi,
8466 *hi_inclusive,
8467 )?
8468 }
8469 _ => self.resolve_condition(c, snapshot)?,
8470 };
8471 sets.push(s);
8472 }
8473 let candidates = RowIdSet::intersect_many(sets);
8474 crate::trace::QueryTrace::record(|t| {
8475 t.survivor_count = Some(candidates.len());
8476 });
8477 if candidates.is_empty() {
8478 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8479 .iter()
8480 .map(|&id| {
8481 (
8482 id,
8483 columnar::null_native(
8484 proj_pairs
8485 .iter()
8486 .find(|(c, _)| c == &id)
8487 .map(|(_, t)| t.clone())
8488 .unwrap_or(TypeId::Bytes),
8489 0,
8490 ),
8491 )
8492 })
8493 .collect();
8494 return Ok(Some(cols));
8495 }
8496 let mut reader = match reader_opt.take() {
8497 Some(r) => r,
8498 None => self.open_reader(self.run_refs[0].run_id)?,
8499 };
8500 let candidate_ids = candidates.into_sorted_vec();
8501 let (positions, fast_rid) = if let Some(positions) =
8502 reader.positions_for_row_ids_fast(&candidate_ids)
8503 {
8504 (positions, true)
8505 } else {
8506 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8507 match col {
8508 columnar::NativeColumn::Int64 { data, .. } => {
8509 let mut p = Vec::with_capacity(candidate_ids.len());
8510 for (index, rid) in candidate_ids.iter().enumerate() {
8511 execution_checkpoint(control, index)?;
8512 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8513 p.push(position);
8514 }
8515 }
8516 p.sort_unstable();
8517 (p, false)
8518 }
8519 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8520 }
8521 };
8522 crate::trace::QueryTrace::record(|t| {
8523 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8524 t.fast_row_id_map = fast_rid;
8525 });
8526 let mut cols = Vec::with_capacity(col_ids.len());
8527 for (index, cid) in col_ids.iter().enumerate() {
8528 execution_checkpoint(control, index)?;
8529 let col = reader.column_native(*cid)?;
8530 cols.push((*cid, col.gather(&positions)));
8531 }
8532 return Ok(Some(cols));
8533 }
8534
8535 if !self.run_refs.is_empty() {
8548 use crate::cursor::{
8549 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8550 };
8551 let remaining: usize;
8552 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
8553 let c = self
8554 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
8555 .expect("single-run cursor should build when run_refs.len() == 1");
8556 remaining = c.remaining_rows();
8557 Box::new(c)
8558 } else {
8559 let c = self
8560 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
8561 .expect("multi-run cursor should build when run_refs.len() >= 1");
8562 remaining = c.remaining_rows();
8563 Box::new(c)
8564 };
8565 crate::trace::QueryTrace::record(|t| {
8566 if t.survivor_count.is_none() {
8567 t.survivor_count = Some(remaining);
8568 }
8569 });
8570 let cols = match control {
8571 Some(control) => {
8572 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
8573 }
8574 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
8575 };
8576 return Ok(Some(cols));
8577 }
8578
8579 crate::trace::QueryTrace::record(|t| {
8584 t.scan_mode = crate::trace::ScanMode::Materialized;
8585 t.row_materialized = true;
8586 });
8587 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8588 for (index, c) in conditions.iter().enumerate() {
8589 execution_checkpoint(control, index)?;
8590 sets.push(self.resolve_condition(c, snapshot)?);
8591 }
8592 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8593 let rows = self.rows_for_rids(&rids, snapshot)?;
8594 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
8595 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
8596 execution_checkpoint(control, index)?;
8597 let vals: Vec<Value> = rows
8598 .iter()
8599 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8600 .collect();
8601 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
8602 }
8603 Ok(Some(cols))
8604 }
8605
8606 pub fn native_page_cursor(
8621 &self,
8622 snapshot: Snapshot,
8623 projection: Vec<(u16, TypeId)>,
8624 conditions: &[crate::query::Condition],
8625 ) -> Result<Option<NativePageCursor>> {
8626 use crate::cursor::build_page_plans;
8627 if self.ttl.is_some() {
8628 return Ok(None);
8629 }
8630 if !conditions.is_empty() && !self.indexes_complete {
8633 return Ok(None);
8634 }
8635 if self.run_refs.len() != 1 {
8636 return Ok(None);
8637 }
8638 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8639 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8640
8641 let overlay_rids: HashSet<u64> = {
8644 let mut s = HashSet::new();
8645 for row in self.memtable.visible_versions(snapshot.epoch) {
8646 s.insert(row.row_id.0);
8647 }
8648 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8649 s.insert(row.row_id.0);
8650 }
8651 s
8652 };
8653
8654 let survivors = if conditions.is_empty() {
8658 None
8659 } else {
8660 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
8661 };
8662
8663 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
8670 survivors.clone()
8671 } else if let Some(s) = &survivors {
8672 let mut run_set = s.clone();
8673 run_set.remove_many(overlay_rids.iter().copied());
8674 Some(run_set)
8675 } else {
8676 Some(RowIdSet::from_unsorted(
8677 rids.iter()
8678 .map(|&r| r as u64)
8679 .filter(|r| !overlay_rids.contains(r))
8680 .collect(),
8681 ))
8682 };
8683
8684 let overlay_rows = if overlay_rids.is_empty() {
8685 Vec::new()
8686 } else {
8687 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8688 self.overlay_visible_rows(snapshot, bound)
8689 };
8690
8691 let plans = if positions.is_empty() {
8693 Vec::new()
8694 } else {
8695 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
8696 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
8697 };
8698
8699 let overlay = if overlay_rows.is_empty() {
8701 None
8702 } else {
8703 let filtered =
8704 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8705 if filtered.is_empty() {
8706 None
8707 } else {
8708 Some(self.materialize_overlay(&filtered, &projection))
8709 }
8710 };
8711
8712 let overlay_row_count = overlay
8713 .as_ref()
8714 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8715 .unwrap_or(0);
8716 crate::trace::QueryTrace::record(|t| {
8717 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
8718 t.run_count = self.run_refs.len();
8719 t.memtable_rows = self.memtable.len();
8720 t.mutable_run_rows = self.mutable_run.len();
8721 t.overlay_rows = overlay_row_count;
8722 t.conditions_pushed = conditions.len();
8723 t.pages_decoded = plans
8724 .iter()
8725 .map(|p| p.positions.len())
8726 .sum::<usize>()
8727 .min(1);
8728 });
8729
8730 Ok(Some(NativePageCursor::new_with_overlay(
8731 reader, projection, plans, overlay,
8732 )))
8733 }
8734 #[allow(clippy::type_complexity)]
8744 pub fn native_multi_run_cursor(
8745 &self,
8746 snapshot: Snapshot,
8747 projection: Vec<(u16, TypeId)>,
8748 conditions: &[crate::query::Condition],
8749 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
8750 use crate::cursor::{MultiRunCursor, RunStream};
8751 use crate::sorted_run::SYS_ROW_ID;
8752 use std::collections::{BinaryHeap, HashMap, HashSet};
8753 if self.ttl.is_some() {
8754 return Ok(None);
8755 }
8756 if !conditions.is_empty() && !self.indexes_complete {
8759 return Ok(None);
8760 }
8761 if self.run_refs.is_empty() {
8762 return Ok(None);
8763 }
8764
8765 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
8767 Vec::with_capacity(self.run_refs.len());
8768 for rr in &self.run_refs {
8769 let mut reader = self.open_reader(rr.run_id)?;
8770 let (rids, eps, del) = reader.system_columns_native()?;
8771 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
8772 run_meta.push((reader, rids, eps, del, page_rows));
8773 }
8774
8775 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
8779 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
8780 for i in 0..rids.len() {
8781 let rid = rids[i] as u64;
8782 let e = eps[i] as u64;
8783 if e > snapshot.epoch.0 {
8784 continue;
8785 }
8786 let is_del = del[i] != 0;
8787 best.entry(rid)
8788 .and_modify(|cur| {
8789 if e > cur.0 {
8790 *cur = (e, run_idx, i, is_del);
8791 }
8792 })
8793 .or_insert((e, run_idx, i, is_del));
8794 }
8795 }
8796
8797 let overlay_rids: HashSet<u64> = {
8799 let mut s = HashSet::new();
8800 for row in self.memtable.visible_versions(snapshot.epoch) {
8801 s.insert(row.row_id.0);
8802 }
8803 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8804 s.insert(row.row_id.0);
8805 }
8806 s
8807 };
8808
8809 let survivors: Option<RowIdSet> = if conditions.is_empty() {
8811 None
8812 } else {
8813 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8814 for c in conditions {
8815 sets.push(self.resolve_condition(c, snapshot)?);
8816 }
8817 Some(RowIdSet::intersect_many(sets))
8818 };
8819
8820 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
8824 for (rid, (_, run_idx, pos, deleted)) in &best {
8825 if *deleted {
8826 continue;
8827 }
8828 if overlay_rids.contains(rid) {
8829 continue;
8830 }
8831 if let Some(s) = &survivors {
8832 if !s.contains(*rid) {
8833 continue;
8834 }
8835 }
8836 per_run[*run_idx].push((*rid, *pos));
8837 }
8838 for v in per_run.iter_mut() {
8839 v.sort_unstable_by_key(|&(rid, _)| rid);
8840 }
8841
8842 let mut streams = Vec::with_capacity(run_meta.len());
8844 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
8845 let mut total = 0usize;
8846 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
8847 let mut starts = Vec::with_capacity(page_rows.len());
8848 let mut acc = 0usize;
8849 for &r in &page_rows {
8850 starts.push(acc);
8851 acc += r;
8852 }
8853 let mut survivors_vec: Vec<(u64, usize, usize)> =
8854 Vec::with_capacity(per_run[run_idx].len());
8855 for &(rid, pos) in &per_run[run_idx] {
8856 let page_seq = match starts.partition_point(|&s| s <= pos) {
8857 0 => continue,
8858 p => p - 1,
8859 };
8860 let within = pos - starts[page_seq];
8861 survivors_vec.push((rid, page_seq, within));
8862 }
8863 total += survivors_vec.len();
8864 if let Some(&(rid, _, _)) = survivors_vec.first() {
8865 heap.push(std::cmp::Reverse((rid, run_idx)));
8866 }
8867 streams.push(RunStream::new(reader, survivors_vec, page_rows));
8868 }
8869
8870 let overlay_rows = if overlay_rids.is_empty() {
8872 Vec::new()
8873 } else {
8874 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8875 self.overlay_visible_rows(snapshot, bound)
8876 };
8877 let overlay = if overlay_rows.is_empty() {
8878 None
8879 } else {
8880 let filtered =
8881 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8882 if filtered.is_empty() {
8883 None
8884 } else {
8885 Some(self.materialize_overlay(&filtered, &projection))
8886 }
8887 };
8888
8889 let overlay_row_count = overlay
8890 .as_ref()
8891 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8892 .unwrap_or(0);
8893 crate::trace::QueryTrace::record(|t| {
8894 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
8895 t.run_count = self.run_refs.len();
8896 t.memtable_rows = self.memtable.len();
8897 t.mutable_run_rows = self.mutable_run.len();
8898 t.overlay_rows = overlay_row_count;
8899 t.conditions_pushed = conditions.len();
8900 t.survivor_count = Some(total);
8901 });
8902
8903 Ok(Some(MultiRunCursor::new(
8904 streams, projection, heap, total, overlay,
8905 )))
8906 }
8907
8908 fn overlay_materialization_bound<'a>(
8920 conditions: &[crate::query::Condition],
8921 survivors: &'a Option<RowIdSet>,
8922 ) -> Option<&'a RowIdSet> {
8923 use crate::query::Condition;
8924 let has_range = conditions
8925 .iter()
8926 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
8927 if has_range {
8928 None
8929 } else {
8930 survivors.as_ref()
8931 }
8932 }
8933
8934 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
8946 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
8947 let mut fold = |row: Row| {
8948 if let Some(b) = bound {
8949 if !b.contains(row.row_id.0) {
8950 return;
8951 }
8952 }
8953 best.entry(row.row_id.0)
8954 .and_modify(|(be, br)| {
8955 if row.committed_epoch > *be {
8956 *be = row.committed_epoch;
8957 *br = row.clone();
8958 }
8959 })
8960 .or_insert_with(|| (row.committed_epoch, row));
8961 };
8962 for row in self.memtable.visible_versions(snapshot.epoch) {
8963 fold(row);
8964 }
8965 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8966 fold(row);
8967 }
8968 let mut out: Vec<Row> = best
8969 .into_values()
8970 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
8971 .collect();
8972 out.sort_by_key(|r| r.row_id);
8973 out
8974 }
8975
8976 fn filter_overlay_rows(
8984 &self,
8985 rows: Vec<Row>,
8986 conditions: &[crate::query::Condition],
8987 survivors: Option<&RowIdSet>,
8988 snapshot: Snapshot,
8989 ) -> Result<Vec<Row>> {
8990 if conditions.is_empty() {
8991 return Ok(rows);
8992 }
8993 use crate::query::Condition;
8994 let all_index_served = !conditions
8998 .iter()
8999 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9000 if all_index_served {
9001 return Ok(rows
9002 .into_iter()
9003 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9004 .collect());
9005 }
9006 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9009 for c in conditions {
9010 let s = match c {
9011 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9012 _ => self.resolve_condition(c, snapshot)?,
9013 };
9014 per_cond_sets.push(s);
9015 }
9016 Ok(rows
9017 .into_iter()
9018 .filter(|row| {
9019 conditions.iter().enumerate().all(|(i, c)| match c {
9020 Condition::Range { column_id, lo, hi } => {
9021 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9022 }
9023 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9024 match row.columns.get(column_id) {
9025 Some(Value::Float64(v)) => {
9026 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9027 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9028 lo_ok && hi_ok
9029 }
9030 _ => false,
9031 }
9032 }
9033 _ => per_cond_sets[i].contains(row.row_id.0),
9034 })
9035 })
9036 .collect())
9037 }
9038
9039 fn materialize_overlay(
9042 &self,
9043 rows: &[Row],
9044 projection: &[(u16, TypeId)],
9045 ) -> Vec<columnar::NativeColumn> {
9046 if projection.is_empty() {
9047 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9048 }
9049 let mut cols = Vec::with_capacity(projection.len());
9050 for (cid, ty) in projection {
9051 let vals: Vec<Value> = rows
9052 .iter()
9053 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9054 .collect();
9055 cols.push(columnar::values_to_native(ty.clone(), &vals));
9056 }
9057 cols
9058 }
9059
9060 fn resolve_survivor_rids(
9065 &self,
9066 conditions: &[crate::query::Condition],
9067 reader: &mut RunReader,
9068 snapshot: Snapshot,
9069 ) -> Result<RowIdSet> {
9070 use crate::query::Condition;
9071 let mut sets: Vec<RowIdSet> = Vec::new();
9072 for c in conditions {
9073 self.validate_condition(c)?;
9074 let s: RowIdSet = match c {
9075 Condition::Pk(key) => {
9076 let lookup = self
9077 .schema
9078 .primary_key()
9079 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9080 .unwrap_or_else(|| key.clone());
9081 self.hot
9082 .get(&lookup)
9083 .map(|r| RowIdSet::one(r.0))
9084 .unwrap_or_else(RowIdSet::empty)
9085 }
9086 Condition::BitmapEq { column_id, value } => {
9087 let lookup = self.index_lookup_key_bytes(*column_id, value);
9088 self.bitmap
9089 .get(column_id)
9090 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9091 .unwrap_or_else(RowIdSet::empty)
9092 }
9093 Condition::BitmapIn { column_id, values } => {
9094 let bm = self.bitmap.get(column_id);
9095 let mut acc = roaring::RoaringBitmap::new();
9096 if let Some(b) = bm {
9097 for v in values {
9098 let lookup = self.index_lookup_key_bytes(*column_id, v);
9099 acc |= b.get(&lookup);
9100 }
9101 }
9102 RowIdSet::from_roaring(acc)
9103 }
9104 Condition::BytesPrefix { column_id, prefix } => {
9105 if let Some(b) = self.bitmap.get(column_id) {
9106 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9107 let mut acc = roaring::RoaringBitmap::new();
9108 for key in b.keys() {
9109 if key.starts_with(&lookup_prefix) {
9110 acc |= b.get(&key);
9111 }
9112 }
9113 RowIdSet::from_roaring(acc)
9114 } else {
9115 RowIdSet::empty()
9116 }
9117 }
9118 Condition::FmContains { column_id, pattern } => self
9119 .fm
9120 .get(column_id)
9121 .map(|f| {
9122 RowIdSet::from_unsorted(
9123 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9124 )
9125 })
9126 .unwrap_or_else(RowIdSet::empty),
9127 Condition::FmContainsAll {
9128 column_id,
9129 patterns,
9130 } => {
9131 if let Some(f) = self.fm.get(column_id) {
9132 let sets: Vec<RowIdSet> = patterns
9133 .iter()
9134 .map(|pat| {
9135 RowIdSet::from_unsorted(
9136 f.locate(pat).into_iter().map(|r| r.0).collect(),
9137 )
9138 })
9139 .collect();
9140 RowIdSet::intersect_many(sets)
9141 } else {
9142 RowIdSet::empty()
9143 }
9144 }
9145 Condition::Ann {
9146 column_id,
9147 query,
9148 k,
9149 } => RowIdSet::from_unsorted(
9150 self.retrieve_filtered(
9151 &crate::query::Retriever::Ann {
9152 column_id: *column_id,
9153 query: query.clone(),
9154 k: *k,
9155 },
9156 snapshot,
9157 None,
9158 None,
9159 None,
9160 None,
9161 )?
9162 .into_iter()
9163 .map(|hit| hit.row_id.0)
9164 .collect(),
9165 ),
9166 Condition::SparseMatch {
9167 column_id,
9168 query,
9169 k,
9170 } => RowIdSet::from_unsorted(
9171 self.retrieve_filtered(
9172 &crate::query::Retriever::Sparse {
9173 column_id: *column_id,
9174 query: query.clone(),
9175 k: *k,
9176 },
9177 snapshot,
9178 None,
9179 None,
9180 None,
9181 None,
9182 )?
9183 .into_iter()
9184 .map(|hit| hit.row_id.0)
9185 .collect(),
9186 ),
9187 Condition::MinHashSimilar {
9188 column_id,
9189 query,
9190 k,
9191 } => match self.minhash.get(column_id) {
9192 Some(index) => {
9193 let candidates = index.candidate_row_ids(query);
9194 let eligible =
9195 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9196 RowIdSet::from_unsorted(
9197 index
9198 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9199 .into_iter()
9200 .map(|(row_id, _)| row_id.0)
9201 .collect(),
9202 )
9203 }
9204 None => RowIdSet::empty(),
9205 },
9206 Condition::Range { column_id, lo, hi } => {
9207 if let Some(li) = self.learned_range.get(column_id) {
9208 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9209 } else {
9210 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9211 }
9212 }
9213 Condition::RangeF64 {
9214 column_id,
9215 lo,
9216 lo_inclusive,
9217 hi,
9218 hi_inclusive,
9219 } => {
9220 if let Some(li) = self.learned_range.get(column_id) {
9221 RowIdSet::from_unsorted(
9222 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9223 .into_iter()
9224 .collect(),
9225 )
9226 } else {
9227 reader.range_row_id_set_f64(
9228 *column_id,
9229 *lo,
9230 *lo_inclusive,
9231 *hi,
9232 *hi_inclusive,
9233 )?
9234 }
9235 }
9236 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9237 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9238 };
9239 sets.push(s);
9240 }
9241 Ok(RowIdSet::intersect_many(sets))
9242 }
9243
9244 pub fn scan_cursor(
9265 &self,
9266 snapshot: Snapshot,
9267 projection: Vec<(u16, TypeId)>,
9268 conditions: &[crate::query::Condition],
9269 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9270 if self.ttl.is_some() {
9271 return Ok(None);
9272 }
9273 if !conditions.is_empty() && !self.indexes_complete {
9279 return Ok(None);
9280 }
9281 if self.run_refs.len() == 1 {
9282 Ok(self
9283 .native_page_cursor(snapshot, projection, conditions)?
9284 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9285 } else {
9286 Ok(self
9287 .native_multi_run_cursor(snapshot, projection, conditions)?
9288 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9289 }
9290 }
9291
9292 pub fn aggregate_native(
9306 &self,
9307 snapshot: Snapshot,
9308 column: Option<u16>,
9309 conditions: &[crate::query::Condition],
9310 agg: NativeAgg,
9311 ) -> Result<Option<NativeAggResult>> {
9312 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9313 }
9314
9315 pub fn aggregate_native_with_control(
9316 &self,
9317 snapshot: Snapshot,
9318 column: Option<u16>,
9319 conditions: &[crate::query::Condition],
9320 agg: NativeAgg,
9321 control: &crate::ExecutionControl,
9322 ) -> Result<Option<NativeAggResult>> {
9323 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9324 }
9325
9326 fn aggregate_native_inner(
9327 &self,
9328 snapshot: Snapshot,
9329 column: Option<u16>,
9330 conditions: &[crate::query::Condition],
9331 agg: NativeAgg,
9332 control: Option<&crate::ExecutionControl>,
9333 ) -> Result<Option<NativeAggResult>> {
9334 execution_checkpoint(control, 0)?;
9335 if self.ttl.is_some() {
9336 return Ok(None);
9337 }
9338 if self.run_refs.len() == 1 && conditions.is_empty() {
9340 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9341 return Ok(Some(res));
9342 }
9343 }
9344 if matches!(agg, NativeAgg::Count) && column.is_none() {
9346 return Ok(self
9347 .scan_cursor(snapshot, Vec::new(), conditions)?
9348 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
9349 }
9350 let cid = match column {
9353 Some(c) => c,
9354 None => return Ok(None),
9355 };
9356 let ty = self.column_type(cid);
9357 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
9358 else {
9359 return Ok(None);
9360 };
9361 execution_checkpoint(control, 0)?;
9362 match ty {
9363 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9364 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9365 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9366 }
9367 TypeId::Float64 => {
9368 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9369 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9370 }
9371 _ => Ok(None),
9372 }
9373 }
9374
9375 fn aggregate_from_stats(
9383 &self,
9384 snapshot: Snapshot,
9385 column: Option<u16>,
9386 agg: NativeAgg,
9387 ) -> Result<Option<NativeAggResult>> {
9388 let cid = match (agg, column) {
9389 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9390 _ => return Ok(None), };
9392 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9393 return Ok(None);
9394 };
9395 let Some(cs) = stats.get(&cid) else {
9396 return Ok(None);
9397 };
9398 match agg {
9399 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9401 self.live_count.saturating_sub(cs.null_count),
9402 ))),
9403 NativeAgg::Min | NativeAgg::Max => {
9404 let bound = if agg == NativeAgg::Min {
9405 &cs.min
9406 } else {
9407 &cs.max
9408 };
9409 match bound {
9410 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9411 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9412 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9417 None => Ok(None),
9418 }
9419 }
9420 _ => Ok(None),
9421 }
9422 }
9423
9424 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9433 if self.ttl.is_some() {
9434 return Ok(None);
9435 }
9436 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9437 return Ok(None);
9438 }
9439 self.ensure_indexes_complete()?;
9442 let reader = self.open_reader(self.run_refs[0].run_id)?;
9443 if self.live_count != reader.row_count() as u64 {
9444 return Ok(None);
9445 }
9446 let Some(bm) = self.bitmap.get(&column_id) else {
9447 return Ok(None); };
9449 let mut distinct = bm.value_count() as u64;
9450 if !bm.get(&Value::Null.encode_key()).is_empty() {
9453 distinct = distinct.saturating_sub(1);
9454 }
9455 Ok(Some(distinct))
9456 }
9457
9458 pub fn aggregate_incremental(
9470 &mut self,
9471 cache_key: u64,
9472 conditions: &[crate::query::Condition],
9473 column: Option<u16>,
9474 agg: NativeAgg,
9475 ) -> Result<IncrementalAggResult> {
9476 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9477 }
9478
9479 pub fn aggregate_incremental_with_control(
9480 &mut self,
9481 cache_key: u64,
9482 conditions: &[crate::query::Condition],
9483 column: Option<u16>,
9484 agg: NativeAgg,
9485 control: &crate::ExecutionControl,
9486 ) -> Result<IncrementalAggResult> {
9487 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9488 }
9489
9490 fn aggregate_incremental_inner(
9491 &mut self,
9492 cache_key: u64,
9493 conditions: &[crate::query::Condition],
9494 column: Option<u16>,
9495 agg: NativeAgg,
9496 control: Option<&crate::ExecutionControl>,
9497 ) -> Result<IncrementalAggResult> {
9498 execution_checkpoint(control, 0)?;
9499 let snap = self.snapshot();
9500 let cur_wm = self.allocator.current().0;
9501 let cur_epoch = snap.epoch.0;
9502 let incremental_ok = self.ttl.is_none()
9509 && !self.had_deletes
9510 && self.memtable.is_empty()
9511 && self.mutable_run.is_empty();
9512
9513 if incremental_ok {
9516 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
9517 if cached.epoch == cur_epoch {
9518 return Ok(IncrementalAggResult {
9519 state: cached.state,
9520 incremental: true,
9521 delta_rows: 0,
9522 });
9523 }
9524 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
9525 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
9526 let mut delta_rids = Vec::with_capacity(delta_len);
9527 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
9528 execution_checkpoint(control, index)?;
9529 delta_rids.push(row_id);
9530 }
9531 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
9532 execution_checkpoint(control, 0)?;
9533 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9534 let delta_state = agg_state_from_rows(
9535 &delta_rows,
9536 conditions,
9537 &index_sets,
9538 column,
9539 agg,
9540 &self.schema,
9541 control,
9542 )?;
9543 let merged = cached.state.merge(delta_state);
9544 let delta_n = delta_rids.len() as u64;
9545 Arc::make_mut(&mut self.agg_cache).insert(
9546 cache_key,
9547 CachedAgg {
9548 state: merged.clone(),
9549 watermark: cur_wm,
9550 epoch: cur_epoch,
9551 },
9552 );
9553 return Ok(IncrementalAggResult {
9554 state: merged,
9555 incremental: true,
9556 delta_rows: delta_n,
9557 });
9558 }
9559 }
9560 }
9561
9562 let cursor_ok =
9567 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9568 let state = if cursor_ok && agg != NativeAgg::Avg {
9569 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
9570 Some(result) => {
9571 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
9572 }
9573 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
9574 }
9575 } else {
9576 self.agg_state_full_scan(conditions, column, agg, snap, control)?
9577 };
9578 if incremental_ok {
9580 Arc::make_mut(&mut self.agg_cache).insert(
9581 cache_key,
9582 CachedAgg {
9583 state: state.clone(),
9584 watermark: cur_wm,
9585 epoch: cur_epoch,
9586 },
9587 );
9588 }
9589 Ok(IncrementalAggResult {
9590 state,
9591 incremental: false,
9592 delta_rows: 0,
9593 })
9594 }
9595
9596 fn agg_state_full_scan(
9599 &self,
9600 conditions: &[crate::query::Condition],
9601 column: Option<u16>,
9602 agg: NativeAgg,
9603 snap: Snapshot,
9604 control: Option<&crate::ExecutionControl>,
9605 ) -> Result<AggState> {
9606 execution_checkpoint(control, 0)?;
9607 let rows = self.visible_rows(snap)?;
9608 execution_checkpoint(control, 0)?;
9609 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9610 agg_state_from_rows(
9611 &rows,
9612 conditions,
9613 &index_sets,
9614 column,
9615 agg,
9616 &self.schema,
9617 control,
9618 )
9619 }
9620
9621 fn resolve_index_conditions(
9624 &self,
9625 conditions: &[crate::query::Condition],
9626 snapshot: Snapshot,
9627 ) -> Result<Vec<RowIdSet>> {
9628 use crate::query::Condition;
9629 let mut sets = Vec::new();
9630 for c in conditions {
9631 if matches!(
9632 c,
9633 Condition::Ann { .. }
9634 | Condition::SparseMatch { .. }
9635 | Condition::MinHashSimilar { .. }
9636 ) {
9637 sets.push(self.resolve_condition(c, snapshot)?);
9638 }
9639 }
9640 Ok(sets)
9641 }
9642
9643 fn column_type(&self, cid: u16) -> TypeId {
9644 self.schema
9645 .columns
9646 .iter()
9647 .find(|c| c.id == cid)
9648 .map(|c| c.ty.clone())
9649 .unwrap_or(TypeId::Bytes)
9650 }
9651
9652 pub fn approx_aggregate(
9661 &mut self,
9662 conditions: &[crate::query::Condition],
9663 column: Option<u16>,
9664 agg: ApproxAgg,
9665 z: f64,
9666 ) -> Result<Option<ApproxResult>> {
9667 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
9668 }
9669
9670 pub fn approx_aggregate_with_candidate_authorization(
9673 &mut self,
9674 conditions: &[crate::query::Condition],
9675 column: Option<u16>,
9676 agg: ApproxAgg,
9677 z: f64,
9678 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9679 ) -> Result<Option<ApproxResult>> {
9680 use crate::query::Condition;
9681 self.ensure_reservoir_complete()?;
9682 let snapshot = self.snapshot();
9683 let n_pop = self.count();
9684 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
9685 if sample_rids.is_empty() {
9686 return Ok(None);
9687 }
9688 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
9690 let s = live_sample.len();
9691 if s == 0 {
9692 return Ok(None);
9693 }
9694 let authorized = authorization
9695 .map(|authorization| {
9696 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
9697 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
9698 })
9699 .transpose()?;
9700
9701 let mut index_sets: Vec<RowIdSet> = Vec::new();
9704 for c in conditions {
9705 if matches!(
9706 c,
9707 Condition::Ann { .. }
9708 | Condition::SparseMatch { .. }
9709 | Condition::MinHashSimilar { .. }
9710 ) {
9711 index_sets.push(self.resolve_condition(c, snapshot)?);
9712 }
9713 }
9714
9715 let cid = match (agg, column) {
9717 (ApproxAgg::Count, _) => None,
9718 (_, Some(c)) => Some(c),
9719 _ => return Ok(None),
9720 };
9721 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
9722 for r in &live_sample {
9723 if authorized
9724 .as_ref()
9725 .is_some_and(|authorized| !authorized.contains(&r.row_id))
9726 {
9727 continue;
9728 }
9729 if !conditions
9731 .iter()
9732 .all(|c| condition_matches_row(c, r, &self.schema))
9733 {
9734 continue;
9735 }
9736 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
9738 continue;
9739 }
9740 if let Some(cid) = cid {
9741 let mut cells = r
9742 .columns
9743 .get(&cid)
9744 .cloned()
9745 .map(|value| vec![(cid, value)])
9746 .unwrap_or_default();
9747 if let Some(authorization) = authorization {
9748 authorization.security.apply_masks_to_cells(
9749 authorization.table,
9750 &mut cells,
9751 authorization.principal,
9752 );
9753 }
9754 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
9755 passing_vals.push(v);
9756 } } else {
9758 passing_vals.push(0.0); }
9760 }
9761 let m = passing_vals.len();
9762
9763 let (point, half) = match agg {
9764 ApproxAgg::Count => {
9765 let p = m as f64 / s as f64;
9767 let point = n_pop as f64 * p;
9768 let var = if s > 1 {
9769 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
9770 * (1.0 - s as f64 / n_pop as f64).max(0.0)
9771 } else {
9772 0.0
9773 };
9774 (point, z * var.sqrt())
9775 }
9776 ApproxAgg::Sum => {
9777 let y: Vec<f64> = live_sample
9779 .iter()
9780 .map(|r| {
9781 let passes_row = authorized
9782 .as_ref()
9783 .is_none_or(|authorized| authorized.contains(&r.row_id))
9784 && conditions
9785 .iter()
9786 .all(|c| condition_matches_row(c, r, &self.schema))
9787 && index_sets.iter().all(|set| set.contains(r.row_id.0));
9788 if passes_row {
9789 cid.and_then(|cid| {
9790 let mut cells = r
9791 .columns
9792 .get(&cid)
9793 .cloned()
9794 .map(|value| vec![(cid, value)])
9795 .unwrap_or_default();
9796 if let Some(authorization) = authorization {
9797 authorization.security.apply_masks_to_cells(
9798 authorization.table,
9799 &mut cells,
9800 authorization.principal,
9801 );
9802 }
9803 as_f64(cells.first().map(|(_, value)| value))
9804 })
9805 .unwrap_or(0.0)
9806 } else {
9807 0.0
9808 }
9809 })
9810 .collect();
9811 let mean_y = y.iter().sum::<f64>() / s as f64;
9812 let point = n_pop as f64 * mean_y;
9813 let var = if s > 1 {
9814 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
9815 let var_y = ss / (s - 1) as f64;
9816 n_pop as f64 * n_pop as f64 * var_y / s as f64
9817 * (1.0 - s as f64 / n_pop as f64).max(0.0)
9818 } else {
9819 0.0
9820 };
9821 (point, z * var.sqrt())
9822 }
9823 ApproxAgg::Avg => {
9824 if m == 0 {
9825 return Ok(Some(ApproxResult {
9826 point: 0.0,
9827 ci_low: 0.0,
9828 ci_high: 0.0,
9829 n_population: n_pop,
9830 n_sample_live: s,
9831 n_passing: 0,
9832 }));
9833 }
9834 let mean = passing_vals.iter().sum::<f64>() / m as f64;
9835 let half = if m > 1 {
9836 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
9837 let sd = (ss / (m - 1) as f64).sqrt();
9838 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
9839 z * sd / (m as f64).sqrt() * fpc.sqrt()
9840 } else {
9841 0.0
9842 };
9843 (mean, half)
9844 }
9845 };
9846
9847 Ok(Some(ApproxResult {
9848 point,
9849 ci_low: point - half,
9850 ci_high: point + half,
9851 n_population: n_pop,
9852 n_sample_live: s,
9853 n_passing: m,
9854 }))
9855 }
9856
9857 pub fn exact_column_stats(
9865 &self,
9866 _snapshot: Snapshot,
9867 projection: &[u16],
9868 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
9869 if self.ttl.is_some()
9870 || !(self.memtable.is_empty()
9871 && self.mutable_run.is_empty()
9872 && self.run_refs.len() == 1)
9873 {
9874 return Ok(None);
9875 }
9876 let reader = self.open_reader(self.run_refs[0].run_id)?;
9877 if self.live_count != reader.row_count() as u64 {
9878 return Ok(None);
9879 }
9880 let mut out = HashMap::new();
9881 for &cid in projection {
9882 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
9883 Some(c) => c,
9884 None => continue,
9885 };
9886 let Some(stats) = reader.column_page_stats(cid) else {
9888 out.insert(
9889 cid,
9890 ColumnStat {
9891 min: None,
9892 max: None,
9893 null_count: self.live_count,
9894 },
9895 );
9896 continue;
9897 };
9898 let stat = match cdef.ty {
9899 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9900 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
9901 min: mn.map(Value::Int64),
9902 max: mx.map(Value::Int64),
9903 null_count: n,
9904 })
9905 }
9906 TypeId::Float64 => {
9907 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
9908 min: mn.map(Value::Float64),
9909 max: mx.map(Value::Float64),
9910 null_count: n,
9911 })
9912 }
9913 _ => None,
9914 };
9915 if let Some(s) = stat {
9916 out.insert(cid, s);
9917 }
9918 }
9919 Ok(Some(out))
9920 }
9921
9922 pub fn dir(&self) -> &Path {
9923 &self.dir
9924 }
9925
9926 pub fn schema(&self) -> &Schema {
9927 &self.schema
9928 }
9929
9930 pub(crate) fn set_catalog_name(&mut self, name: String) {
9931 self.name = name;
9932 }
9933
9934 pub(crate) fn prepare_alter_column(
9935 &mut self,
9936 column_name: &str,
9937 change: &AlterColumn,
9938 ) -> Result<(ColumnDef, Option<Schema>)> {
9939 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
9940 return Err(MongrelError::InvalidArgument(
9941 "ALTER COLUMN requires committing staged writes first".into(),
9942 ));
9943 }
9944 let old = self
9945 .schema
9946 .columns
9947 .iter()
9948 .find(|c| c.name == column_name)
9949 .cloned()
9950 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
9951 let mut next = old.clone();
9952
9953 if let Some(name) = &change.name {
9954 let trimmed = name.trim();
9955 if trimmed.is_empty() {
9956 return Err(MongrelError::InvalidArgument(
9957 "ALTER COLUMN name must not be empty".into(),
9958 ));
9959 }
9960 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
9961 return Err(MongrelError::Schema(format!(
9962 "column {trimmed} already exists"
9963 )));
9964 }
9965 next.name = trimmed.to_string();
9966 }
9967
9968 if let Some(ty) = &change.ty {
9969 next.ty = ty.clone();
9970 }
9971 if let Some(flags) = change.flags {
9972 validate_alter_column_flags(old.flags, flags)?;
9973 next.flags = flags;
9974 }
9975
9976 if let Some(default_change) = &change.default_value {
9977 next.default_value = default_change.clone();
9978 }
9979
9980 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
9981 if old.flags.contains(ColumnFlags::NULLABLE)
9982 && !next.flags.contains(ColumnFlags::NULLABLE)
9983 && self.column_has_nulls(old.id)?
9984 {
9985 return Err(MongrelError::InvalidArgument(format!(
9986 "column '{}' contains NULL values",
9987 old.name
9988 )));
9989 }
9990 if next == old {
9991 return Ok((next, None));
9992 }
9993 let mut schema = self.schema.clone();
9994 let index = schema
9995 .columns
9996 .iter()
9997 .position(|column| column.id == next.id)
9998 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
9999 schema.columns[index] = next.clone();
10000 schema.schema_id = schema
10001 .schema_id
10002 .checked_add(1)
10003 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10004 schema.validate_auto_increment()?;
10005 schema.validate_defaults()?;
10006 Ok((next, Some(schema)))
10007 }
10008
10009 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10010 self.schema = schema;
10011 self.auto_inc = resolve_auto_inc(&self.schema);
10012 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10013 self.clear_result_cache();
10014 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10015 }
10016
10017 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10018 checkpoint_current_schema(self)
10019 }
10020
10021 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10022 self.ensure_writable()?;
10023 let previous_schema = self.schema.clone();
10024 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10025 if let Some(schema) = schema {
10026 self.apply_altered_schema_prepared(schema);
10027 self.checkpoint_standalone_schema_change(previous_schema)?;
10028 }
10029 Ok(column)
10030 }
10031
10032 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10033 if self.live_count == 0 {
10034 return Ok(false);
10035 }
10036 let snap = self.snapshot();
10037 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10038 Ok(columns
10039 .first()
10040 .map(|(_, col)| col.null_count(col.len()) != 0)
10041 .unwrap_or(true))
10042 }
10043
10044 fn has_stored_versions(&self) -> bool {
10045 !self.memtable.is_empty()
10046 || !self.mutable_run.is_empty()
10047 || self.run_refs.iter().any(|r| r.row_count > 0)
10048 || !self.retiring.is_empty()
10049 }
10050
10051 pub fn add_column(
10056 &mut self,
10057 name: &str,
10058 ty: TypeId,
10059 flags: ColumnFlags,
10060 default_value: Option<crate::schema::DefaultExpr>,
10061 ) -> Result<u16> {
10062 self.add_column_with_id(name, ty, flags, default_value, None)
10063 }
10064
10065 pub fn add_column_with_id(
10066 &mut self,
10067 name: &str,
10068 ty: TypeId,
10069 flags: ColumnFlags,
10070 default_value: Option<crate::schema::DefaultExpr>,
10071 requested_id: Option<u16>,
10072 ) -> Result<u16> {
10073 self.ensure_writable()?;
10074 if self.schema.columns.iter().any(|c| c.name == name) {
10075 return Err(MongrelError::Schema(format!(
10076 "column {name} already exists"
10077 )));
10078 }
10079 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10080 if self.schema.columns.iter().any(|c| c.id == id) {
10081 return Err(MongrelError::Schema(format!(
10082 "column id {id} already exists"
10083 )));
10084 }
10085 id
10086 } else {
10087 self.schema
10088 .columns
10089 .iter()
10090 .map(|c| c.id)
10091 .max()
10092 .unwrap_or(0)
10093 .checked_add(1)
10094 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10095 };
10096 let previous_schema = self.schema.clone();
10097 let mut next_schema = previous_schema.clone();
10098 next_schema.columns.push(ColumnDef {
10099 id,
10100 name: name.to_string(),
10101 ty,
10102 flags,
10103 default_value,
10104 });
10105 next_schema.schema_id = next_schema
10106 .schema_id
10107 .checked_add(1)
10108 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10109 next_schema.validate_auto_increment()?;
10110 next_schema.validate_defaults()?;
10111 self.apply_altered_schema_prepared(next_schema);
10112 self.checkpoint_standalone_schema_change(previous_schema)?;
10113 Ok(id)
10114 }
10115
10116 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10125 self.ensure_writable()?;
10126 let cid = self
10127 .schema
10128 .columns
10129 .iter()
10130 .find(|c| c.name == column_name)
10131 .map(|c| c.id)
10132 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10133 let ty = self
10134 .schema
10135 .columns
10136 .iter()
10137 .find(|c| c.id == cid)
10138 .map(|c| c.ty.clone())
10139 .unwrap_or(TypeId::Int64);
10140 if !matches!(
10141 ty,
10142 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10143 ) {
10144 return Err(MongrelError::Schema(format!(
10145 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10146 )));
10147 }
10148 if self
10149 .schema
10150 .indexes
10151 .iter()
10152 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10153 {
10154 return Ok(()); }
10156 let previous_schema = self.schema.clone();
10157 let previous_learned_range = Arc::clone(&self.learned_range);
10158 let mut next_schema = previous_schema.clone();
10159 next_schema.indexes.push(IndexDef {
10160 name: format!("{}_learned_range", column_name),
10161 column_id: cid,
10162 kind: IndexKind::LearnedRange,
10163 predicate: None,
10164 options: Default::default(),
10165 });
10166 next_schema.schema_id = next_schema
10167 .schema_id
10168 .checked_add(1)
10169 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10170 self.apply_altered_schema_prepared(next_schema);
10171 if let Err(error) = self.build_learned_ranges() {
10172 self.apply_altered_schema_prepared(previous_schema);
10173 self.learned_range = previous_learned_range;
10174 return Err(error);
10175 }
10176 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10177 if !matches!(
10178 &error,
10179 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10180 ) {
10181 self.learned_range = previous_learned_range;
10182 }
10183 return Err(error);
10184 }
10185 Ok(())
10186 }
10187
10188 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10189 let mut schema_published = false;
10190 let schema_result = match self._root_guard.as_deref() {
10191 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10192 schema_published = true;
10193 }),
10194 None => write_schema_with_after(&self.dir, &self.schema, || {
10195 schema_published = true;
10196 }),
10197 };
10198 if schema_result.is_err() && !schema_published {
10199 self.apply_altered_schema_prepared(previous_schema);
10200 return schema_result;
10201 }
10202
10203 let manifest_result = self.persist_manifest(self.current_epoch());
10204 match (schema_result, manifest_result) {
10205 (_, Ok(())) => Ok(()),
10206 (Ok(()), Err(error)) => {
10207 self.poison_after_maintenance_publish_failure();
10208 Err(MongrelError::DurableCommit {
10209 epoch: self.current_epoch().0,
10210 message: format!(
10211 "schema is durable but matching manifest publication failed: {error}"
10212 ),
10213 })
10214 }
10215 (Err(schema_error), Err(manifest_error)) => {
10216 self.poison_after_maintenance_publish_failure();
10217 Err(MongrelError::CommitOutcomeUnknown {
10218 epoch: self.current_epoch().0,
10219 message: format!(
10220 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10221 ),
10222 })
10223 }
10224 }
10225 }
10226
10227 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10230 self.sync_byte_threshold = threshold;
10231 if let WalSink::Private(w) = &mut self.wal {
10232 w.set_sync_byte_threshold(threshold);
10233 }
10234 }
10235
10236 pub fn page_cache_flush(&self) {
10240 self.page_cache.flush_to_disk();
10241 }
10242
10243 pub fn page_cache_len(&self) -> usize {
10245 self.page_cache.len()
10246 }
10247
10248 pub fn decoded_cache_len(&self) -> usize {
10251 self.decoded_cache.len()
10252 }
10253
10254 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10257 self.memtable.drain_sorted()
10258 }
10259
10260 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10261 self.runs_dir().join(format!("r-{run_id}.sr"))
10262 }
10263
10264 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10265 match self.runs_root.as_deref() {
10266 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10267 None => Ok(None),
10268 }
10269 }
10270
10271 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10272 match self.runs_root.as_deref() {
10273 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10274 None => Ok(None),
10275 }
10276 }
10277
10278 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10279 match self.runs_root.as_deref() {
10280 Some(root) => match root.remove_file(name) {
10281 Ok(()) => Ok(()),
10282 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10283 Err(error) => Err(error.into()),
10284 },
10285 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10286 Ok(()) => Ok(()),
10287 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10288 Err(error) => Err(error.into()),
10289 },
10290 }
10291 }
10292
10293 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10294 match self.runs_root.as_deref() {
10295 Some(root) => root
10296 .rename_file_new(source, destination)
10297 .map_err(Into::into),
10298 None => crate::durable_file::rename(
10299 &self.runs_dir().join(source),
10300 &self.runs_dir().join(destination),
10301 )
10302 .map_err(Into::into),
10303 }
10304 }
10305
10306 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10307 self.run_refs.iter().map(|run| run.run_id)
10308 }
10309
10310 pub(crate) fn table_dir(&self) -> &Path {
10311 &self.dir
10312 }
10313
10314 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10315 &self.schema
10316 }
10317
10318 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10319 let id = self.next_run_id;
10320 self.next_run_id = self
10321 .next_run_id
10322 .checked_add(1)
10323 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10324 Ok(id)
10325 }
10326
10327 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10328 self.run_refs.push(run_ref);
10329 }
10330
10331 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10341 self.retiring.push(crate::manifest::RetiredRun {
10342 run_id,
10343 retire_epoch,
10344 });
10345 }
10346
10347 pub(crate) fn reap_retiring(
10351 &mut self,
10352 min_active: Epoch,
10353 backup_pinned: &std::collections::HashSet<u128>,
10354 ) -> Result<usize> {
10355 if self.retiring.is_empty() {
10356 return Ok(0);
10357 }
10358 let mut reaped = 0;
10359 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
10360 for r in std::mem::take(&mut self.retiring) {
10366 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
10367 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
10368 reaped += 1;
10369 } else {
10370 kept.push(r);
10371 }
10372 }
10373 self.retiring = kept;
10374 if reaped > 0 {
10375 self.persist_manifest(self.current_epoch())?;
10376 }
10377 Ok(reaped)
10378 }
10379
10380 pub(crate) fn has_reapable_retiring(
10381 &self,
10382 min_active: Epoch,
10383 backup_pinned: &std::collections::HashSet<u128>,
10384 ) -> bool {
10385 self.retiring
10386 .iter()
10387 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
10388 }
10389
10390 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
10391 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
10392 return false;
10393 }
10394 self.live_count = self.live_count.saturating_add(run_ref.row_count);
10395 self.run_refs.push(run_ref);
10396 self.indexes_complete = false;
10397 true
10398 }
10399
10400 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
10401 self.kek.as_ref()
10402 }
10403
10404 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
10405 let mut reader = match self.runs_root.as_deref() {
10406 Some(root) => RunReader::open_file_with_cache(
10407 root.open_regular(format!("r-{run_id}.sr"))?,
10408 self.schema.clone(),
10409 self.kek.clone(),
10410 Some(self.page_cache.clone()),
10411 Some(self.decoded_cache.clone()),
10412 self.table_id,
10413 Some(&self.verified_runs),
10414 None,
10415 )?,
10416 None => RunReader::open_with_cache(
10417 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
10418 self.schema.clone(),
10419 self.kek.clone(),
10420 Some(self.page_cache.clone()),
10421 Some(self.decoded_cache.clone()),
10422 self.table_id,
10423 Some(&self.verified_runs),
10424 )?,
10425 };
10426 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
10430 reader.set_uniform_epoch(Epoch(rr.epoch_created));
10431 }
10432 Ok(reader)
10433 }
10434
10435 pub(crate) fn run_refs(&self) -> &[RunRef] {
10436 &self.run_refs
10437 }
10438
10439 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10440 self.retiring.iter().map(|run| run.run_id)
10441 }
10442
10443 pub(crate) fn runs_dir(&self) -> PathBuf {
10444 self.runs_root
10445 .as_deref()
10446 .and_then(|root| root.io_path().ok())
10447 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
10448 }
10449
10450 pub(crate) fn wal_dir(&self) -> PathBuf {
10451 self.dir.join(WAL_DIR)
10452 }
10453
10454 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
10455 self.run_refs = refs;
10456 }
10457
10458 pub(crate) fn compaction_zstd_level(&self) -> i32 {
10459 self.compaction_zstd_level
10460 }
10461
10462 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
10463 self.kek.clone()
10464 }
10465
10466 #[cfg(feature = "encryption")]
10470 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10471 self.kek.as_ref().map(|k| k.derive_idx_key())
10472 }
10473
10474 #[cfg(not(feature = "encryption"))]
10475 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10476 None
10477 }
10478
10479 #[cfg(feature = "encryption")]
10483 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10484 self.kek.as_ref().map(|k| *k.derive_meta_key())
10485 }
10486
10487 #[cfg(not(feature = "encryption"))]
10488 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10489 None
10490 }
10491
10492 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
10495 self.column_keys
10496 .iter()
10497 .map(|(&id, &(_, scheme))| (id, scheme))
10498 .collect()
10499 }
10500
10501 #[cfg(feature = "encryption")]
10506 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
10507 self.tokenize_value_enc(column_id, v)
10508 }
10509
10510 #[cfg(feature = "encryption")]
10511 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
10512 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10513 let (key, scheme) = self.column_keys.get(&column_id)?;
10514 let token: Vec<u8> = match (*scheme, v) {
10515 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
10516 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
10517 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
10518 _ => hmac_token(key, &v.encode_key()).to_vec(),
10519 };
10520 Some(Value::Bytes(token))
10521 }
10522
10523 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
10525 self.index_lookup_key_bytes(column_id, &v.encode_key())
10526 }
10527
10528 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
10531 #[cfg(feature = "encryption")]
10532 {
10533 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
10534 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
10535 if *scheme == SCHEME_HMAC_EQ {
10536 return hmac_token(key, encoded).to_vec();
10537 }
10538 }
10539 }
10540 let _ = column_id;
10541 encoded.to_vec()
10542 }
10543}
10544
10545fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
10546 let columnar::NativeColumn::Int64 { data, validity } = col else {
10547 return false;
10548 };
10549 if data.len() < n || !columnar::all_non_null(validity, n) {
10550 return false;
10551 }
10552 data.iter()
10553 .take(n)
10554 .zip(data.iter().skip(1))
10555 .all(|(a, b)| a < b)
10556}
10557
10558#[derive(Debug, Clone)]
10562pub struct ColumnStat {
10563 pub min: Option<Value>,
10564 pub max: Option<Value>,
10565 pub null_count: u64,
10566}
10567
10568#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10570pub enum NativeAgg {
10571 Count,
10572 Sum,
10573 Min,
10574 Max,
10575 Avg,
10576}
10577
10578#[derive(Debug, Clone, PartialEq)]
10580pub enum NativeAggResult {
10581 Count(u64),
10582 Int(i64),
10583 Float(f64),
10584 Null,
10586}
10587
10588#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10590pub enum ApproxAgg {
10591 Count,
10592 Sum,
10593 Avg,
10594}
10595
10596#[derive(Debug, Clone)]
10600pub struct ApproxResult {
10601 pub point: f64,
10603 pub ci_low: f64,
10605 pub ci_high: f64,
10607 pub n_population: u64,
10609 pub n_sample_live: usize,
10611 pub n_passing: usize,
10613}
10614
10615#[derive(Debug, Clone, PartialEq)]
10620pub enum AggState {
10621 Count(u64),
10623 SumI {
10625 sum: i128,
10626 count: u64,
10627 },
10628 SumF {
10630 sum: f64,
10631 count: u64,
10632 },
10633 AvgI {
10635 sum: i128,
10636 count: u64,
10637 },
10638 AvgF {
10640 sum: f64,
10641 count: u64,
10642 },
10643 MinI(i64),
10645 MaxI(i64),
10646 MinF(f64),
10648 MaxF(f64),
10649 Empty,
10651}
10652
10653impl AggState {
10654 pub fn merge(self, other: AggState) -> AggState {
10656 use AggState::*;
10657 match (self, other) {
10658 (Empty, x) | (x, Empty) => x,
10659 (Count(a), Count(b)) => Count(a + b),
10660 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
10661 sum: sa + sb,
10662 count: ca + cb,
10663 },
10664 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
10665 sum: sa + sb,
10666 count: ca + cb,
10667 },
10668 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
10669 sum: sa + sb,
10670 count: ca + cb,
10671 },
10672 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
10673 sum: sa + sb,
10674 count: ca + cb,
10675 },
10676 (MinI(a), MinI(b)) => MinI(a.min(b)),
10677 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
10678 (MinF(a), MinF(b)) => MinF(a.min(b)),
10679 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
10680 _ => Empty, }
10682 }
10683
10684 pub fn point(&self) -> Option<f64> {
10686 match self {
10687 AggState::Count(n) => Some(*n as f64),
10688 AggState::SumI { sum, .. } => Some(*sum as f64),
10689 AggState::SumF { sum, .. } => Some(*sum),
10690 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
10691 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
10692 AggState::MinI(n) => Some(*n as f64),
10693 AggState::MaxI(n) => Some(*n as f64),
10694 AggState::MinF(n) => Some(*n),
10695 AggState::MaxF(n) => Some(*n),
10696 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
10697 }
10698 }
10699
10700 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
10704 let is_float = matches!(ty, Some(TypeId::Float64));
10705 match (agg, result) {
10706 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
10707 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
10708 sum: x as i128,
10709 count: 1, },
10711 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
10712 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
10713 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
10714 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
10715 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
10716 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
10717 (NativeAgg::Count, _) => AggState::Empty,
10718 (_, NativeAggResult::Null) => AggState::Empty,
10719 _ => {
10720 let _ = is_float;
10721 AggState::Empty
10722 }
10723 }
10724 }
10725}
10726
10727#[derive(Debug, Clone)]
10730pub struct CachedAgg {
10731 pub state: AggState,
10732 pub watermark: u64,
10733 pub epoch: u64,
10734}
10735
10736#[derive(Debug, Clone)]
10738pub struct IncrementalAggResult {
10739 pub state: AggState,
10741 pub incremental: bool,
10744 pub delta_rows: u64,
10746}
10747
10748fn agg_state_from_rows(
10752 rows: &[Row],
10753 conditions: &[crate::query::Condition],
10754 index_sets: &[RowIdSet],
10755 column: Option<u16>,
10756 agg: NativeAgg,
10757 schema: &Schema,
10758 control: Option<&crate::ExecutionControl>,
10759) -> Result<AggState> {
10760 let mut count: u64 = 0;
10761 let mut sum_i: i128 = 0;
10762 let mut sum_f: f64 = 0.0;
10763 let mut mn_i: i64 = i64::MAX;
10764 let mut mx_i: i64 = i64::MIN;
10765 let mut mn_f: f64 = f64::INFINITY;
10766 let mut mx_f: f64 = f64::NEG_INFINITY;
10767 let mut saw_int = false;
10768 let mut saw_float = false;
10769 for (index, r) in rows.iter().enumerate() {
10770 execution_checkpoint(control, index)?;
10771 if !conditions
10772 .iter()
10773 .all(|c| condition_matches_row(c, r, schema))
10774 {
10775 continue;
10776 }
10777 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
10778 continue;
10779 }
10780 match agg {
10781 NativeAgg::Count => match column {
10782 None => count += 1,
10784 Some(cid) => match r.columns.get(&cid) {
10787 None | Some(Value::Null) => {}
10788 Some(_) => count += 1,
10789 },
10790 },
10791 _ => match column.and_then(|cid| r.columns.get(&cid)) {
10792 Some(Value::Int64(n)) => {
10793 count += 1;
10794 sum_i += *n as i128;
10795 mn_i = mn_i.min(*n);
10796 mx_i = mx_i.max(*n);
10797 saw_int = true;
10798 }
10799 Some(Value::Float64(f)) => {
10800 count += 1;
10801 sum_f += f;
10802 mn_f = mn_f.min(*f);
10803 mx_f = mx_f.max(*f);
10804 saw_float = true;
10805 }
10806 _ => {}
10807 },
10808 }
10809 }
10810 Ok(match agg {
10811 NativeAgg::Count => {
10812 if count == 0 {
10813 AggState::Empty
10814 } else {
10815 AggState::Count(count)
10816 }
10817 }
10818 NativeAgg::Sum => {
10819 if count == 0 {
10820 AggState::Empty
10821 } else if saw_int {
10822 AggState::SumI { sum: sum_i, count }
10823 } else {
10824 AggState::SumF { sum: sum_f, count }
10825 }
10826 }
10827 NativeAgg::Avg => {
10828 if count == 0 {
10829 AggState::Empty
10830 } else if saw_int {
10831 AggState::AvgI { sum: sum_i, count }
10832 } else {
10833 AggState::AvgF { sum: sum_f, count }
10834 }
10835 }
10836 NativeAgg::Min => {
10837 if !saw_int && !saw_float {
10838 AggState::Empty
10839 } else if saw_int {
10840 AggState::MinI(mn_i)
10841 } else {
10842 AggState::MinF(mn_f)
10843 }
10844 }
10845 NativeAgg::Max => {
10846 if !saw_int && !saw_float {
10847 AggState::Empty
10848 } else if saw_int {
10849 AggState::MaxI(mx_i)
10850 } else {
10851 AggState::MaxF(mx_f)
10852 }
10853 }
10854 })
10855}
10856
10857fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
10861 use crate::query::Condition;
10862 match c {
10863 Condition::Pk(key) => match schema.primary_key() {
10864 Some(pk) => row
10865 .columns
10866 .get(&pk.id)
10867 .map(|v| v.encode_key() == *key)
10868 .unwrap_or(false),
10869 None => false,
10870 },
10871 Condition::BitmapEq { column_id, value } => row
10872 .columns
10873 .get(column_id)
10874 .map(|v| v.encode_key() == *value)
10875 .unwrap_or(false),
10876 Condition::BitmapIn { column_id, values } => {
10877 let key = row.columns.get(column_id).map(|v| v.encode_key());
10878 match key {
10879 Some(k) => values.contains(&k),
10880 None => false,
10881 }
10882 }
10883 Condition::BytesPrefix { column_id, prefix } => row
10884 .columns
10885 .get(column_id)
10886 .map(|v| v.encode_key().starts_with(prefix))
10887 .unwrap_or(false),
10888 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
10889 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
10890 _ => false,
10891 },
10892 Condition::RangeF64 {
10893 column_id,
10894 lo,
10895 lo_inclusive,
10896 hi,
10897 hi_inclusive,
10898 } => match row.columns.get(column_id) {
10899 Some(Value::Float64(n)) => {
10900 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
10901 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
10902 lo_ok && hi_ok
10903 }
10904 _ => false,
10905 },
10906 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
10907 Some(Value::Bytes(b)) => {
10908 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
10909 }
10910 _ => false,
10911 },
10912 Condition::FmContainsAll {
10913 column_id,
10914 patterns,
10915 } => match row.columns.get(column_id) {
10916 Some(Value::Bytes(b)) => patterns
10917 .iter()
10918 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
10919 _ => false,
10920 },
10921 Condition::Ann { .. }
10922 | Condition::SparseMatch { .. }
10923 | Condition::MinHashSimilar { .. } => true,
10924 Condition::IsNull { column_id } => {
10925 matches!(row.columns.get(column_id), Some(Value::Null) | None)
10926 }
10927 Condition::IsNotNull { column_id } => {
10928 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
10929 }
10930 }
10931}
10932
10933fn as_f64(v: Option<&Value>) -> Option<f64> {
10935 match v {
10936 Some(Value::Int64(n)) => Some(*n as f64),
10937 Some(Value::Float64(f)) => Some(*f),
10938 _ => None,
10939 }
10940}
10941
10942fn accumulate_int(
10946 cursor: &mut dyn crate::cursor::Cursor,
10947 control: Option<&crate::ExecutionControl>,
10948) -> Result<(u64, i128, i64, i64)> {
10949 let mut count: u64 = 0;
10950 let mut sum: i128 = 0;
10951 let mut mn: i64 = i64::MAX;
10952 let mut mx: i64 = i64::MIN;
10953 while let Some(cols) = cursor.next_batch()? {
10954 execution_checkpoint(control, 0)?;
10955 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
10956 if crate::columnar::all_non_null(validity, data.len()) {
10957 count += data.len() as u64;
10959 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
10960 execution_checkpoint(control, chunk_index * 1024)?;
10961 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
10962 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
10963 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
10964 }
10965 } else {
10966 for (i, &v) in data.iter().enumerate() {
10967 execution_checkpoint(control, i)?;
10968 if crate::columnar::validity_bit(validity, i) {
10969 count += 1;
10970 sum += v as i128;
10971 mn = mn.min(v);
10972 mx = mx.max(v);
10973 }
10974 }
10975 }
10976 }
10977 }
10978 Ok((count, sum, mn, mx))
10979}
10980
10981fn accumulate_float(
10983 cursor: &mut dyn crate::cursor::Cursor,
10984 control: Option<&crate::ExecutionControl>,
10985) -> Result<(u64, f64, f64, f64)> {
10986 let mut count: u64 = 0;
10987 let mut sum: f64 = 0.0;
10988 let mut mn: f64 = f64::INFINITY;
10989 let mut mx: f64 = f64::NEG_INFINITY;
10990 while let Some(cols) = cursor.next_batch()? {
10991 execution_checkpoint(control, 0)?;
10992 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
10993 if crate::columnar::all_non_null(validity, data.len()) {
10994 count += data.len() as u64;
10995 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
10996 execution_checkpoint(control, chunk_index * 1024)?;
10997 sum += chunk.iter().sum::<f64>();
10998 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
10999 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11000 }
11001 } else {
11002 for (i, &v) in data.iter().enumerate() {
11003 execution_checkpoint(control, i)?;
11004 if crate::columnar::validity_bit(validity, i) {
11005 count += 1;
11006 sum += v;
11007 mn = mn.min(v);
11008 mx = mx.max(v);
11009 }
11010 }
11011 }
11012 }
11013 }
11014 Ok((count, sum, mn, mx))
11015}
11016
11017#[inline]
11018fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11019 if index.is_multiple_of(256) {
11020 control
11021 .map(crate::ExecutionControl::checkpoint)
11022 .transpose()?;
11023 }
11024 Ok(())
11025}
11026
11027fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11028 if count == 0 && !matches!(agg, NativeAgg::Count) {
11029 return NativeAggResult::Null;
11030 }
11031 match agg {
11032 NativeAgg::Count => NativeAggResult::Count(count),
11033 NativeAgg::Sum => match sum.try_into() {
11036 Ok(v) => NativeAggResult::Int(v),
11037 Err(_) => NativeAggResult::Null,
11038 },
11039 NativeAgg::Min => NativeAggResult::Int(mn),
11040 NativeAgg::Max => NativeAggResult::Int(mx),
11041 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11042 }
11043}
11044
11045fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11046 if count == 0 && !matches!(agg, NativeAgg::Count) {
11047 return NativeAggResult::Null;
11048 }
11049 match agg {
11050 NativeAgg::Count => NativeAggResult::Count(count),
11051 NativeAgg::Sum => NativeAggResult::Float(sum),
11052 NativeAgg::Min => NativeAggResult::Float(mn),
11053 NativeAgg::Max => NativeAggResult::Float(mx),
11054 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11055 }
11056}
11057
11058fn agg_int(
11061 stats: &[crate::page::PageStat],
11062 decode: fn(Option<&[u8]>) -> Option<i64>,
11063) -> Option<(Option<i64>, Option<i64>, u64)> {
11064 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11065 let mut any = false;
11066 for s in stats {
11067 if let Some(v) = decode(s.min.as_deref()) {
11068 mn = mn.min(v);
11069 any = true;
11070 }
11071 if let Some(v) = decode(s.max.as_deref()) {
11072 mx = mx.max(v);
11073 any = true;
11074 }
11075 nulls += s.null_count;
11076 }
11077 any.then_some((Some(mn), Some(mx), nulls))
11078}
11079
11080fn agg_float(
11082 stats: &[crate::page::PageStat],
11083 decode: fn(Option<&[u8]>) -> Option<f64>,
11084) -> Option<(Option<f64>, Option<f64>, u64)> {
11085 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11086 let mut any = false;
11087 for s in stats {
11088 if let Some(v) = decode(s.min.as_deref()) {
11089 mn = mn.min(v);
11090 any = true;
11091 }
11092 if let Some(v) = decode(s.max.as_deref()) {
11093 mx = mx.max(v);
11094 any = true;
11095 }
11096 nulls += s.null_count;
11097 }
11098 any.then_some((Some(mn), Some(mx), nulls))
11099}
11100
11101type SecondaryIndexes = (
11103 HashMap<u16, BitmapIndex>,
11104 HashMap<u16, AnnIndex>,
11105 HashMap<u16, FmIndex>,
11106 HashMap<u16, SparseIndex>,
11107 HashMap<u16, MinHashIndex>,
11108);
11109
11110fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11111 let mut bitmap = HashMap::new();
11112 let mut ann = HashMap::new();
11113 let mut fm = HashMap::new();
11114 let mut sparse = HashMap::new();
11115 let mut minhash = HashMap::new();
11116 for idef in &schema.indexes {
11117 match idef.kind {
11118 IndexKind::Bitmap => {
11119 bitmap.insert(idef.column_id, BitmapIndex::new());
11120 }
11121 IndexKind::Ann => {
11122 let dim = schema
11123 .columns
11124 .iter()
11125 .find(|c| c.id == idef.column_id)
11126 .and_then(|c| match c.ty {
11127 TypeId::Embedding { dim } => Some(dim as usize),
11128 _ => None,
11129 })
11130 .unwrap_or(0);
11131 let options = idef.options.ann.clone().unwrap_or_default();
11132 ann.insert(
11133 idef.column_id,
11134 AnnIndex::with_options(
11135 dim,
11136 options.m,
11137 options.ef_construction,
11138 options.ef_search,
11139 ),
11140 );
11141 }
11142 IndexKind::FmIndex => {
11143 fm.insert(idef.column_id, FmIndex::new());
11144 }
11145 IndexKind::Sparse => {
11146 sparse.insert(idef.column_id, SparseIndex::new());
11147 }
11148 IndexKind::MinHash => {
11149 let options = idef.options.minhash.clone().unwrap_or_default();
11150 minhash.insert(
11151 idef.column_id,
11152 MinHashIndex::with_options(options.permutations, options.bands),
11153 );
11154 }
11155 _ => {}
11156 }
11157 }
11158 (bitmap, ann, fm, sparse, minhash)
11159}
11160
11161const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11162 | ColumnFlags::AUTO_INCREMENT
11163 | ColumnFlags::ENCRYPTED
11164 | ColumnFlags::ENCRYPTED_INDEXABLE
11165 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11166
11167fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11168 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11169 return Err(MongrelError::Schema(
11170 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11171 ));
11172 }
11173 Ok(())
11174}
11175
11176fn validate_alter_column_type(
11177 schema: &Schema,
11178 old: &ColumnDef,
11179 next: &ColumnDef,
11180 has_stored_versions: bool,
11181) -> Result<()> {
11182 if old.ty == next.ty {
11183 return Ok(());
11184 }
11185 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11186 return Err(MongrelError::Schema(format!(
11187 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11188 old.name
11189 )));
11190 }
11191 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11192 return Ok(());
11193 }
11194 Err(MongrelError::Schema(format!(
11195 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11196 old.ty, next.ty
11197 )))
11198}
11199
11200fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11201 matches!(
11202 (old, new),
11203 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11204 )
11205}
11206
11207fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11213 let mut prev: Option<i64> = None;
11214 for r in rows {
11215 match r.columns.get(&pk_id) {
11216 Some(Value::Int64(v)) => {
11217 if prev.is_some_and(|p| p >= *v) {
11218 return false;
11219 }
11220 prev = Some(*v);
11221 }
11222 _ => return false,
11223 }
11224 }
11225 true
11226}
11227
11228#[allow(clippy::too_many_arguments)]
11229fn index_into(
11230 schema: &Schema,
11231 row: &Row,
11232 hot: &mut HotIndex,
11233 bitmap: &mut HashMap<u16, BitmapIndex>,
11234 ann: &mut HashMap<u16, AnnIndex>,
11235 fm: &mut HashMap<u16, FmIndex>,
11236 sparse: &mut HashMap<u16, SparseIndex>,
11237 minhash: &mut HashMap<u16, MinHashIndex>,
11238) {
11239 for idef in &schema.indexes {
11240 let Some(val) = row.columns.get(&idef.column_id) else {
11241 continue;
11242 };
11243 match idef.kind {
11244 IndexKind::Bitmap => {
11245 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11246 b.insert(val.encode_key(), row.row_id);
11247 }
11248 }
11249 IndexKind::Ann => {
11250 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11251 a.insert_validated(v, row.row_id);
11252 }
11253 }
11254 IndexKind::FmIndex => {
11255 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11256 f.insert(b.clone(), row.row_id);
11257 }
11258 }
11259 IndexKind::Sparse => {
11260 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11261 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11264 s.insert(&terms, row.row_id);
11265 }
11266 }
11267 }
11268 IndexKind::MinHash => {
11269 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11270 let tokens = crate::index::token_hashes_from_bytes(b);
11273 mh.insert(&tokens, row.row_id);
11274 }
11275 }
11276 _ => {}
11277 }
11278 }
11279 if let Some(pk_col) = schema.primary_key() {
11280 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11281 hot.insert(pk_val.encode_key(), row.row_id);
11282 }
11283 }
11284}
11285
11286#[allow(clippy::too_many_arguments)]
11289fn index_into_single(
11290 idef: &IndexDef,
11291 _schema: &Schema,
11292 row: &Row,
11293 _hot: &mut HotIndex,
11294 bitmap: &mut HashMap<u16, BitmapIndex>,
11295 ann: &mut HashMap<u16, AnnIndex>,
11296 fm: &mut HashMap<u16, FmIndex>,
11297 sparse: &mut HashMap<u16, SparseIndex>,
11298 minhash: &mut HashMap<u16, MinHashIndex>,
11299) {
11300 let Some(val) = row.columns.get(&idef.column_id) else {
11301 return;
11302 };
11303 match idef.kind {
11304 IndexKind::Bitmap => {
11305 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11306 b.insert(val.encode_key(), row.row_id);
11307 }
11308 }
11309 IndexKind::Ann => {
11310 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11311 a.insert_validated(v, row.row_id);
11312 }
11313 }
11314 IndexKind::FmIndex => {
11315 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11316 f.insert(b.clone(), row.row_id);
11317 }
11318 }
11319 IndexKind::Sparse => {
11320 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11321 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11322 s.insert(&terms, row.row_id);
11323 }
11324 }
11325 }
11326 IndexKind::MinHash => {
11327 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11328 let tokens = crate::index::token_hashes_from_bytes(b);
11329 mh.insert(&tokens, row.row_id);
11330 }
11331 }
11332 _ => {}
11333 }
11334}
11335
11336fn eval_partial_predicate(
11342 pred: &str,
11343 columns_map: &HashMap<u16, &Value>,
11344 name_to_id: &HashMap<&str, u16>,
11345) -> bool {
11346 let lower = pred.trim().to_ascii_lowercase();
11347 if let Some(rest) = lower.strip_suffix(" is not null") {
11349 let col_name = rest.trim();
11350 if let Some(col_id) = name_to_id.get(col_name) {
11351 return columns_map
11352 .get(col_id)
11353 .is_some_and(|v| !matches!(v, Value::Null));
11354 }
11355 }
11356 if let Some(rest) = lower.strip_suffix(" is null") {
11358 let col_name = rest.trim();
11359 if let Some(col_id) = name_to_id.get(col_name) {
11360 return columns_map
11361 .get(col_id)
11362 .is_none_or(|v| matches!(v, Value::Null));
11363 }
11364 }
11365 true
11368}
11369
11370#[allow(dead_code)]
11376fn bulk_index_key(
11377 column_keys: &HashMap<u16, ([u8; 32], u8)>,
11378 column_id: u16,
11379 ty: TypeId,
11380 col: &columnar::NativeColumn,
11381 i: usize,
11382) -> Option<Vec<u8>> {
11383 let encoded = columnar::encode_key_native(ty, col, i)?;
11384 #[cfg(feature = "encryption")]
11385 {
11386 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11387 if let Some((key, scheme)) = column_keys.get(&column_id) {
11388 return Some(match (*scheme, col) {
11389 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
11390 (_, columnar::NativeColumn::Int64 { data, .. }) => {
11391 ope_token_i64(key, data[i]).to_vec()
11392 }
11393 (_, columnar::NativeColumn::Float64 { data, .. }) => {
11394 ope_token_f64(key, data[i]).to_vec()
11395 }
11396 _ => hmac_token(key, &encoded).to_vec(),
11397 });
11398 }
11399 }
11400 #[cfg(not(feature = "encryption"))]
11401 {
11402 let _ = (column_id, column_keys, col);
11403 }
11404 Some(encoded)
11405}
11406
11407pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
11408 write_schema_with_after(dir, schema, || {})
11409}
11410
11411pub(crate) fn write_schema_durable(
11412 root: &crate::durable_file::DurableRoot,
11413 schema: &Schema,
11414) -> Result<()> {
11415 write_schema_durable_with_after(root, schema, || {})
11416}
11417
11418fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
11419where
11420 F: FnOnce(),
11421{
11422 let json = serde_json::to_string_pretty(schema)
11423 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
11424 crate::durable_file::write_atomic_with_after(
11425 &dir.join(SCHEMA_FILENAME),
11426 json.as_bytes(),
11427 after_publish,
11428 )?;
11429 Ok(())
11430}
11431
11432fn write_schema_durable_with_after<F>(
11433 root: &crate::durable_file::DurableRoot,
11434 schema: &Schema,
11435 after_publish: F,
11436) -> Result<()>
11437where
11438 F: FnOnce(),
11439{
11440 let json = serde_json::to_string_pretty(schema)
11441 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
11442 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
11443 Ok(())
11444}
11445
11446fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
11447 let mut schema_published = false;
11448 let schema_result = match table._root_guard.as_deref() {
11449 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
11450 schema_published = true;
11451 }),
11452 None => write_schema_with_after(&table.dir, &table.schema, || {
11453 schema_published = true;
11454 }),
11455 };
11456 if schema_result.is_err() && !schema_published {
11457 return schema_result;
11458 }
11459 match table.persist_manifest(table.current_epoch()) {
11460 Ok(()) => Ok(()),
11461 Err(manifest_error) => Err(match schema_result {
11462 Ok(()) => manifest_error,
11463 Err(schema_error) => MongrelError::Other(format!(
11464 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11465 )),
11466 }),
11467 }
11468}
11469
11470fn read_schema(dir: &Path) -> Result<Schema> {
11471 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
11472 read_schema_file(file)
11473}
11474
11475fn read_schema_file(file: std::fs::File) -> Result<Schema> {
11476 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
11477 use std::io::Read;
11478
11479 let length = file.metadata()?.len();
11480 if length > MAX_SCHEMA_BYTES {
11481 return Err(MongrelError::ResourceLimitExceeded {
11482 resource: "schema bytes",
11483 requested: usize::try_from(length).unwrap_or(usize::MAX),
11484 limit: MAX_SCHEMA_BYTES as usize,
11485 });
11486 }
11487 let mut bytes = Vec::with_capacity(length as usize);
11488 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
11489 if bytes.len() as u64 != length {
11490 return Err(MongrelError::Schema(
11491 "schema length changed while reading".into(),
11492 ));
11493 }
11494 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
11495}
11496
11497fn preflight_standalone_open(
11498 dir: &Path,
11499 runs_root: Option<&crate::durable_file::DurableRoot>,
11500 idx_root: Option<&crate::durable_file::DurableRoot>,
11501 manifest: &Manifest,
11502 schema: &Schema,
11503 records: &[crate::wal::Record],
11504 kek: Option<Arc<Kek>>,
11505) -> Result<()> {
11506 crate::wal::validate_shared_transaction_framing(records)?;
11507 if manifest.schema_id > schema.schema_id
11508 || manifest.flushed_epoch > manifest.current_epoch
11509 || manifest.global_idx_epoch > manifest.current_epoch
11510 || manifest.next_row_id == u64::MAX
11511 || manifest.auto_inc_next < 0
11512 || manifest.auto_inc_next == i64::MAX
11513 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
11514 {
11515 return Err(MongrelError::InvalidArgument(
11516 "manifest counters or schema identity are invalid".into(),
11517 ));
11518 }
11519 let mut run_ids = HashSet::new();
11520 let mut maximum_row_id = None::<u64>;
11521 for run in &manifest.runs {
11522 if run.run_id >= u64::MAX as u128
11523 || !run_ids.insert(run.run_id)
11524 || run.epoch_created > manifest.current_epoch
11525 {
11526 return Err(MongrelError::InvalidArgument(
11527 "manifest contains an invalid or duplicate active run".into(),
11528 ));
11529 }
11530 let mut reader = match runs_root {
11531 Some(root) => RunReader::open_file(
11532 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
11533 schema.clone(),
11534 kek.clone(),
11535 )?,
11536 None => RunReader::open(
11537 dir.join(RUNS_DIR)
11538 .join(format!("r-{}.sr", run.run_id as u64)),
11539 schema.clone(),
11540 kek.clone(),
11541 )?,
11542 };
11543 let header = reader.header();
11544 if header.run_id != run.run_id
11545 || header.level != run.level
11546 || header.row_count != run.row_count
11547 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
11548 || header.is_uniform_epoch() && header.epoch_created != 0
11549 || header.schema_id > schema.schema_id
11550 {
11551 return Err(MongrelError::InvalidArgument(format!(
11552 "run {} differs from its manifest",
11553 run.run_id
11554 )));
11555 }
11556 if header.row_count != 0 {
11557 maximum_row_id = Some(
11558 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
11559 );
11560 }
11561 reader.validate_all_pages()?;
11562 }
11563 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
11564 return Err(MongrelError::InvalidArgument(
11565 "manifest next_row_id does not advance beyond persisted rows".into(),
11566 ));
11567 }
11568 for run in &manifest.retiring {
11569 if run.run_id >= u64::MAX as u128
11570 || run.retire_epoch > manifest.current_epoch
11571 || !run_ids.insert(run.run_id)
11572 {
11573 return Err(MongrelError::InvalidArgument(
11574 "manifest contains an invalid or duplicate retired run".into(),
11575 ));
11576 }
11577 }
11578 #[cfg(feature = "encryption")]
11579 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
11580 #[cfg(not(feature = "encryption"))]
11581 let idx_dek: Option<Zeroizing<[u8; DEK_LEN]>> = None;
11582 match idx_root {
11583 Some(root) => {
11584 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
11585 }
11586 None => {
11587 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
11588 }
11589 }
11590
11591 let committed = records
11592 .iter()
11593 .filter_map(|record| match record.op {
11594 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
11595 _ => None,
11596 })
11597 .collect::<HashMap<_, _>>();
11598 for record in records {
11599 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
11600 continue;
11601 };
11602 match &record.op {
11603 Op::Put { table_id, rows } => {
11604 if *table_id != manifest.table_id {
11605 return Err(MongrelError::CorruptWal {
11606 offset: record.seq.0,
11607 reason: format!(
11608 "private WAL record references table {table_id}, expected {}",
11609 manifest.table_id
11610 ),
11611 });
11612 }
11613 let rows: Vec<Row> =
11614 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
11615 offset: record.seq.0,
11616 reason: format!("committed Put payload could not be decoded: {error}"),
11617 })?;
11618 for row in rows {
11619 if row.deleted || row.row_id.0 == u64::MAX {
11620 return Err(MongrelError::CorruptWal {
11621 offset: record.seq.0,
11622 reason: "committed Put contains an invalid row identity".into(),
11623 });
11624 }
11625 let cells = row.columns.into_iter().collect::<Vec<_>>();
11626 schema
11627 .validate_values(&cells)
11628 .map_err(|error| MongrelError::CorruptWal {
11629 offset: record.seq.0,
11630 reason: format!("committed Put violates table schema: {error}"),
11631 })?;
11632 if schema.auto_increment_column().is_some_and(|column| {
11633 matches!(
11634 cells.iter().find(|(id, _)| *id == column.id),
11635 Some((_, Value::Int64(value))) if *value == i64::MAX
11636 )
11637 }) {
11638 return Err(MongrelError::CorruptWal {
11639 offset: record.seq.0,
11640 reason: "committed Put exhausts AUTO_INCREMENT".into(),
11641 });
11642 }
11643 }
11644 }
11645 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
11646 if *table_id != manifest.table_id =>
11647 {
11648 return Err(MongrelError::CorruptWal {
11649 offset: record.seq.0,
11650 reason: format!(
11651 "private WAL record references table {table_id}, expected {}",
11652 manifest.table_id
11653 ),
11654 });
11655 }
11656 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
11657 return Err(MongrelError::CorruptWal {
11658 offset: record.seq.0,
11659 reason: "private WAL contains shared spilled-run metadata".into(),
11660 });
11661 }
11662 _ => {}
11663 }
11664 }
11665 Ok(())
11666}
11667
11668fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
11669 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
11670}
11671
11672fn wal_segment_number(path: &Path) -> Option<u64> {
11673 path.file_stem()
11674 .and_then(|stem| stem.to_str())
11675 .and_then(|stem| stem.strip_prefix("seg-"))
11676 .and_then(|number| number.parse().ok())
11677}
11678
11679fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
11680 let n = list_wal_numbers(wal_dir)?;
11681 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
11682}
11683
11684fn next_wal_number(wal_dir: &Path) -> Result<u32> {
11685 list_wal_numbers(wal_dir)?
11686 .map(|maximum| {
11687 maximum
11688 .checked_add(1)
11689 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
11690 })
11691 .unwrap_or(Ok(0))
11692}
11693
11694fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
11695 let mut max_n = None;
11696 let entries = match std::fs::read_dir(wal_dir) {
11697 Ok(entries) => entries,
11698 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
11699 Err(error) => return Err(error.into()),
11700 };
11701 for entry in entries {
11702 let entry = entry?;
11703 let fname = entry.file_name();
11704 let Some(s) = fname.to_str() else {
11705 continue;
11706 };
11707 let Some(stripped) = s.strip_prefix("seg-") else {
11708 continue;
11709 };
11710 let Some(number) = stripped.strip_suffix(".wal") else {
11711 return Err(MongrelError::CorruptWal {
11712 offset: 0,
11713 reason: format!("malformed WAL segment name {s:?}"),
11714 });
11715 };
11716 let n = number
11717 .parse::<u32>()
11718 .map_err(|_| MongrelError::CorruptWal {
11719 offset: 0,
11720 reason: format!("malformed WAL segment name {s:?}"),
11721 })?;
11722 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
11723 return Err(MongrelError::CorruptWal {
11724 offset: n as u64,
11725 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
11726 });
11727 }
11728 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
11729 }
11730 Ok(max_n)
11731}