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, IndexGeneration, MinHashIndex,
18 SparseIndex,
19};
20use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
21use crate::memtable::{Memtable, Row, Value};
22use crate::mutable_run::MutableRun;
23use crate::row_id_set::RowIdSet;
24use crate::rowid::{RowId, RowIdAllocator};
25use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
26use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
27use crate::txn::{GroupCommit, OwnedRow};
28use crate::wal::{Op, SharedWal, Wal};
29use crate::{MongrelError, Result};
30use arc_swap::ArcSwap;
31use std::cmp::Reverse;
32use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
33use std::path::{Path, PathBuf};
34use std::sync::atomic::AtomicBool;
35use std::sync::Arc;
36use zeroize::Zeroizing;
37
38pub const WAL_DIR: &str = "_wal";
39pub const RUNS_DIR: &str = "_runs";
40pub const CACHE_DIR: &str = "_cache";
41pub const META_DIR: &str = "_meta";
42pub const RCACHE_DIR: &str = "_rcache";
43pub const KEYS_FILENAME: &str = "keys";
44pub const SCHEMA_FILENAME: &str = "schema.json";
45
46fn derive_next_run_id(
47 dir: &Path,
48 runs_root: Option<&crate::durable_file::DurableRoot>,
49 active: &[RunRef],
50 retiring: &[crate::manifest::RetiredRun],
51) -> Result<u64> {
52 let mut maximum = 0_u64;
53 for run_id in active
54 .iter()
55 .map(|run| run.run_id)
56 .chain(retiring.iter().map(|run| run.run_id))
57 {
58 let run_id = u64::try_from(run_id)
59 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
60 maximum = maximum.max(run_id);
61 }
62 let names = match runs_root {
63 Some(root) => root.list_regular_files(".")?,
64 None => std::fs::read_dir(dir.join(RUNS_DIR))?
65 .map(|entry| entry.map(|entry| entry.file_name()))
66 .collect::<std::io::Result<Vec<_>>>()?,
67 };
68 for name in names {
69 let Some(name) = name.to_str() else {
70 continue;
71 };
72 let Some(digits) = name
73 .strip_prefix("r-")
74 .and_then(|name| name.strip_suffix(".sr"))
75 else {
76 continue;
77 };
78 let Ok(run_id) = digits.parse::<u64>() else {
79 continue;
80 };
81 if name == format!("r-{run_id}.sr") {
82 maximum = maximum.max(run_id);
83 }
84 }
85 maximum
86 .checked_add(1)
87 .map(|next| next.max(1))
88 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
89}
90
91enum ControlledVisibleCandidate {
92 Memory(Row),
93 Run(RunVisibleVersion),
94}
95
96impl ControlledVisibleCandidate {
97 fn row_id(&self) -> RowId {
98 match self {
99 Self::Memory(row) => row.row_id,
100 Self::Run(version) => version.row_id,
101 }
102 }
103
104 fn committed_epoch(&self) -> Epoch {
105 match self {
106 Self::Memory(row) => row.committed_epoch,
107 Self::Run(version) => version.committed_epoch,
108 }
109 }
110
111 fn deleted(&self) -> bool {
112 match self {
113 Self::Memory(row) => row.deleted,
114 Self::Run(version) => version.deleted,
115 }
116 }
117}
118
119enum ControlledVisibleCursor {
120 Memory(std::vec::IntoIter<Row>),
121 Run(Box<RunVisibleVersionCursor>),
122 #[cfg(test)]
123 Synthetic {
124 next: u64,
125 end: u64,
126 },
127}
128
129struct ControlledVisibleSource {
130 cursor: ControlledVisibleCursor,
131 current: Option<ControlledVisibleCandidate>,
132}
133
134impl ControlledVisibleSource {
135 fn memory(rows: Vec<Row>) -> Self {
136 Self {
137 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
138 current: None,
139 }
140 }
141
142 fn run(cursor: RunVisibleVersionCursor) -> Self {
143 Self {
144 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
145 current: None,
146 }
147 }
148
149 #[cfg(test)]
150 fn synthetic(end: u64) -> Self {
151 Self {
152 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
153 current: None,
154 }
155 }
156
157 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
158 self.current = match &mut self.cursor {
159 ControlledVisibleCursor::Memory(rows) => {
160 rows.next().map(ControlledVisibleCandidate::Memory)
161 }
162 ControlledVisibleCursor::Run(cursor) => cursor
163 .next_visible_version(control)?
164 .map(ControlledVisibleCandidate::Run),
165 #[cfg(test)]
166 ControlledVisibleCursor::Synthetic { next, end } => {
167 if *next > *end {
168 None
169 } else {
170 let row = Row::new(RowId(*next), Epoch(1));
171 *next += 1;
172 Some(ControlledVisibleCandidate::Memory(row))
173 }
174 }
175 };
176 Ok(())
177 }
178
179 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate> {
180 let current = self.current.take().ok_or_else(|| {
181 MongrelError::Other("controlled visible source was not primed".into())
182 })?;
183 self.advance(control)?;
184 Ok(current)
185 }
186
187 fn materialize(
188 &mut self,
189 candidate: ControlledVisibleCandidate,
190 control: &crate::ExecutionControl,
191 ) -> Result<Row> {
192 match candidate {
193 ControlledVisibleCandidate::Memory(row) => Ok(row),
194 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
195 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
196 _ => Err(MongrelError::Other(
197 "run candidate escaped its controlled cursor".into(),
198 )),
199 },
200 }
201 }
202}
203
204fn merge_controlled_visible_sources(
205 sources: &mut [ControlledVisibleSource],
206 control: &crate::ExecutionControl,
207 mut expired: impl FnMut(&Row) -> bool,
208 mut visit: impl FnMut(Row) -> Result<()>,
209) -> Result<()> {
210 let mut heap = BinaryHeap::new();
211 for (source_index, source) in sources.iter_mut().enumerate() {
212 source.advance(control)?;
213 if let Some(candidate) = &source.current {
214 heap.push(Reverse((candidate.row_id(), source_index)));
215 }
216 }
217 let mut merged = 0_usize;
218 while let Some(Reverse((row_id, source_index))) = heap.pop() {
219 if merged.is_multiple_of(256) {
220 control.checkpoint()?;
221 }
222 merged += 1;
223 let mut best_source = source_index;
224 let mut best = sources[source_index].pop(control)?;
225 if let Some(next) = &sources[source_index].current {
226 heap.push(Reverse((next.row_id(), source_index)));
227 }
228 while heap
229 .peek()
230 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
231 {
232 let Some(Reverse((_, source_index))) = heap.pop() else {
233 break;
234 };
235 let candidate = sources[source_index].pop(control)?;
236 if candidate.committed_epoch() > best.committed_epoch() {
237 best = candidate;
238 best_source = source_index;
239 }
240 if let Some(next) = &sources[source_index].current {
241 heap.push(Reverse((next.row_id(), source_index)));
242 }
243 }
244 if best.deleted() {
245 continue;
246 }
247 let row = sources[best_source].materialize(best, control)?;
248 if !expired(&row) {
249 visit(row)?;
250 }
251 }
252 control.checkpoint()
253}
254
255#[cfg(test)]
256mod controlled_visible_cursor_tests {
257 use super::*;
258
259 #[test]
260 fn streams_more_than_one_million_rows_without_a_source_cap() {
261 let control = crate::ExecutionControl::new(None);
262 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
263 let mut count = 0_u64;
264 let mut last = 0_u64;
265 merge_controlled_visible_sources(
266 &mut sources,
267 &control,
268 |_| false,
269 |row| {
270 count += 1;
271 assert!(row.row_id.0 > last);
272 last = row.row_id.0;
273 Ok(())
274 },
275 )
276 .unwrap();
277 assert_eq!(count, 1_000_001);
278 assert_eq!(last, 1_000_001);
279 }
280
281 #[test]
282 fn merge_orders_rows_and_honors_newest_tombstones() {
283 let control = crate::ExecutionControl::new(None);
284 let older = vec![
285 Row::new(RowId(1), Epoch(1)),
286 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
287 Row::new(RowId(4), Epoch(1)),
288 ];
289 let mut deleted = Row::new(RowId(1), Epoch(2));
290 deleted.deleted = true;
291 let newer = vec![
292 deleted,
293 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
294 Row::new(RowId(3), Epoch(2)),
295 ];
296 let mut sources = vec![
297 ControlledVisibleSource::memory(older),
298 ControlledVisibleSource::memory(newer),
299 ];
300 let mut rows = Vec::new();
301 merge_controlled_visible_sources(
302 &mut sources,
303 &control,
304 |_| false,
305 |row| {
306 rows.push(row);
307 Ok(())
308 },
309 )
310 .unwrap();
311 assert_eq!(
312 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
313 vec![2, 3, 4]
314 );
315 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
316 }
317}
318
319fn iso_now_bytes() -> Vec<u8> {
322 let secs = std::time::SystemTime::now()
323 .duration_since(std::time::UNIX_EPOCH)
324 .map(|d| d.as_secs() as i64)
325 .unwrap_or(0);
326 let days = secs.div_euclid(86_400);
327 let rem = secs.rem_euclid(86_400);
328 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
329 let (year, month, day) = civil_from_days(days);
330 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
331}
332
333pub(crate) fn unix_nanos_now() -> i64 {
334 std::time::SystemTime::now()
335 .duration_since(std::time::UNIX_EPOCH)
336 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
337 .unwrap_or(0)
338}
339
340fn ann_candidate_cap(
341 index_len: usize,
342 context: Option<&crate::query::AiExecutionContext>,
343) -> usize {
344 index_len
345 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
346 .min(context.map_or(
347 crate::query::MAX_RAW_INDEX_CANDIDATES,
348 crate::query::AiExecutionContext::max_fused_candidates,
349 ))
350}
351
352#[cfg(test)]
353mod ann_candidate_cap_tests {
354 use super::*;
355
356 #[test]
357 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
358 assert_eq!(
359 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
360 crate::query::MAX_RAW_INDEX_CANDIDATES,
361 );
362 let context = crate::query::AiExecutionContext::with_limits(
363 std::time::Duration::from_secs(1),
364 usize::MAX,
365 17,
366 );
367 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
368 }
369}
370
371fn civil_from_days(z: i64) -> (i64, u32, u32) {
372 let z = z + 719_468;
373 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
374 let doe = z - era * 146_097;
375 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
376 let y = yoe + era * 400;
377 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
378 let mp = (5 * doy + 2) / 153;
379 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
380 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
381 (if m <= 2 { y + 1 } else { y }, m, d)
382}
383
384pub fn clustered_row_id(primary_key: &Value) -> RowId {
391 let mut hash: u64 = 0xcbf29ce484222325;
392 for byte in primary_key.encode_key() {
393 hash ^= u64::from(byte);
394 hash = hash.wrapping_mul(0x100000001b3);
395 }
396 RowId(hash.max(1))
397}
398
399const 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;
406
407#[derive(Clone, Copy, Debug)]
422struct AutoIncState {
423 column_id: u16,
424 next: i64,
425 seeded: bool,
426}
427
428pub(crate) struct RecoveryMetadataPlan {
429 live_count: u64,
430 auto_inc: Option<AutoIncState>,
431 changed: bool,
432}
433
434type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
435
436fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
439 schema.auto_increment_column().map(|c| AutoIncState {
440 column_id: c.id,
441 next: 0,
442 seeded: false,
443 })
444}
445
446#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
459pub enum IndexBuildPolicy {
460 #[default]
462 Deferred,
463 Eager,
465}
466
467#[derive(Clone)]
468struct ReversePkSegment {
469 values: HashMap<RowId, Vec<u8>>,
470 removed: HashSet<RowId>,
471}
472
473#[derive(Clone)]
474struct ReversePkMap {
475 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
476 active: ReversePkSegment,
477}
478
479impl ReversePkMap {
480 fn new() -> Self {
481 Self {
482 frozen: Arc::new(Vec::new()),
483 active: ReversePkSegment {
484 values: HashMap::new(),
485 removed: HashSet::new(),
486 },
487 }
488 }
489
490 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
491 let mut map = Self::new();
492 map.active.values.extend(entries);
493 map
494 }
495
496 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
497 self.active.removed.remove(&row_id);
498 self.active.values.insert(row_id, key);
499 }
500
501 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
502 if let Some(key) = self.active.values.get(row_id) {
503 return Some(key);
504 }
505 if self.active.removed.contains(row_id) {
506 return None;
507 }
508 for segment in self.frozen.iter().rev() {
509 if let Some(key) = segment.values.get(row_id) {
510 return Some(key);
511 }
512 if segment.removed.contains(row_id) {
513 return None;
514 }
515 }
516 None
517 }
518
519 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
520 let previous = self.get(row_id).cloned();
521 self.active.values.remove(row_id);
522 self.active.removed.insert(*row_id);
523 previous
524 }
525
526 fn clear(&mut self) {
527 *self = Self::new();
528 }
529
530 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
531 let mut entries = HashMap::new();
532 for segment in self
533 .frozen
534 .iter()
535 .map(Arc::as_ref)
536 .chain(std::iter::once(&self.active))
537 {
538 for row_id in &segment.removed {
539 entries.remove(row_id);
540 }
541 entries.extend(
542 segment
543 .values
544 .iter()
545 .map(|(row_id, key)| (*row_id, key.clone())),
546 );
547 }
548 entries
549 }
550
551 fn seal(&mut self) {
552 if self.active.values.is_empty() && self.active.removed.is_empty() {
553 return;
554 }
555 let active = std::mem::replace(
556 &mut self.active,
557 ReversePkSegment {
558 values: HashMap::new(),
559 removed: HashSet::new(),
560 },
561 );
562 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
563 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
564 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
565 values: self.entries(),
566 removed: HashSet::new(),
567 })]);
568 }
569 }
570}
571
572#[derive(Clone)]
585pub struct ReadGeneration {
586 schema: Arc<Schema>,
587 base_runs: Arc<Vec<RunRef>>,
588 deltas: TableDeltas,
589 indexes: Arc<IndexGeneration>,
590 visible_through: Epoch,
591}
592
593pub(crate) enum SecondaryIndexArtifact {
597 Bitmap(u16, BitmapIndex),
598 LearnedRange(u16, ColumnLearnedRange),
599 Fm(u16, Box<FmIndex>),
600 Ann(u16, Box<AnnIndex>),
601 Sparse(u16, SparseIndex),
602 MinHash(u16, MinHashIndex),
603}
604
605#[derive(Clone)]
611pub struct TableDeltas {
612 memtable: Memtable,
613 mutable_run: MutableRun,
614 hot: HotIndex,
615 pk_by_row: ReversePkMap,
616}
617
618impl ReadGeneration {
619 fn empty(schema: &Schema) -> Self {
622 Self {
623 schema: Arc::new(schema.clone()),
624 base_runs: Arc::new(Vec::new()),
625 deltas: TableDeltas {
626 memtable: Memtable::new(),
627 mutable_run: MutableRun::new(),
628 hot: HotIndex::new(),
629 pk_by_row: ReversePkMap::new(),
630 },
631 indexes: Arc::new(IndexGeneration::default()),
632 visible_through: Epoch(0),
633 }
634 }
635
636 pub fn schema(&self) -> &Arc<Schema> {
638 &self.schema
639 }
640
641 pub fn base_runs(&self) -> &[RunRef] {
643 &self.base_runs
644 }
645
646 pub fn indexes(&self) -> &Arc<IndexGeneration> {
648 &self.indexes
649 }
650
651 pub fn visible_through(&self) -> Epoch {
653 self.visible_through
654 }
655
656 pub fn deltas(&self) -> &TableDeltas {
660 &self.deltas
661 }
662}
663
664impl TableDeltas {
665 pub fn approx_bytes(&self) -> u64 {
668 self.memtable
669 .approx_bytes()
670 .saturating_add(self.mutable_run.approx_bytes())
671 }
672
673 pub fn memtable_len(&self) -> usize {
675 self.memtable.len()
676 }
677
678 pub fn mutable_run_len(&self) -> usize {
680 self.mutable_run.len()
681 }
682
683 pub fn hot_len(&self) -> usize {
685 self.hot.len()
686 }
687
688 pub fn reverse_pk_len(&self) -> usize {
690 self.pk_by_row.entries().len()
691 }
692}
693
694#[derive(Clone)]
696pub struct Table {
697 dir: PathBuf,
698 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
699 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
700 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
701 table_id: u64,
702 name: String,
706 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
711 read_only: bool,
714 durable_commit_failed: bool,
718 wal: WalSink,
719 memtable: Memtable,
720 mutable_run: MutableRun,
725 mutable_run_spill_bytes: u64,
727 compaction_zstd_level: i32,
730 allocator: RowIdAllocator,
731 epoch: Arc<EpochAuthority>,
732 data_generation: u64,
735 schema: Schema,
736 hot: HotIndex,
737 kek: Option<Arc<Kek>>,
740 column_keys: HashMap<u16, ([u8; 32], u8)>,
744 run_refs: Vec<RunRef>,
745 retiring: Vec<crate::manifest::RetiredRun>,
748 next_run_id: u64,
749 sync_byte_threshold: u64,
750 current_txn_id: u64,
755 pending_private_mutations: bool,
759 bitmap: HashMap<u16, BitmapIndex>,
760 ann: HashMap<u16, AnnIndex>,
761 fm: HashMap<u16, FmIndex>,
762 sparse: HashMap<u16, SparseIndex>,
763 minhash: HashMap<u16, MinHashIndex>,
764 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
767 pk_by_row: ReversePkMap,
769 pinned: BTreeMap<Epoch, usize>,
772 pub(crate) live_count: u64,
775 reservoir: crate::reservoir::Reservoir,
778 reservoir_complete: bool,
786 had_deletes: bool,
790 agg_cache: Arc<HashMap<u64, CachedAgg>>,
794 global_idx_epoch: u64,
798 indexes_complete: bool,
803 index_build_policy: IndexBuildPolicy,
805 pk_by_row_complete: bool,
812 flushed_epoch: u64,
815 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
818 snapshots: Arc<crate::retention::SnapshotRegistry>,
821 commit_lock: Arc<parking_lot::Mutex<()>>,
823 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
826 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
836 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
845 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
847 pending_delete_rids: roaring::RoaringBitmap,
850 pending_put_cols: std::collections::HashSet<u16>,
853 pending_rows: Vec<Row>,
859 pending_rows_auto_inc: Vec<bool>,
860 pending_dels: Vec<RowId>,
863 pending_truncate: Option<Epoch>,
867 auto_inc: Option<AutoIncState>,
870 ttl: Option<TtlPolicy>,
873 pins: Arc<crate::retention::PinRegistry>,
880 published: Arc<ArcSwap<ReadGeneration>>,
885 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
890}
891
892const _: () = {
899 const fn assert_sync<T: ?Sized + Sync>() {}
900 assert_sync::<Table>();
901};
902
903enum CachedData {
909 Rows(Arc<Vec<Row>>),
910 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
911}
912
913impl CachedData {
914 fn approx_bytes(&self) -> u64 {
915 match self {
916 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
917 CachedData::Columns(c) => c
918 .iter()
919 .map(|(_, c)| c.approx_bytes())
920 .sum::<u64>()
921 .saturating_add(c.len() as u64 * 16),
922 }
923 }
924}
925
926struct CachedEntry {
930 data: CachedData,
931 footprint: roaring::RoaringBitmap,
932 condition_cols: Vec<u16>,
933}
934
935struct ResultCache {
946 entries: std::collections::HashMap<u64, CachedEntry>,
947 order: BTreeSet<(u64, u64)>,
948 generations: HashMap<u64, u64>,
949 next_generation: u64,
950 condition_index: HashMap<u16, HashSet<u64>>,
953 empty_footprint_entries: HashSet<u64>,
956 bytes: u64,
957 max_bytes: u64,
958 dir: Option<std::path::PathBuf>,
959 #[allow(dead_code)]
960 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
961}
962
963#[derive(serde::Serialize, serde::Deserialize)]
965struct SerializedEntry {
966 condition_cols: Vec<u16>,
967 footprint_bits: Vec<u32>,
968 data: SerializedData,
969}
970
971#[derive(serde::Serialize, serde::Deserialize)]
972enum SerializedData {
973 Rows(Vec<Row>),
974 Columns(Vec<(u16, columnar::NativeColumn)>),
975}
976
977#[derive(serde::Serialize)]
978struct SerializedEntryRef<'a> {
979 condition_cols: &'a [u16],
980 footprint_bits: Vec<u32>,
981 data: SerializedDataRef<'a>,
982}
983
984#[derive(serde::Serialize)]
985enum SerializedDataRef<'a> {
986 Rows(&'a [Row]),
987 Columns(&'a [(u16, columnar::NativeColumn)]),
988}
989
990impl<'a> SerializedEntryRef<'a> {
991 fn from_entry(entry: &'a CachedEntry) -> Self {
992 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
993 let data = match &entry.data {
994 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
995 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
996 };
997 Self {
998 condition_cols: &entry.condition_cols,
999 footprint_bits,
1000 data,
1001 }
1002 }
1003}
1004
1005impl SerializedEntry {
1006 fn into_entry(self) -> Option<CachedEntry> {
1007 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1008 let data = match self.data {
1009 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1010 SerializedData::Columns(c) => {
1011 if !c.iter().all(|(_, col)| col.validate()) {
1014 return None;
1015 }
1016 CachedData::Columns(Arc::new(c))
1017 }
1018 };
1019 Some(CachedEntry {
1020 data,
1021 footprint,
1022 condition_cols: self.condition_cols,
1023 })
1024 }
1025}
1026
1027impl ResultCache {
1028 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1029
1030 fn new() -> Self {
1031 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1032 }
1033
1034 fn with_max_bytes(max_bytes: u64) -> Self {
1035 Self {
1036 entries: std::collections::HashMap::new(),
1037 order: BTreeSet::new(),
1038 generations: HashMap::new(),
1039 next_generation: 0,
1040 condition_index: HashMap::new(),
1041 empty_footprint_entries: HashSet::new(),
1042 bytes: 0,
1043 max_bytes,
1044 dir: None,
1045 cache_dek: None,
1046 }
1047 }
1048
1049 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1050 let _ = std::fs::create_dir_all(&dir);
1051 self.dir = Some(dir);
1052 self
1053 }
1054
1055 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1056 self.cache_dek = dek;
1057 self
1058 }
1059
1060 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1061 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1062 }
1063
1064 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1068 let Some(path) = self.disk_path(key) else {
1069 return;
1070 };
1071 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1072 Ok(s) => s,
1073 Err(_) => return,
1074 };
1075 let on_disk = if let Some(dek) = &self.cache_dek {
1077 match self.encrypt_cache(&serialized, dek) {
1078 Some(b) => b,
1079 None => return,
1080 }
1081 } else {
1082 serialized
1083 };
1084 let tmp = path.with_extension("tmp");
1085 use std::io::Write;
1086 let write = || -> std::io::Result<()> {
1087 let mut f = std::fs::File::create(&tmp)?;
1088 f.write_all(&on_disk)?;
1089 f.flush()?;
1090 Ok(())
1091 };
1092 if write().is_err() {
1093 let _ = std::fs::remove_file(&tmp);
1094 return;
1095 }
1096 let _ = std::fs::rename(&tmp, &path);
1097 }
1098
1099 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1101 let path = self.disk_path(key)?;
1102 let bytes = std::fs::read(&path).ok()?;
1103 let plaintext = if let Some(dek) = &self.cache_dek {
1104 self.decrypt_cache(&bytes, dek)?
1105 } else {
1106 bytes
1107 };
1108 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1109 serialized.into_entry()
1110 }
1111
1112 fn remove_from_disk(&self, key: u64) {
1114 if let Some(path) = self.disk_path(key) {
1115 let _ = std::fs::remove_file(&path);
1116 }
1117 }
1118
1119 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1121 use crate::encryption::Cipher;
1122 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1123 let mut nonce = [0u8; 12];
1124 crate::encryption::fill_random(&mut nonce).ok()?;
1125 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1126 let mut out = Vec::with_capacity(12 + ct.len());
1127 out.extend_from_slice(&nonce);
1128 out.extend_from_slice(&ct);
1129 Some(out)
1130 }
1131
1132 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1134 use crate::encryption::Cipher;
1135 if bytes.len() < 28 {
1136 return None;
1137 }
1138 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1139 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1140 let ct = &bytes[12..];
1141 cipher.decrypt_page(&nonce, ct).ok()
1142 }
1143
1144 fn load_persistent(&mut self) {
1147 let Some(dir) = self.dir.as_ref().cloned() else {
1148 return;
1149 };
1150 let entries = match std::fs::read_dir(&dir) {
1151 Ok(e) => e,
1152 Err(_) => return,
1153 };
1154 for entry in entries.flatten() {
1155 let path = entry.path();
1156 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1158 let _ = std::fs::remove_file(&path);
1159 continue;
1160 }
1161 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1162 continue;
1163 }
1164 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1165 Some(s) => s,
1166 None => continue,
1167 };
1168 let key = match u64::from_str_radix(stem, 16) {
1169 Ok(k) => k,
1170 Err(_) => continue,
1171 };
1172 let bytes = match std::fs::read(&path) {
1173 Ok(b) => b,
1174 Err(_) => continue,
1175 };
1176 let plaintext = if let Some(dek) = &self.cache_dek {
1178 match self.decrypt_cache(&bytes, dek) {
1179 Some(p) => p,
1180 None => {
1181 let _ = std::fs::remove_file(&path);
1182 continue;
1183 }
1184 }
1185 } else {
1186 bytes
1187 };
1188 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1189 Ok(serialized) => {
1190 if let Some(entry) = serialized.into_entry() {
1191 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1192 self.index_entry(key, &entry);
1193 self.entries.insert(key, entry);
1194 self.touch(key);
1195 } else {
1196 let _ = std::fs::remove_file(&path);
1197 }
1198 }
1199 Err(_) => {
1200 let _ = std::fs::remove_file(&path);
1201 }
1202 }
1203 }
1204 self.evict();
1205 }
1206
1207 fn set_max_bytes(&mut self, max_bytes: u64) {
1208 self.max_bytes = max_bytes;
1209 self.evict();
1210 }
1211
1212 fn touch(&mut self, key: u64) {
1215 if !self.entries.contains_key(&key) {
1216 return;
1217 }
1218 if let Some(previous) = self.generations.remove(&key) {
1219 self.order.remove(&(previous, key));
1220 }
1221 let generation = self.allocate_generation();
1222 self.generations.insert(key, generation);
1223 self.order.insert((generation, key));
1224 }
1225
1226 fn untrack(&mut self, key: u64) {
1227 if let Some(generation) = self.generations.remove(&key) {
1228 self.order.remove(&(generation, key));
1229 }
1230 }
1231
1232 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
1233 for column in &entry.condition_cols {
1234 self.condition_index.entry(*column).or_default().insert(key);
1235 }
1236 if entry.footprint.is_empty() {
1237 self.empty_footprint_entries.insert(key);
1238 }
1239 }
1240
1241 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
1242 for column in &entry.condition_cols {
1243 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
1244 keys.remove(&key);
1245 keys.is_empty()
1246 });
1247 if remove_column {
1248 self.condition_index.remove(column);
1249 }
1250 }
1251 if entry.footprint.is_empty() {
1252 self.empty_footprint_entries.remove(&key);
1253 }
1254 }
1255
1256 fn allocate_generation(&mut self) -> u64 {
1257 if self.next_generation == u64::MAX {
1258 self.rebase_generations();
1259 }
1260 let generation = self.next_generation;
1261 self.next_generation += 1;
1262 generation
1263 }
1264
1265 fn rebase_generations(&mut self) {
1266 let ordered = std::mem::take(&mut self.order);
1267 self.generations.clear();
1268 for (generation, (_, key)) in ordered.into_iter().enumerate() {
1269 let generation = generation as u64;
1270 self.generations.insert(key, generation);
1271 self.order.insert((generation, key));
1272 }
1273 self.next_generation = self.order.len() as u64;
1274 }
1275
1276 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1277 let res = self.entries.get(&key).and_then(|e| match &e.data {
1278 CachedData::Rows(r) => Some(r.clone()),
1279 CachedData::Columns(_) => None,
1280 });
1281 if res.is_some() {
1282 self.touch(key);
1283 return res;
1284 }
1285 if let Some(entry) = self.load_from_disk(key) {
1287 let res = match &entry.data {
1288 CachedData::Rows(r) => Some(r.clone()),
1289 CachedData::Columns(_) => None,
1290 };
1291 if res.is_some() {
1292 let approx = entry.data.approx_bytes();
1293 self.bytes = self.bytes.saturating_add(approx);
1294 self.index_entry(key, &entry);
1295 self.entries.insert(key, entry);
1296 self.touch(key);
1297 self.evict();
1298 return res;
1299 }
1300 }
1301 None
1302 }
1303
1304 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1305 let res = self.entries.get(&key).and_then(|e| match &e.data {
1306 CachedData::Columns(c) => Some(c.clone()),
1307 CachedData::Rows(_) => None,
1308 });
1309 if res.is_some() {
1310 self.touch(key);
1311 return res;
1312 }
1313 if let Some(entry) = self.load_from_disk(key) {
1315 let res = match &entry.data {
1316 CachedData::Columns(c) => Some(c.clone()),
1317 CachedData::Rows(_) => None,
1318 };
1319 if res.is_some() {
1320 let approx = entry.data.approx_bytes();
1321 self.bytes = self.bytes.saturating_add(approx);
1322 self.index_entry(key, &entry);
1323 self.entries.insert(key, entry);
1324 self.touch(key);
1325 self.evict();
1326 return res;
1327 }
1328 }
1329 None
1330 }
1331
1332 fn insert(&mut self, key: u64, entry: CachedEntry) {
1333 let approx = entry.data.approx_bytes();
1334 if let Some(previous) = self.entries.remove(&key) {
1335 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
1336 self.unindex_entry(key, &previous);
1337 self.untrack(key);
1338 }
1339 self.store_to_disk(key, &entry);
1341 self.bytes = self.bytes.saturating_add(approx);
1342 self.index_entry(key, &entry);
1343 self.entries.insert(key, entry);
1344 self.touch(key);
1345 self.evict();
1346 }
1347
1348 fn invalidate(
1357 &mut self,
1358 delete_rids: &roaring::RoaringBitmap,
1359 put_cols: &std::collections::HashSet<u16>,
1360 ) {
1361 if self.entries.is_empty() {
1362 return;
1363 }
1364 let has_deletes = !delete_rids.is_empty();
1365 let mut to_remove = HashSet::new();
1366
1367 for column in put_cols {
1371 if let Some(keys) = self.condition_index.get(column) {
1372 to_remove.extend(keys.iter().copied());
1373 }
1374 }
1375
1376 if has_deletes {
1377 to_remove.extend(self.empty_footprint_entries.iter().copied());
1381 for (&key, entry) in &self.entries {
1382 if !to_remove.contains(&key)
1383 && !entry.footprint.is_empty()
1384 && entry.footprint.intersection_len(delete_rids) > 0
1385 {
1386 to_remove.insert(key);
1387 }
1388 }
1389 }
1390
1391 for key in to_remove {
1392 if let Some(entry) = self.entries.remove(&key) {
1393 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
1394 self.unindex_entry(key, &entry);
1395 }
1396 self.remove_from_disk(key);
1397 self.untrack(key);
1398 }
1399 }
1400
1401 fn clear(&mut self) {
1402 if let Some(dir) = &self.dir {
1404 if let Ok(entries) = std::fs::read_dir(dir) {
1405 for entry in entries.flatten() {
1406 let path = entry.path();
1407 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1408 let _ = std::fs::remove_file(&path);
1409 }
1410 }
1411 }
1412 }
1413 self.entries.clear();
1414 self.order.clear();
1415 self.generations.clear();
1416 self.next_generation = 0;
1417 self.condition_index.clear();
1418 self.empty_footprint_entries.clear();
1419 self.bytes = 0;
1420 }
1421
1422 fn evict(&mut self) {
1423 while self.bytes > self.max_bytes {
1424 let Some((_, key)) = self.order.pop_first() else {
1425 break;
1426 };
1427 self.generations.remove(&key);
1428 if let Some(entry) = self.entries.remove(&key) {
1429 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
1430 self.unindex_entry(key, &entry);
1431 self.remove_from_disk(key);
1435 }
1436 }
1437 }
1438}
1439
1440#[cfg(test)]
1441mod result_cache_lru_tests {
1442 use super::*;
1443
1444 fn row_entry(row_id: u64) -> CachedEntry {
1445 CachedEntry {
1446 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
1447 footprint: std::iter::once(row_id as u32).collect(),
1448 condition_cols: vec![1],
1449 }
1450 }
1451
1452 #[test]
1453 fn hits_update_recency_without_growing_an_order_queue() {
1454 let mut cache = ResultCache::with_max_bytes(u64::MAX);
1455 cache.insert(1, row_entry(1));
1456 cache.insert(2, row_entry(2));
1457 assert_eq!(cache.order.len(), cache.entries.len());
1458
1459 for _ in 0..1_000 {
1460 assert!(cache.get_rows(1).is_some());
1461 }
1462
1463 assert_eq!(cache.order.len(), cache.entries.len());
1464 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
1465
1466 let one_entry = cache.entries[&1].data.approx_bytes();
1467 cache.set_max_bytes(one_entry);
1468 assert!(cache.entries.contains_key(&1));
1469 assert!(!cache.entries.contains_key(&2));
1470 assert_eq!(cache.order.len(), cache.entries.len());
1471 }
1472
1473 #[test]
1474 fn insert_invalidation_uses_the_condition_reverse_index() {
1475 let mut cache = ResultCache::with_max_bytes(u64::MAX);
1476 let mut first = row_entry(1);
1477 first.condition_cols = vec![1];
1478 let mut second = row_entry(2);
1479 second.condition_cols = vec![2];
1480 cache.insert(1, first);
1481 cache.insert(2, second);
1482
1483 let put_cols = std::iter::once(1_u16).collect();
1484 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
1485
1486 assert!(!cache.entries.contains_key(&1));
1487 assert!(cache.entries.contains_key(&2));
1488 assert!(!cache
1489 .condition_index
1490 .get(&1)
1491 .is_some_and(|keys| keys.contains(&1)));
1492 assert!(cache
1493 .condition_index
1494 .get(&2)
1495 .is_some_and(|keys| keys.contains(&2)));
1496 assert_eq!(cache.order.len(), cache.entries.len());
1497 }
1498
1499 #[test]
1500 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
1501 let mut cache = ResultCache::with_max_bytes(u64::MAX);
1502 cache.insert(1, row_entry(1));
1503 cache.insert(2, row_entry(2));
1504 let mut unknown = row_entry(3);
1505 unknown.footprint.clear();
1506 cache.insert(3, unknown);
1507
1508 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
1509 cache.invalidate(&delete_rids, &HashSet::new());
1510
1511 assert!(!cache.entries.contains_key(&1));
1512 assert!(cache.entries.contains_key(&2));
1513 assert!(!cache.entries.contains_key(&3));
1514 assert!(cache.empty_footprint_entries.is_empty());
1515 assert_eq!(cache.order.len(), cache.entries.len());
1516 }
1517
1518 #[test]
1519 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
1520 let entry = row_entry(7);
1521 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
1522 let owned = bincode::serialize(&SerializedEntry {
1523 condition_cols: entry.condition_cols.clone(),
1524 footprint_bits: entry.footprint.iter().collect(),
1525 data: match &entry.data {
1526 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
1527 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
1528 },
1529 })
1530 .unwrap();
1531 assert_eq!(borrowed, owned);
1532 }
1533}
1534
1535type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1542
1543fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1544 let _ = kek;
1545 {
1546 if let Some(k) = kek {
1547 return (
1548 Some(k.derive_table_wal_key(_table_id)),
1549 Some(k.derive_cache_key()),
1550 );
1551 }
1552 }
1553 (None, None)
1554}
1555
1556fn read_table_encryption_salt_root(
1557 root: &crate::durable_file::DurableRoot,
1558) -> Result<[u8; crate::encryption::SALT_LEN]> {
1559 use std::io::Read;
1560
1561 let mut file = root
1562 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1563 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1564 let length = file.metadata()?.len();
1565 if length != crate::encryption::SALT_LEN as u64 {
1566 return Err(MongrelError::InvalidArgument(format!(
1567 "salt file is {length} bytes, expected {}",
1568 crate::encryption::SALT_LEN
1569 )));
1570 }
1571 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1572 file.read_exact(&mut salt)?;
1573 Ok(salt)
1574}
1575
1576fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1578 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1579}
1580
1581fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1582 let Some(kek) = kek else {
1583 return HashMap::new();
1584 };
1585 {
1586 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1587 schema
1588 .columns
1589 .iter()
1590 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1591 .map(|c| {
1592 let scheme = if schema
1593 .indexes
1594 .iter()
1595 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1596 {
1597 SCHEME_OPE_RANGE
1598 } else {
1599 SCHEME_HMAC_EQ
1600 };
1601 let key: [u8; 32] = *kek.derive_column_key(c.id);
1602 (c.id, (key, scheme))
1603 })
1604 .collect()
1605 }
1606}
1607
1608pub(crate) struct SharedCtx {
1613 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1614 pub epoch: Arc<EpochAuthority>,
1615 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1616 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1617 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1618 pub kek: Option<Arc<Kek>>,
1619 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1625 pub shared: Option<SharedWalCtx>,
1629 pub table_name: Option<String>,
1632 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1635 pub read_only: bool,
1637}
1638
1639#[derive(Clone)]
1645pub(crate) struct SharedWalCtx {
1646 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1647 pub group: Arc<GroupCommit>,
1648 pub poisoned: Arc<AtomicBool>,
1649 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1650 pub change_wake: tokio::sync::broadcast::Sender<()>,
1651 pub lifecycle: Arc<crate::core::LifecycleController>,
1654 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
1656}
1657
1658enum WalSink {
1661 Private(Wal),
1662 Shared(SharedWalCtx),
1663 ReadOnly,
1664}
1665
1666impl Clone for WalSink {
1667 fn clone(&self) -> Self {
1668 match self {
1669 Self::Shared(shared) => Self::Shared(shared.clone()),
1670 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1671 }
1672 }
1673}
1674
1675impl SharedCtx {
1676 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1680 let n_shards = if cache_dir.is_some() {
1684 1
1685 } else {
1686 crate::cache::CACHE_SHARDS
1687 };
1688 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1689 let page_cache = if let Some(d) = cache_dir {
1690 Arc::new(crate::cache::Sharded::new(1, || {
1691 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1692 }))
1693 } else {
1694 Arc::new(crate::cache::Sharded::new(n_shards, || {
1695 crate::cache::PageCache::new(per_shard)
1696 }))
1697 };
1698 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1699 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1700 crate::cache::CACHE_SHARDS,
1701 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1702 ));
1703 Self {
1704 root_guard: None,
1705 epoch: Arc::new(EpochAuthority::new(0)),
1706 page_cache,
1707 decoded_cache,
1708 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1709 kek,
1710 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1711 shared: None,
1712 table_name: None,
1713 auth: None,
1714 read_only: false,
1715 }
1716 }
1717}
1718
1719fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1723 use crate::query::Condition;
1724 match c {
1725 Condition::Pk(_)
1727 | Condition::BitmapEq { .. }
1728 | Condition::BitmapIn { .. }
1729 | Condition::BytesPrefix { .. }
1730 | Condition::IsNull { .. }
1731 | Condition::IsNotNull { .. } => 0,
1732 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1734 1
1735 }
1736 Condition::FmContains { .. }
1738 | Condition::FmContainsAll { .. }
1739 | Condition::Ann { .. }
1740 | Condition::SparseMatch { .. } => 2,
1741 }
1742}
1743
1744impl Table {
1745 pub(crate) fn build_secondary_index_artifact<F>(
1751 &self,
1752 definition: &IndexDef,
1753 rows: &[Row],
1754 mut checkpoint: F,
1755 ) -> Result<SecondaryIndexArtifact>
1756 where
1757 F: FnMut(usize, usize) -> Result<()>,
1758 {
1759 let mut build_schema = self.schema.clone();
1760 build_schema.indexes = vec![definition.clone()];
1761 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
1762 let name_to_id: HashMap<&str, u16> = self
1763 .schema
1764 .columns
1765 .iter()
1766 .map(|column| (column.name.as_str(), column.id))
1767 .collect();
1768 let mut hot = HotIndex::new();
1769 let mut accepted = Vec::with_capacity(rows.len());
1770
1771 for (offset, row) in rows.iter().enumerate() {
1772 checkpoint(offset, rows.len())?;
1773 if row.deleted {
1774 continue;
1775 }
1776 if let Some(predicate) = &definition.predicate {
1777 let columns: HashMap<u16, &Value> =
1778 row.columns.iter().map(|(id, value)| (*id, value)).collect();
1779 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
1780 continue;
1781 }
1782 }
1783 accepted.push(row);
1784 if definition.kind == IndexKind::LearnedRange {
1785 continue;
1786 }
1787 let effective = if self.column_keys.is_empty() {
1788 row.clone()
1789 } else {
1790 self.tokenized_for_indexes(row)
1791 };
1792 if definition.kind == IndexKind::Ann {
1793 if let (Some(index), Some(vector)) = (
1794 ann.get_mut(&definition.column_id),
1795 effective
1796 .columns
1797 .get(&definition.column_id)
1798 .and_then(Value::as_embedding),
1799 ) {
1800 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
1801 checkpoint(offset, rows.len())
1802 })?;
1803 }
1804 } else {
1805 index_into_single(
1806 definition,
1807 &build_schema,
1808 &effective,
1809 &mut hot,
1810 &mut bitmap,
1811 &mut ann,
1812 &mut fm,
1813 &mut sparse,
1814 &mut minhash,
1815 );
1816 }
1817 }
1818 checkpoint(rows.len(), rows.len())?;
1819
1820 if let Some(index) = ann.get_mut(&definition.column_id) {
1821 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
1822 }
1823
1824 let column_id = definition.column_id;
1825 match definition.kind {
1826 IndexKind::Bitmap => bitmap
1827 .remove(&column_id)
1828 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
1829 IndexKind::Ann => ann
1830 .remove(&column_id)
1831 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
1832 IndexKind::FmIndex => fm
1833 .remove(&column_id)
1834 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
1835 IndexKind::Sparse => sparse
1836 .remove(&column_id)
1837 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
1838 IndexKind::MinHash => minhash
1839 .remove(&column_id)
1840 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
1841 IndexKind::LearnedRange => {
1842 let epsilon = definition
1843 .options
1844 .learned_range
1845 .as_ref()
1846 .map(|options| options.epsilon)
1847 .unwrap_or(16);
1848 let column = self
1849 .schema
1850 .columns
1851 .iter()
1852 .find(|column| column.id == column_id)
1853 .ok_or_else(|| {
1854 MongrelError::Schema(format!(
1855 "index {} references unknown column {column_id}",
1856 definition.name
1857 ))
1858 })?;
1859 let index = match column.ty {
1860 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1861 let pairs: Vec<_> = accepted
1862 .iter()
1863 .filter_map(|row| match row.columns.get(&column_id) {
1864 Some(Value::Int64(value)) if !row.deleted => {
1865 Some((*value, row.row_id.0))
1866 }
1867 _ => None,
1868 })
1869 .collect();
1870 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
1871 }
1872 TypeId::Float32 | TypeId::Float64 => {
1873 let pairs: Vec<_> = accepted
1874 .iter()
1875 .filter_map(|row| match row.columns.get(&column_id) {
1876 Some(Value::Float64(value)) if !row.deleted => {
1877 Some((*value, row.row_id.0))
1878 }
1879 _ => None,
1880 })
1881 .collect();
1882 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
1883 }
1884 ref ty => {
1885 return Err(MongrelError::Schema(format!(
1886 "LearnedRange index {} does not support {ty:?}",
1887 definition.name
1888 )));
1889 }
1890 };
1891 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
1892 }
1893 }
1894 .ok_or_else(|| {
1895 MongrelError::Other(format!(
1896 "failed to construct hidden index {}",
1897 definition.name
1898 ))
1899 })
1900 }
1901
1902 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1903 let dir = dir.as_ref().to_path_buf();
1904 crate::durable_file::create_directory_all(&dir)?;
1905 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1906 let pinned = root.io_path()?;
1907 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1908 ctx.root_guard = Some(root);
1909 Self::create_in(&pinned, schema, table_id, ctx)
1910 }
1911
1912 pub fn create_encrypted(
1923 dir: impl AsRef<Path>,
1924 schema: Schema,
1925 table_id: u64,
1926 passphrase: &str,
1927 ) -> Result<Self> {
1928 let dir = dir.as_ref().to_path_buf();
1929 crate::durable_file::create_directory_all(&dir)?;
1930 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1931 root.create_directory_all(META_DIR)?;
1932 let salt = crate::encryption::random_salt()?;
1933 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1934 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1935 let pinned = root.io_path()?;
1936 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1937 ctx.root_guard = Some(root);
1938 Self::create_in(&pinned, schema, table_id, ctx)
1939 }
1940
1941 pub fn create_with_key(
1946 dir: impl AsRef<Path>,
1947 schema: Schema,
1948 table_id: u64,
1949 key: &[u8],
1950 ) -> Result<Self> {
1951 let dir = dir.as_ref().to_path_buf();
1952 crate::durable_file::create_directory_all(&dir)?;
1953 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1954 root.create_directory_all(META_DIR)?;
1955 let salt = crate::encryption::random_salt()?;
1956 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1957 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1958 let pinned = root.io_path()?;
1959 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1960 ctx.root_guard = Some(root);
1961 Self::create_in(&pinned, schema, table_id, ctx)
1962 }
1963
1964 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1966 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1967 let salt = read_table_encryption_salt_root(&root)?;
1968 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1969 let pinned = root.io_path()?;
1970 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1971 ctx.root_guard = Some(root);
1972 Self::open_in(&pinned, ctx)
1973 }
1974
1975 pub(crate) fn create_in(
1976 dir: impl AsRef<Path>,
1977 schema: Schema,
1978 table_id: u64,
1979 ctx: SharedCtx,
1980 ) -> Result<Self> {
1981 schema.validate_auto_increment()?;
1982 schema.validate_defaults()?;
1983 schema.validate_ai()?;
1984 for index in &schema.indexes {
1985 index.validate_options()?;
1986 }
1987 let dir = dir.as_ref().to_path_buf();
1988 let runs_root = match ctx.root_guard.as_ref() {
1989 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1990 None => {
1991 crate::durable_file::create_directory_all(&dir)?;
1992 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1993 None
1994 }
1995 };
1996 match ctx.root_guard.as_deref() {
1997 Some(root) => write_schema_durable(root, &schema)?,
1998 None => write_schema(&dir, &schema)?,
1999 }
2000 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
2001 let (wal, current_txn_id) = match ctx.shared.clone() {
2004 Some(s) => (WalSink::Shared(s), 0),
2005 None => {
2006 let pinned_wal_root = match ctx.root_guard.as_deref() {
2007 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
2008 None => None,
2009 };
2010 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
2011 root.io_path()?
2012 } else {
2013 let wal_dir = dir.join(WAL_DIR);
2014 crate::durable_file::create_directory_all(&wal_dir)?;
2015 wal_dir
2016 };
2017 let mut w = if let Some(ref dk) = wal_dek {
2018 Wal::create_with_cipher(
2019 wal_dir.join("seg-000000.wal"),
2020 Epoch(0),
2021 Some(make_cipher(dk)),
2022 0,
2023 )?
2024 } else {
2025 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
2026 };
2027 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2028 (WalSink::Private(w), 1)
2029 }
2030 };
2031 let mut manifest = Manifest::new(table_id, schema.schema_id);
2032 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2037 match ctx.root_guard.as_deref() {
2038 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
2039 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
2040 }
2041 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
2042 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2043 let auto_inc = resolve_auto_inc(&schema);
2044 let rcache_dir = dir.join(RCACHE_DIR);
2045 let initial_view = ReadGeneration::empty(&schema);
2046 Ok(Self {
2047 dir,
2048 _root_guard: ctx.root_guard,
2049 runs_root,
2050 idx_root: None,
2051 table_id,
2052 name: ctx.table_name.unwrap_or_default(),
2053 auth: ctx.auth,
2054 read_only: ctx.read_only,
2055 durable_commit_failed: false,
2056 wal,
2057 memtable: Memtable::new(),
2058 mutable_run: MutableRun::new(),
2059 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2060 compaction_zstd_level: 3,
2061 allocator: RowIdAllocator::new(0),
2062 epoch: ctx.epoch,
2063 data_generation: 0,
2064 schema,
2065 hot: HotIndex::new(),
2066 kek: ctx.kek,
2067 column_keys,
2068 run_refs: Vec::new(),
2069 retiring: Vec::new(),
2070 next_run_id: 1,
2071 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2072 current_txn_id,
2073 pending_private_mutations: false,
2074 bitmap,
2075 ann,
2076 fm,
2077 sparse,
2078 minhash,
2079 learned_range: Arc::new(HashMap::new()),
2080 pk_by_row: ReversePkMap::new(),
2081 pinned: BTreeMap::new(),
2082 live_count: 0,
2083 reservoir: crate::reservoir::Reservoir::default(),
2084 reservoir_complete: true,
2085 had_deletes: false,
2086 agg_cache: Arc::new(HashMap::new()),
2087 global_idx_epoch: 0,
2088 indexes_complete: true,
2089 index_build_policy: IndexBuildPolicy::default(),
2090 pk_by_row_complete: false,
2091 flushed_epoch: 0,
2092 page_cache: ctx.page_cache,
2093 decoded_cache: ctx.decoded_cache,
2094 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2095 snapshots: ctx.snapshots,
2096 commit_lock: ctx.commit_lock,
2097 result_cache: Arc::new(parking_lot::Mutex::new(
2098 ResultCache::new()
2099 .with_dir(rcache_dir)
2100 .with_cache_dek(cache_dek.clone()),
2101 )),
2102 pending_delete_rids: roaring::RoaringBitmap::new(),
2103 pending_put_cols: std::collections::HashSet::new(),
2104 pending_rows: Vec::new(),
2105 pending_rows_auto_inc: Vec::new(),
2106 pending_dels: Vec::new(),
2107 pending_truncate: None,
2108 wal_dek,
2109 auto_inc,
2110 ttl: None,
2111 pins: Arc::new(crate::retention::PinRegistry::new()),
2112 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2113 read_generation_pin: None,
2114 })
2115 }
2116
2117 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
2121 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2122 let pinned = root.io_path()?;
2123 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2124 ctx.root_guard = Some(root);
2125 Self::open_in(&pinned, ctx)
2126 }
2127
2128 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
2131 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2132 let salt = read_table_encryption_salt_root(&root)?;
2133 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2134 let pinned = root.io_path()?;
2135 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2136 ctx.root_guard = Some(root);
2137 let t = Self::open_in(&pinned, ctx)?;
2138 Ok(t)
2139 }
2140
2141 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
2142 let dir = dir.as_ref().to_path_buf();
2143 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2144 let mut manifest = match ctx.root_guard.as_ref() {
2145 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
2146 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
2147 };
2148 let schema: Schema = match ctx.root_guard.as_ref() {
2149 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
2150 None => read_schema(&dir)?,
2151 };
2152 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
2158 let runs_root = match ctx.root_guard.as_ref() {
2159 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
2160 None => None,
2161 };
2162 let idx_root = match ctx.root_guard.as_ref() {
2163 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
2164 Ok(root) => Some(Arc::new(root)),
2165 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2166 Err(error) => return Err(error.into()),
2167 },
2168 None => None,
2169 };
2170 schema.validate_auto_increment()?;
2171 schema.validate_defaults()?;
2172 schema.validate_ai()?;
2173 for index in &schema.indexes {
2174 index.validate_options()?;
2175 }
2176 let replay_epoch = Epoch(manifest.current_epoch);
2177 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
2178 let private_replayed = if ctx.shared.is_none() {
2179 match latest_wal_segment(&dir.join(WAL_DIR))? {
2180 Some(path) => {
2181 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
2182 crate::wal::replay_with_cipher(path, cipher)?
2183 }
2184 None => Vec::new(),
2185 }
2186 } else {
2187 Vec::new()
2188 };
2189 if ctx.shared.is_none() {
2190 preflight_standalone_open(
2191 &dir,
2192 runs_root.as_deref(),
2193 idx_root.as_deref(),
2194 &manifest,
2195 &schema,
2196 &private_replayed,
2197 ctx.kek.clone(),
2198 )?;
2199 }
2200 let next_run_id = derive_next_run_id(
2201 &dir,
2202 runs_root.as_deref(),
2203 &manifest.runs,
2204 &manifest.retiring,
2205 )?;
2206 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2210 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2211 None => {
2212 let replayed = private_replayed;
2213 let wal_dir = dir.join(WAL_DIR);
2219 crate::durable_file::create_directory_all(&wal_dir)?;
2220 let segment = next_wal_segment(&wal_dir)?;
2221 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2222 let temporary = wal_dir.join(format!(
2223 ".recovery-{}-{}-{segment_no:06}.tmp",
2224 std::process::id(),
2225 std::time::SystemTime::now()
2226 .duration_since(std::time::UNIX_EPOCH)
2227 .unwrap_or_default()
2228 .as_nanos()
2229 ));
2230 let mut w = Wal::create_with_cipher(
2231 &temporary,
2232 replay_epoch,
2233 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2234 segment_no,
2235 )?;
2236 for record in &replayed {
2237 w.append_txn(record.txn_id, record.op.clone())?;
2238 }
2239 let mut w = w.publish_as(segment)?;
2240 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2241 let next_txn_id = replayed
2242 .iter()
2243 .map(|record| record.txn_id)
2244 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2245 .max()
2246 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2247 .unwrap_or(1);
2248 (WalSink::Private(w), replayed, next_txn_id)
2249 }
2250 };
2251
2252 let mut memtable = Memtable::new();
2253 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2254 let persisted_epoch = manifest.current_epoch;
2255 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2262 s.next = manifest.auto_inc_next;
2263 s.seeded = manifest.auto_inc_next > 0;
2264 s
2265 });
2266
2267 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2274 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2275 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2276 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2277 std::collections::BTreeMap::new();
2278 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2279 let mut recovered_epoch = manifest.current_epoch;
2280 let mut recovered_manifest_dirty = schema_manifest_repair;
2281 let mut saw_delete = false;
2282 for record in replayed {
2283 let txn_id = record.txn_id;
2284 match record.op {
2285 Op::Put { rows, .. } => {
2286 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2287 for row in &rows {
2288 allocator.advance_to(row.row_id)?;
2289 if let Some(ai) = auto_inc.as_mut() {
2290 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2291 let next = n.checked_add(1).ok_or_else(|| {
2292 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2293 })?;
2294 if next > ai.next {
2295 ai.next = next;
2296 }
2297 }
2298 }
2299 }
2300 staged_puts.entry(txn_id).or_default().extend(rows);
2301 }
2302 Op::Delete { row_ids, .. } => {
2303 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2304 }
2305 Op::TxnCommit { epoch, .. } => {
2306 let commit_epoch = Epoch(epoch);
2307 recovered_epoch = recovered_epoch.max(epoch);
2308 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2309 memtable = Memtable::new();
2310 replayed_puts.clear();
2311 replayed_deletes.clear();
2312 manifest.runs.clear();
2313 manifest.retiring.clear();
2314 manifest.live_count = 0;
2315 manifest.global_idx_epoch = 0;
2316 manifest.current_epoch = manifest.current_epoch.max(epoch);
2317 recovered_manifest_dirty = true;
2318 saw_delete = true;
2319 }
2320 if let Some(puts) = staged_puts.remove(&txn_id) {
2321 if commit_epoch.0 > manifest.flushed_epoch {
2322 for row in &puts {
2323 memtable.upsert(row.clone());
2324 }
2325 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2326 }
2327 }
2328 if let Some(dels) = staged_deletes.remove(&txn_id) {
2329 saw_delete = true;
2330 if commit_epoch.0 > manifest.flushed_epoch {
2331 for rid in dels {
2332 memtable.tombstone(rid, commit_epoch);
2333 replayed_deletes.push((rid, commit_epoch));
2334 }
2335 }
2336 }
2337 }
2338 Op::TxnAbort => {
2339 staged_puts.remove(&txn_id);
2340 staged_deletes.remove(&txn_id);
2341 staged_truncates.remove(&txn_id);
2342 }
2343 Op::TruncateTable { .. } => {
2344 staged_puts.remove(&txn_id);
2345 staged_deletes.remove(&txn_id);
2346 staged_truncates.insert(txn_id);
2347 }
2348 Op::ExternalTableState { .. }
2349 | Op::Flush { .. }
2350 | Op::Ddl(_)
2351 | Op::BeforeImage { .. }
2352 | Op::CommitTimestamp { .. }
2353 | Op::SpilledRows { .. } => {}
2354 }
2355 }
2356
2357 let rcache_dir = dir.join(RCACHE_DIR);
2358 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2359 let initial_view = ReadGeneration::empty(&schema);
2360 let mut db = Self {
2361 dir,
2362 _root_guard: ctx.root_guard,
2363 runs_root,
2364 idx_root,
2365 table_id: manifest.table_id,
2366 name: ctx.table_name.unwrap_or_default(),
2367 auth: ctx.auth,
2368 read_only: ctx.read_only,
2369 durable_commit_failed: false,
2370 wal,
2371 memtable,
2372 mutable_run: MutableRun::new(),
2373 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2374 compaction_zstd_level: 3,
2375 allocator,
2376 epoch: ctx.epoch,
2377 data_generation: persisted_epoch,
2378 schema,
2379 hot: HotIndex::new(),
2380 kek: ctx.kek,
2381 column_keys,
2382 run_refs: manifest.runs.clone(),
2383 retiring: manifest.retiring.clone(),
2384 next_run_id,
2385 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2386 current_txn_id,
2387 pending_private_mutations: false,
2388 bitmap: HashMap::new(),
2389 ann: HashMap::new(),
2390 fm: HashMap::new(),
2391 sparse: HashMap::new(),
2392 minhash: HashMap::new(),
2393 learned_range: Arc::new(HashMap::new()),
2394 pk_by_row: ReversePkMap::new(),
2395 pinned: BTreeMap::new(),
2396 live_count: manifest.live_count,
2397 reservoir: crate::reservoir::Reservoir::default(),
2398 reservoir_complete: false,
2399 had_deletes: saw_delete
2400 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2401 != manifest.live_count,
2402 agg_cache: Arc::new(HashMap::new()),
2403 global_idx_epoch: manifest.global_idx_epoch,
2404 indexes_complete: true,
2405 index_build_policy: IndexBuildPolicy::default(),
2406 pk_by_row_complete: false,
2407 flushed_epoch: manifest.flushed_epoch,
2408 page_cache: ctx.page_cache,
2409 decoded_cache: ctx.decoded_cache,
2410 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2411 snapshots: ctx.snapshots,
2412 commit_lock: ctx.commit_lock,
2413 result_cache: Arc::new(parking_lot::Mutex::new(
2414 ResultCache::new()
2415 .with_dir(rcache_dir)
2416 .with_cache_dek(cache_dek.clone()),
2417 )),
2418 pending_delete_rids: roaring::RoaringBitmap::new(),
2419 pending_put_cols: std::collections::HashSet::new(),
2420 pending_rows: Vec::new(),
2421 pending_rows_auto_inc: Vec::new(),
2422 pending_dels: Vec::new(),
2423 pending_truncate: None,
2424 wal_dek,
2425 auto_inc,
2426 ttl: manifest.ttl,
2427 pins: Arc::new(crate::retention::PinRegistry::new()),
2428 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2429 read_generation_pin: None,
2430 };
2431
2432 db.epoch.advance_recovered(Epoch(recovered_epoch));
2435
2436 let checkpoint = match db.idx_root.as_deref() {
2441 Some(root) => {
2442 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2443 }
2444 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2445 };
2446 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2447 c.epoch_built == manifest.global_idx_epoch
2448 && manifest.global_idx_epoch > 0
2449 && manifest
2450 .runs
2451 .iter()
2452 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2453 });
2454 if let Some(loaded) = checkpoint {
2455 if checkpoint_valid {
2456 db.hot = loaded.hot;
2457 db.bitmap = loaded.bitmap;
2458 db.ann = loaded.ann;
2459 db.fm = loaded.fm;
2460 db.sparse = loaded.sparse;
2461 db.minhash = loaded.minhash;
2462 db.learned_range = Arc::new(loaded.learned_range);
2463 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2468 for (cid, idx) in bitmap0 {
2469 db.bitmap.entry(cid).or_insert(idx);
2470 }
2471 for (cid, idx) in ann0 {
2472 db.ann.entry(cid).or_insert(idx);
2473 }
2474 for (cid, idx) in fm0 {
2475 db.fm.entry(cid).or_insert(idx);
2476 }
2477 for (cid, idx) in sparse0 {
2478 db.sparse.entry(cid).or_insert(idx);
2479 }
2480 for (cid, idx) in minhash0 {
2481 db.minhash.entry(cid).or_insert(idx);
2482 }
2483 }
2486 }
2487 if !checkpoint_valid {
2488 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2489 db.bitmap = bitmap;
2490 db.ann = ann;
2491 db.fm = fm;
2492 db.sparse = sparse;
2493 db.minhash = minhash;
2494 db.rebuild_indexes_from_runs()?;
2495 db.build_learned_ranges()?;
2496 }
2497
2498 for (epoch, group) in replayed_puts {
2503 let (losers, winner_pks) = db.partition_pk_winners(&group);
2504 for (key, &row_id) in &winner_pks {
2505 if let Some(old_rid) = db.hot.get(key) {
2506 if old_rid != row_id {
2507 db.tombstone_row(old_rid, epoch, None, false);
2508 }
2509 }
2510 }
2511 for &loser_rid in &losers {
2512 db.tombstone_row(loser_rid, epoch, None, false);
2513 }
2514 for (key, row_id) in winner_pks {
2515 db.insert_hot_pk(key, row_id);
2516 }
2517 if db.schema.primary_key().is_none() {
2518 for r in &group {
2519 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2520 }
2521 }
2522 for r in &group {
2523 if !losers.contains(&r.row_id) {
2524 db.index_row(r);
2525 }
2526 }
2527 }
2528 for (rid, epoch) in &replayed_deletes {
2532 db.remove_hot_for_row(*rid, *epoch);
2533 }
2534
2535 if recovered_manifest_dirty {
2536 let rows = db.visible_rows(Snapshot::unbounded())?;
2537 db.live_count = rows.len() as u64;
2538 db.persist_manifest(Epoch(recovered_epoch))?;
2539 }
2540
2541 db.result_cache.lock().load_persistent();
2548 Ok(db)
2549 }
2550
2551 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2557 if self.reservoir_complete {
2558 return Ok(());
2559 }
2560 self.rebuild_reservoir()?;
2561 self.reservoir_complete = true;
2562 Ok(())
2563 }
2564
2565 fn rebuild_reservoir(&mut self) -> Result<()> {
2568 let snap = self.snapshot();
2569 let rows = self.visible_rows(snap)?;
2570 self.reservoir.reset();
2571 for r in rows {
2572 self.reservoir.offer(r.row_id.0);
2573 }
2574 Ok(())
2575 }
2576
2577 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2578 self.rebuild_indexes_from_runs_inner(None)
2579 }
2580
2581 fn rebuild_indexes_from_runs_inner(
2582 &mut self,
2583 control: Option<&crate::ExecutionControl>,
2584 ) -> Result<()> {
2585 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2588 crate::retention::PinSource::OnlineIndexBuild,
2589 self.current_epoch(),
2590 );
2591 self.hot = HotIndex::new();
2592 self.pk_by_row.clear();
2593 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2594 self.bitmap = bitmap;
2595 self.ann = ann;
2596 self.fm = fm;
2597 self.sparse = sparse;
2598 self.minhash = minhash;
2599 let snapshot = Epoch(u64::MAX);
2600 let ttl_now = unix_nanos_now();
2601 let mut scanned = 0_usize;
2602 for rr in self.run_refs.clone() {
2603 if let Some(control) = control {
2604 control.checkpoint()?;
2605 }
2606 let mut reader = self.open_reader(rr.run_id)?;
2607 for row in reader.visible_rows(snapshot)? {
2608 if scanned.is_multiple_of(256) {
2609 if let Some(control) = control {
2610 control.checkpoint()?;
2611 }
2612 }
2613 scanned += 1;
2614 if self.row_expired_at(&row, ttl_now) {
2615 continue;
2616 }
2617 let tok_row = self.tokenized_for_indexes(&row);
2618 index_into(
2619 &self.schema,
2620 &tok_row,
2621 &mut self.hot,
2622 &mut self.bitmap,
2623 &mut self.ann,
2624 &mut self.fm,
2625 &mut self.sparse,
2626 &mut self.minhash,
2627 );
2628 }
2629 }
2630 for row in self.mutable_run.visible_versions(snapshot) {
2631 if scanned.is_multiple_of(256) {
2632 if let Some(control) = control {
2633 control.checkpoint()?;
2634 }
2635 }
2636 scanned += 1;
2637 if row.deleted {
2638 self.remove_hot_for_row(row.row_id, snapshot);
2639 } else if !self.row_expired_at(&row, ttl_now) {
2640 self.index_row(&row);
2641 }
2642 }
2643 for row in self.memtable.visible_versions(snapshot) {
2644 if scanned.is_multiple_of(256) {
2645 if let Some(control) = control {
2646 control.checkpoint()?;
2647 }
2648 }
2649 scanned += 1;
2650 if row.deleted {
2651 self.remove_hot_for_row(row.row_id, snapshot);
2652 } else if !self.row_expired_at(&row, ttl_now) {
2653 self.index_row(&row);
2654 }
2655 }
2656 self.refresh_pk_by_row_from_hot();
2657 Ok(())
2658 }
2659
2660 fn refresh_pk_by_row_from_hot(&mut self) {
2661 self.pk_by_row_complete = true;
2662 if self.schema.primary_key().is_none() {
2663 self.pk_by_row.clear();
2664 return;
2665 }
2666 self.pk_by_row = ReversePkMap::from_entries(
2672 self.hot
2673 .entries()
2674 .into_iter()
2675 .map(|(key, row_id)| (row_id, key)),
2676 );
2677 }
2678
2679 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2680 if self.schema.primary_key().is_some() {
2681 self.pk_by_row.insert(row_id, key.clone());
2682 }
2683 self.hot.insert(key, row_id);
2684 }
2685
2686 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2690 self.build_learned_ranges_inner(None)
2691 }
2692
2693 fn build_learned_ranges_inner(
2694 &mut self,
2695 control: Option<&crate::ExecutionControl>,
2696 ) -> Result<()> {
2697 self.learned_range = Arc::new(HashMap::new());
2698 if self.run_refs.len() != 1 {
2699 return Ok(());
2700 }
2701 let cols: Vec<(u16, usize)> = self
2702 .schema
2703 .indexes
2704 .iter()
2705 .filter(|i| i.kind == IndexKind::LearnedRange)
2706 .map(|i| {
2707 (
2708 i.column_id,
2709 i.options
2710 .learned_range
2711 .as_ref()
2712 .map(|options| options.epsilon)
2713 .unwrap_or(16),
2714 )
2715 })
2716 .collect();
2717 if cols.is_empty() {
2718 return Ok(());
2719 }
2720 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2721 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2722 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2723 _ => return Ok(()),
2724 };
2725 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2726 if column_index % 256 == 0 {
2727 if let Some(control) = control {
2728 control.checkpoint()?;
2729 }
2730 }
2731 let ty = self
2732 .schema
2733 .columns
2734 .iter()
2735 .find(|c| c.id == cid)
2736 .map(|c| c.ty.clone())
2737 .unwrap_or(TypeId::Int64);
2738 match ty {
2739 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2740 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2741 let pairs: Vec<(i64, u64)> = data
2742 .iter()
2743 .zip(row_ids.iter())
2744 .map(|(v, r)| (*v, *r))
2745 .collect();
2746 Arc::make_mut(&mut self.learned_range).insert(
2747 cid,
2748 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2749 );
2750 }
2751 }
2752 TypeId::Float64 => {
2753 if let columnar::NativeColumn::Float64 { data, .. } =
2754 reader.column_native(cid)?
2755 {
2756 let pairs: Vec<(f64, u64)> = data
2757 .iter()
2758 .zip(row_ids.iter())
2759 .map(|(v, r)| (*v, *r))
2760 .collect();
2761 Arc::make_mut(&mut self.learned_range).insert(
2762 cid,
2763 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2764 );
2765 }
2766 }
2767 _ => {}
2768 }
2769 }
2770 Ok(())
2771 }
2772
2773 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2780 if self.indexes_complete {
2781 crate::trace::QueryTrace::record(|t| {
2782 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2783 });
2784 return Ok(());
2785 }
2786 crate::trace::QueryTrace::record(|t| {
2787 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2788 });
2789 self.rebuild_indexes_from_runs()?;
2790 self.build_learned_ranges()?;
2791 self.indexes_complete = true;
2792 let epoch = self.current_epoch();
2793 self.checkpoint_indexes(epoch);
2794 Ok(())
2795 }
2796
2797 #[doc(hidden)]
2800 pub fn ensure_indexes_complete_controlled<F>(
2801 &mut self,
2802 control: &crate::ExecutionControl,
2803 before_publish: F,
2804 ) -> Result<bool>
2805 where
2806 F: FnOnce() -> bool,
2807 {
2808 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2809 .map(|(changed, _)| changed)
2810 }
2811
2812 #[doc(hidden)]
2815 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2816 &mut self,
2817 control: &crate::ExecutionControl,
2818 before_publish: F,
2819 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2820 where
2821 F: FnOnce() -> bool,
2822 {
2823 if self.indexes_complete {
2824 crate::trace::QueryTrace::record(|trace| {
2825 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2826 });
2827 return Ok((false, None));
2828 }
2829 crate::trace::QueryTrace::record(|trace| {
2830 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2831 });
2832 control.checkpoint()?;
2833 let maintenance_epoch = self.current_epoch();
2834 self.rebuild_indexes_from_runs_inner(Some(control))?;
2835 self.build_learned_ranges_inner(Some(control))?;
2836 control.checkpoint()?;
2837 if !before_publish() {
2838 return Err(MongrelError::Cancelled);
2839 }
2840 self.indexes_complete = true;
2841 self.checkpoint_indexes(maintenance_epoch);
2842 Ok((
2843 true,
2844 Some(MaintenanceReceipt {
2845 epoch: maintenance_epoch,
2846 }),
2847 ))
2848 }
2849
2850 fn pending_epoch(&self) -> Epoch {
2851 Epoch(self.epoch.visible().0 + 1)
2852 }
2853
2854 fn is_shared(&self) -> bool {
2857 matches!(self.wal, WalSink::Shared(_))
2858 }
2859
2860 fn ensure_txn_id(&mut self) -> Result<u64> {
2864 if self.current_txn_id == 0 {
2865 let id = match &self.wal {
2866 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2867 WalSink::Private(_) => {
2868 return Err(MongrelError::Full(
2869 "standalone transaction id namespace exhausted".into(),
2870 ))
2871 }
2872 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2873 };
2874 self.current_txn_id = id;
2875 }
2876 Ok(self.current_txn_id)
2877 }
2878
2879 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2882 self.ensure_writable()?;
2883 let txn_id = self.ensure_txn_id()?;
2884 let table_id = self.table_id;
2885 match &mut self.wal {
2886 WalSink::Private(w) => {
2887 w.append_txn(txn_id, op)?;
2888 self.pending_private_mutations = true;
2889 }
2890 WalSink::Shared(s) => {
2891 s.wal.lock().append(txn_id, table_id, op)?;
2892 }
2893 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2894 }
2895 Ok(())
2896 }
2897
2898 fn ensure_writable(&self) -> Result<()> {
2899 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2900 return Err(MongrelError::ReadOnlyReplica);
2901 }
2902 if self.durable_commit_failed {
2903 return Err(MongrelError::Other(
2904 "table poisoned by post-commit failure; reopen required".into(),
2905 ));
2906 }
2907 Ok(())
2908 }
2909
2910 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2921 match &self.auth {
2922 Some(checker) => checker.check(&self.name, perm),
2923 None => Ok(()),
2924 }
2925 }
2926 pub fn require_select(&self) -> Result<()> {
2931 self.require(crate::auth_state::RequiredPermission::Select)
2932 }
2933 fn require_insert(&self) -> Result<()> {
2934 self.require(crate::auth_state::RequiredPermission::Insert)
2935 }
2936 #[allow(dead_code)]
2940 fn require_update(&self) -> Result<()> {
2941 self.require(crate::auth_state::RequiredPermission::Update)
2942 }
2943 fn require_delete(&self) -> Result<()> {
2944 self.require(crate::auth_state::RequiredPermission::Delete)
2945 }
2946
2947 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2950 self.require_insert()?;
2951 Ok(self.put_returning(columns)?.0)
2952 }
2953
2954 pub fn put_returning(
2959 &mut self,
2960 mut columns: Vec<(u16, Value)>,
2961 ) -> Result<(RowId, Option<i64>)> {
2962 self.require_insert()?;
2963 let assigned = self.fill_auto_inc(&mut columns)?;
2964 self.apply_defaults(&mut columns)?;
2965 self.schema.validate_values(&columns)?;
2966 let row_id = if self.schema.clustered {
2971 self.derive_clustered_row_id(&columns)?
2972 } else {
2973 self.allocator.alloc()?
2974 };
2975 let epoch = self.pending_epoch();
2976 let mut row = Row::new(row_id, epoch);
2977 for (col_id, val) in columns {
2978 row.columns.insert(col_id, val);
2979 }
2980 self.commit_rows(vec![row], assigned.is_some())?;
2981 Ok((row_id, assigned))
2982 }
2983
2984 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2987 self.require_insert()?;
2988 Ok(self
2989 .put_batch_returning(batch)?
2990 .into_iter()
2991 .map(|(r, _)| r)
2992 .collect())
2993 }
2994
2995 pub fn put_batch_returning(
2998 &mut self,
2999 batch: Vec<Vec<(u16, Value)>>,
3000 ) -> Result<Vec<(RowId, Option<i64>)>> {
3001 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
3002 for mut cols in batch {
3003 let assigned = self.fill_auto_inc(&mut cols)?;
3004 self.apply_defaults(&mut cols)?;
3005 filled.push((cols, assigned));
3006 }
3007 for (cols, _) in &filled {
3008 self.schema.validate_values(cols)?;
3009 }
3010 let epoch = self.pending_epoch();
3011 let mut rows = Vec::with_capacity(filled.len());
3012 let mut ids = Vec::with_capacity(filled.len());
3013 let first_row_id = if self.schema.clustered {
3014 None
3015 } else {
3016 let count = u64::try_from(filled.len())
3017 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
3018 Some(self.allocator.alloc_range(count)?.0)
3019 };
3020 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
3021 let row_id = match first_row_id {
3022 Some(first) => RowId(first + row_index as u64),
3023 None => self.derive_clustered_row_id(&cols)?,
3024 };
3025 let mut row = Row::new(row_id, epoch);
3026 for (c, v) in cols {
3027 row.columns.insert(c, v);
3028 }
3029 ids.push((row_id, assigned));
3030 rows.push(row);
3031 }
3032 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
3033 self.commit_rows(rows, all_auto_generated)?;
3034 Ok(ids)
3035 }
3036
3037 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
3043 self.ensure_writable()?;
3044 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3045 return Ok(None);
3046 };
3047 let pos = columns.iter().position(|(c, _)| *c == cid);
3048 let assigned = match pos {
3049 Some(i) => match &columns[i].1 {
3050 Value::Null => {
3051 let next = self.alloc_auto_inc_value()?;
3052 columns[i].1 = Value::Int64(next);
3053 Some(next)
3054 }
3055 Value::Int64(n) => {
3056 self.advance_auto_inc_past(*n)?;
3057 None
3058 }
3059 other => {
3060 return Err(MongrelError::InvalidArgument(format!(
3061 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
3062 other
3063 )))
3064 }
3065 },
3066 None => {
3067 let next = self.alloc_auto_inc_value()?;
3068 columns.push((cid, Value::Int64(next)));
3069 Some(next)
3070 }
3071 };
3072 Ok(assigned)
3073 }
3074
3075 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
3081 for col in &self.schema.columns {
3082 let Some(expr) = &col.default_value else {
3083 continue;
3084 };
3085 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
3087 continue;
3088 }
3089 let pos = columns.iter().position(|(c, _)| *c == col.id);
3090 let needs_default = match pos {
3091 None => true,
3092 Some(i) => matches!(columns[i].1, Value::Null),
3093 };
3094 if !needs_default {
3095 continue;
3096 }
3097 let v = match expr {
3098 crate::schema::DefaultExpr::Static(v) => v.clone(),
3099 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
3100 crate::schema::DefaultExpr::Uuid => {
3101 let mut buf = [0u8; 16];
3102 getrandom::getrandom(&mut buf)
3103 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
3104 Value::Uuid(buf)
3105 }
3106 };
3107 match pos {
3108 None => columns.push((col.id, v)),
3109 Some(i) => columns[i].1 = v,
3110 }
3111 }
3112 Ok(())
3113 }
3114
3115 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
3117 self.ensure_auto_inc_seeded()?;
3118 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3120 let v = ai.next;
3121 ai.next = ai
3122 .next
3123 .checked_add(1)
3124 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3125 Ok(v)
3126 }
3127
3128 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
3131 self.ensure_auto_inc_seeded()?;
3132 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3133 let floor = used
3134 .checked_add(1)
3135 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3136 .max(1);
3137 if ai.next < floor {
3138 ai.next = floor;
3139 }
3140 Ok(())
3141 }
3142
3143 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
3148 let needs_seed = match self.auto_inc {
3149 Some(ai) => !ai.seeded,
3150 None => return Ok(()),
3151 };
3152 if !needs_seed {
3153 return Ok(());
3154 }
3155 if self.seed_empty_auto_inc() {
3156 return Ok(());
3157 }
3158 let cid = self
3159 .auto_inc
3160 .as_ref()
3161 .expect("auto-inc column present")
3162 .column_id;
3163 let max = self.scan_max_int64(cid)?;
3164 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3165 let floor = max
3166 .checked_add(1)
3167 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3168 .max(1);
3169 if ai.next < floor {
3170 ai.next = floor;
3171 }
3172 ai.seeded = true;
3173 Ok(())
3174 }
3175
3176 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
3177 if n == 0 || self.auto_inc.is_none() {
3178 return Ok(None);
3179 }
3180 self.ensure_auto_inc_seeded()?;
3181 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3182 let start = ai.next;
3183 let count = i64::try_from(n)
3184 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
3185 ai.next = ai
3186 .next
3187 .checked_add(count)
3188 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3189 Ok(Some(start))
3190 }
3191
3192 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
3196 let mut max: i64 = 0;
3197 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
3198 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3199 if *n > max {
3200 max = *n;
3201 }
3202 }
3203 }
3204 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3205 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3206 if *n > max {
3207 max = *n;
3208 }
3209 }
3210 }
3211 for rr in self.run_refs.clone() {
3212 let reader = self.open_reader(rr.run_id)?;
3213 if let Some(stats) = reader.column_page_stats(column_id) {
3214 for s in stats {
3215 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3216 if n > max {
3217 max = n;
3218 }
3219 }
3220 }
3221 } else if reader.has_column(column_id) {
3222 if let columnar::NativeColumn::Int64 { data, validity } =
3223 reader.column_native_shared(column_id)?
3224 {
3225 for (i, n) in data.iter().enumerate() {
3226 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3227 {
3228 max = *n;
3229 }
3230 }
3231 }
3232 }
3233 }
3234 Ok(max)
3235 }
3236
3237 fn seed_empty_auto_inc(&mut self) -> bool {
3238 let Some(ai) = self.auto_inc.as_mut() else {
3239 return false;
3240 };
3241 if ai.seeded || self.live_count != 0 {
3242 return false;
3243 }
3244 if ai.next < 1 {
3245 ai.next = 1;
3246 }
3247 ai.seeded = true;
3248 true
3249 }
3250
3251 fn advance_auto_inc_from_native_columns(
3252 &mut self,
3253 columns: &[(u16, columnar::NativeColumn)],
3254 n: usize,
3255 live_before: u64,
3256 ) -> Result<()> {
3257 let Some(ai) = self.auto_inc.as_mut() else {
3258 return Ok(());
3259 };
3260 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3261 return Ok(());
3262 };
3263 let columnar::NativeColumn::Int64 { data, validity } = col else {
3264 return Err(MongrelError::InvalidArgument(format!(
3265 "AUTO_INCREMENT column {} must be Int64",
3266 ai.column_id
3267 )));
3268 };
3269 let max = if native_int64_strictly_increasing(col, n) {
3270 data.get(n.saturating_sub(1)).copied()
3271 } else {
3272 data.iter()
3273 .take(n)
3274 .enumerate()
3275 .filter_map(|(i, v)| {
3276 if validity.is_empty() || columnar::validity_bit(validity, i) {
3277 Some(*v)
3278 } else {
3279 None
3280 }
3281 })
3282 .max()
3283 };
3284 if let Some(max) = max {
3285 let floor = max
3286 .checked_add(1)
3287 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3288 .max(1);
3289 if ai.next < floor {
3290 ai.next = floor;
3291 }
3292 if ai.seeded || live_before == 0 {
3293 ai.seeded = true;
3294 }
3295 }
3296 Ok(())
3297 }
3298
3299 fn fill_auto_inc_native_columns(
3300 &mut self,
3301 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3302 n: usize,
3303 ) -> Result<()> {
3304 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3305 return Ok(());
3306 };
3307 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3308 if let Some(start) = self.alloc_auto_inc_range(n)? {
3309 columns.push((
3310 cid,
3311 columnar::NativeColumn::Int64 {
3312 data: (start..start.saturating_add(n as i64)).collect(),
3313 validity: vec![0xFF; n.div_ceil(8)],
3314 },
3315 ));
3316 }
3317 return Ok(());
3318 };
3319
3320 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3321 return Err(MongrelError::InvalidArgument(format!(
3322 "AUTO_INCREMENT column {cid} must be Int64"
3323 )));
3324 };
3325 if data.len() < n {
3326 return Err(MongrelError::InvalidArgument(format!(
3327 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3328 data.len()
3329 )));
3330 }
3331 if columnar::all_non_null(validity, n) {
3332 return Ok(());
3333 }
3334 if validity.iter().all(|b| *b == 0) {
3335 if let Some(start) = self.alloc_auto_inc_range(n)? {
3336 for (i, slot) in data.iter_mut().take(n).enumerate() {
3337 *slot = start.saturating_add(i as i64);
3338 }
3339 *validity = vec![0xFF; n.div_ceil(8)];
3340 }
3341 return Ok(());
3342 }
3343
3344 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3345 for (i, slot) in data.iter_mut().enumerate().take(n) {
3346 if columnar::validity_bit(validity, i) {
3347 self.advance_auto_inc_past(*slot)?;
3348 } else {
3349 *slot = self.alloc_auto_inc_value()?;
3350 }
3351 }
3352 *validity = new_validity;
3353 Ok(())
3354 }
3355
3356 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3370 self.ensure_writable()?;
3371 if self.auto_inc.is_none() {
3372 return Ok(None);
3373 }
3374 Ok(Some(self.alloc_auto_inc_value()?))
3375 }
3376
3377 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3383 let payload = bincode::serialize(&rows)?;
3384 self.wal_append_data(Op::Put {
3385 table_id: self.table_id,
3386 rows: payload,
3387 })?;
3388 if self.is_shared() {
3389 self.pending_rows_auto_inc
3390 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3391 self.pending_rows.extend(rows);
3392 } else {
3393 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3394 }
3395 Ok(())
3396 }
3397
3398 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3401 self.ensure_indexes_complete()
3402 }
3403
3404 pub(crate) fn prepare_durable_publish_controlled(
3405 &mut self,
3406 control: &crate::ExecutionControl,
3407 ) -> Result<()> {
3408 self.ensure_indexes_complete_controlled(control, || true)?;
3409 Ok(())
3410 }
3411
3412 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3413 self.apply_put_rows_inner_prepared(rows, true);
3414 }
3415
3416 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3417 if check_existing_pk {
3418 self.ensure_indexes_complete()?;
3419 }
3420 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3421 Ok(())
3422 }
3423
3424 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3428 if rows.len() == 1 {
3432 let row = rows.into_iter().next().expect("len checked");
3433 self.apply_put_row_single(row, check_existing_pk);
3434 return;
3435 }
3436 let pk_id = self.schema.primary_key().map(|c| c.id);
3453 let probe = match pk_id {
3454 Some(pid) => {
3455 check_existing_pk
3456 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3457 }
3458 None => false,
3459 };
3460 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3463 for r in rows {
3464 for &cid in r.columns.keys() {
3465 self.pending_put_cols.insert(cid);
3466 }
3467 match pk_id {
3468 Some(pid) if probe || maintain_pk_by_row => {
3469 if let Some(pk_val) = r.columns.get(&pid) {
3470 let key = self.index_lookup_key(pid, pk_val);
3471 if probe {
3472 if let Some(old_rid) = self.hot.get(&key) {
3473 if old_rid != r.row_id {
3474 self.tombstone_row(
3475 old_rid,
3476 r.committed_epoch,
3477 r.commit_ts,
3478 true,
3479 );
3480 }
3481 }
3482 }
3483 if maintain_pk_by_row {
3484 self.pk_by_row.insert(r.row_id, key);
3485 }
3486 }
3487 }
3488 Some(_) => {}
3489 None => {
3490 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3491 }
3492 }
3493 self.index_row(&r);
3494 self.reservoir.offer(r.row_id.0);
3495 self.memtable.upsert(r);
3496 self.live_count = self.live_count.saturating_add(1);
3499 }
3500 self.data_generation = self.data_generation.wrapping_add(1);
3501 }
3502
3503 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3507 for &cid in row.columns.keys() {
3508 self.pending_put_cols.insert(cid);
3509 }
3510 let epoch = row.committed_epoch;
3511 if let Some(pk_col) = self.schema.primary_key() {
3512 let pk_id = pk_col.id;
3513 if let Some(pk_val) = row.columns.get(&pk_id) {
3514 let maintain_pk_by_row = self.pk_by_row_complete;
3518 if check_existing_pk || maintain_pk_by_row {
3519 let key = self.index_lookup_key(pk_id, pk_val);
3520 if check_existing_pk {
3521 if let Some(old_rid) = self.hot.get(&key) {
3522 if old_rid != row.row_id {
3523 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
3524 }
3525 }
3526 }
3527 if maintain_pk_by_row {
3528 self.pk_by_row.insert(row.row_id, key);
3529 }
3530 }
3531 }
3532 } else {
3533 self.hot
3534 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3535 }
3536 self.index_row(&row);
3537 self.reservoir.offer(row.row_id.0);
3538 self.memtable.upsert(row);
3539 self.live_count = self.live_count.saturating_add(1);
3540 self.data_generation = self.data_generation.wrapping_add(1);
3541 }
3542
3543 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3546 self.allocator.alloc()
3547 }
3548
3549 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3555 let pk = self.schema.primary_key().ok_or_else(|| {
3556 MongrelError::Schema("clustered table requires a single-column primary key".into())
3557 })?;
3558 let pk_val = columns
3559 .iter()
3560 .find(|(id, _)| *id == pk.id)
3561 .map(|(_, v)| v)
3562 .ok_or_else(|| {
3563 MongrelError::Schema(format!(
3564 "clustered table missing primary key column {} ({})",
3565 pk.id, pk.name
3566 ))
3567 })?;
3568 Ok(clustered_row_id(pk_val))
3569 }
3570
3571 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3579 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3580 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3581 }
3582 let n = rows.len();
3583 for r in rows {
3584 for &cid in r.columns.keys() {
3585 self.pending_put_cols.insert(cid);
3586 }
3587 }
3588 let (losers, winner_pks) = self.partition_pk_winners(rows);
3589 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3590 let group_ts = rows.first().and_then(|r| r.commit_ts);
3592 for (key, &row_id) in &winner_pks {
3593 if let Some(old_rid) = self.hot.get(key) {
3594 if old_rid != row_id {
3595 self.tombstone_row(old_rid, epoch, group_ts, true);
3596 }
3597 }
3598 }
3599 for &loser_rid in &losers {
3602 self.tombstone_row(loser_rid, epoch, group_ts, false);
3603 }
3604 for (key, row_id) in winner_pks {
3606 self.insert_hot_pk(key, row_id);
3607 }
3608 if self.schema.primary_key().is_none() {
3609 for r in rows {
3610 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3611 }
3612 }
3613 for r in rows {
3614 self.allocator.advance_to(r.row_id)?;
3615 if !losers.contains(&r.row_id) {
3616 self.index_row(r);
3617 }
3618 }
3619 for r in rows {
3620 if !losers.contains(&r.row_id) {
3621 self.reservoir.offer(r.row_id.0);
3622 }
3623 }
3624 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3625 self.data_generation = self.data_generation.wrapping_add(1);
3626 Ok(())
3627 }
3628
3629 pub(crate) fn recover_apply(
3634 &mut self,
3635 rows: Vec<Row>,
3636 deletes: Vec<(RowId, Epoch)>,
3637 ) -> Result<()> {
3638 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3642 std::collections::BTreeMap::new();
3643 for row in rows {
3644 if row.row_id.0 >= u64::MAX - 1 {
3645 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3646 }
3647 self.allocator.advance_to(row.row_id)?;
3648 if let Some(ai) = self.auto_inc.as_mut() {
3653 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3654 let next = n.checked_add(1).ok_or_else(|| {
3655 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3656 })?;
3657 if next > ai.next {
3658 ai.next = next;
3659 }
3660 }
3661 }
3662 by_epoch.entry(row.committed_epoch).or_default().push(row);
3663 }
3664 for (epoch, group) in by_epoch {
3665 let (losers, winner_pks) = self.partition_pk_winners(&group);
3666 let group_ts = group.first().and_then(|r| r.commit_ts);
3668 for (key, &row_id) in &winner_pks {
3669 if let Some(old_rid) = self.hot.get(key) {
3670 if old_rid != row_id {
3671 self.tombstone_row(old_rid, epoch, group_ts, false);
3672 }
3673 }
3674 }
3675 for (key, row_id) in winner_pks {
3676 self.insert_hot_pk(key, row_id);
3677 }
3678 if self.schema.primary_key().is_none() {
3679 for r in &group {
3680 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3681 }
3682 }
3683 for r in &group {
3684 if !losers.contains(&r.row_id) {
3685 self.memtable.upsert(r.clone());
3686 self.index_row(r);
3687 }
3688 }
3689 }
3690 for (rid, epoch) in deletes {
3691 self.memtable.tombstone(rid, epoch);
3692 self.remove_hot_for_row(rid, epoch);
3693 }
3694 self.reservoir_complete = false;
3697 Ok(())
3698 }
3699
3700 pub(crate) fn flushed_epoch(&self) -> u64 {
3702 self.flushed_epoch
3703 }
3704
3705 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3706 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3707 }
3708
3709 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3711 self.schema.validate_values(cells)
3712 }
3713
3714 fn validate_columns_not_null(
3718 &self,
3719 columns: &[(u16, columnar::NativeColumn)],
3720 n: usize,
3721 ) -> Result<()> {
3722 let by_id: HashMap<u16, &columnar::NativeColumn> =
3723 columns.iter().map(|(id, c)| (*id, c)).collect();
3724 for col in &self.schema.columns {
3725 if !col.flags.contains(ColumnFlags::NULLABLE) {
3726 match by_id.get(&col.id) {
3727 None => {
3728 return Err(MongrelError::InvalidArgument(format!(
3729 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3730 col.name, col.id
3731 )));
3732 }
3733 Some(c) => {
3734 if c.null_count(n) != 0 {
3735 return Err(MongrelError::InvalidArgument(format!(
3736 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3737 col.name, col.id
3738 )));
3739 }
3740 }
3741 }
3742 }
3743 if let TypeId::Enum { variants } = &col.ty {
3744 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3745 if by_id.contains_key(&col.id) {
3746 return Err(MongrelError::InvalidArgument(format!(
3747 "column '{}' ({}) enum requires a bytes column",
3748 col.name, col.id
3749 )));
3750 }
3751 continue;
3752 };
3753 for index in 0..n {
3754 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3755 continue;
3756 };
3757 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3758 return Err(MongrelError::InvalidArgument(format!(
3759 "column '{}' ({}) enum value {:?} is not one of {:?}",
3760 col.name,
3761 col.id,
3762 String::from_utf8_lossy(value),
3763 variants
3764 )));
3765 }
3766 }
3767 }
3768 }
3769 Ok(())
3770 }
3771
3772 fn bulk_pk_winner_indices(
3777 &self,
3778 columns: &[(u16, columnar::NativeColumn)],
3779 n: usize,
3780 ) -> Option<Vec<usize>> {
3781 let pk_col = self.schema.primary_key()?;
3782 let pk_id = pk_col.id;
3783 let pk_ty = pk_col.ty.clone();
3784 let by_id: HashMap<u16, &columnar::NativeColumn> =
3785 columns.iter().map(|(id, c)| (*id, c)).collect();
3786 let pk_native = by_id.get(&pk_id)?;
3787 if native_int64_strictly_increasing(pk_native, n) {
3788 return None;
3789 }
3790 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3792 let mut null_pk_rows: Vec<usize> = Vec::new();
3793 for i in 0..n {
3794 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3795 Some(key) => {
3796 last.insert(key, i);
3797 }
3798 None => null_pk_rows.push(i),
3799 }
3800 }
3801 let mut winners: HashSet<usize> = last.values().copied().collect();
3802 for i in null_pk_rows {
3803 winners.insert(i);
3804 }
3805 Some((0..n).filter(|i| winners.contains(i)).collect())
3806 }
3807
3808 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3810 self.require_delete()?;
3811 let epoch = self.pending_epoch();
3812 self.wal_append_data(Op::Delete {
3813 table_id: self.table_id,
3814 row_ids: vec![row_id],
3815 })?;
3816 if self.is_shared() {
3817 self.pending_dels.push(row_id);
3818 } else {
3819 self.apply_delete(row_id, epoch);
3820 }
3821 Ok(())
3822 }
3823
3824 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3825 let pre = self.get(row_id, self.snapshot());
3826 self.delete(row_id)?;
3827 Ok(pre.map(|row| {
3828 let mut columns: Vec<_> = row.columns.into_iter().collect();
3829 columns.sort_by_key(|(id, _)| *id);
3830 OwnedRow { columns }
3831 }))
3832 }
3833
3834 pub fn truncate(&mut self) -> Result<()> {
3836 self.require_delete()?;
3837 let epoch = self.pending_epoch();
3838 self.wal_append_data(Op::TruncateTable {
3839 table_id: self.table_id,
3840 })?;
3841 self.pending_rows.clear();
3842 self.pending_rows_auto_inc.clear();
3843 self.pending_dels.clear();
3844 self.pending_truncate = Some(epoch);
3845 Ok(())
3846 }
3847
3848 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3850 self.run_refs.clear();
3855 self.retiring.clear();
3856 self.memtable = Memtable::new();
3857 self.mutable_run = MutableRun::new();
3858 self.hot = HotIndex::new();
3859 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3860 self.bitmap = bitmap;
3861 self.ann = ann;
3862 self.fm = fm;
3863 self.sparse = sparse;
3864 self.minhash = minhash;
3865 self.learned_range = Arc::new(HashMap::new());
3866 self.pk_by_row.clear();
3867 self.pk_by_row_complete = false;
3868 self.live_count = 0;
3869 self.reservoir = crate::reservoir::Reservoir::default();
3870 self.reservoir_complete = true;
3871 self.had_deletes = true;
3872 self.agg_cache = Arc::new(HashMap::new());
3873 self.global_idx_epoch = 0;
3874 self.indexes_complete = true;
3875 self.pending_delete_rids.clear();
3876 self.pending_put_cols.clear();
3877 self.pending_rows.clear();
3878 self.pending_rows_auto_inc.clear();
3879 self.pending_dels.clear();
3880 self.clear_result_cache();
3881 self.invalidate_index_checkpoint();
3882 self.data_generation = self.data_generation.wrapping_add(1);
3883 }
3884
3885 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3888 self.apply_delete_at(row_id, epoch, None);
3889 }
3890
3891 pub(crate) fn apply_delete_at(
3893 &mut self,
3894 row_id: RowId,
3895 epoch: Epoch,
3896 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
3897 ) {
3898 self.remove_hot_for_row(row_id, epoch);
3899 self.tombstone_row(row_id, epoch, commit_ts, true);
3900 self.data_generation = self.data_generation.wrapping_add(1);
3901 }
3902
3903 fn tombstone_row(
3907 &mut self,
3908 row_id: RowId,
3909 epoch: Epoch,
3910 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
3911 adjust_live_count: bool,
3912 ) {
3913 let tombstone = Row {
3914 row_id,
3915 committed_epoch: epoch,
3916 columns: std::collections::HashMap::new(),
3917 deleted: true,
3918 commit_ts,
3919 };
3920 self.memtable.upsert(tombstone);
3921 self.pk_by_row.remove(&row_id);
3922 if adjust_live_count {
3923 self.live_count = self.live_count.saturating_sub(1);
3924 }
3925 self.pending_delete_rids.insert(row_id.0 as u32);
3927 self.had_deletes = true;
3930 self.agg_cache = Arc::new(HashMap::new());
3931 }
3932
3933 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3937 let Some(pk_col) = self.schema.primary_key() else {
3938 return;
3939 };
3940 if self.pk_by_row_complete {
3943 if let Some(key) = self.pk_by_row.remove(&row_id) {
3944 if self.hot.get(&key) == Some(row_id) {
3945 self.hot.remove(&key);
3946 }
3947 }
3948 return;
3949 }
3950 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3969 if self.indexes_complete {
3970 let pk_val = self
3971 .memtable
3972 .get_version(row_id, lookup_epoch)
3973 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3974 .or_else(|| {
3975 self.mutable_run
3976 .get_version(row_id, lookup_epoch)
3977 .filter(|(_, r)| !r.deleted)
3978 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3979 })
3980 .or_else(|| {
3981 self.run_refs.iter().find_map(|rr| {
3982 let mut reader = self.open_reader(rr.run_id).ok()?;
3983 let (_, deleted, val) = reader
3984 .get_version_column(row_id, lookup_epoch, pk_col.id)
3985 .ok()??;
3986 if deleted {
3987 return None;
3988 }
3989 val
3990 })
3991 });
3992 if let Some(pk_val) = pk_val {
3993 let key = self.index_lookup_key(pk_col.id, &pk_val);
3994 if self.hot.get(&key) == Some(row_id) {
3995 self.hot.remove(&key);
3996 }
3997 return;
3998 }
3999 }
4000 self.refresh_pk_by_row_from_hot();
4005 if let Some(key) = self.pk_by_row.remove(&row_id) {
4006 if self.hot.get(&key) == Some(row_id) {
4007 self.hot.remove(&key);
4008 }
4009 }
4010 }
4011
4012 fn partition_pk_winners(
4017 &self,
4018 rows: &[Row],
4019 ) -> (
4020 std::collections::HashSet<RowId>,
4021 std::collections::HashMap<Vec<u8>, RowId>,
4022 ) {
4023 let mut losers = std::collections::HashSet::new();
4024 let Some(pk_col) = self.schema.primary_key() else {
4025 return (losers, std::collections::HashMap::new());
4026 };
4027 let pk_id = pk_col.id;
4028 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
4029 std::collections::HashMap::new();
4030 for r in rows {
4031 let Some(pk_val) = r.columns.get(&pk_id) else {
4032 continue;
4033 };
4034 let key = self.index_lookup_key(pk_id, pk_val);
4035 if let Some(&old_rid) = winners.get(&key) {
4036 losers.insert(old_rid);
4037 }
4038 winners.insert(key, r.row_id);
4039 }
4040 (losers, winners)
4041 }
4042
4043 fn index_row(&mut self, row: &Row) {
4044 if row.deleted {
4045 return;
4046 }
4047 let any_predicate = self
4055 .schema
4056 .indexes
4057 .iter()
4058 .any(|idx| idx.predicate.is_some());
4059 if any_predicate {
4060 let columns_map: HashMap<u16, &Value> =
4061 row.columns.iter().map(|(k, v)| (*k, v)).collect();
4062 let name_to_id: HashMap<&str, u16> = self
4063 .schema
4064 .columns
4065 .iter()
4066 .map(|c| (c.name.as_str(), c.id))
4067 .collect();
4068 for idx in &self.schema.indexes {
4069 if let Some(pred) = &idx.predicate {
4070 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4071 continue; }
4073 }
4074 index_into_single(
4076 idx,
4077 &self.schema,
4078 row,
4079 &mut self.hot,
4080 &mut self.bitmap,
4081 &mut self.ann,
4082 &mut self.fm,
4083 &mut self.sparse,
4084 &mut self.minhash,
4085 );
4086 }
4087 return;
4088 }
4089 if self.column_keys.is_empty() {
4093 index_into(
4094 &self.schema,
4095 row,
4096 &mut self.hot,
4097 &mut self.bitmap,
4098 &mut self.ann,
4099 &mut self.fm,
4100 &mut self.sparse,
4101 &mut self.minhash,
4102 );
4103 return;
4104 }
4105 let effective_row = self.tokenized_for_indexes(row);
4106 index_into(
4107 &self.schema,
4108 &effective_row,
4109 &mut self.hot,
4110 &mut self.bitmap,
4111 &mut self.ann,
4112 &mut self.fm,
4113 &mut self.sparse,
4114 &mut self.minhash,
4115 );
4116 }
4117
4118 fn tokenized_for_indexes(&self, row: &Row) -> Row {
4124 if self.column_keys.is_empty() {
4125 return row.clone();
4126 }
4127 {
4128 use crate::encryption::SCHEME_HMAC_EQ;
4129 let mut tok = row.clone();
4130 for (&cid, &(_, scheme)) in &self.column_keys {
4131 if scheme != SCHEME_HMAC_EQ {
4132 continue;
4133 }
4134 if let Some(v) = tok.columns.get(&cid).cloned() {
4135 if let Some(t) = self.tokenize_value(cid, &v) {
4136 tok.columns.insert(cid, t);
4137 }
4138 }
4139 }
4140 tok
4141 }
4142 }
4143
4144 pub fn commit(&mut self) -> Result<Epoch> {
4149 self.commit_inner(None)
4150 }
4151
4152 #[doc(hidden)]
4155 pub fn commit_controlled<F>(
4156 &mut self,
4157 control: &crate::ExecutionControl,
4158 mut before_commit: F,
4159 ) -> Result<Epoch>
4160 where
4161 F: FnMut() -> Result<()>,
4162 {
4163 self.commit_inner(Some((control, &mut before_commit)))
4164 }
4165
4166 fn commit_inner(
4167 &mut self,
4168 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4169 ) -> Result<Epoch> {
4170 self.ensure_writable()?;
4171 if !self.has_pending_mutations() {
4172 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
4173 return Err(MongrelError::Full(
4174 "standalone transaction id namespace exhausted".into(),
4175 ));
4176 }
4177 return Ok(self.epoch.visible());
4178 }
4179 self.commit_new_epoch_inner(controlled)
4180 }
4181
4182 fn commit_new_epoch(&mut self) -> Result<Epoch> {
4186 self.commit_new_epoch_inner(None)
4187 }
4188
4189 fn commit_new_epoch_inner(
4190 &mut self,
4191 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4192 ) -> Result<Epoch> {
4193 self.ensure_writable()?;
4194 if self.is_shared() {
4195 self.commit_shared(controlled)
4196 } else {
4197 self.commit_private(controlled)
4198 }
4199 }
4200
4201 fn commit_private(
4203 &mut self,
4204 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4205 ) -> Result<Epoch> {
4206 let commit_lock = Arc::clone(&self.commit_lock);
4210 let _g = commit_lock.lock();
4211 let txn_id = self.ensure_txn_id()?;
4214 if let Some((control, before_commit)) = controlled {
4215 control.checkpoint()?;
4216 before_commit()?;
4217 }
4218 let new_epoch = self.epoch.bump_assigned();
4219 let epoch_authority = Arc::clone(&self.epoch);
4220 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4221 let wal_result = match &mut self.wal {
4225 WalSink::Private(w) => w
4226 .append_txn(
4227 txn_id,
4228 Op::TxnCommit {
4229 epoch: new_epoch.0,
4230 added_runs: Vec::new(),
4231 },
4232 )
4233 .and_then(|_| w.sync()),
4234 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4235 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4236 };
4237 if let Err(error) = wal_result {
4238 self.durable_commit_failed = true;
4239 return Err(MongrelError::CommitOutcomeUnknown {
4240 epoch: new_epoch.0,
4241 message: error.to_string(),
4242 });
4243 }
4244 if let Some(epoch) = self.pending_truncate.take() {
4247 self.apply_truncate(epoch);
4248 }
4249 self.invalidate_pending_cache();
4250 let publish_result = self.persist_manifest(new_epoch);
4251 self.epoch.publish_in_order(new_epoch);
4255 epoch_guard.disarm();
4256 if let Err(error) = publish_result {
4257 self.durable_commit_failed = true;
4258 return Err(MongrelError::DurableCommit {
4259 epoch: new_epoch.0,
4260 message: error.to_string(),
4261 });
4262 }
4263 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4264 self.pending_private_mutations = false;
4265 self.data_generation = self.data_generation.wrapping_add(1);
4266 Ok(new_epoch)
4267 }
4268
4269 fn commit_shared(
4275 &mut self,
4276 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4277 ) -> Result<Epoch> {
4278 use std::sync::atomic::Ordering;
4279 let s = match &self.wal {
4280 WalSink::Shared(s) => s.clone(),
4281 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4282 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4283 };
4284 if s.poisoned.load(Ordering::Relaxed) {
4285 return Err(MongrelError::Other(
4286 "database poisoned by fsync error".into(),
4287 ));
4288 }
4289 let commit_lock = Arc::clone(&self.commit_lock);
4296 let _g = commit_lock.lock();
4297 if !self.pending_rows.is_empty() {
4298 match controlled.as_ref() {
4299 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4300 None => self.prepare_durable_publish()?,
4301 }
4302 }
4303 let txn_id = self.ensure_txn_id()?;
4306 let mut wal = s.wal.lock();
4307 if let Some((control, before_commit)) = controlled {
4308 control.checkpoint()?;
4309 before_commit()?;
4310 }
4311 let new_epoch = self.epoch.bump_assigned();
4312 let epoch_authority = Arc::clone(&self.epoch);
4313 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4314 let commit_ts = s.hlc.now().map_err(|skew| {
4317 MongrelError::Other(format!(
4318 "clock skew rejected commit timestamp allocation: {skew}"
4319 ))
4320 })?;
4321 let commit_seq = match wal.append_commit_at(
4322 txn_id,
4323 new_epoch,
4324 &[],
4325 commit_ts.physical_micros.saturating_mul(1_000),
4326 ) {
4327 Ok(commit_seq) => commit_seq,
4328 Err(error) => {
4329 s.poisoned.store(true, Ordering::Relaxed);
4330 s.lifecycle.poison();
4331 return Err(MongrelError::CommitOutcomeUnknown {
4332 epoch: new_epoch.0,
4333 message: error.to_string(),
4334 });
4335 }
4336 };
4337 drop(wal);
4338 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4339 s.poisoned.store(true, Ordering::Relaxed);
4340 s.lifecycle.poison();
4341 return Err(MongrelError::CommitOutcomeUnknown {
4342 epoch: new_epoch.0,
4343 message: error.to_string(),
4344 });
4345 }
4346
4347 if self.pending_truncate.take().is_some() {
4350 self.apply_truncate(new_epoch);
4351 }
4352 let mut rows = std::mem::take(&mut self.pending_rows);
4353 if !rows.is_empty() {
4354 for r in &mut rows {
4355 r.committed_epoch = new_epoch;
4356 r.commit_ts = Some(commit_ts);
4357 }
4358 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4359 let all_auto_generated =
4360 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4361 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4362 } else {
4363 self.pending_rows_auto_inc.clear();
4364 }
4365 let dels = std::mem::take(&mut self.pending_dels);
4366 for rid in dels {
4367 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
4368 }
4369
4370 self.invalidate_pending_cache();
4371 let publish_result = self.persist_manifest(new_epoch);
4372 self.epoch.publish_in_order(new_epoch);
4373 epoch_guard.disarm();
4374 let _ = s.change_wake.send(());
4375 if let Err(error) = publish_result {
4376 self.durable_commit_failed = true;
4377 s.poisoned.store(true, Ordering::Relaxed);
4378 s.lifecycle.poison();
4379 return Err(MongrelError::DurableCommit {
4380 epoch: new_epoch.0,
4381 message: error.to_string(),
4382 });
4383 }
4384 self.current_txn_id = 0;
4386 self.data_generation = self.data_generation.wrapping_add(1);
4387 Ok(new_epoch)
4388 }
4389
4390 pub fn flush(&mut self) -> Result<Epoch> {
4398 self.flush_with_outcome().map(|(epoch, _)| epoch)
4399 }
4400
4401 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4403 self.flush_with_outcome_inner(None)
4404 }
4405
4406 #[doc(hidden)]
4409 pub fn flush_with_outcome_controlled<F>(
4410 &mut self,
4411 control: &crate::ExecutionControl,
4412 mut before_commit: F,
4413 ) -> Result<(Epoch, bool)>
4414 where
4415 F: FnMut() -> Result<()>,
4416 {
4417 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4418 }
4419
4420 fn flush_with_outcome_inner(
4421 &mut self,
4422 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4423 ) -> Result<(Epoch, bool)> {
4424 match controlled.as_ref() {
4425 Some((control, _)) => {
4426 self.ensure_indexes_complete_controlled(control, || true)?;
4427 }
4428 None => self.ensure_indexes_complete()?,
4429 }
4430 let committed = self.has_pending_mutations();
4431 let epoch = self.commit_inner(controlled)?;
4432 let finish: Result<(Epoch, bool)> = (|| {
4433 let rows = self.memtable.drain_sorted();
4434 if !rows.is_empty() {
4435 self.mutable_run.insert_many(rows);
4436 }
4437 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4438 self.spill_mutable_run(epoch)?;
4439 self.mark_flushed(epoch)?;
4443 self.persist_manifest(epoch)?;
4444 self.build_learned_ranges()?;
4445 self.checkpoint_indexes(epoch);
4448 }
4449 Ok((epoch, committed))
4452 })();
4453 let outcome = match finish {
4454 Err(error) if committed => Err(MongrelError::DurableCommit {
4455 epoch: epoch.0,
4456 message: error.to_string(),
4457 }),
4458 result => result,
4459 };
4460 if outcome.is_ok() {
4461 let _ = self.publish_read_generation();
4467 }
4468 outcome
4469 }
4470
4471 fn has_pending_mutations(&self) -> bool {
4472 self.pending_private_mutations
4473 || !self.pending_rows.is_empty()
4474 || !self.pending_dels.is_empty()
4475 || self.pending_truncate.is_some()
4476 }
4477
4478 pub fn has_pending_writes(&self) -> bool {
4479 self.has_pending_mutations()
4480 }
4481
4482 pub fn force_flush(&mut self) -> Result<Epoch> {
4491 let saved = self.mutable_run_spill_bytes;
4492 self.mutable_run_spill_bytes = 1;
4493 let result = self.flush();
4494 self.mutable_run_spill_bytes = saved;
4495 result
4496 }
4497
4498 pub fn close(&mut self) -> Result<()> {
4505 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4506 self.force_flush()?;
4507 }
4508 Ok(())
4509 }
4510
4511 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4518 let op = Op::Flush {
4519 table_id: self.table_id,
4520 flushed_epoch: epoch.0,
4521 };
4522 match &mut self.wal {
4523 WalSink::Private(w) => {
4524 w.append_system(op)?;
4525 w.sync()?;
4526 }
4527 WalSink::Shared(s) => {
4528 s.wal.lock().append_system(op)?;
4533 }
4534 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4535 }
4536 self.flushed_epoch = epoch.0;
4537 if matches!(self.wal, WalSink::Private(_)) {
4538 self.rotate_wal(epoch)?;
4539 }
4540 Ok(())
4541 }
4542
4543 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4547 if self.mutable_run.is_empty() {
4548 return Ok(());
4549 }
4550 let run_id = self.alloc_run_id()?;
4551 let rows = self.mutable_run.drain_sorted();
4552 if rows.is_empty() {
4553 return Ok(());
4554 }
4555 let path = self.run_path(run_id);
4556 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4557 if let Some(kek) = &self.kek {
4558 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4559 }
4560 let header = match self.create_run_file(run_id)? {
4561 Some(file) => writer.write_file(file, &rows)?,
4562 None => writer.write(&path, &rows)?,
4563 };
4564 self.run_refs.push(RunRef {
4565 run_id: run_id as u128,
4566 level: 0,
4567 epoch_created: epoch.0,
4568 row_count: header.row_count,
4569 });
4570 Ok(())
4571 }
4572
4573 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4577 self.mutable_run_spill_bytes = bytes.max(1);
4578 }
4579
4580 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4584 self.compaction_zstd_level = level;
4585 }
4586
4587 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4591 self.result_cache.lock().set_max_bytes(max_bytes);
4592 }
4593
4594 pub(crate) fn clear_result_cache(&mut self) {
4598 self.result_cache.lock().clear();
4599 }
4600
4601 pub fn mutable_run_len(&self) -> usize {
4603 self.mutable_run.len()
4604 }
4605
4606 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4609 self.mutable_run.drain_sorted()
4610 }
4611
4612 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4614 let mut snapshot = self.mutable_run.clone();
4615 snapshot.drain_sorted()
4616 }
4617
4618 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4623 self.ensure_writable()?;
4624 let n = batch.len();
4625 if n == 0 {
4626 return Ok(self.current_epoch());
4627 }
4628 for row in &batch {
4629 self.schema.validate_values(row)?;
4630 }
4631 let epoch = self.commit_new_epoch()?;
4632 let live_before = self.live_count;
4633 self.spill_mutable_run(epoch)?;
4637 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4638 && self.indexes_complete
4639 && self.run_refs.is_empty()
4640 && self.memtable.is_empty()
4641 && self.mutable_run.is_empty();
4642 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4648 use rayon::prelude::*;
4649 self.schema
4650 .columns
4651 .par_iter()
4652 .map(|cdef| {
4653 (
4654 cdef.id,
4655 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4656 )
4657 })
4658 .collect::<Vec<_>>()
4659 };
4660 drop(batch);
4661 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4666 self.validate_columns_not_null(&user_columns, n)?;
4667 let winner_idx = self
4668 .bulk_pk_winner_indices(&user_columns, n)
4669 .filter(|idx| idx.len() != n);
4670 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4671 match winner_idx.as_deref() {
4672 Some(idx) => {
4673 let compacted = user_columns
4674 .iter()
4675 .map(|(id, c)| (*id, c.gather(idx)))
4676 .collect();
4677 (compacted, idx.len())
4678 }
4679 None => (std::mem::take(&mut user_columns), n),
4680 };
4681 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4682 let first = self.allocator.alloc_range(write_n as u64)?.0;
4683 for rid in first..first + write_n as u64 {
4684 self.reservoir.offer(rid);
4685 }
4686 let run_id = self.alloc_run_id()?;
4687 let path = self.run_path(run_id);
4688 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4689 .clean(true)
4690 .with_lz4()
4691 .with_native_endian();
4692 if let Some(kek) = &self.kek {
4693 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4694 }
4695 let header = match self.create_run_file(run_id)? {
4696 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4697 None => writer.write_native(&path, &write_columns, write_n, first)?,
4698 };
4699 self.run_refs.push(RunRef {
4700 run_id: run_id as u128,
4701 level: 0,
4702 epoch_created: epoch.0,
4703 row_count: header.row_count,
4704 });
4705 self.live_count = self.live_count.saturating_add(write_n as u64);
4706 if eager_index_build {
4707 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4708 self.index_columns_bulk(&write_columns, &row_ids);
4709 self.indexes_complete = true;
4710 self.build_learned_ranges()?;
4711 } else {
4712 self.indexes_complete = false;
4713 }
4714 self.mark_flushed(epoch)?;
4715 self.persist_manifest(epoch)?;
4716 if eager_index_build {
4717 self.checkpoint_indexes(epoch);
4718 }
4719 self.clear_result_cache();
4720 Ok(epoch)
4721 }
4722
4723 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4726 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4727 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4728 let segment_no = segment
4731 .file_stem()
4732 .and_then(|s| s.to_str())
4733 .and_then(|s| s.strip_prefix("seg-"))
4734 .and_then(|s| s.parse::<u64>().ok())
4735 .unwrap_or(0);
4736 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4737 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4738 wal.sync()?;
4739 self.wal = WalSink::Private(wal);
4740 Ok(())
4741 }
4742
4743 pub(crate) fn invalidate_pending_cache(&mut self) {
4748 self.result_cache
4749 .lock()
4750 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4751 self.pending_delete_rids.clear();
4752 self.pending_put_cols.clear();
4753 }
4754
4755 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4756 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4757 m.current_epoch = epoch.0;
4758 m.next_row_id = self.allocator.current().0;
4759 m.runs = self.run_refs.clone();
4760 m.live_count = self.live_count;
4761 m.global_idx_epoch = self.global_idx_epoch;
4762 m.flushed_epoch = self.flushed_epoch;
4763 m.retiring = self.retiring.clone();
4764 m.auto_inc_next = match self.auto_inc {
4768 Some(ai) if ai.seeded => ai.next,
4769 _ => 0,
4770 };
4771 m.ttl = self.ttl;
4772 let meta_dek = self.manifest_meta_dek();
4773 match self._root_guard.as_deref() {
4774 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4775 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4776 }
4777 Ok(())
4778 }
4779
4780 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4781 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
4785 let live_count = u64::try_from(rows.len())
4786 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4787 let auto_inc = match self.auto_inc {
4788 Some(mut state) => {
4789 let maximum = self.scan_max_int64(state.column_id)?;
4790 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4791 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4792 })?;
4793 state.next = state.next.max(after_maximum).max(1);
4794 state.seeded = true;
4795 Some(state)
4796 }
4797 None => None,
4798 };
4799 Ok(RecoveryMetadataPlan {
4800 live_count,
4801 auto_inc,
4802 changed: live_count != self.live_count
4803 || auto_inc.is_some_and(|planned| {
4804 self.auto_inc.is_none_or(|current| {
4805 current.next != planned.next || current.seeded != planned.seeded
4806 })
4807 }),
4808 })
4809 }
4810
4811 pub(crate) fn apply_recovered_metadata(
4812 &mut self,
4813 plan: RecoveryMetadataPlan,
4814 epoch: Epoch,
4815 ) -> Result<()> {
4816 if !plan.changed {
4817 return Ok(());
4818 }
4819 self.live_count = plan.live_count;
4820 self.auto_inc = plan.auto_inc;
4821 self.persist_manifest(epoch)
4822 }
4823
4824 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4830 if !self.indexes_complete {
4833 return;
4834 }
4835 if crate::catalog::inject_hook("index.publish.before").is_err() {
4838 return;
4839 }
4840 if self.idx_root.is_none() {
4841 if let Some(root) = self._root_guard.as_ref() {
4842 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4843 return;
4844 };
4845 self.idx_root = Some(Arc::new(idx_root));
4846 }
4847 }
4848 let snap = global_idx::IndexSnapshot {
4849 hot: &self.hot,
4850 bitmap: &self.bitmap,
4851 ann: &self.ann,
4852 fm: &self.fm,
4853 sparse: &self.sparse,
4854 minhash: &self.minhash,
4855 learned_range: &self.learned_range,
4856 };
4857 let idx_dek = self.idx_dek();
4859 let written = match self.idx_root.as_deref() {
4860 Some(root) => global_idx::write_atomic_root(
4861 root,
4862 self.table_id,
4863 epoch.0,
4864 snap,
4865 idx_dek.as_deref(),
4866 ),
4867 None => global_idx::write_atomic(
4868 &self.dir,
4869 self.table_id,
4870 epoch.0,
4871 snap,
4872 idx_dek.as_deref(),
4873 ),
4874 };
4875 if written.is_ok() {
4876 self.global_idx_epoch = epoch.0;
4877 let _ = self.persist_manifest(epoch);
4878 let _ = crate::catalog::inject_hook("index.publish.after");
4880 }
4881 }
4882
4883 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4886 self.global_idx_epoch = 0;
4887 if let Some(root) = self.idx_root.as_deref() {
4888 let _ = root.remove_file(global_idx::IDX_FILENAME);
4889 } else {
4890 global_idx::remove(&self.dir);
4891 }
4892 let _ = self.persist_manifest(self.epoch.visible());
4893 }
4894
4895 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4900 self.indexes_complete = false;
4901 self.global_idx_epoch = 0;
4902 if let Some(root) = self.idx_root.as_deref() {
4903 let _ = root.remove_file(global_idx::IDX_FILENAME);
4904 } else {
4905 global_idx::remove(&self.dir);
4906 }
4907 }
4908
4909 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4910 self.indexes_complete = true;
4911 }
4912
4913 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4919 self.durable_commit_failed = true;
4920 if let WalSink::Shared(shared) = &self.wal {
4921 shared
4922 .poisoned
4923 .store(true, std::sync::atomic::Ordering::Relaxed);
4924 }
4925 }
4926
4927 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4931 self.durable_commit_failed = true;
4932 }
4933
4934 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4943 let mut best: Option<Row> = None;
4944 let mut consider = |row: Row| {
4945 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
4946 return;
4947 }
4948 if best.as_ref().is_none_or(|current| {
4949 Snapshot::version_is_newer(
4950 row.committed_epoch,
4951 row.commit_ts,
4952 current.committed_epoch,
4953 current.commit_ts,
4954 )
4955 }) {
4956 best = Some(row);
4957 }
4958 };
4959 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
4960 consider(row);
4961 }
4962 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
4963 consider(row);
4964 }
4965 for rr in &self.run_refs {
4966 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4967 continue;
4968 };
4969 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
4972 continue;
4973 };
4974 consider(row);
4975 }
4976 let now_nanos = unix_nanos_now();
4977 match best {
4978 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4979 Some(r) => Some(r),
4980 None => None,
4981 }
4982 }
4983
4984 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4988 self.visible_rows_at_time(snapshot, unix_nanos_now())
4989 }
4990
4991 #[doc(hidden)]
4994 pub fn visible_rows_controlled(
4995 &self,
4996 snapshot: Snapshot,
4997 control: &crate::ExecutionControl,
4998 ) -> Result<Vec<Row>> {
4999 let mut out = Vec::new();
5000 self.for_each_visible_row_controlled(snapshot, control, |row| {
5001 out.push(row);
5002 Ok(())
5003 })?;
5004 Ok(out)
5005 }
5006
5007 #[doc(hidden)]
5010 pub fn for_each_visible_row_controlled<F>(
5011 &self,
5012 snapshot: Snapshot,
5013 control: &crate::ExecutionControl,
5014 mut visit: F,
5015 ) -> Result<()>
5016 where
5017 F: FnMut(Row) -> Result<()>,
5018 {
5019 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
5020 control.checkpoint()?;
5021 let memtable = self.memtable.visible_versions_at(snapshot);
5023 if !memtable.is_empty() {
5024 sources.push(ControlledVisibleSource::memory(memtable));
5025 }
5026 control.checkpoint()?;
5027 let mutable = self.mutable_run.visible_versions_at(snapshot);
5030 if !mutable.is_empty() {
5031 sources.push(ControlledVisibleSource::memory(mutable));
5032 }
5033 for run in &self.run_refs {
5034 control.checkpoint()?;
5035 let reader = self.open_reader(run.run_id)?;
5036 sources.push(ControlledVisibleSource::run(
5038 reader.into_visible_version_cursor(snapshot.epoch)?,
5039 ));
5040 }
5041 let now_nanos = unix_nanos_now();
5042 merge_controlled_visible_sources(
5043 &mut sources,
5044 control,
5045 |row| self.row_expired_at(row, now_nanos),
5046 |row| {
5047 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5050 return Ok(());
5051 }
5052 visit(row)
5053 },
5054 )
5055 }
5056
5057 #[doc(hidden)]
5058 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
5059 let mut best: HashMap<u64, Row> = HashMap::new();
5060 let mut fold = |row: Row| {
5061 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5062 return;
5063 }
5064 best.entry(row.row_id.0)
5065 .and_modify(|existing| {
5066 if Snapshot::version_is_newer(
5067 row.committed_epoch,
5068 row.commit_ts,
5069 existing.committed_epoch,
5070 existing.commit_ts,
5071 ) {
5072 *existing = row.clone();
5073 }
5074 })
5075 .or_insert(row);
5076 };
5077 for row in self.memtable.visible_versions_at(snapshot) {
5078 fold(row);
5079 }
5080 for row in self.mutable_run.visible_versions_at(snapshot) {
5081 fold(row);
5082 }
5083 for rr in &self.run_refs {
5084 let mut reader = self.open_reader(rr.run_id)?;
5085 for row in reader.visible_versions(snapshot.epoch)? {
5087 fold(row);
5088 }
5089 }
5090 let mut out: Vec<Row> = best
5091 .into_values()
5092 .filter_map(|r| {
5093 if r.deleted || self.row_expired_at(&r, now_nanos) {
5094 None
5095 } else {
5096 Some(r)
5097 }
5098 })
5099 .collect();
5100 out.sort_by_key(|r| r.row_id);
5101 Ok(out)
5102 }
5103
5104 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
5111 if self.ttl.is_none()
5112 && self.memtable.is_empty()
5113 && self.mutable_run.is_empty()
5114 && self.run_refs.len() == 1
5115 {
5116 let rr = self.run_refs[0].clone();
5117 let mut reader = self.open_reader(rr.run_id)?;
5118 let idxs = reader.visible_indices(snapshot.epoch)?;
5119 let mut cols = Vec::with_capacity(self.schema.columns.len());
5120 for cdef in &self.schema.columns {
5121 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
5122 }
5123 return Ok(cols);
5124 }
5125 let rows = self.visible_rows(snapshot)?;
5127 let mut cols: Vec<(u16, Vec<Value>)> = self
5128 .schema
5129 .columns
5130 .iter()
5131 .map(|c| (c.id, Vec::with_capacity(rows.len())))
5132 .collect();
5133 for r in &rows {
5134 for (cid, vec) in cols.iter_mut() {
5135 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
5136 }
5137 }
5138 Ok(cols)
5139 }
5140
5141 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
5143 let row_id = self.hot.get(key)?;
5144 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
5145 Some(row_id)
5146 } else {
5147 None
5148 }
5149 }
5150
5151 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
5156 self.query_at_with_allowed(q, self.snapshot(), None)
5157 }
5158
5159 pub fn query_controlled(
5162 &mut self,
5163 q: &crate::query::Query,
5164 control: &crate::ExecutionControl,
5165 ) -> Result<Vec<Row>> {
5166 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
5167 }
5168
5169 pub fn query_at_with_allowed(
5172 &mut self,
5173 q: &crate::query::Query,
5174 snapshot: Snapshot,
5175 allowed: Option<&std::collections::HashSet<RowId>>,
5176 ) -> Result<Vec<Row>> {
5177 self.query_at_with_allowed_after(q, snapshot, allowed, None)
5178 }
5179
5180 #[doc(hidden)]
5181 pub fn query_at_with_allowed_controlled(
5182 &mut self,
5183 q: &crate::query::Query,
5184 snapshot: Snapshot,
5185 allowed: Option<&std::collections::HashSet<RowId>>,
5186 control: &crate::ExecutionControl,
5187 ) -> Result<Vec<Row>> {
5188 self.require_select()?;
5189 self.ensure_indexes_complete_controlled(control, || true)?;
5190 self.validate_native_query(q)?;
5191 self.query_conditions_at(
5192 &q.conditions,
5193 snapshot,
5194 allowed,
5195 q.limit,
5196 q.offset,
5197 None,
5198 unix_nanos_now(),
5199 Some(control),
5200 )
5201 }
5202
5203 #[doc(hidden)]
5204 pub fn query_at_with_allowed_after(
5205 &mut self,
5206 q: &crate::query::Query,
5207 snapshot: Snapshot,
5208 allowed: Option<&std::collections::HashSet<RowId>>,
5209 after_row_id: Option<RowId>,
5210 ) -> Result<Vec<Row>> {
5211 self.query_at_with_allowed_after_at_time(
5212 q,
5213 snapshot,
5214 allowed,
5215 after_row_id,
5216 unix_nanos_now(),
5217 )
5218 }
5219
5220 #[doc(hidden)]
5221 pub fn query_at_with_allowed_after_at_time(
5222 &mut self,
5223 q: &crate::query::Query,
5224 snapshot: Snapshot,
5225 allowed: Option<&std::collections::HashSet<RowId>>,
5226 after_row_id: Option<RowId>,
5227 query_time_nanos: i64,
5228 ) -> Result<Vec<Row>> {
5229 self.require_select()?;
5230 self.ensure_indexes_complete()?;
5231 self.validate_native_query(q)?;
5232 self.query_conditions_at(
5233 &q.conditions,
5234 snapshot,
5235 allowed,
5236 q.limit,
5237 q.offset,
5238 after_row_id,
5239 query_time_nanos,
5240 None,
5241 )
5242 }
5243
5244 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
5245 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5246 return Err(MongrelError::InvalidArgument(format!(
5247 "query exceeds {} conditions",
5248 crate::query::MAX_HARD_CONDITIONS
5249 )));
5250 }
5251 if let Some(limit) = q.limit {
5252 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
5253 return Err(MongrelError::InvalidArgument(format!(
5254 "query limit must be between 1 and {}",
5255 crate::query::MAX_FINAL_LIMIT
5256 )));
5257 }
5258 }
5259 if q.offset > crate::query::MAX_QUERY_OFFSET {
5260 return Err(MongrelError::InvalidArgument(format!(
5261 "query offset exceeds {}",
5262 crate::query::MAX_QUERY_OFFSET
5263 )));
5264 }
5265 Ok(())
5266 }
5267
5268 #[doc(hidden)]
5271 pub fn query_all_at(
5272 &mut self,
5273 conditions: &[crate::query::Condition],
5274 snapshot: Snapshot,
5275 ) -> Result<Vec<Row>> {
5276 self.require_select()?;
5277 self.ensure_indexes_complete()?;
5278 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5279 return Err(MongrelError::InvalidArgument(format!(
5280 "query exceeds {} conditions",
5281 crate::query::MAX_HARD_CONDITIONS
5282 )));
5283 }
5284 self.query_conditions_at(
5285 conditions,
5286 snapshot,
5287 None,
5288 None,
5289 0,
5290 None,
5291 unix_nanos_now(),
5292 None,
5293 )
5294 }
5295
5296 #[allow(clippy::too_many_arguments)]
5297 fn query_conditions_at(
5298 &self,
5299 conditions: &[crate::query::Condition],
5300 snapshot: Snapshot,
5301 allowed: Option<&std::collections::HashSet<RowId>>,
5302 limit: Option<usize>,
5303 offset: usize,
5304 after_row_id: Option<RowId>,
5305 query_time_nanos: i64,
5306 control: Option<&crate::ExecutionControl>,
5307 ) -> Result<Vec<Row>> {
5308 control
5309 .map(crate::ExecutionControl::checkpoint)
5310 .transpose()?;
5311 crate::trace::QueryTrace::record(|t| {
5312 t.run_count = self.run_refs.len();
5313 t.memtable_rows = self.memtable.len();
5314 t.mutable_run_rows = self.mutable_run.len();
5315 });
5316 if conditions.is_empty() {
5320 crate::trace::QueryTrace::record(|t| {
5321 t.scan_mode = crate::trace::ScanMode::Materialized;
5322 t.row_materialized = true;
5323 });
5324 let mut rows = match control {
5325 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5326 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5327 };
5328 if let Some(allowed) = allowed {
5329 let mut filtered = Vec::with_capacity(rows.len());
5330 for (index, row) in rows.into_iter().enumerate() {
5331 if index & 255 == 0 {
5332 control
5333 .map(crate::ExecutionControl::checkpoint)
5334 .transpose()?;
5335 }
5336 if allowed.contains(&row.row_id) {
5337 filtered.push(row);
5338 }
5339 }
5340 rows = filtered;
5341 }
5342 if let Some(after_row_id) = after_row_id {
5343 rows.retain(|row| row.row_id > after_row_id);
5344 }
5345 rows.drain(..offset.min(rows.len()));
5346 if let Some(limit) = limit {
5347 rows.truncate(limit);
5348 }
5349 return Ok(rows);
5350 }
5351 crate::trace::QueryTrace::record(|t| {
5352 t.conditions_pushed = conditions.len();
5353 t.scan_mode = crate::trace::ScanMode::Materialized;
5354 t.row_materialized = true;
5355 });
5356 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5363 ordered.sort_by_key(|c| condition_cost_rank(c));
5364 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5365 for c in &ordered {
5366 control
5367 .map(crate::ExecutionControl::checkpoint)
5368 .transpose()?;
5369 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5370 let empty = s.is_empty();
5371 sets.push(s);
5372 if empty {
5373 break;
5374 }
5375 }
5376 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5377 if let Some(allowed) = allowed {
5378 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5379 }
5380 if let Some(after_row_id) = after_row_id {
5381 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5382 rids.drain(..first);
5383 }
5384 rids.drain(..offset.min(rids.len()));
5385 if let Some(limit) = limit {
5386 rids.truncate(limit);
5387 }
5388 control
5389 .map(crate::ExecutionControl::checkpoint)
5390 .transpose()?;
5391 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5392 }
5393
5394 pub fn retrieve(
5396 &mut self,
5397 retriever: &crate::query::Retriever,
5398 ) -> Result<Vec<crate::query::RetrieverHit>> {
5399 self.retrieve_with_allowed(retriever, None)
5400 }
5401
5402 pub fn retrieve_at(
5403 &mut self,
5404 retriever: &crate::query::Retriever,
5405 snapshot: Snapshot,
5406 allowed: Option<&std::collections::HashSet<RowId>>,
5407 ) -> Result<Vec<crate::query::RetrieverHit>> {
5408 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5409 }
5410
5411 pub fn retrieve_with_allowed(
5414 &mut self,
5415 retriever: &crate::query::Retriever,
5416 allowed: Option<&std::collections::HashSet<RowId>>,
5417 ) -> Result<Vec<crate::query::RetrieverHit>> {
5418 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5419 }
5420
5421 pub fn retrieve_at_with_allowed(
5422 &mut self,
5423 retriever: &crate::query::Retriever,
5424 snapshot: Snapshot,
5425 allowed: Option<&std::collections::HashSet<RowId>>,
5426 ) -> Result<Vec<crate::query::RetrieverHit>> {
5427 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5428 }
5429
5430 pub fn retrieve_at_with_allowed_and_context(
5431 &mut self,
5432 retriever: &crate::query::Retriever,
5433 snapshot: Snapshot,
5434 allowed: Option<&std::collections::HashSet<RowId>>,
5435 context: Option<&crate::query::AiExecutionContext>,
5436 ) -> Result<Vec<crate::query::RetrieverHit>> {
5437 self.require_select()?;
5438 self.ensure_indexes_complete()?;
5439 self.validate_retriever(retriever)?;
5440 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5441 }
5442
5443 pub fn retrieve_at_with_candidate_authorization_and_context(
5444 &mut self,
5445 retriever: &crate::query::Retriever,
5446 snapshot: Snapshot,
5447 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5448 context: Option<&crate::query::AiExecutionContext>,
5449 ) -> Result<Vec<crate::query::RetrieverHit>> {
5450 self.require_select()?;
5451 self.ensure_indexes_complete()?;
5452 self.retrieve_at_with_candidate_authorization_on_generation(
5453 retriever,
5454 snapshot,
5455 authorization,
5456 context,
5457 )
5458 }
5459
5460 #[doc(hidden)]
5461 pub fn retrieve_at_with_candidate_authorization_on_generation(
5462 &self,
5463 retriever: &crate::query::Retriever,
5464 snapshot: Snapshot,
5465 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5466 context: Option<&crate::query::AiExecutionContext>,
5467 ) -> Result<Vec<crate::query::RetrieverHit>> {
5468 self.require_select()?;
5469 self.validate_retriever(retriever)?;
5470 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5471 }
5472
5473 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5474 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5475 let (column_id, k) = match retriever {
5476 Retriever::Ann {
5477 column_id,
5478 query,
5479 k,
5480 } => {
5481 let index = self.ann.get(column_id).ok_or_else(|| {
5482 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5483 })?;
5484 if query.len() != index.dim() {
5485 return Err(MongrelError::InvalidArgument(format!(
5486 "ANN query dimension must be {}, got {}",
5487 index.dim(),
5488 query.len()
5489 )));
5490 }
5491 if query.iter().any(|value| !value.is_finite()) {
5492 return Err(MongrelError::InvalidArgument(
5493 "ANN query values must be finite".into(),
5494 ));
5495 }
5496 (*column_id, *k)
5497 }
5498 Retriever::Sparse {
5499 column_id,
5500 query,
5501 k,
5502 } => {
5503 if !self.sparse.contains_key(column_id) {
5504 return Err(MongrelError::InvalidArgument(format!(
5505 "column {column_id} has no Sparse index"
5506 )));
5507 }
5508 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5509 return Err(MongrelError::InvalidArgument(
5510 "Sparse query must be non-empty with finite weights".into(),
5511 ));
5512 }
5513 if query.len() > MAX_SPARSE_TERMS {
5514 return Err(MongrelError::InvalidArgument(format!(
5515 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5516 )));
5517 }
5518 (*column_id, *k)
5519 }
5520 Retriever::MinHash {
5521 column_id,
5522 members,
5523 k,
5524 } => {
5525 if !self.minhash.contains_key(column_id) {
5526 return Err(MongrelError::InvalidArgument(format!(
5527 "column {column_id} has no MinHash index"
5528 )));
5529 }
5530 if members.is_empty() {
5531 return Err(MongrelError::InvalidArgument(
5532 "MinHash members must not be empty".into(),
5533 ));
5534 }
5535 if members.len() > MAX_SET_MEMBERS {
5536 return Err(MongrelError::InvalidArgument(format!(
5537 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5538 )));
5539 }
5540 let mut total_bytes = 0usize;
5541 for member in members {
5542 let bytes = member.encoded_len();
5543 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5544 return Err(MongrelError::InvalidArgument(format!(
5545 "MinHash member exceeds {} bytes",
5546 crate::query::MAX_SET_MEMBER_BYTES
5547 )));
5548 }
5549 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5550 MongrelError::InvalidArgument("MinHash input size overflow".into())
5551 })?;
5552 }
5553 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5554 return Err(MongrelError::InvalidArgument(format!(
5555 "MinHash input exceeds {} bytes",
5556 crate::query::MAX_SET_INPUT_BYTES
5557 )));
5558 }
5559 (*column_id, *k)
5560 }
5561 };
5562 if k == 0 {
5563 return Err(MongrelError::InvalidArgument(
5564 "retriever k must be > 0".into(),
5565 ));
5566 }
5567 if k > MAX_RETRIEVER_K {
5568 return Err(MongrelError::InvalidArgument(format!(
5569 "retriever k exceeds {MAX_RETRIEVER_K}"
5570 )));
5571 }
5572 debug_assert!(self
5573 .schema
5574 .columns
5575 .iter()
5576 .any(|column| column.id == column_id));
5577 Ok(())
5578 }
5579
5580 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5581 use crate::query::Condition;
5582 match condition {
5583 Condition::Ann {
5584 column_id,
5585 query,
5586 k,
5587 } => self.validate_retriever(&crate::query::Retriever::Ann {
5588 column_id: *column_id,
5589 query: query.clone(),
5590 k: *k,
5591 }),
5592 Condition::SparseMatch {
5593 column_id,
5594 query,
5595 k,
5596 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5597 column_id: *column_id,
5598 query: query.clone(),
5599 k: *k,
5600 }),
5601 Condition::MinHashSimilar {
5602 column_id,
5603 query,
5604 k,
5605 } => {
5606 if !self.minhash.contains_key(column_id) {
5607 return Err(MongrelError::InvalidArgument(format!(
5608 "column {column_id} has no MinHash index"
5609 )));
5610 }
5611 if query.is_empty() || *k == 0 {
5612 return Err(MongrelError::InvalidArgument(
5613 "MinHash query must be non-empty and k must be > 0".into(),
5614 ));
5615 }
5616 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5617 {
5618 return Err(MongrelError::InvalidArgument(format!(
5619 "MinHash query must have <= {} members and k <= {}",
5620 crate::query::MAX_SET_MEMBERS,
5621 crate::query::MAX_RETRIEVER_K
5622 )));
5623 }
5624 Ok(())
5625 }
5626 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5627 Err(MongrelError::InvalidArgument(format!(
5628 "bitmap IN exceeds {} values",
5629 crate::query::MAX_SET_MEMBERS
5630 )))
5631 }
5632 Condition::FmContainsAll { patterns, .. }
5633 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5634 {
5635 Err(MongrelError::InvalidArgument(format!(
5636 "FM query exceeds {} patterns",
5637 crate::query::MAX_HARD_CONDITIONS
5638 )))
5639 }
5640 _ => Ok(()),
5641 }
5642 }
5643
5644 fn retrieve_filtered(
5645 &self,
5646 retriever: &crate::query::Retriever,
5647 snapshot: Snapshot,
5648 hard_filter: Option<&RowIdSet>,
5649 allowed: Option<&std::collections::HashSet<RowId>>,
5650 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5651 context: Option<&crate::query::AiExecutionContext>,
5652 ) -> Result<Vec<crate::query::RetrieverHit>> {
5653 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5654 let started = std::time::Instant::now();
5655 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5656 Retriever::Ann {
5657 column_id,
5658 query,
5659 k,
5660 } => {
5661 let Some(index) = self.ann.get(column_id) else {
5662 return Ok(Vec::new());
5663 };
5664 let cap = ann_candidate_cap(index.len(), context);
5665 if cap == 0 {
5666 return Ok(Vec::new());
5667 }
5668 let mut breadth = (*k).max(1).min(cap);
5669 let mut eligibility = std::collections::HashMap::new();
5670 let mut filtered = loop {
5671 let mut seen = std::collections::HashSet::new();
5672 if let Some(context) = context {
5673 context.checkpoint()?;
5674 }
5675 let raw = index.search_with_context(query, breadth, context)?;
5676 let unchecked: Vec<_> = raw
5677 .iter()
5678 .map(|(row_id, _)| *row_id)
5679 .filter(|row_id| !eligibility.contains_key(row_id))
5680 .filter(|row_id| {
5681 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5682 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5683 })
5684 .collect();
5685 let eligible = self.eligible_and_authorized_candidate_ids(
5686 &unchecked,
5687 *column_id,
5688 snapshot,
5689 candidate_authorization,
5690 context,
5691 )?;
5692 for row_id in unchecked {
5693 eligibility.insert(row_id, eligible.contains(&row_id));
5694 }
5695 let filtered: Vec<_> = raw
5696 .into_iter()
5697 .filter(|(row_id, _)| {
5698 seen.insert(*row_id)
5699 && eligibility.get(row_id).copied().unwrap_or(false)
5700 })
5701 .map(|(row_id, score)| {
5702 let score = match score {
5703 crate::index::AnnDistance::Hamming(d) => {
5704 RetrieverScore::AnnHammingDistance(d)
5705 }
5706 crate::index::AnnDistance::Cosine(d) => {
5707 RetrieverScore::AnnCosineDistance(d)
5708 }
5709 };
5710 (row_id, score)
5711 })
5712 .collect();
5713 if filtered.len() >= *k || breadth >= cap {
5714 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5715 crate::trace::QueryTrace::record(|trace| {
5716 trace.ann_candidate_cap_hit = true;
5717 });
5718 }
5719 break filtered;
5720 }
5721 breadth = breadth.saturating_mul(2).min(cap);
5722 };
5723 filtered.truncate(*k);
5724 filtered
5725 }
5726 Retriever::Sparse {
5727 column_id,
5728 query,
5729 k,
5730 } => self
5731 .sparse
5732 .get(column_id)
5733 .map(|index| -> Result<Vec<_>> {
5734 let mut breadth = (*k).max(1);
5735 let mut eligibility = std::collections::HashMap::new();
5736 loop {
5737 if let Some(context) = context {
5738 context.checkpoint()?;
5739 }
5740 let raw = index.search_with_context(query, breadth, context)?;
5741 let unchecked: Vec<_> = raw
5742 .iter()
5743 .map(|(row_id, _)| *row_id)
5744 .filter(|row_id| !eligibility.contains_key(row_id))
5745 .filter(|row_id| {
5746 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5747 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5748 })
5749 .collect();
5750 let eligible = self.eligible_and_authorized_candidate_ids(
5751 &unchecked,
5752 *column_id,
5753 snapshot,
5754 candidate_authorization,
5755 context,
5756 )?;
5757 for row_id in unchecked {
5758 eligibility.insert(row_id, eligible.contains(&row_id));
5759 }
5760 let filtered: Vec<_> = raw
5761 .iter()
5762 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5763 .take(*k)
5764 .map(|(row_id, score)| {
5765 (*row_id, RetrieverScore::SparseDotProduct(*score))
5766 })
5767 .collect();
5768 if filtered.len() >= *k || raw.len() < breadth {
5769 break Ok(filtered);
5770 }
5771 let next = breadth.saturating_mul(2);
5772 if next == breadth {
5773 break Ok(filtered);
5774 }
5775 breadth = next;
5776 }
5777 })
5778 .transpose()?
5779 .unwrap_or_default(),
5780 Retriever::MinHash {
5781 column_id,
5782 members,
5783 k,
5784 } => self
5785 .minhash
5786 .get(column_id)
5787 .map(|index| -> Result<Vec<_>> {
5788 let mut hashes = Vec::with_capacity(members.len());
5789 for member in members {
5790 if let Some(context) = context {
5791 context.consume(crate::query::work_units(
5792 member.encoded_len(),
5793 crate::query::PARSE_WORK_QUANTUM,
5794 ))?;
5795 }
5796 hashes.push(member.hash_v1());
5797 }
5798 let mut breadth = (*k).max(1);
5799 let mut eligibility = std::collections::HashMap::new();
5800 loop {
5801 if let Some(context) = context {
5802 context.checkpoint()?;
5803 }
5804 let raw = index.search_with_context(&hashes, breadth, context)?;
5805 let unchecked: Vec<_> = raw
5806 .iter()
5807 .map(|(row_id, _)| *row_id)
5808 .filter(|row_id| !eligibility.contains_key(row_id))
5809 .filter(|row_id| {
5810 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5811 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5812 })
5813 .collect();
5814 let eligible = self.eligible_and_authorized_candidate_ids(
5815 &unchecked,
5816 *column_id,
5817 snapshot,
5818 candidate_authorization,
5819 context,
5820 )?;
5821 for row_id in unchecked {
5822 eligibility.insert(row_id, eligible.contains(&row_id));
5823 }
5824 let filtered: Vec<_> = raw
5825 .iter()
5826 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5827 .take(*k)
5828 .map(|(row_id, score)| {
5829 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5830 })
5831 .collect();
5832 if filtered.len() >= *k || raw.len() < breadth {
5833 break Ok(filtered);
5834 }
5835 let next = breadth.saturating_mul(2);
5836 if next == breadth {
5837 break Ok(filtered);
5838 }
5839 breadth = next;
5840 }
5841 })
5842 .transpose()?
5843 .unwrap_or_default(),
5844 };
5845 let elapsed = started.elapsed().as_nanos() as u64;
5846 crate::trace::QueryTrace::record(|trace| {
5847 match retriever {
5848 Retriever::Ann { .. } => {
5849 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
5850 if let Retriever::Ann { column_id, .. } = retriever {
5851 if let Some(index) = self.ann.get(column_id) {
5852 trace.ann_algorithm = Some(index.algorithm());
5853 trace.ann_quantization = Some(index.quantization());
5854 trace.ann_backend = Some(index.backend_name());
5855 }
5856 }
5857 }
5858 Retriever::Sparse { .. } => {
5859 trace.sparse_candidate_nanos =
5860 trace.sparse_candidate_nanos.saturating_add(elapsed)
5861 }
5862 Retriever::MinHash { .. } => {
5863 trace.minhash_candidate_nanos =
5864 trace.minhash_candidate_nanos.saturating_add(elapsed)
5865 }
5866 }
5867 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5868 });
5869 Ok(scored
5870 .into_iter()
5871 .enumerate()
5872 .map(|(rank, (row_id, score))| RetrieverHit {
5873 row_id,
5874 rank: rank + 1,
5875 score,
5876 })
5877 .collect())
5878 }
5879
5880 fn eligible_candidate_ids(
5881 &self,
5882 candidates: &[RowId],
5883 _column_id: u16,
5884 snapshot: Snapshot,
5885 context: Option<&crate::query::AiExecutionContext>,
5886 ) -> Result<std::collections::HashSet<RowId>> {
5887 if !self.had_deletes
5888 && self.ttl.is_none()
5889 && self.pending_put_cols.is_empty()
5890 && snapshot.epoch == self.snapshot().epoch
5891 {
5892 return Ok(candidates.iter().copied().collect());
5893 }
5894 let mut readers: Vec<_> = self
5895 .run_refs
5896 .iter()
5897 .map(|run| self.open_reader(run.run_id))
5898 .collect::<Result<_>>()?;
5899 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5900 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5901 for &row_id in candidates {
5902 if let Some(context) = context {
5903 context.consume(1)?;
5904 }
5905 let mem = self.memtable.get_version_at(row_id, snapshot);
5906 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
5907 let overlay = match (mem, mutable) {
5908 (Some(left), Some(right)) => Some(
5909 if Snapshot::version_is_newer(
5910 left.1.committed_epoch,
5911 left.1.commit_ts,
5912 right.1.committed_epoch,
5913 right.1.commit_ts,
5914 ) {
5915 left
5916 } else {
5917 right
5918 },
5919 ),
5920 (Some(value), None) | (None, Some(value)) => Some(value),
5921 (None, None) => None,
5922 };
5923 if let Some((_, row)) = overlay {
5924 if !row.deleted && !self.row_expired_at(&row, now) {
5925 eligible.insert(row_id);
5926 }
5927 continue;
5928 }
5929 let mut best: Option<(Epoch, bool, usize)> = None;
5930 for (index, reader) in readers.iter_mut().enumerate() {
5931 if let Some((epoch, deleted)) =
5932 reader.get_version_visibility(row_id, snapshot.epoch)?
5933 {
5934 if best
5935 .as_ref()
5936 .map(|(best_epoch, ..)| epoch > *best_epoch)
5937 .unwrap_or(true)
5938 {
5939 best = Some((epoch, deleted, index));
5940 }
5941 }
5942 }
5943 let Some((_, false, reader_index)) = best else {
5944 continue;
5945 };
5946 if let Some(ttl) = self.ttl {
5947 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5948 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5949 {
5950 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5951 continue;
5952 }
5953 }
5954 }
5955 eligible.insert(row_id);
5956 }
5957 Ok(eligible)
5958 }
5959
5960 fn eligible_and_authorized_candidate_ids(
5961 &self,
5962 candidates: &[RowId],
5963 column_id: u16,
5964 snapshot: Snapshot,
5965 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5966 context: Option<&crate::query::AiExecutionContext>,
5967 ) -> Result<std::collections::HashSet<RowId>> {
5968 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5969 let Some(authorization) = authorization else {
5970 return Ok(eligible);
5971 };
5972 let candidates: Vec<_> = eligible.into_iter().collect();
5973 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5974 }
5975
5976 fn policy_allowed_candidate_ids(
5977 &self,
5978 candidates: &[RowId],
5979 snapshot: Snapshot,
5980 authorization: &crate::security::CandidateAuthorization<'_>,
5981 context: Option<&crate::query::AiExecutionContext>,
5982 ) -> Result<std::collections::HashSet<RowId>> {
5983 let started = std::time::Instant::now();
5984 if candidates.is_empty()
5985 || authorization.principal.is_admin
5986 || !authorization.security.rls_enabled(authorization.table)
5987 {
5988 return Ok(candidates.iter().copied().collect());
5989 }
5990 if let Some(context) = context {
5991 context.checkpoint()?;
5992 }
5993 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5994 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5995 .iter()
5996 .map(|row_id| {
5997 (
5998 *row_id,
5999 Row {
6000 row_id: *row_id,
6001 committed_epoch: snapshot.epoch,
6002 columns: std::collections::HashMap::new(),
6003 deleted: false,
6004 commit_ts: None,
6005 },
6006 )
6007 })
6008 .collect();
6009 let columns = authorization
6010 .security
6011 .select_policy_columns(authorization.table, authorization.principal);
6012 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6013 let mut decoded = 0usize;
6014 for column_id in &columns {
6015 if let Some(context) = context {
6016 context.checkpoint()?;
6017 }
6018 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6019 &row_ids, *column_id, snapshot, query_now, context,
6020 )? {
6021 if let Some(row) = rows.get_mut(&row_id) {
6022 row.columns.insert(*column_id, value);
6023 decoded = decoded.saturating_add(1);
6024 }
6025 }
6026 }
6027 if let Some(context) = context {
6028 context.consume(candidates.len().saturating_add(decoded))?;
6029 }
6030 let allowed = rows
6031 .into_values()
6032 .filter_map(|row| {
6033 authorization
6034 .security
6035 .row_allowed(
6036 authorization.table,
6037 crate::security::PolicyCommand::Select,
6038 &row,
6039 authorization.principal,
6040 false,
6041 )
6042 .then_some(row.row_id)
6043 })
6044 .collect();
6045 crate::trace::QueryTrace::record(|trace| {
6046 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
6047 trace.rls_policy_columns_decoded =
6048 trace.rls_policy_columns_decoded.saturating_add(decoded);
6049 trace.authorization_nanos = trace
6050 .authorization_nanos
6051 .saturating_add(started.elapsed().as_nanos() as u64);
6052 });
6053 Ok(allowed)
6054 }
6055
6056 pub fn search(
6058 &mut self,
6059 request: &crate::query::SearchRequest,
6060 ) -> Result<Vec<crate::query::SearchHit>> {
6061 self.search_with_allowed(request, None)
6062 }
6063
6064 pub fn search_at(
6065 &mut self,
6066 request: &crate::query::SearchRequest,
6067 snapshot: Snapshot,
6068 authorized: Option<&std::collections::HashSet<RowId>>,
6069 ) -> Result<Vec<crate::query::SearchHit>> {
6070 self.search_at_with_allowed(request, snapshot, authorized)
6071 }
6072
6073 pub fn search_with_allowed(
6074 &mut self,
6075 request: &crate::query::SearchRequest,
6076 authorized: Option<&std::collections::HashSet<RowId>>,
6077 ) -> Result<Vec<crate::query::SearchHit>> {
6078 self.search_at_with_allowed(request, self.snapshot(), authorized)
6079 }
6080
6081 pub fn search_at_with_allowed(
6082 &mut self,
6083 request: &crate::query::SearchRequest,
6084 snapshot: Snapshot,
6085 authorized: Option<&std::collections::HashSet<RowId>>,
6086 ) -> Result<Vec<crate::query::SearchHit>> {
6087 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
6088 }
6089
6090 pub fn search_at_with_allowed_and_context(
6091 &mut self,
6092 request: &crate::query::SearchRequest,
6093 snapshot: Snapshot,
6094 authorized: Option<&std::collections::HashSet<RowId>>,
6095 context: Option<&crate::query::AiExecutionContext>,
6096 ) -> Result<Vec<crate::query::SearchHit>> {
6097 self.ensure_indexes_complete()?;
6098 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
6099 }
6100
6101 pub fn search_at_with_candidate_authorization_and_context(
6102 &mut self,
6103 request: &crate::query::SearchRequest,
6104 snapshot: Snapshot,
6105 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6106 context: Option<&crate::query::AiExecutionContext>,
6107 ) -> Result<Vec<crate::query::SearchHit>> {
6108 self.ensure_indexes_complete()?;
6109 self.search_at_with_filters_and_context(
6110 request,
6111 snapshot,
6112 None,
6113 authorization,
6114 context,
6115 None,
6116 )
6117 }
6118
6119 #[doc(hidden)]
6120 pub fn search_at_with_candidate_authorization_on_generation(
6121 &self,
6122 request: &crate::query::SearchRequest,
6123 snapshot: Snapshot,
6124 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6125 context: Option<&crate::query::AiExecutionContext>,
6126 ) -> Result<Vec<crate::query::SearchHit>> {
6127 self.search_at_with_filters_and_context(
6128 request,
6129 snapshot,
6130 None,
6131 authorization,
6132 context,
6133 None,
6134 )
6135 }
6136
6137 #[doc(hidden)]
6138 pub fn search_at_with_candidate_authorization_on_generation_after(
6139 &self,
6140 request: &crate::query::SearchRequest,
6141 snapshot: Snapshot,
6142 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6143 context: Option<&crate::query::AiExecutionContext>,
6144 after: Option<crate::query::SearchAfter>,
6145 ) -> Result<Vec<crate::query::SearchHit>> {
6146 self.search_at_with_filters_and_context(
6147 request,
6148 snapshot,
6149 None,
6150 authorization,
6151 context,
6152 after,
6153 )
6154 }
6155
6156 fn search_at_with_filters_and_context(
6157 &self,
6158 request: &crate::query::SearchRequest,
6159 snapshot: Snapshot,
6160 authorized: Option<&std::collections::HashSet<RowId>>,
6161 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6162 context: Option<&crate::query::AiExecutionContext>,
6163 after: Option<crate::query::SearchAfter>,
6164 ) -> Result<Vec<crate::query::SearchHit>> {
6165 use crate::query::{
6166 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
6167 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
6168 };
6169 let total_started = std::time::Instant::now();
6170 let rank_offset = after.map_or(0, |after| after.returned_count);
6171 self.require_select()?;
6172 if request.limit == 0 {
6173 return Err(MongrelError::InvalidArgument(
6174 "search limit must be > 0".into(),
6175 ));
6176 }
6177 if request.limit > MAX_FINAL_LIMIT {
6178 return Err(MongrelError::InvalidArgument(format!(
6179 "search limit exceeds {MAX_FINAL_LIMIT}"
6180 )));
6181 }
6182 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
6183 return Err(MongrelError::InvalidArgument(
6184 "search-after score must be finite".into(),
6185 ));
6186 }
6187 if request.retrievers.is_empty() {
6188 return Err(MongrelError::InvalidArgument(
6189 "search requires at least one retriever".into(),
6190 ));
6191 }
6192 if request.retrievers.len() > MAX_RETRIEVERS {
6193 return Err(MongrelError::InvalidArgument(format!(
6194 "search exceeds {MAX_RETRIEVERS} retrievers"
6195 )));
6196 }
6197 if request.must.len() > MAX_HARD_CONDITIONS {
6198 return Err(MongrelError::InvalidArgument(format!(
6199 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
6200 )));
6201 }
6202 for condition in &request.must {
6203 self.validate_condition(condition)?;
6204 }
6205 if request.must.iter().any(|condition| {
6206 matches!(
6207 condition,
6208 Condition::Ann { .. }
6209 | Condition::SparseMatch { .. }
6210 | Condition::MinHashSimilar { .. }
6211 )
6212 }) {
6213 return Err(MongrelError::InvalidArgument(
6214 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
6215 .into(),
6216 ));
6217 }
6218 let mut names = std::collections::HashSet::new();
6219 for named in &request.retrievers {
6220 if named.name.is_empty()
6221 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
6222 || !names.insert(named.name.as_str())
6223 {
6224 return Err(MongrelError::InvalidArgument(format!(
6225 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
6226 crate::query::MAX_RETRIEVER_NAME_BYTES
6227 )));
6228 }
6229 if !named.weight.is_finite()
6230 || named.weight < 0.0
6231 || named.weight > MAX_RETRIEVER_WEIGHT
6232 {
6233 return Err(MongrelError::InvalidArgument(format!(
6234 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6235 )));
6236 }
6237 self.validate_retriever(&named.retriever)?;
6238 }
6239 let projection = request
6240 .projection
6241 .clone()
6242 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
6243 if projection.len() > MAX_PROJECTION_COLUMNS {
6244 return Err(MongrelError::InvalidArgument(format!(
6245 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
6246 )));
6247 }
6248 for column_id in &projection {
6249 if !self
6250 .schema
6251 .columns
6252 .iter()
6253 .any(|column| column.id == *column_id)
6254 {
6255 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
6256 }
6257 }
6258 if let Some(crate::query::Rerank::ExactVector {
6259 embedding_column,
6260 query,
6261 candidate_limit,
6262 weight,
6263 ..
6264 }) = &request.rerank
6265 {
6266 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
6267 {
6268 return Err(MongrelError::InvalidArgument(format!(
6269 "rerank candidate_limit must be between search limit and {}",
6270 crate::query::MAX_RETRIEVER_K
6271 )));
6272 }
6273 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
6274 return Err(MongrelError::InvalidArgument(format!(
6275 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6276 )));
6277 }
6278 let column = self
6279 .schema
6280 .columns
6281 .iter()
6282 .find(|column| column.id == *embedding_column)
6283 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
6284 let crate::schema::TypeId::Embedding { dim } = column.ty else {
6285 return Err(MongrelError::InvalidArgument(format!(
6286 "rerank column {embedding_column} is not an embedding"
6287 )));
6288 };
6289 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
6290 return Err(MongrelError::InvalidArgument(format!(
6291 "rerank query must contain {dim} finite values"
6292 )));
6293 }
6294 }
6295
6296 let hard_filter_started = std::time::Instant::now();
6297 let hard_filter = if request.must.is_empty() {
6298 None
6299 } else {
6300 let mut sets = Vec::with_capacity(request.must.len());
6301 for condition in &request.must {
6302 if let Some(context) = context {
6303 context.checkpoint()?;
6304 }
6305 sets.push(self.resolve_condition(condition, snapshot)?);
6306 }
6307 Some(RowIdSet::intersect_many(sets))
6308 };
6309 crate::trace::QueryTrace::record(|trace| {
6310 trace.hard_filter_nanos = trace
6311 .hard_filter_nanos
6312 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6313 });
6314 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6315 return Ok(Vec::new());
6316 }
6317
6318 let constant = match request.fusion {
6319 Fusion::ReciprocalRank { constant } => constant,
6320 };
6321 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6322 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6323 let mut fusion_nanos = 0u64;
6324 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6325 std::collections::HashMap::new();
6326 for named in retrievers {
6327 if named.weight == 0.0 {
6328 continue;
6329 }
6330 if let Some(context) = context {
6331 context.checkpoint()?;
6332 }
6333 let hits = self.retrieve_filtered(
6334 &named.retriever,
6335 snapshot,
6336 hard_filter.as_ref(),
6337 authorized,
6338 candidate_authorization,
6339 context,
6340 )?;
6341 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6342 let fusion_started = std::time::Instant::now();
6343 for hit in hits {
6344 if let Some(context) = context {
6345 context.consume(1)?;
6346 }
6347 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6348 if !contribution.is_finite() {
6349 return Err(MongrelError::InvalidArgument(
6350 "retriever contribution must be finite".into(),
6351 ));
6352 }
6353 let max_fused_candidates = context.map_or(
6354 crate::query::MAX_FUSED_CANDIDATES,
6355 crate::query::AiExecutionContext::max_fused_candidates,
6356 );
6357 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6358 return Err(MongrelError::WorkBudgetExceeded);
6359 }
6360 let entry = fused.entry(hit.row_id).or_default();
6361 entry.0 += contribution;
6362 if !entry.0.is_finite() {
6363 return Err(MongrelError::InvalidArgument(
6364 "fused score must be finite".into(),
6365 ));
6366 }
6367 entry.1.push(ComponentScore {
6368 retriever_name: retriever_name.clone(),
6369 rank: hit.rank,
6370 raw_score: hit.score,
6371 contribution,
6372 });
6373 }
6374 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6375 }
6376 let union_size = fused.len();
6377 let mut ranked: Vec<_> = fused
6378 .into_iter()
6379 .map(|(row_id, (fused_score, components))| {
6380 (row_id, fused_score, components, None, fused_score)
6381 })
6382 .collect();
6383 let order = |(a_row, _, _, _, a_score): &(
6384 RowId,
6385 f64,
6386 Vec<ComponentScore>,
6387 Option<f32>,
6388 f64,
6389 ),
6390 (b_row, _, _, _, b_score): &(
6391 RowId,
6392 f64,
6393 Vec<ComponentScore>,
6394 Option<f32>,
6395 f64,
6396 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6397 if let Some(crate::query::Rerank::ExactVector {
6398 embedding_column,
6399 query,
6400 metric,
6401 candidate_limit,
6402 weight,
6403 }) = &request.rerank
6404 {
6405 let fused_order = |(a_row, a_score, ..): &(
6406 RowId,
6407 f64,
6408 Vec<ComponentScore>,
6409 Option<f32>,
6410 f64,
6411 ),
6412 (b_row, b_score, ..): &(
6413 RowId,
6414 f64,
6415 Vec<ComponentScore>,
6416 Option<f32>,
6417 f64,
6418 )| {
6419 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6420 };
6421 let selection_started = std::time::Instant::now();
6422 if let Some(context) = context {
6423 context.consume(ranked.len())?;
6424 }
6425 if ranked.len() > *candidate_limit {
6426 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6427 ranked.truncate(*candidate_limit);
6428 }
6429 ranked.sort_by(fused_order);
6430 fusion_nanos =
6431 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6432 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6433 if let Some(context) = context {
6434 context.consume(row_ids.len())?;
6435 }
6436 let query_now =
6437 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6438 let gather_started = std::time::Instant::now();
6439 let vectors = self.values_for_rids_batch_at_with_context(
6440 &row_ids,
6441 *embedding_column,
6442 snapshot,
6443 query_now,
6444 context,
6445 )?;
6446 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6447 let vector_work =
6448 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6449 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6450 if let Some(context) = context {
6451 context.consume(vector_work)?;
6452 }
6453 query
6454 .iter()
6455 .map(|value| f64::from(*value).powi(2))
6456 .sum::<f64>()
6457 .sqrt()
6458 } else {
6459 0.0
6460 };
6461 let score_started = std::time::Instant::now();
6462 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6463 for (row_id, value) in vectors {
6464 if let Some(meta) = value.generated_embedding_metadata() {
6465 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6466 continue;
6467 }
6468 }
6469 let Some(vector) = value.as_embedding() else {
6470 continue;
6471 };
6472 let score = match metric {
6473 crate::query::VectorMetric::DotProduct => {
6474 if let Some(context) = context {
6475 context.consume(vector_work)?;
6476 }
6477 query
6478 .iter()
6479 .zip(vector)
6480 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6481 .sum::<f64>()
6482 }
6483 crate::query::VectorMetric::Cosine => {
6484 if let Some(context) = context {
6485 context.consume(vector_work.saturating_mul(2))?;
6486 }
6487 let dot = query
6488 .iter()
6489 .zip(vector)
6490 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6491 .sum::<f64>();
6492 let norm = vector
6493 .iter()
6494 .map(|value| f64::from(*value).powi(2))
6495 .sum::<f64>()
6496 .sqrt();
6497 if query_norm == 0.0 || norm == 0.0 {
6498 0.0
6499 } else {
6500 dot / (query_norm * norm)
6501 }
6502 }
6503 crate::query::VectorMetric::Euclidean => {
6504 if let Some(context) = context {
6505 context.consume(vector_work)?;
6506 }
6507 query
6508 .iter()
6509 .zip(vector)
6510 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6511 .sum::<f64>()
6512 .sqrt()
6513 }
6514 };
6515 if !score.is_finite() {
6516 return Err(MongrelError::InvalidArgument(
6517 "exact rerank score must be finite".into(),
6518 ));
6519 }
6520 scores.insert(row_id, score as f32);
6521 }
6522 let mut reranked = Vec::with_capacity(ranked.len());
6523 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6524 let Some(score) = scores.get(&row_id).copied() else {
6525 continue;
6526 };
6527 let ordering_score = match metric {
6528 crate::query::VectorMetric::Euclidean => -f64::from(score),
6529 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6530 f64::from(score)
6531 }
6532 };
6533 let final_score = fused_score + *weight * ordering_score;
6534 if !final_score.is_finite() {
6535 return Err(MongrelError::InvalidArgument(
6536 "final rerank score must be finite".into(),
6537 ));
6538 }
6539 reranked.push((row_id, fused_score, components, Some(score), final_score));
6540 }
6541 ranked = reranked;
6542 ranked.sort_by(order);
6543 crate::trace::QueryTrace::record(|trace| {
6544 trace.exact_vector_gather_nanos =
6545 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6546 trace.exact_vector_score_nanos = trace
6547 .exact_vector_score_nanos
6548 .saturating_add(score_started.elapsed().as_nanos() as u64);
6549 });
6550 }
6551 if let Some(after) = after {
6552 ranked.retain(|(row_id, _, _, _, final_score)| {
6553 final_score.total_cmp(&after.final_score).is_lt()
6554 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6555 });
6556 }
6557 let projection_started = std::time::Instant::now();
6558 let sentinel = projection
6559 .first()
6560 .copied()
6561 .or_else(|| self.schema.columns.first().map(|column| column.id));
6562 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6563 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6564 let mut projection_rows = 0usize;
6565 let mut projection_cells = 0usize;
6566 while out.len() < request.limit && !ranked.is_empty() {
6567 if let Some(context) = context {
6568 context.checkpoint()?;
6569 context.consume(ranked.len())?;
6570 }
6571 let needed = request.limit - out.len();
6572 let window_size = ranked
6573 .len()
6574 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6575 let selection_started = std::time::Instant::now();
6576 let mut remainder = if ranked.len() > window_size {
6577 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6578 ranked.split_off(window_size)
6579 } else {
6580 Vec::new()
6581 };
6582 ranked.sort_by(order);
6583 fusion_nanos =
6584 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6585 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6586 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6587 if let Some(context) = context {
6588 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6589 }
6590 projection_rows = projection_rows.saturating_add(row_ids.len());
6591 projection_cells =
6592 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6593 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6594 std::collections::HashMap::new();
6595 if let Some(column_id) = sentinel {
6596 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6597 &row_ids, column_id, snapshot, query_now, context,
6598 )? {
6599 cells.entry(row_id).or_default().insert(column_id, value);
6600 }
6601 }
6602 for &column_id in &projection {
6603 if Some(column_id) == sentinel {
6604 continue;
6605 }
6606 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6607 &row_ids, column_id, snapshot, query_now, context,
6608 )? {
6609 cells.entry(row_id).or_default().insert(column_id, value);
6610 }
6611 }
6612 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6613 ranked.drain(..)
6614 {
6615 let Some(row_cells) = cells.remove(&row_id) else {
6616 continue;
6617 };
6618 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6619 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6620 out.push(SearchHit {
6621 row_id,
6622 cells: projection
6623 .iter()
6624 .filter_map(|column_id| {
6625 row_cells
6626 .get(column_id)
6627 .cloned()
6628 .map(|value| (*column_id, value))
6629 })
6630 .collect(),
6631 components,
6632 fused_score,
6633 exact_rerank_score,
6634 final_score,
6635 final_rank,
6636 });
6637 if out.len() == request.limit {
6638 break;
6639 }
6640 }
6641 ranked.append(&mut remainder);
6642 }
6643 crate::trace::QueryTrace::record(|trace| {
6644 trace.union_size = union_size;
6645 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6646 trace.projection_nanos = trace
6647 .projection_nanos
6648 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6649 trace.total_nanos = trace
6650 .total_nanos
6651 .saturating_add(total_started.elapsed().as_nanos() as u64);
6652 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6653 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6654 if let Some(context) = context {
6655 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6656 }
6657 });
6658 Ok(out)
6659 }
6660
6661 pub fn set_similarity(
6664 &mut self,
6665 request: &crate::query::SetSimilarityRequest,
6666 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6667 self.set_similarity_with_allowed(request, None)
6668 }
6669
6670 pub fn set_similarity_at(
6671 &mut self,
6672 request: &crate::query::SetSimilarityRequest,
6673 snapshot: Snapshot,
6674 allowed: Option<&std::collections::HashSet<RowId>>,
6675 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6676 self.set_similarity_explained_at(request, snapshot, allowed)
6677 .map(|(hits, _)| hits)
6678 }
6679
6680 pub fn ann_rerank(
6682 &mut self,
6683 request: &crate::query::AnnRerankRequest,
6684 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6685 self.ann_rerank_with_allowed(request, None)
6686 }
6687
6688 pub fn ann_rerank_with_allowed(
6689 &mut self,
6690 request: &crate::query::AnnRerankRequest,
6691 allowed: Option<&std::collections::HashSet<RowId>>,
6692 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6693 self.ann_rerank_at(request, self.snapshot(), allowed)
6694 }
6695
6696 pub fn ann_rerank_at(
6697 &mut self,
6698 request: &crate::query::AnnRerankRequest,
6699 snapshot: Snapshot,
6700 allowed: Option<&std::collections::HashSet<RowId>>,
6701 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6702 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6703 }
6704
6705 pub fn ann_rerank_at_with_context(
6706 &mut self,
6707 request: &crate::query::AnnRerankRequest,
6708 snapshot: Snapshot,
6709 allowed: Option<&std::collections::HashSet<RowId>>,
6710 context: Option<&crate::query::AiExecutionContext>,
6711 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6712 self.ensure_indexes_complete()?;
6713 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6714 }
6715
6716 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6717 &mut self,
6718 request: &crate::query::AnnRerankRequest,
6719 snapshot: Snapshot,
6720 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6721 context: Option<&crate::query::AiExecutionContext>,
6722 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6723 self.ensure_indexes_complete()?;
6724 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6725 }
6726
6727 #[doc(hidden)]
6728 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6729 &self,
6730 request: &crate::query::AnnRerankRequest,
6731 snapshot: Snapshot,
6732 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6733 context: Option<&crate::query::AiExecutionContext>,
6734 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6735 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6736 }
6737
6738 fn ann_rerank_at_with_filters_and_context(
6739 &self,
6740 request: &crate::query::AnnRerankRequest,
6741 snapshot: Snapshot,
6742 allowed: Option<&std::collections::HashSet<RowId>>,
6743 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6744 context: Option<&crate::query::AiExecutionContext>,
6745 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6746 use crate::query::{
6747 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6748 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6749 };
6750 if request.candidate_k == 0 || request.limit == 0 {
6751 return Err(MongrelError::InvalidArgument(
6752 "candidate_k and limit must be > 0".into(),
6753 ));
6754 }
6755 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6756 return Err(MongrelError::InvalidArgument(format!(
6757 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6758 )));
6759 }
6760 let retriever = Retriever::Ann {
6761 column_id: request.column_id,
6762 query: request.query.clone(),
6763 k: request.candidate_k,
6764 };
6765 self.require_select()?;
6766 self.validate_retriever(&retriever)?;
6767 let hits = self.retrieve_filtered(
6768 &retriever,
6769 snapshot,
6770 None,
6771 allowed,
6772 candidate_authorization,
6773 context,
6774 )?;
6775 let distances: std::collections::HashMap<_, _> = hits
6776 .iter()
6777 .filter_map(|hit| match hit.score {
6778 RetrieverScore::AnnHammingDistance(distance) => {
6779 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
6780 }
6781 RetrieverScore::AnnCosineDistance(distance) => {
6782 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
6783 }
6784 _ => None,
6785 })
6786 .collect();
6787 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6788 if let Some(context) = context {
6789 context.consume(row_ids.len())?;
6790 }
6791 let gather_started = std::time::Instant::now();
6792 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6793 let values = self.values_for_rids_batch_at_with_context(
6794 &row_ids,
6795 request.column_id,
6796 snapshot,
6797 query_now,
6798 context,
6799 )?;
6800 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6801 let score_started = std::time::Instant::now();
6802 let vector_work =
6803 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6804 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6805 if let Some(context) = context {
6806 context.consume(vector_work)?;
6807 }
6808 request
6809 .query
6810 .iter()
6811 .map(|value| f64::from(*value).powi(2))
6812 .sum::<f64>()
6813 .sqrt()
6814 } else {
6815 0.0
6816 };
6817 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6818 for (row_id, value) in values {
6819 if let Some(meta) = value.generated_embedding_metadata() {
6820 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6821 continue;
6822 }
6823 }
6824 let Some(vector) = value.as_embedding() else {
6825 continue;
6826 };
6827 let exact_score = match request.metric {
6828 VectorMetric::DotProduct => {
6829 if let Some(context) = context {
6830 context.consume(vector_work)?;
6831 }
6832 request
6833 .query
6834 .iter()
6835 .zip(vector)
6836 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6837 .sum::<f64>()
6838 }
6839 VectorMetric::Cosine => {
6840 if let Some(context) = context {
6841 context.consume(vector_work.saturating_mul(2))?;
6842 }
6843 let dot = request
6844 .query
6845 .iter()
6846 .zip(vector)
6847 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6848 .sum::<f64>();
6849 let norm = vector
6850 .iter()
6851 .map(|value| f64::from(*value).powi(2))
6852 .sum::<f64>()
6853 .sqrt();
6854 if query_norm == 0.0 || norm == 0.0 {
6855 0.0
6856 } else {
6857 dot / (query_norm * norm)
6858 }
6859 }
6860 VectorMetric::Euclidean => {
6861 if let Some(context) = context {
6862 context.consume(vector_work)?;
6863 }
6864 request
6865 .query
6866 .iter()
6867 .zip(vector)
6868 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6869 .sum::<f64>()
6870 .sqrt()
6871 }
6872 };
6873 let exact_score = exact_score as f32;
6874 if !exact_score.is_finite() {
6875 return Err(MongrelError::InvalidArgument(
6876 "exact ANN score must be finite".into(),
6877 ));
6878 }
6879 let Some(candidate_distance) = distances.get(&row_id).copied() else {
6880 continue;
6881 };
6882 reranked.push(AnnRerankHit {
6883 row_id,
6884 candidate_distance,
6885 exact_score,
6886 });
6887 }
6888 reranked.sort_by(|left, right| {
6889 let score = match request.metric {
6890 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6891 VectorMetric::Cosine | VectorMetric::DotProduct => {
6892 right.exact_score.total_cmp(&left.exact_score)
6893 }
6894 };
6895 score.then_with(|| left.row_id.cmp(&right.row_id))
6896 });
6897 reranked.truncate(request.limit);
6898 crate::trace::QueryTrace::record(|trace| {
6899 trace.exact_vector_gather_nanos =
6900 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6901 trace.exact_vector_score_nanos = trace
6902 .exact_vector_score_nanos
6903 .saturating_add(score_started.elapsed().as_nanos() as u64);
6904 });
6905 Ok(reranked)
6906 }
6907
6908 pub fn set_similarity_with_allowed(
6909 &mut self,
6910 request: &crate::query::SetSimilarityRequest,
6911 allowed: Option<&std::collections::HashSet<RowId>>,
6912 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6913 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6914 .map(|(hits, _)| hits)
6915 }
6916
6917 pub fn set_similarity_explained(
6918 &mut self,
6919 request: &crate::query::SetSimilarityRequest,
6920 ) -> Result<(
6921 Vec<crate::query::SetSimilarityHit>,
6922 crate::query::SetSimilarityTrace,
6923 )> {
6924 self.set_similarity_explained_at(request, self.snapshot(), None)
6925 }
6926
6927 fn set_similarity_explained_at(
6928 &mut self,
6929 request: &crate::query::SetSimilarityRequest,
6930 snapshot: Snapshot,
6931 allowed: Option<&std::collections::HashSet<RowId>>,
6932 ) -> Result<(
6933 Vec<crate::query::SetSimilarityHit>,
6934 crate::query::SetSimilarityTrace,
6935 )> {
6936 self.ensure_indexes_complete()?;
6937 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6938 }
6939
6940 pub fn set_similarity_at_with_context(
6941 &mut self,
6942 request: &crate::query::SetSimilarityRequest,
6943 snapshot: Snapshot,
6944 allowed: Option<&std::collections::HashSet<RowId>>,
6945 context: Option<&crate::query::AiExecutionContext>,
6946 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6947 self.ensure_indexes_complete()?;
6948 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6949 .map(|(hits, _)| hits)
6950 }
6951
6952 pub fn set_similarity_at_with_candidate_authorization_and_context(
6953 &mut self,
6954 request: &crate::query::SetSimilarityRequest,
6955 snapshot: Snapshot,
6956 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6957 context: Option<&crate::query::AiExecutionContext>,
6958 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6959 self.ensure_indexes_complete()?;
6960 self.set_similarity_explained_at_with_context(
6961 request,
6962 snapshot,
6963 None,
6964 authorization,
6965 context,
6966 )
6967 .map(|(hits, _)| hits)
6968 }
6969
6970 #[doc(hidden)]
6971 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6972 &self,
6973 request: &crate::query::SetSimilarityRequest,
6974 snapshot: Snapshot,
6975 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6976 context: Option<&crate::query::AiExecutionContext>,
6977 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6978 self.set_similarity_explained_at_with_context(
6979 request,
6980 snapshot,
6981 None,
6982 authorization,
6983 context,
6984 )
6985 .map(|(hits, _)| hits)
6986 }
6987
6988 fn set_similarity_explained_at_with_context(
6989 &self,
6990 request: &crate::query::SetSimilarityRequest,
6991 snapshot: Snapshot,
6992 allowed: Option<&std::collections::HashSet<RowId>>,
6993 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6994 context: Option<&crate::query::AiExecutionContext>,
6995 ) -> Result<(
6996 Vec<crate::query::SetSimilarityHit>,
6997 crate::query::SetSimilarityTrace,
6998 )> {
6999 use crate::query::{
7000 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
7001 MAX_SET_MEMBERS,
7002 };
7003 let mut trace = crate::query::SetSimilarityTrace::default();
7004 if request.members.is_empty() {
7005 return Ok((Vec::new(), trace));
7006 }
7007 if request.candidate_k == 0 || request.limit == 0 {
7008 return Err(MongrelError::InvalidArgument(
7009 "candidate_k and limit must be > 0".into(),
7010 ));
7011 }
7012 if request.candidate_k > MAX_RETRIEVER_K
7013 || request.limit > MAX_FINAL_LIMIT
7014 || request.members.len() > MAX_SET_MEMBERS
7015 {
7016 return Err(MongrelError::InvalidArgument(format!(
7017 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
7018 )));
7019 }
7020 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
7021 return Err(MongrelError::InvalidArgument(
7022 "min_jaccard must be finite and between 0 and 1".into(),
7023 ));
7024 }
7025 let started = std::time::Instant::now();
7026 let retriever = Retriever::MinHash {
7027 column_id: request.column_id,
7028 members: request.members.clone(),
7029 k: request.candidate_k,
7030 };
7031 self.require_select()?;
7032 self.validate_retriever(&retriever)?;
7033 let hits = self.retrieve_filtered(
7034 &retriever,
7035 snapshot,
7036 None,
7037 allowed,
7038 candidate_authorization,
7039 context,
7040 )?;
7041 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
7042 trace.candidate_count = hits.len();
7043 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7044 if let Some(context) = context {
7045 context.consume(row_ids.len())?;
7046 }
7047 let started = std::time::Instant::now();
7048 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7049 let values = self.values_for_rids_batch_at_with_context(
7050 &row_ids,
7051 request.column_id,
7052 snapshot,
7053 query_now,
7054 context,
7055 )?;
7056 trace.gather_us = started.elapsed().as_micros() as u64;
7057 if let Some(context) = context {
7058 context.consume(request.members.len())?;
7059 }
7060 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
7061 let estimates: std::collections::HashMap<_, _> = hits
7062 .into_iter()
7063 .filter_map(|hit| match hit.score {
7064 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
7065 _ => None,
7066 })
7067 .collect();
7068 let started = std::time::Instant::now();
7069 let mut parsed = Vec::with_capacity(values.len());
7070 for (row_id, value) in values {
7071 let Value::Bytes(bytes) = value else {
7072 continue;
7073 };
7074 if let Some(context) = context {
7075 context.consume(crate::query::work_units(
7076 bytes.len(),
7077 crate::query::PARSE_WORK_QUANTUM,
7078 ))?;
7079 }
7080 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
7081 continue;
7082 };
7083 if let Some(context) = context {
7084 context.consume(members.len())?;
7085 }
7086 let stored = members
7087 .into_iter()
7088 .filter_map(|member| match member {
7089 serde_json::Value::String(value) => {
7090 Some(crate::query::SetMember::String(value))
7091 }
7092 serde_json::Value::Number(value) => {
7093 Some(crate::query::SetMember::Number(value))
7094 }
7095 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
7096 _ => None,
7097 })
7098 .collect::<std::collections::HashSet<_>>();
7099 parsed.push((row_id, stored));
7100 }
7101 trace.parse_us = started.elapsed().as_micros() as u64;
7102 trace.verified_count = parsed.len();
7103 let started = std::time::Instant::now();
7104 let mut exact = Vec::new();
7105 for (row_id, stored) in parsed {
7106 if let Some(context) = context {
7107 context.consume(query.len().saturating_add(stored.len()))?;
7108 }
7109 let union = query.union(&stored).count();
7110 let score = if union == 0 {
7111 1.0
7112 } else {
7113 query.intersection(&stored).count() as f32 / union as f32
7114 };
7115 if score >= request.min_jaccard {
7116 exact.push(SetSimilarityHit {
7117 row_id,
7118 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
7119 exact_jaccard: score,
7120 });
7121 }
7122 }
7123 exact.sort_by(|a, b| {
7124 b.exact_jaccard
7125 .total_cmp(&a.exact_jaccard)
7126 .then_with(|| a.row_id.cmp(&b.row_id))
7127 });
7128 exact.truncate(request.limit);
7129 trace.score_us = started.elapsed().as_micros() as u64;
7130 crate::trace::QueryTrace::record(|query_trace| {
7131 query_trace.exact_set_gather_nanos = query_trace
7132 .exact_set_gather_nanos
7133 .saturating_add(trace.gather_us.saturating_mul(1_000));
7134 query_trace.exact_set_parse_nanos = query_trace
7135 .exact_set_parse_nanos
7136 .saturating_add(trace.parse_us.saturating_mul(1_000));
7137 query_trace.exact_set_score_nanos = query_trace
7138 .exact_set_score_nanos
7139 .saturating_add(trace.score_us.saturating_mul(1_000));
7140 });
7141 Ok((exact, trace))
7142 }
7143
7144 fn values_for_rids_batch_at(
7146 &self,
7147 row_ids: &[u64],
7148 column_id: u16,
7149 snapshot: Snapshot,
7150 now: i64,
7151 ) -> Result<Vec<(RowId, Value)>> {
7152 if self.ttl.is_none()
7153 && self.memtable.is_empty()
7154 && self.mutable_run.is_empty()
7155 && self.run_refs.len() == 1
7156 {
7157 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7158 if row_ids.len().saturating_mul(24) < reader.row_count() {
7163 let mut values = Vec::with_capacity(row_ids.len());
7164 for &raw_row_id in row_ids {
7165 let row_id = RowId(raw_row_id);
7166 if let Some((_, false, Some(value))) =
7167 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7168 {
7169 values.push((row_id, value));
7170 }
7171 }
7172 return Ok(values);
7173 }
7174 let (positions, visible_row_ids) =
7175 reader.visible_positions_with_rids(snapshot.epoch)?;
7176 let requested: Vec<(RowId, usize)> = row_ids
7177 .iter()
7178 .filter_map(|raw| {
7179 visible_row_ids
7180 .binary_search(&(*raw as i64))
7181 .ok()
7182 .map(|index| (RowId(*raw), positions[index]))
7183 })
7184 .collect();
7185 let values = reader.gather_column(
7186 column_id,
7187 &requested
7188 .iter()
7189 .map(|(_, position)| *position)
7190 .collect::<Vec<_>>(),
7191 )?;
7192 return Ok(requested
7193 .into_iter()
7194 .zip(values)
7195 .map(|((row_id, _), value)| (row_id, value))
7196 .collect());
7197 }
7198 self.values_for_rids_at(row_ids, column_id, snapshot, now)
7199 }
7200
7201 fn values_for_rids_batch_at_with_context(
7202 &self,
7203 row_ids: &[u64],
7204 column_id: u16,
7205 snapshot: Snapshot,
7206 now: i64,
7207 context: Option<&crate::query::AiExecutionContext>,
7208 ) -> Result<Vec<(RowId, Value)>> {
7209 let Some(context) = context else {
7210 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
7211 };
7212 let mut values = Vec::with_capacity(row_ids.len());
7213 for chunk in row_ids.chunks(256) {
7214 context.checkpoint()?;
7215 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
7216 }
7217 Ok(values)
7218 }
7219
7220 fn values_for_rids_at(
7222 &self,
7223 row_ids: &[u64],
7224 column_id: u16,
7225 snapshot: Snapshot,
7226 now: i64,
7227 ) -> Result<Vec<(RowId, Value)>> {
7228 let mut readers: Vec<_> = self
7229 .run_refs
7230 .iter()
7231 .map(|run| self.open_reader(run.run_id))
7232 .collect::<Result<_>>()?;
7233 let mut out = Vec::with_capacity(row_ids.len());
7234 for &raw_row_id in row_ids {
7235 let row_id = RowId(raw_row_id);
7236 let mem = self.memtable.get_version_at(row_id, snapshot);
7237 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
7238 let overlay = match (mem, mutable) {
7239 (Some((_, a)), Some((_, b))) => Some(
7240 if Snapshot::version_is_newer(
7241 a.committed_epoch,
7242 a.commit_ts,
7243 b.committed_epoch,
7244 b.commit_ts,
7245 ) {
7246 a
7247 } else {
7248 b
7249 },
7250 ),
7251 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
7252 (None, None) => None,
7253 };
7254 if let Some(row) = overlay {
7255 if !row.deleted && !self.row_expired_at(&row, now) {
7256 if let Some(value) = row.columns.get(&column_id) {
7257 out.push((row_id, value.clone()));
7258 }
7259 }
7260 continue;
7261 }
7262
7263 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
7264 for (index, reader) in readers.iter_mut().enumerate() {
7265 if let Some((epoch, deleted, value)) =
7266 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7267 {
7268 if best
7269 .as_ref()
7270 .map(|(best_epoch, ..)| epoch > *best_epoch)
7271 .unwrap_or(true)
7272 {
7273 best = Some((epoch, deleted, value, index));
7274 }
7275 }
7276 }
7277 let Some((_, false, Some(value), reader_index)) = best else {
7278 continue;
7279 };
7280 if let Some(ttl) = self.ttl {
7281 if ttl.column_id != column_id {
7282 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7283 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
7284 {
7285 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7286 continue;
7287 }
7288 }
7289 } else if let Value::Int64(timestamp) = value {
7290 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7291 continue;
7292 }
7293 }
7294 }
7295 out.push((row_id, value));
7296 }
7297 Ok(out)
7298 }
7299
7300 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
7305 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
7306 }
7307
7308 pub fn rows_for_rids_with_context(
7309 &self,
7310 rids: &[u64],
7311 snapshot: Snapshot,
7312 context: &crate::query::AiExecutionContext,
7313 ) -> Result<Vec<Row>> {
7314 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
7315 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
7316 }
7317
7318 fn rows_for_rids_at_time(
7319 &self,
7320 rids: &[u64],
7321 snapshot: Snapshot,
7322 ttl_now: i64,
7323 control: Option<&crate::ExecutionControl>,
7324 ) -> Result<Vec<Row>> {
7325 use std::collections::HashMap;
7326 let mut rows = Vec::with_capacity(rids.len());
7327 let tier_size = self.memtable.len() + self.mutable_run.len();
7342 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7343 if rids.len().saturating_mul(24) < tier_size {
7344 for &rid in rids {
7345 if overlay.len() & 255 == 0 {
7346 control
7347 .map(crate::ExecutionControl::checkpoint)
7348 .transpose()?;
7349 }
7350 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
7351 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
7352 let newest = match (mem, mrun) {
7353 (Some((_, mr)), Some((_, rr))) => Some(
7354 if Snapshot::version_is_newer(
7355 mr.committed_epoch,
7356 mr.commit_ts,
7357 rr.committed_epoch,
7358 rr.commit_ts,
7359 ) {
7360 mr
7361 } else {
7362 rr
7363 },
7364 ),
7365 (Some((_, mr)), None) => Some(mr),
7366 (None, Some((_, rr))) => Some(rr),
7367 (None, None) => None,
7368 };
7369 if let Some(row) = newest {
7370 overlay.insert(rid, row);
7371 }
7372 }
7373 } else {
7374 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7375 overlay
7376 .entry(row.row_id.0)
7377 .and_modify(|e| {
7378 if Snapshot::version_is_newer(
7379 row.committed_epoch,
7380 row.commit_ts,
7381 e.committed_epoch,
7382 e.commit_ts,
7383 ) {
7384 *e = row.clone();
7385 }
7386 })
7387 .or_insert(row);
7388 };
7389 for (index, row) in self
7390 .memtable
7391 .visible_versions_at(snapshot)
7392 .into_iter()
7393 .enumerate()
7394 {
7395 if index & 255 == 0 {
7396 control
7397 .map(crate::ExecutionControl::checkpoint)
7398 .transpose()?;
7399 }
7400 fold_newest(row, &mut overlay);
7401 }
7402 for (index, row) in self
7403 .mutable_run
7404 .visible_versions_at(snapshot)
7405 .into_iter()
7406 .enumerate()
7407 {
7408 if index & 255 == 0 {
7409 control
7410 .map(crate::ExecutionControl::checkpoint)
7411 .transpose()?;
7412 }
7413 fold_newest(row, &mut overlay);
7414 }
7415 }
7416 if self.run_refs.len() == 1 {
7417 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7418 if rids.len().saturating_mul(24) < reader.row_count() {
7426 for (index, &rid) in rids.iter().enumerate() {
7427 if index & 255 == 0 {
7428 control
7429 .map(crate::ExecutionControl::checkpoint)
7430 .transpose()?;
7431 }
7432 if let Some(r) = overlay.get(&rid) {
7433 if !r.deleted {
7434 rows.push(r.clone());
7435 }
7436 continue;
7437 }
7438 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7439 if !row.deleted {
7440 rows.push(row);
7441 }
7442 }
7443 }
7444 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7445 return Ok(rows);
7446 }
7447 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7456 enum Src {
7459 Overlay,
7460 Run,
7461 }
7462 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7463 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7464 for (index, rid) in rids.iter().enumerate() {
7465 if index & 255 == 0 {
7466 control
7467 .map(crate::ExecutionControl::checkpoint)
7468 .transpose()?;
7469 }
7470 if overlay.contains_key(rid) {
7471 plan.push(Src::Overlay);
7472 continue;
7473 }
7474 match vis_rids.binary_search(&(*rid as i64)) {
7475 Ok(i) => {
7476 plan.push(Src::Run);
7477 fetch.push(positions[i]);
7478 }
7479 Err(_) => { }
7480 }
7481 }
7482 let fetched = reader.materialize_batch(&fetch)?;
7483 let mut fetched_iter = fetched.into_iter();
7484 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7485 if index & 255 == 0 {
7486 control
7487 .map(crate::ExecutionControl::checkpoint)
7488 .transpose()?;
7489 }
7490 match src {
7491 Src::Overlay => {
7492 if let Some(r) = overlay.get(rid) {
7493 if !r.deleted {
7494 rows.push(r.clone());
7495 }
7496 }
7497 }
7498 Src::Run => {
7499 if let Some(row) = fetched_iter.next() {
7500 if !row.deleted {
7501 rows.push(row);
7502 }
7503 }
7504 }
7505 }
7506 }
7507 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7508 return Ok(rows);
7509 }
7510 let mut readers: Vec<_> = self
7514 .run_refs
7515 .iter()
7516 .map(|rr| self.open_reader(rr.run_id))
7517 .collect::<Result<Vec<_>>>()?;
7518 for (index, rid) in rids.iter().enumerate() {
7519 if index & 255 == 0 {
7520 control
7521 .map(crate::ExecutionControl::checkpoint)
7522 .transpose()?;
7523 }
7524 if let Some(r) = overlay.get(rid) {
7525 if !r.deleted {
7526 rows.push(r.clone());
7527 }
7528 continue;
7529 }
7530 let mut best: Option<(Epoch, Row)> = None;
7531 for reader in readers.iter_mut() {
7532 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7533 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7534 best = Some((epoch, row));
7535 }
7536 }
7537 }
7538 if let Some((_, r)) = best {
7539 if !r.deleted {
7540 rows.push(r);
7541 }
7542 }
7543 }
7544 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7545 Ok(rows)
7546 }
7547
7548 pub fn indexes_complete(&self) -> bool {
7558 self.indexes_complete
7559 }
7560
7561 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7563 self.index_build_policy
7564 }
7565
7566 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7570 self.index_build_policy = policy;
7571 }
7572
7573 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7578 if !self.indexes_complete {
7582 return None;
7583 }
7584 let b = self.bitmap.get(&column_id)?;
7585 let result: Vec<Vec<u8>> = b
7586 .keys()
7587 .into_iter()
7588 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7589 .collect();
7590 Some(result)
7591 }
7592
7593 pub fn fk_join_row_ids(
7594 &self,
7595 fk_column_id: u16,
7596 pk_values: &[Vec<u8>],
7597 fk_conditions: &[crate::query::Condition],
7598 snapshot: Snapshot,
7599 ) -> Result<Vec<u64>> {
7600 let Some(b) = self.bitmap.get(&fk_column_id) else {
7601 return Ok(Vec::new());
7602 };
7603 let mut join_set = {
7604 let mut acc = roaring::RoaringBitmap::new();
7605 for v in pk_values {
7606 acc |= b.get(v);
7607 }
7608 RowIdSet::from_roaring(acc)
7609 };
7610 if !fk_conditions.is_empty() {
7611 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7612 sets.push(join_set);
7613 for c in fk_conditions {
7614 sets.push(self.resolve_condition(c, snapshot)?);
7615 }
7616 join_set = RowIdSet::intersect_many(sets);
7617 }
7618 Ok(join_set.into_sorted_vec())
7619 }
7620
7621 pub fn fk_join_count(
7627 &self,
7628 fk_column_id: u16,
7629 pk_values: &[Vec<u8>],
7630 fk_conditions: &[crate::query::Condition],
7631 snapshot: Snapshot,
7632 ) -> Result<u64> {
7633 let Some(b) = self.bitmap.get(&fk_column_id) else {
7634 return Ok(0);
7635 };
7636 let mut acc = roaring::RoaringBitmap::new();
7637 for v in pk_values {
7638 acc |= b.get(v);
7639 }
7640 if fk_conditions.is_empty() {
7641 return Ok(acc.len());
7642 }
7643 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7644 sets.push(RowIdSet::from_roaring(acc));
7645 for c in fk_conditions {
7646 sets.push(self.resolve_condition(c, snapshot)?);
7647 }
7648 Ok(RowIdSet::intersect_many(sets).len() as u64)
7649 }
7650
7651 fn resolve_condition(
7656 &self,
7657 c: &crate::query::Condition,
7658 snapshot: Snapshot,
7659 ) -> Result<RowIdSet> {
7660 self.resolve_condition_with_allowed(c, snapshot, None)
7661 }
7662
7663 fn resolve_condition_with_allowed(
7664 &self,
7665 c: &crate::query::Condition,
7666 snapshot: Snapshot,
7667 allowed: Option<&std::collections::HashSet<RowId>>,
7668 ) -> Result<RowIdSet> {
7669 use crate::query::Condition;
7670 self.validate_condition(c)?;
7671 Ok(match c {
7672 Condition::Pk(key) => {
7673 let lookup = self
7674 .schema
7675 .primary_key()
7676 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7677 .unwrap_or_else(|| key.clone());
7678 self.hot
7679 .get(&lookup)
7680 .map(|r| RowIdSet::one(r.0))
7681 .unwrap_or_else(RowIdSet::empty)
7682 }
7683 Condition::BitmapEq { column_id, value } => {
7684 let lookup = self.index_lookup_key_bytes(*column_id, value);
7685 self.bitmap
7686 .get(column_id)
7687 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7688 .unwrap_or_else(RowIdSet::empty)
7689 }
7690 Condition::BitmapIn { column_id, values } => {
7691 let bm = self.bitmap.get(column_id);
7692 let mut acc = roaring::RoaringBitmap::new();
7693 if let Some(b) = bm {
7694 for v in values {
7695 let lookup = self.index_lookup_key_bytes(*column_id, v);
7696 acc |= b.get(&lookup);
7697 }
7698 }
7699 RowIdSet::from_roaring(acc)
7700 }
7701 Condition::BytesPrefix { column_id, prefix } => {
7702 if let Some(b) = self.bitmap.get(column_id) {
7707 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7708 let mut acc = roaring::RoaringBitmap::new();
7709 for key in b.keys() {
7710 if key.starts_with(&lookup_prefix) {
7711 acc |= b.get(&key);
7712 }
7713 }
7714 RowIdSet::from_roaring(acc)
7715 } else {
7716 RowIdSet::empty()
7717 }
7718 }
7719 Condition::FmContains { column_id, pattern } => self
7720 .fm
7721 .get(column_id)
7722 .map(|f| {
7723 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7724 })
7725 .unwrap_or_else(RowIdSet::empty),
7726 Condition::FmContainsAll {
7727 column_id,
7728 patterns,
7729 } => {
7730 if let Some(f) = self.fm.get(column_id) {
7733 let sets: Vec<RowIdSet> = patterns
7734 .iter()
7735 .map(|pat| {
7736 RowIdSet::from_unsorted(
7737 f.locate(pat).into_iter().map(|r| r.0).collect(),
7738 )
7739 })
7740 .collect();
7741 RowIdSet::intersect_many(sets)
7742 } else {
7743 RowIdSet::empty()
7744 }
7745 }
7746 Condition::Ann {
7747 column_id,
7748 query,
7749 k,
7750 } => RowIdSet::from_unsorted(
7751 self.retrieve_filtered(
7752 &crate::query::Retriever::Ann {
7753 column_id: *column_id,
7754 query: query.clone(),
7755 k: *k,
7756 },
7757 snapshot,
7758 None,
7759 allowed,
7760 None,
7761 None,
7762 )?
7763 .into_iter()
7764 .map(|hit| hit.row_id.0)
7765 .collect(),
7766 ),
7767 Condition::SparseMatch {
7768 column_id,
7769 query,
7770 k,
7771 } => RowIdSet::from_unsorted(
7772 self.retrieve_filtered(
7773 &crate::query::Retriever::Sparse {
7774 column_id: *column_id,
7775 query: query.clone(),
7776 k: *k,
7777 },
7778 snapshot,
7779 None,
7780 allowed,
7781 None,
7782 None,
7783 )?
7784 .into_iter()
7785 .map(|hit| hit.row_id.0)
7786 .collect(),
7787 ),
7788 Condition::MinHashSimilar {
7789 column_id,
7790 query,
7791 k,
7792 } => match self.minhash.get(column_id) {
7793 Some(index) => {
7794 let candidates = index.candidate_row_ids(query);
7795 let eligible =
7796 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7797 RowIdSet::from_unsorted(
7798 index
7799 .search_filtered(query, *k, |row_id| {
7800 eligible.contains(&row_id)
7801 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7802 })
7803 .into_iter()
7804 .map(|(row_id, _)| row_id.0)
7805 .collect(),
7806 )
7807 }
7808 None => RowIdSet::empty(),
7809 },
7810 Condition::Range { column_id, lo, hi } => {
7811 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7820 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7821 } else if self.run_refs.len() == 1 {
7822 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7823 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7824 } else {
7825 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7826 };
7827 set.remove_many(self.overlay_rid_set(snapshot));
7828 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7829 set
7830 }
7831 Condition::RangeF64 {
7832 column_id,
7833 lo,
7834 lo_inclusive,
7835 hi,
7836 hi_inclusive,
7837 } => {
7838 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7841 RowIdSet::from_unsorted(
7842 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7843 .into_iter()
7844 .collect(),
7845 )
7846 } else if self.run_refs.len() == 1 {
7847 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7848 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7849 } else {
7850 return self.range_scan_f64(
7851 *column_id,
7852 *lo,
7853 *lo_inclusive,
7854 *hi,
7855 *hi_inclusive,
7856 snapshot,
7857 );
7858 };
7859 set.remove_many(self.overlay_rid_set(snapshot));
7860 self.range_scan_overlay_f64(
7861 &mut set,
7862 *column_id,
7863 *lo,
7864 *lo_inclusive,
7865 *hi,
7866 *hi_inclusive,
7867 snapshot,
7868 );
7869 set
7870 }
7871 Condition::IsNull { column_id } => {
7872 let mut set = if self.run_refs.len() == 1 {
7873 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7874 r.null_row_id_set(*column_id, true)?
7875 } else {
7876 return self.null_scan(*column_id, true, snapshot);
7877 };
7878 set.remove_many(self.overlay_rid_set(snapshot));
7879 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7880 set
7881 }
7882 Condition::IsNotNull { column_id } => {
7883 let mut set = if self.run_refs.len() == 1 {
7884 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7885 r.null_row_id_set(*column_id, false)?
7886 } else {
7887 return self.null_scan(*column_id, false, snapshot);
7888 };
7889 set.remove_many(self.overlay_rid_set(snapshot));
7890 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7891 set
7892 }
7893 })
7894 }
7895
7896 fn range_scan_i64(
7904 &self,
7905 column_id: u16,
7906 lo: i64,
7907 hi: i64,
7908 snapshot: Snapshot,
7909 ) -> Result<RowIdSet> {
7910 let mut row_ids = Vec::new();
7911 let overlay_rids = self.overlay_rid_set(snapshot);
7912 for rr in &self.run_refs {
7913 let mut reader = self.open_reader(rr.run_id)?;
7914 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7915 for rid in matched {
7916 if !overlay_rids.contains(&rid) {
7917 row_ids.push(rid);
7918 }
7919 }
7920 }
7921 let mut s = RowIdSet::from_unsorted(row_ids);
7922 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7923 Ok(s)
7924 }
7925
7926 fn range_scan_f64(
7929 &self,
7930 column_id: u16,
7931 lo: f64,
7932 lo_inclusive: bool,
7933 hi: f64,
7934 hi_inclusive: bool,
7935 snapshot: Snapshot,
7936 ) -> Result<RowIdSet> {
7937 let mut row_ids = Vec::new();
7938 let overlay_rids = self.overlay_rid_set(snapshot);
7939 for rr in &self.run_refs {
7940 let mut reader = self.open_reader(rr.run_id)?;
7941 let matched = reader.range_row_ids_visible_f64(
7942 column_id,
7943 lo,
7944 lo_inclusive,
7945 hi,
7946 hi_inclusive,
7947 snapshot.epoch,
7948 )?;
7949 for rid in matched {
7950 if !overlay_rids.contains(&rid) {
7951 row_ids.push(rid);
7952 }
7953 }
7954 }
7955 let mut s = RowIdSet::from_unsorted(row_ids);
7956 self.range_scan_overlay_f64(
7957 &mut s,
7958 column_id,
7959 lo,
7960 lo_inclusive,
7961 hi,
7962 hi_inclusive,
7963 snapshot,
7964 );
7965 Ok(s)
7966 }
7967
7968 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7970 let mut s = HashSet::new();
7971 for row in self.memtable.visible_versions_at(snapshot) {
7972 s.insert(row.row_id.0);
7973 }
7974 for row in self.mutable_run.visible_versions_at(snapshot) {
7975 s.insert(row.row_id.0);
7976 }
7977 s
7978 }
7979
7980 fn range_scan_overlay_i64(
7981 &self,
7982 s: &mut RowIdSet,
7983 column_id: u16,
7984 lo: i64,
7985 hi: i64,
7986 snapshot: Snapshot,
7987 ) {
7988 let mut newest: HashMap<u64, Row> = HashMap::new();
7995 for r in self.mutable_run.visible_versions_at(snapshot) {
7996 newest.insert(r.row_id.0, r);
7997 }
7998 for r in self.memtable.visible_versions_at(snapshot) {
7999 newest
8000 .entry(r.row_id.0)
8001 .and_modify(|cur| {
8002 if Snapshot::version_is_newer(
8003 r.committed_epoch,
8004 r.commit_ts,
8005 cur.committed_epoch,
8006 cur.commit_ts,
8007 ) {
8008 *cur = r.clone();
8009 }
8010 })
8011 .or_insert(r);
8012 }
8013 for row in newest.values() {
8014 if !row.deleted {
8015 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
8016 if *v >= lo && *v <= hi {
8017 s.insert(row.row_id.0);
8018 }
8019 }
8020 }
8021 }
8022 }
8023
8024 #[allow(clippy::too_many_arguments)]
8025 fn range_scan_overlay_f64(
8026 &self,
8027 s: &mut RowIdSet,
8028 column_id: u16,
8029 lo: f64,
8030 lo_inclusive: bool,
8031 hi: f64,
8032 hi_inclusive: bool,
8033 snapshot: Snapshot,
8034 ) {
8035 let mut newest: HashMap<u64, Row> = HashMap::new();
8038 for r in self.mutable_run.visible_versions_at(snapshot) {
8039 newest.insert(r.row_id.0, r);
8040 }
8041 for r in self.memtable.visible_versions_at(snapshot) {
8042 newest
8043 .entry(r.row_id.0)
8044 .and_modify(|cur| {
8045 if Snapshot::version_is_newer(
8046 r.committed_epoch,
8047 r.commit_ts,
8048 cur.committed_epoch,
8049 cur.commit_ts,
8050 ) {
8051 *cur = r.clone();
8052 }
8053 })
8054 .or_insert(r);
8055 }
8056 for row in newest.values() {
8057 if !row.deleted {
8058 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
8059 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
8060 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
8061 if ok_lo && ok_hi {
8062 s.insert(row.row_id.0);
8063 }
8064 }
8065 }
8066 }
8067 }
8068
8069 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
8072 let mut row_ids = Vec::new();
8073 let overlay_rids = self.overlay_rid_set(snapshot);
8074 for rr in &self.run_refs {
8075 let mut reader = self.open_reader(rr.run_id)?;
8076 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
8077 for rid in matched {
8078 if !overlay_rids.contains(&rid) {
8079 row_ids.push(rid);
8080 }
8081 }
8082 }
8083 let mut s = RowIdSet::from_unsorted(row_ids);
8084 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
8085 Ok(s)
8086 }
8087
8088 fn null_scan_overlay(
8092 &self,
8093 s: &mut RowIdSet,
8094 column_id: u16,
8095 want_nulls: bool,
8096 snapshot: Snapshot,
8097 ) {
8098 let mut newest: HashMap<u64, Row> = HashMap::new();
8099 for r in self.mutable_run.visible_versions_at(snapshot) {
8100 newest.insert(r.row_id.0, r);
8101 }
8102 for r in self.memtable.visible_versions_at(snapshot) {
8103 newest
8104 .entry(r.row_id.0)
8105 .and_modify(|cur| {
8106 if Snapshot::version_is_newer(
8107 r.committed_epoch,
8108 r.commit_ts,
8109 cur.committed_epoch,
8110 cur.commit_ts,
8111 ) {
8112 *cur = r.clone();
8113 }
8114 })
8115 .or_insert(r);
8116 }
8117 for row in newest.values() {
8118 if row.deleted {
8119 continue;
8120 }
8121 let is_null = !row.columns.contains_key(&column_id)
8122 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
8123 if is_null == want_nulls {
8124 s.insert(row.row_id.0);
8125 }
8126 }
8127 }
8128
8129 pub fn snapshot(&self) -> Snapshot {
8130 let epoch = self.epoch.visible();
8131 match &self.wal {
8135 WalSink::Shared(shared) => match shared.hlc.now() {
8136 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
8137 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
8139 },
8140 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
8141 }
8142 }
8143
8144 pub fn data_generation(&self) -> u64 {
8146 self.data_generation
8147 }
8148
8149 pub(crate) fn bump_data_generation(&mut self) {
8150 self.data_generation = self.data_generation.wrapping_add(1);
8151 }
8152
8153 pub fn table_id(&self) -> u64 {
8155 self.table_id
8156 }
8157
8158 fn seal_generations(&mut self) {
8164 self.memtable.seal();
8165 self.mutable_run.seal();
8166 self.hot.seal();
8167 for index in self.bitmap.values_mut() {
8168 index.seal();
8169 }
8170 for index in self.ann.values_mut() {
8171 index.seal();
8172 }
8173 for index in self.fm.values_mut() {
8174 index.seal();
8175 }
8176 for index in self.sparse.values_mut() {
8177 index.seal();
8178 }
8179 for index in self.minhash.values_mut() {
8180 index.seal();
8181 }
8182 self.pk_by_row.seal();
8183 }
8184
8185 fn capture_read_generation(&self) -> ReadGeneration {
8190 let visible_through = self.current_epoch();
8191 ReadGeneration {
8192 schema: Arc::new(self.schema.clone()),
8193 base_runs: Arc::new(self.run_refs.clone()),
8194 deltas: TableDeltas {
8195 memtable: self.memtable.clone(),
8196 mutable_run: self.mutable_run.clone(),
8197 hot: self.hot.clone(),
8198 pk_by_row: self.pk_by_row.clone(),
8199 },
8200 indexes: Arc::new(IndexGeneration::capture(
8201 &self.bitmap,
8202 &self.learned_range,
8203 &self.fm,
8204 &self.ann,
8205 &self.sparse,
8206 &self.minhash,
8207 visible_through,
8208 match &self.wal {
8210 WalSink::Shared(shared) => shared
8211 .hlc
8212 .now()
8213 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
8214 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
8215 },
8216 )),
8217 visible_through,
8218 }
8219 }
8220
8221 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
8226 self.ensure_indexes_complete()?;
8227 self.seal_generations();
8228 let view = Arc::new(self.capture_read_generation());
8229 self.published.store(Arc::clone(&view));
8230 Ok(view)
8231 }
8232
8233 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
8239 self.published.load_full()
8240 }
8241
8242 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
8244 &self.pins
8245 }
8246
8247 pub fn version_gc_floor(&self) -> Epoch {
8252 self.min_active_snapshot()
8253 .unwrap_or_else(|| self.current_epoch())
8254 }
8255
8256 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
8263 let mut report = self.pins.report();
8264 let transaction_floor = [
8265 self.pinned.keys().next().copied(),
8266 self.snapshots.min_pinned(),
8267 ]
8268 .into_iter()
8269 .flatten()
8270 .min();
8271 if let Some(epoch) = transaction_floor {
8272 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
8273 }
8274 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
8275 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
8276 }
8277 report
8278 }
8279
8280 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
8281 self.publish_read_generation()?;
8282 let mut generation = self.clone();
8283 generation.read_only = true;
8284 generation.wal = WalSink::ReadOnly;
8285 generation.pending_delete_rids.clear();
8286 generation.pending_put_cols.clear();
8287 generation.pending_rows.clear();
8288 generation.pending_rows_auto_inc.clear();
8289 generation.pending_dels.clear();
8290 generation.pending_truncate = None;
8291 generation.agg_cache = Arc::new(HashMap::new());
8292 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
8295 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
8298 crate::retention::PinSource::ReadGeneration,
8299 self.current_epoch(),
8300 )));
8301 Ok(generation)
8302 }
8303
8304 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
8305 (std::mem::size_of::<Self>() as u64)
8306 .saturating_add(self.memtable.approx_bytes())
8307 .saturating_add(self.mutable_run.approx_bytes())
8308 .saturating_add(self.live_count.saturating_mul(64))
8309 }
8310
8311 pub fn pin_snapshot(&mut self) -> Snapshot {
8318 let snap = self.snapshot();
8319 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
8320 snap
8321 }
8322
8323 pub fn hlc_gc_floor(
8333 &self,
8334 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
8335 ) -> crate::epoch::GcFloor {
8336 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
8337 epoch
8338 .and_then(&mut project_epoch)
8339 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
8340 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
8341 };
8342 let transaction = [
8343 self.pinned.keys().next().copied(),
8344 self.snapshots.min_pinned(),
8345 ]
8346 .into_iter()
8347 .flatten()
8348 .min();
8349 let history = self.snapshots.history_floor(self.current_epoch());
8350 crate::epoch::GcFloor {
8351 transaction_snapshot: project(transaction),
8352 history_retention: project(history),
8353 backup_pitr: project(
8354 self.pins
8355 .oldest_for(crate::retention::PinSource::BackupPitr),
8356 ),
8357 replication: project(
8358 self.pins
8359 .oldest_for(crate::retention::PinSource::Replication),
8360 ),
8361 read_generation: project(
8362 self.pins
8363 .oldest_for(crate::retention::PinSource::ReadGeneration),
8364 ),
8365 online_index_build: project(
8366 self.pins
8367 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
8368 ),
8369 }
8370 }
8371
8372 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
8374 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
8375 *count -= 1;
8376 if *count == 0 {
8377 self.pinned.remove(&snap.epoch);
8378 }
8379 }
8380 }
8381
8382 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
8394 let local = self.pinned.keys().next().copied();
8395 let global = self.snapshots.min_pinned();
8396 let history = self.snapshots.history_floor(self.current_epoch());
8397 let pinned = self.pins.oldest_pinned();
8398 [local, global, history, pinned].into_iter().flatten().min()
8399 }
8400
8401 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
8405 self.ensure_writable()?;
8406 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
8407 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
8408 }
8409
8410 pub fn clear_ttl(&mut self) -> Result<()> {
8411 self.ensure_writable()?;
8412 self.apply_ttl_policy_at(None, self.current_epoch())
8413 }
8414
8415 pub fn ttl(&self) -> Option<TtlPolicy> {
8416 self.ttl
8417 }
8418
8419 pub(crate) fn prepare_ttl_policy(
8420 &self,
8421 column_name: &str,
8422 duration_nanos: u64,
8423 ) -> Result<TtlPolicy> {
8424 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
8425 return Err(MongrelError::InvalidArgument(
8426 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
8427 ));
8428 }
8429 let column = self
8430 .schema
8431 .columns
8432 .iter()
8433 .find(|column| column.name == column_name)
8434 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
8435 if column.ty != TypeId::TimestampNanos {
8436 return Err(MongrelError::Schema(format!(
8437 "TTL column {column_name} must be TimestampNanos, is {:?}",
8438 column.ty
8439 )));
8440 }
8441 Ok(TtlPolicy {
8442 column_id: column.id,
8443 duration_nanos,
8444 })
8445 }
8446
8447 pub(crate) fn apply_ttl_policy_at(
8448 &mut self,
8449 policy: Option<TtlPolicy>,
8450 epoch: Epoch,
8451 ) -> Result<()> {
8452 if let Some(policy) = policy {
8453 let column = self
8454 .schema
8455 .columns
8456 .iter()
8457 .find(|column| column.id == policy.column_id)
8458 .ok_or_else(|| {
8459 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8460 })?;
8461 if column.ty != TypeId::TimestampNanos
8462 || policy.duration_nanos == 0
8463 || policy.duration_nanos > i64::MAX as u64
8464 {
8465 return Err(MongrelError::Schema("invalid TTL policy".into()));
8466 }
8467 }
8468 self.ttl = policy;
8469 self.agg_cache = Arc::new(HashMap::new());
8470 self.clear_result_cache();
8471 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8472 self.persist_manifest(epoch)
8473 }
8474
8475 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8476 let Some(policy) = self.ttl else {
8477 return false;
8478 };
8479 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8480 return false;
8481 };
8482 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8483 }
8484
8485 pub fn current_epoch(&self) -> Epoch {
8486 self.epoch.visible()
8487 }
8488
8489 pub fn memtable_len(&self) -> usize {
8490 self.memtable.len()
8491 }
8492
8493 pub fn count(&self) -> u64 {
8496 if self.ttl.is_none()
8497 && self.pending_put_cols.is_empty()
8498 && self.pending_delete_rids.is_empty()
8499 && self.pending_rows.is_empty()
8500 && self.pending_dels.is_empty()
8501 && self.pending_truncate.is_none()
8502 {
8503 self.live_count
8504 } else {
8505 self.visible_rows(self.snapshot())
8506 .map(|rows| rows.len() as u64)
8507 .unwrap_or(self.live_count)
8508 }
8509 }
8510
8511 pub fn count_conditions(
8515 &mut self,
8516 conditions: &[crate::query::Condition],
8517 snapshot: Snapshot,
8518 ) -> Result<Option<u64>> {
8519 use crate::query::Condition;
8520 if self.ttl.is_some() {
8521 if conditions.is_empty() {
8522 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8523 }
8524 let mut sets = Vec::with_capacity(conditions.len());
8525 for condition in conditions {
8526 sets.push(self.resolve_condition(condition, snapshot)?);
8527 }
8528 let survivors = RowIdSet::intersect_many(sets);
8529 let rows = self.visible_rows(snapshot)?;
8530 return Ok(Some(
8531 rows.into_iter()
8532 .filter(|row| survivors.contains(row.row_id.0))
8533 .count() as u64,
8534 ));
8535 }
8536 if conditions.is_empty() {
8537 return Ok(Some(self.count()));
8538 }
8539 let served = |c: &Condition| {
8540 matches!(
8541 c,
8542 Condition::Pk(_)
8543 | Condition::BitmapEq { .. }
8544 | Condition::BitmapIn { .. }
8545 | Condition::BytesPrefix { .. }
8546 | Condition::FmContains { .. }
8547 | Condition::FmContainsAll { .. }
8548 | Condition::Ann { .. }
8549 | Condition::Range { .. }
8550 | Condition::RangeF64 { .. }
8551 | Condition::SparseMatch { .. }
8552 | Condition::MinHashSimilar { .. }
8553 | Condition::IsNull { .. }
8554 | Condition::IsNotNull { .. }
8555 )
8556 };
8557 if !conditions.iter().all(served) {
8558 return Ok(None);
8559 }
8560 self.ensure_indexes_complete()?;
8561 if !self.pending_put_cols.is_empty()
8562 || !self.pending_delete_rids.is_empty()
8563 || !self.pending_rows.is_empty()
8564 || !self.pending_dels.is_empty()
8565 || self.pending_truncate.is_some()
8566 {
8567 let mut sets = Vec::with_capacity(conditions.len());
8568 for condition in conditions {
8569 sets.push(self.resolve_condition(condition, snapshot)?);
8570 }
8571 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8572 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8573 }
8574 let mut sets = Vec::with_capacity(conditions.len());
8575 for condition in conditions {
8576 sets.push(self.resolve_condition(condition, snapshot)?);
8577 }
8578 let mut rids = RowIdSet::intersect_many(sets);
8579 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8589 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8590 }
8591 let count = rids.len() as u64;
8592 crate::trace::QueryTrace::record(|t| {
8593 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8594 t.survivor_count = Some(count as usize);
8595 t.conditions_pushed = conditions.len();
8596 });
8597 Ok(Some(count))
8598 }
8599
8600 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8605 let mut out = Vec::new();
8606 for row in self.memtable.visible_versions(snapshot.epoch) {
8607 if row.deleted {
8608 out.push(row.row_id.0);
8609 }
8610 }
8611 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8612 if row.deleted {
8613 out.push(row.row_id.0);
8614 }
8615 }
8616 out
8617 }
8618
8619 pub fn bulk_load_columns(
8628 &mut self,
8629 user_columns: Vec<(u16, columnar::NativeColumn)>,
8630 ) -> Result<Epoch> {
8631 self.bulk_load_columns_with(user_columns, 3, false, true)
8632 }
8633
8634 pub fn bulk_load_fast(
8641 &mut self,
8642 user_columns: Vec<(u16, columnar::NativeColumn)>,
8643 ) -> Result<Epoch> {
8644 self.bulk_load_columns_with(user_columns, -1, true, false)
8645 }
8646
8647 fn bulk_load_columns_with(
8648 &mut self,
8649 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8650 zstd_level: i32,
8651 force_plain: bool,
8652 lz4: bool,
8653 ) -> Result<Epoch> {
8654 self.ensure_writable()?;
8655 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8656 if n == 0 {
8657 return Ok(self.current_epoch());
8658 }
8659 let epoch = self.commit_new_epoch()?;
8660 let live_before = self.live_count;
8661 self.spill_mutable_run(epoch)?;
8663 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8664 && self.indexes_complete
8665 && self.run_refs.is_empty()
8666 && self.memtable.is_empty()
8667 && self.mutable_run.is_empty();
8668 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8671 self.validate_columns_not_null(&user_columns, n)?;
8672 let winner_idx = self
8673 .bulk_pk_winner_indices(&user_columns, n)
8674 .filter(|idx| idx.len() != n);
8675 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8676 match winner_idx.as_deref() {
8677 Some(idx) => {
8678 let compacted = user_columns
8679 .iter()
8680 .map(|(id, c)| (*id, c.gather(idx)))
8681 .collect();
8682 (compacted, idx.len())
8683 }
8684 None => (user_columns, n),
8685 };
8686 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8687 let first = self.allocator.alloc_range(write_n as u64)?.0;
8688 for rid in first..first + write_n as u64 {
8689 self.reservoir.offer(rid);
8690 }
8691 let run_id = self.alloc_run_id()?;
8692 let path = self.run_path(run_id);
8693 let mut writer =
8694 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8695 if force_plain {
8696 writer = writer.with_plain();
8697 } else if lz4 {
8698 writer = writer.with_lz4();
8701 } else {
8702 writer = writer.with_zstd_level(zstd_level);
8703 }
8704 if let Some(kek) = &self.kek {
8705 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8706 }
8707 let header = match self.create_run_file(run_id)? {
8708 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8709 None => writer.write_native(&path, &write_columns, write_n, first)?,
8710 };
8711 self.run_refs.push(RunRef {
8712 run_id: run_id as u128,
8713 level: 0,
8714 epoch_created: epoch.0,
8715 row_count: header.row_count,
8716 });
8717 self.live_count = self.live_count.saturating_add(write_n as u64);
8718 if eager_index_build {
8719 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8720 self.index_columns_bulk(&write_columns, &row_ids);
8721 self.indexes_complete = true;
8722 self.build_learned_ranges()?;
8723 } else {
8724 self.indexes_complete = false;
8728 }
8729 self.mark_flushed(epoch)?;
8730 self.persist_manifest(epoch)?;
8731 if eager_index_build {
8732 self.checkpoint_indexes(epoch);
8733 }
8734 self.clear_result_cache();
8735 self.data_generation = self.data_generation.wrapping_add(1);
8736 Ok(epoch)
8737 }
8738
8739 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8757 let n = row_ids.len();
8758 if n == 0 {
8759 return;
8760 }
8761 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8762 columns.iter().map(|(id, c)| (*id, c)).collect();
8763 let ty_of: std::collections::HashMap<u16, TypeId> = self
8764 .schema
8765 .columns
8766 .iter()
8767 .map(|c| (c.id, c.ty.clone()))
8768 .collect();
8769 let pk_id = self.schema.primary_key().map(|c| c.id);
8770
8771 for (i, &rid) in row_ids.iter().enumerate() {
8772 let row_id = RowId(rid);
8773 if let Some(pid) = pk_id {
8774 if let Some(col) = by_id.get(&pid) {
8775 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8776 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8777 self.insert_hot_pk(key, row_id);
8778 }
8779 }
8780 }
8781 for idef in &self.schema.indexes {
8782 let Some(col) = by_id.get(&idef.column_id) else {
8783 continue;
8784 };
8785 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8786 match idef.kind {
8787 IndexKind::Bitmap => {
8788 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8789 if let Some(key) =
8790 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8791 {
8792 b.insert(key, row_id);
8793 }
8794 }
8795 }
8796 IndexKind::FmIndex => {
8797 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8798 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8799 f.insert(bytes.to_vec(), row_id);
8800 }
8801 }
8802 }
8803 IndexKind::Sparse => {
8804 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8805 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8806 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8807 s.insert(&terms, row_id);
8808 }
8809 }
8810 }
8811 }
8812 IndexKind::MinHash => {
8813 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8814 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8815 let tokens = crate::index::token_hashes_from_bytes(bytes);
8816 mh.insert(&tokens, row_id);
8817 }
8818 }
8819 }
8820 _ => {}
8821 }
8822 }
8823 }
8824 }
8825
8826 pub fn visible_columns_native(
8831 &self,
8832 snapshot: Snapshot,
8833 projection: Option<&[u16]>,
8834 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8835 self.visible_columns_native_inner(snapshot, projection, None)
8836 }
8837
8838 pub fn visible_columns_native_with_control(
8839 &self,
8840 snapshot: Snapshot,
8841 projection: Option<&[u16]>,
8842 control: &crate::ExecutionControl,
8843 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8844 self.visible_columns_native_inner(snapshot, projection, Some(control))
8845 }
8846
8847 fn visible_columns_native_inner(
8848 &self,
8849 snapshot: Snapshot,
8850 projection: Option<&[u16]>,
8851 control: Option<&crate::ExecutionControl>,
8852 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8853 execution_checkpoint(control, 0)?;
8854 let wanted: Vec<u16> = match projection {
8855 Some(p) => p.to_vec(),
8856 None => self.schema.columns.iter().map(|c| c.id).collect(),
8857 };
8858 if self.ttl.is_none()
8859 && self.memtable.is_empty()
8860 && self.mutable_run.is_empty()
8861 && self.run_refs.len() == 1
8862 {
8863 let rr = self.run_refs[0].clone();
8864 let mut reader = self.open_reader(rr.run_id)?;
8865 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8866 execution_checkpoint(control, 0)?;
8867 let all_visible = idxs.len() == reader.row_count();
8868 if reader.has_mmap() && control.is_none() {
8874 use rayon::prelude::*;
8875 let valid: Vec<u16> = wanted
8878 .iter()
8879 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8880 .copied()
8881 .collect();
8882 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8884 .par_iter()
8885 .filter_map(|cid| {
8886 reader
8887 .column_native_shared(*cid)
8888 .ok()
8889 .map(|col| (*cid, col))
8890 })
8891 .collect();
8892 let cols = decoded
8893 .into_iter()
8894 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8895 .collect();
8896 return Ok(cols);
8897 }
8898 let mut cols = Vec::with_capacity(wanted.len());
8899 for (index, cid) in wanted.iter().enumerate() {
8900 execution_checkpoint(control, index)?;
8901 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8902 Some(c) => c,
8903 None => continue,
8904 };
8905 let col = reader.column_native(cdef.id)?;
8906 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8907 }
8908 return Ok(cols);
8909 }
8910 let vcols = self.visible_columns(snapshot)?;
8911 execution_checkpoint(control, 0)?;
8912 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8913 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8914 .into_iter()
8915 .filter(|(id, _)| want_set.contains(id))
8916 .map(|(id, vals)| {
8917 let ty = self
8918 .schema
8919 .columns
8920 .iter()
8921 .find(|c| c.id == id)
8922 .map(|c| c.ty.clone())
8923 .unwrap_or(TypeId::Bytes);
8924 (id, columnar::values_to_native(ty, &vals))
8925 })
8926 .collect();
8927 Ok(out)
8928 }
8929
8930 pub fn run_count(&self) -> usize {
8931 self.run_refs.len()
8932 }
8933
8934 pub fn memtable_is_empty(&self) -> bool {
8936 self.memtable.is_empty()
8937 }
8938
8939 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8943 self.page_cache.stats()
8944 }
8945
8946 pub fn reset_page_cache_stats(&self) {
8948 self.page_cache.reset_stats();
8949 }
8950
8951 pub fn run_ids(&self) -> Vec<u128> {
8954 self.run_refs.iter().map(|r| r.run_id).collect()
8955 }
8956
8957 pub fn single_run_is_clean(&self) -> bool {
8961 if self.ttl.is_some() || self.run_refs.len() != 1 {
8962 return false;
8963 }
8964 self.open_reader(self.run_refs[0].run_id)
8965 .map(|r| r.is_clean())
8966 .unwrap_or(false)
8967 }
8968
8969 fn resolve_footprint(
8976 &self,
8977 conditions: &[crate::query::Condition],
8978 snapshot: Snapshot,
8979 ) -> roaring::RoaringBitmap {
8980 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8981 return roaring::RoaringBitmap::new();
8982 }
8983 if self.run_refs.is_empty() {
8984 return roaring::RoaringBitmap::new();
8985 }
8986 if self.run_refs.len() == 1 {
8988 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8989 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8990 return rids.to_roaring_lossy();
8991 }
8992 }
8993 }
8994 roaring::RoaringBitmap::new()
8995 }
8996
8997 pub fn query_columns_native_cached(
9008 &mut self,
9009 conditions: &[crate::query::Condition],
9010 projection: Option<&[u16]>,
9011 snapshot: Snapshot,
9012 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9013 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
9014 }
9015
9016 pub fn query_columns_native_cached_with_control(
9017 &mut self,
9018 conditions: &[crate::query::Condition],
9019 projection: Option<&[u16]>,
9020 snapshot: Snapshot,
9021 control: &crate::ExecutionControl,
9022 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9023 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
9024 }
9025
9026 fn query_columns_native_cached_inner(
9027 &mut self,
9028 conditions: &[crate::query::Condition],
9029 projection: Option<&[u16]>,
9030 snapshot: Snapshot,
9031 control: Option<&crate::ExecutionControl>,
9032 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9033 execution_checkpoint(control, 0)?;
9034 if self.ttl.is_some() {
9037 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9038 }
9039 if conditions.is_empty() {
9040 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9041 }
9042 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
9046 if let Some(hit) = self.result_cache.lock().get_columns(key) {
9047 crate::trace::QueryTrace::record(|t| {
9048 t.result_cache_hit = true;
9049 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9050 });
9051 return Ok(Some((*hit).clone()));
9052 }
9053 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
9054 execution_checkpoint(control, 0)?;
9055 if let Some(cols) = &res {
9056 let footprint = self.resolve_footprint(conditions, snapshot);
9057 let condition_cols = crate::query::condition_columns(conditions);
9058 execution_checkpoint(control, 0)?;
9059 self.result_cache.lock().insert(
9060 key,
9061 CachedEntry {
9062 data: CachedData::Columns(Arc::new(cols.clone())),
9063 footprint,
9064 condition_cols,
9065 },
9066 );
9067 }
9068 Ok(res)
9069 }
9070
9071 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
9076 if self.ttl.is_some() {
9077 return self.query(q);
9078 }
9079 if q.conditions.is_empty() {
9080 return self.query(q);
9081 }
9082 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
9083 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
9084 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
9085 if let Some(hit) = self.result_cache.lock().get_rows(key) {
9086 crate::trace::QueryTrace::record(|t| {
9087 t.result_cache_hit = true;
9088 t.scan_mode = crate::trace::ScanMode::Materialized;
9089 });
9090 return Ok((*hit).clone());
9091 }
9092 let rows = self.query(q)?;
9093 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
9094 let condition_cols = crate::query::condition_columns(&q.conditions);
9095 self.result_cache.lock().insert(
9096 key,
9097 CachedEntry {
9098 data: CachedData::Rows(Arc::new(rows.clone())),
9099 footprint,
9100 condition_cols,
9101 },
9102 );
9103 Ok(rows)
9104 }
9105
9106 #[allow(clippy::type_complexity)]
9121 pub fn query_columns_native_traced(
9122 &mut self,
9123 conditions: &[crate::query::Condition],
9124 projection: Option<&[u16]>,
9125 snapshot: Snapshot,
9126 ) -> Result<(
9127 Option<Vec<(u16, columnar::NativeColumn)>>,
9128 crate::trace::QueryTrace,
9129 )> {
9130 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9131 self.query_columns_native(conditions, projection, snapshot)
9132 });
9133 Ok((result?, trace))
9134 }
9135
9136 #[allow(clippy::type_complexity)]
9139 pub fn query_columns_native_cached_traced(
9140 &mut self,
9141 conditions: &[crate::query::Condition],
9142 projection: Option<&[u16]>,
9143 snapshot: Snapshot,
9144 ) -> Result<(
9145 Option<Vec<(u16, columnar::NativeColumn)>>,
9146 crate::trace::QueryTrace,
9147 )> {
9148 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9149 self.query_columns_native_cached(conditions, projection, snapshot)
9150 });
9151 Ok((result?, trace))
9152 }
9153
9154 pub fn native_page_cursor_traced(
9156 &self,
9157 snapshot: Snapshot,
9158 projection: Vec<(u16, TypeId)>,
9159 conditions: &[crate::query::Condition],
9160 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
9161 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9162 self.native_page_cursor(snapshot, projection, conditions)
9163 });
9164 Ok((result?, trace))
9165 }
9166
9167 pub fn native_multi_run_cursor_traced(
9169 &self,
9170 snapshot: Snapshot,
9171 projection: Vec<(u16, TypeId)>,
9172 conditions: &[crate::query::Condition],
9173 ) -> Result<(
9174 Option<crate::cursor::MultiRunCursor>,
9175 crate::trace::QueryTrace,
9176 )> {
9177 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9178 self.native_multi_run_cursor(snapshot, projection, conditions)
9179 });
9180 Ok((result?, trace))
9181 }
9182
9183 pub fn count_conditions_traced(
9185 &mut self,
9186 conditions: &[crate::query::Condition],
9187 snapshot: Snapshot,
9188 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
9189 let (result, trace) =
9190 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
9191 Ok((result?, trace))
9192 }
9193
9194 pub fn query_traced(
9196 &mut self,
9197 q: &crate::query::Query,
9198 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
9199 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
9200 Ok((result?, trace))
9201 }
9202
9203 pub fn query_columns_native(
9208 &mut self,
9209 conditions: &[crate::query::Condition],
9210 projection: Option<&[u16]>,
9211 snapshot: Snapshot,
9212 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9213 self.query_columns_native_inner(conditions, projection, snapshot, None)
9214 }
9215
9216 pub fn query_columns_native_with_control(
9217 &mut self,
9218 conditions: &[crate::query::Condition],
9219 projection: Option<&[u16]>,
9220 snapshot: Snapshot,
9221 control: &crate::ExecutionControl,
9222 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9223 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
9224 }
9225
9226 fn query_columns_native_inner(
9227 &mut self,
9228 conditions: &[crate::query::Condition],
9229 projection: Option<&[u16]>,
9230 snapshot: Snapshot,
9231 control: Option<&crate::ExecutionControl>,
9232 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9233 use crate::query::Condition;
9234 execution_checkpoint(control, 0)?;
9235 if self.ttl.is_some() {
9238 return Ok(None);
9239 }
9240 if conditions.is_empty() {
9241 return Ok(None);
9242 }
9243 self.ensure_indexes_complete()?;
9244
9245 let served = |c: &Condition| {
9250 matches!(
9251 c,
9252 Condition::Pk(_)
9253 | Condition::BitmapEq { .. }
9254 | Condition::BitmapIn { .. }
9255 | Condition::BytesPrefix { .. }
9256 | Condition::FmContains { .. }
9257 | Condition::FmContainsAll { .. }
9258 | Condition::Ann { .. }
9259 | Condition::Range { .. }
9260 | Condition::RangeF64 { .. }
9261 | Condition::SparseMatch { .. }
9262 | Condition::MinHashSimilar { .. }
9263 | Condition::IsNull { .. }
9264 | Condition::IsNotNull { .. }
9265 )
9266 };
9267 if !conditions.iter().all(served) {
9268 return Ok(None);
9269 }
9270 let fast_path =
9271 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9272 crate::trace::QueryTrace::record(|t| {
9273 t.run_count = self.run_refs.len();
9274 t.memtable_rows = self.memtable.len();
9275 t.mutable_run_rows = self.mutable_run.len();
9276 t.conditions_pushed = conditions.len();
9277 t.learned_range_used = conditions.iter().any(|c| match c {
9278 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
9279 self.learned_range.contains_key(column_id)
9280 }
9281 _ => false,
9282 });
9283 });
9284 let col_ids: Vec<u16> = projection
9286 .map(|p| p.to_vec())
9287 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
9288 let proj_pairs: Vec<(u16, TypeId)> = col_ids
9289 .iter()
9290 .map(|&cid| {
9291 let ty = self
9292 .schema
9293 .columns
9294 .iter()
9295 .find(|c| c.id == cid)
9296 .map(|c| c.ty.clone())
9297 .unwrap_or(TypeId::Bytes);
9298 (cid, ty)
9299 })
9300 .collect();
9301
9302 if fast_path {
9308 let needs_column = conditions.iter().any(|c| match c {
9311 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
9312 Condition::RangeF64 { column_id, .. } => {
9313 !self.learned_range.contains_key(column_id)
9314 }
9315 _ => false,
9316 });
9317 let mut reader_opt: Option<RunReader> = if needs_column {
9318 Some(self.open_reader(self.run_refs[0].run_id)?)
9319 } else {
9320 None
9321 };
9322 let mut sets: Vec<RowIdSet> = Vec::new();
9323 for (index, c) in conditions.iter().enumerate() {
9324 execution_checkpoint(control, index)?;
9325 let s = match c {
9326 Condition::Range { column_id, lo, hi }
9327 if !self.learned_range.contains_key(column_id) =>
9328 {
9329 if reader_opt.is_none() {
9330 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9331 }
9332 reader_opt
9333 .as_mut()
9334 .expect("reader opened for range")
9335 .range_row_id_set_i64(*column_id, *lo, *hi)?
9336 }
9337 Condition::RangeF64 {
9338 column_id,
9339 lo,
9340 lo_inclusive,
9341 hi,
9342 hi_inclusive,
9343 } if !self.learned_range.contains_key(column_id) => {
9344 if reader_opt.is_none() {
9345 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9346 }
9347 reader_opt
9348 .as_mut()
9349 .expect("reader opened for range")
9350 .range_row_id_set_f64(
9351 *column_id,
9352 *lo,
9353 *lo_inclusive,
9354 *hi,
9355 *hi_inclusive,
9356 )?
9357 }
9358 _ => self.resolve_condition(c, snapshot)?,
9359 };
9360 sets.push(s);
9361 }
9362 let candidates = RowIdSet::intersect_many(sets);
9363 crate::trace::QueryTrace::record(|t| {
9364 t.survivor_count = Some(candidates.len());
9365 });
9366 if candidates.is_empty() {
9367 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
9368 .iter()
9369 .map(|&id| {
9370 (
9371 id,
9372 columnar::null_native(
9373 proj_pairs
9374 .iter()
9375 .find(|(c, _)| c == &id)
9376 .map(|(_, t)| t.clone())
9377 .unwrap_or(TypeId::Bytes),
9378 0,
9379 ),
9380 )
9381 })
9382 .collect();
9383 return Ok(Some(cols));
9384 }
9385 let mut reader = match reader_opt.take() {
9386 Some(r) => r,
9387 None => self.open_reader(self.run_refs[0].run_id)?,
9388 };
9389 let candidate_ids = candidates.into_sorted_vec();
9390 let (positions, fast_rid) = if let Some(positions) =
9391 reader.positions_for_row_ids_fast(&candidate_ids)
9392 {
9393 (positions, true)
9394 } else {
9395 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
9396 match col {
9397 columnar::NativeColumn::Int64 { data, .. } => {
9398 let mut p = Vec::with_capacity(candidate_ids.len());
9399 for (index, rid) in candidate_ids.iter().enumerate() {
9400 execution_checkpoint(control, index)?;
9401 if let Ok(position) = data.binary_search(&(*rid as i64)) {
9402 p.push(position);
9403 }
9404 }
9405 p.sort_unstable();
9406 (p, false)
9407 }
9408 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
9409 }
9410 };
9411 crate::trace::QueryTrace::record(|t| {
9412 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9413 t.fast_row_id_map = fast_rid;
9414 });
9415 let mut cols = Vec::with_capacity(col_ids.len());
9416 for (index, cid) in col_ids.iter().enumerate() {
9417 execution_checkpoint(control, index)?;
9418 let col = reader.column_native(*cid)?;
9419 cols.push((*cid, col.gather(&positions)));
9420 }
9421 return Ok(Some(cols));
9422 }
9423
9424 if !self.run_refs.is_empty() {
9437 use crate::cursor::{
9438 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
9439 };
9440 let remaining: usize;
9441 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9442 let c = self
9443 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9444 .expect("single-run cursor should build when run_refs.len() == 1");
9445 remaining = c.remaining_rows();
9446 Box::new(c)
9447 } else {
9448 let c = self
9449 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9450 .expect("multi-run cursor should build when run_refs.len() >= 1");
9451 remaining = c.remaining_rows();
9452 Box::new(c)
9453 };
9454 crate::trace::QueryTrace::record(|t| {
9455 if t.survivor_count.is_none() {
9456 t.survivor_count = Some(remaining);
9457 }
9458 });
9459 let cols = match control {
9460 Some(control) => {
9461 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9462 }
9463 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9464 };
9465 return Ok(Some(cols));
9466 }
9467
9468 crate::trace::QueryTrace::record(|t| {
9473 t.scan_mode = crate::trace::ScanMode::Materialized;
9474 t.row_materialized = true;
9475 });
9476 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9477 for (index, c) in conditions.iter().enumerate() {
9478 execution_checkpoint(control, index)?;
9479 sets.push(self.resolve_condition(c, snapshot)?);
9480 }
9481 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9482 let rows = self.rows_for_rids(&rids, snapshot)?;
9483 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9484 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9485 execution_checkpoint(control, index)?;
9486 let vals: Vec<Value> = rows
9487 .iter()
9488 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9489 .collect();
9490 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9491 }
9492 Ok(Some(cols))
9493 }
9494
9495 pub fn native_page_cursor(
9510 &self,
9511 snapshot: Snapshot,
9512 projection: Vec<(u16, TypeId)>,
9513 conditions: &[crate::query::Condition],
9514 ) -> Result<Option<NativePageCursor>> {
9515 use crate::cursor::build_page_plans;
9516 if self.ttl.is_some() {
9517 return Ok(None);
9518 }
9519 if !conditions.is_empty() && !self.indexes_complete {
9522 return Ok(None);
9523 }
9524 if self.run_refs.len() != 1 {
9525 return Ok(None);
9526 }
9527 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9528 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9529
9530 let overlay_rids: HashSet<u64> = {
9533 let mut s = HashSet::new();
9534 for row in self.memtable.visible_versions(snapshot.epoch) {
9535 s.insert(row.row_id.0);
9536 }
9537 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9538 s.insert(row.row_id.0);
9539 }
9540 s
9541 };
9542
9543 let survivors = if conditions.is_empty() {
9547 None
9548 } else {
9549 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9550 };
9551
9552 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9559 survivors.clone()
9560 } else if let Some(s) = &survivors {
9561 let mut run_set = s.clone();
9562 run_set.remove_many(overlay_rids.iter().copied());
9563 Some(run_set)
9564 } else {
9565 Some(RowIdSet::from_unsorted(
9566 rids.iter()
9567 .map(|&r| r as u64)
9568 .filter(|r| !overlay_rids.contains(r))
9569 .collect(),
9570 ))
9571 };
9572
9573 let overlay_rows = if overlay_rids.is_empty() {
9574 Vec::new()
9575 } else {
9576 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9577 self.overlay_visible_rows(snapshot, bound)
9578 };
9579
9580 let plans = if positions.is_empty() {
9582 Vec::new()
9583 } else {
9584 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9585 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9586 };
9587
9588 let overlay = if overlay_rows.is_empty() {
9590 None
9591 } else {
9592 let filtered =
9593 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9594 if filtered.is_empty() {
9595 None
9596 } else {
9597 Some(self.materialize_overlay(&filtered, &projection))
9598 }
9599 };
9600
9601 let overlay_row_count = overlay
9602 .as_ref()
9603 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9604 .unwrap_or(0);
9605 crate::trace::QueryTrace::record(|t| {
9606 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9607 t.run_count = self.run_refs.len();
9608 t.memtable_rows = self.memtable.len();
9609 t.mutable_run_rows = self.mutable_run.len();
9610 t.overlay_rows = overlay_row_count;
9611 t.conditions_pushed = conditions.len();
9612 t.pages_decoded = plans
9613 .iter()
9614 .map(|p| p.positions.len())
9615 .sum::<usize>()
9616 .min(1);
9617 });
9618
9619 Ok(Some(NativePageCursor::new_with_overlay(
9620 reader, projection, plans, overlay,
9621 )))
9622 }
9623 #[allow(clippy::type_complexity)]
9633 pub fn native_multi_run_cursor(
9634 &self,
9635 snapshot: Snapshot,
9636 projection: Vec<(u16, TypeId)>,
9637 conditions: &[crate::query::Condition],
9638 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9639 use crate::cursor::{MultiRunCursor, RunStream};
9640 use crate::sorted_run::SYS_ROW_ID;
9641 use std::collections::{BinaryHeap, HashMap, HashSet};
9642 if self.ttl.is_some() {
9643 return Ok(None);
9644 }
9645 if !conditions.is_empty() && !self.indexes_complete {
9648 return Ok(None);
9649 }
9650 if self.run_refs.is_empty() {
9651 return Ok(None);
9652 }
9653
9654 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9656 Vec::with_capacity(self.run_refs.len());
9657 for rr in &self.run_refs {
9658 let mut reader = self.open_reader(rr.run_id)?;
9659 let (rids, eps, del) = reader.system_columns_native()?;
9660 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9661 run_meta.push((reader, rids, eps, del, page_rows));
9662 }
9663
9664 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9668 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9669 for i in 0..rids.len() {
9670 let rid = rids[i] as u64;
9671 let e = eps[i] as u64;
9672 if e > snapshot.epoch.0 {
9673 continue;
9674 }
9675 let is_del = del[i] != 0;
9676 best.entry(rid)
9677 .and_modify(|cur| {
9678 if e > cur.0 {
9679 *cur = (e, run_idx, i, is_del);
9680 }
9681 })
9682 .or_insert((e, run_idx, i, is_del));
9683 }
9684 }
9685
9686 let overlay_rids: HashSet<u64> = {
9688 let mut s = HashSet::new();
9689 for row in self.memtable.visible_versions(snapshot.epoch) {
9690 s.insert(row.row_id.0);
9691 }
9692 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9693 s.insert(row.row_id.0);
9694 }
9695 s
9696 };
9697
9698 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9700 None
9701 } else {
9702 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9703 for c in conditions {
9704 sets.push(self.resolve_condition(c, snapshot)?);
9705 }
9706 Some(RowIdSet::intersect_many(sets))
9707 };
9708
9709 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9713 for (rid, (_, run_idx, pos, deleted)) in &best {
9714 if *deleted {
9715 continue;
9716 }
9717 if overlay_rids.contains(rid) {
9718 continue;
9719 }
9720 if let Some(s) = &survivors {
9721 if !s.contains(*rid) {
9722 continue;
9723 }
9724 }
9725 per_run[*run_idx].push((*rid, *pos));
9726 }
9727 for v in per_run.iter_mut() {
9728 v.sort_unstable_by_key(|&(rid, _)| rid);
9729 }
9730
9731 let mut streams = Vec::with_capacity(run_meta.len());
9733 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9734 let mut total = 0usize;
9735 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9736 let mut starts = Vec::with_capacity(page_rows.len());
9737 let mut acc = 0usize;
9738 for &r in &page_rows {
9739 starts.push(acc);
9740 acc += r;
9741 }
9742 let mut survivors_vec: Vec<(u64, usize, usize)> =
9743 Vec::with_capacity(per_run[run_idx].len());
9744 for &(rid, pos) in &per_run[run_idx] {
9745 let page_seq = match starts.partition_point(|&s| s <= pos) {
9746 0 => continue,
9747 p => p - 1,
9748 };
9749 let within = pos - starts[page_seq];
9750 survivors_vec.push((rid, page_seq, within));
9751 }
9752 total += survivors_vec.len();
9753 if let Some(&(rid, _, _)) = survivors_vec.first() {
9754 heap.push(std::cmp::Reverse((rid, run_idx)));
9755 }
9756 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9757 }
9758
9759 let overlay_rows = if overlay_rids.is_empty() {
9761 Vec::new()
9762 } else {
9763 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9764 self.overlay_visible_rows(snapshot, bound)
9765 };
9766 let overlay = if overlay_rows.is_empty() {
9767 None
9768 } else {
9769 let filtered =
9770 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9771 if filtered.is_empty() {
9772 None
9773 } else {
9774 Some(self.materialize_overlay(&filtered, &projection))
9775 }
9776 };
9777
9778 let overlay_row_count = overlay
9779 .as_ref()
9780 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9781 .unwrap_or(0);
9782 crate::trace::QueryTrace::record(|t| {
9783 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9784 t.run_count = self.run_refs.len();
9785 t.memtable_rows = self.memtable.len();
9786 t.mutable_run_rows = self.mutable_run.len();
9787 t.overlay_rows = overlay_row_count;
9788 t.conditions_pushed = conditions.len();
9789 t.survivor_count = Some(total);
9790 });
9791
9792 Ok(Some(MultiRunCursor::new(
9793 streams, projection, heap, total, overlay,
9794 )))
9795 }
9796
9797 fn overlay_materialization_bound<'a>(
9809 conditions: &[crate::query::Condition],
9810 survivors: &'a Option<RowIdSet>,
9811 ) -> Option<&'a RowIdSet> {
9812 use crate::query::Condition;
9813 let has_range = conditions
9814 .iter()
9815 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9816 if has_range {
9817 None
9818 } else {
9819 survivors.as_ref()
9820 }
9821 }
9822
9823 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9835 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9836 let mut fold = |row: Row| {
9837 if let Some(b) = bound {
9838 if !b.contains(row.row_id.0) {
9839 return;
9840 }
9841 }
9842 best.entry(row.row_id.0)
9843 .and_modify(|(be, br)| {
9844 if row.committed_epoch > *be {
9845 *be = row.committed_epoch;
9846 *br = row.clone();
9847 }
9848 })
9849 .or_insert_with(|| (row.committed_epoch, row));
9850 };
9851 for row in self.memtable.visible_versions(snapshot.epoch) {
9852 fold(row);
9853 }
9854 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9855 fold(row);
9856 }
9857 let mut out: Vec<Row> = best
9858 .into_values()
9859 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9860 .collect();
9861 out.sort_by_key(|r| r.row_id);
9862 out
9863 }
9864
9865 fn filter_overlay_rows(
9873 &self,
9874 rows: Vec<Row>,
9875 conditions: &[crate::query::Condition],
9876 survivors: Option<&RowIdSet>,
9877 snapshot: Snapshot,
9878 ) -> Result<Vec<Row>> {
9879 if conditions.is_empty() {
9880 return Ok(rows);
9881 }
9882 use crate::query::Condition;
9883 let all_index_served = !conditions
9887 .iter()
9888 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9889 if all_index_served {
9890 return Ok(rows
9891 .into_iter()
9892 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9893 .collect());
9894 }
9895 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9898 for c in conditions {
9899 let s = match c {
9900 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9901 _ => self.resolve_condition(c, snapshot)?,
9902 };
9903 per_cond_sets.push(s);
9904 }
9905 Ok(rows
9906 .into_iter()
9907 .filter(|row| {
9908 conditions.iter().enumerate().all(|(i, c)| match c {
9909 Condition::Range { column_id, lo, hi } => {
9910 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9911 }
9912 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9913 match row.columns.get(column_id) {
9914 Some(Value::Float64(v)) => {
9915 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9916 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9917 lo_ok && hi_ok
9918 }
9919 _ => false,
9920 }
9921 }
9922 _ => per_cond_sets[i].contains(row.row_id.0),
9923 })
9924 })
9925 .collect())
9926 }
9927
9928 fn materialize_overlay(
9931 &self,
9932 rows: &[Row],
9933 projection: &[(u16, TypeId)],
9934 ) -> Vec<columnar::NativeColumn> {
9935 if projection.is_empty() {
9936 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9937 }
9938 let mut cols = Vec::with_capacity(projection.len());
9939 for (cid, ty) in projection {
9940 let vals: Vec<Value> = rows
9941 .iter()
9942 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9943 .collect();
9944 cols.push(columnar::values_to_native(ty.clone(), &vals));
9945 }
9946 cols
9947 }
9948
9949 fn resolve_survivor_rids(
9954 &self,
9955 conditions: &[crate::query::Condition],
9956 reader: &mut RunReader,
9957 snapshot: Snapshot,
9958 ) -> Result<RowIdSet> {
9959 use crate::query::Condition;
9960 let mut sets: Vec<RowIdSet> = Vec::new();
9961 for c in conditions {
9962 self.validate_condition(c)?;
9963 let s: RowIdSet = match c {
9964 Condition::Pk(key) => {
9965 let lookup = self
9966 .schema
9967 .primary_key()
9968 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9969 .unwrap_or_else(|| key.clone());
9970 self.hot
9971 .get(&lookup)
9972 .map(|r| RowIdSet::one(r.0))
9973 .unwrap_or_else(RowIdSet::empty)
9974 }
9975 Condition::BitmapEq { column_id, value } => {
9976 let lookup = self.index_lookup_key_bytes(*column_id, value);
9977 self.bitmap
9978 .get(column_id)
9979 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9980 .unwrap_or_else(RowIdSet::empty)
9981 }
9982 Condition::BitmapIn { column_id, values } => {
9983 let bm = self.bitmap.get(column_id);
9984 let mut acc = roaring::RoaringBitmap::new();
9985 if let Some(b) = bm {
9986 for v in values {
9987 let lookup = self.index_lookup_key_bytes(*column_id, v);
9988 acc |= b.get(&lookup);
9989 }
9990 }
9991 RowIdSet::from_roaring(acc)
9992 }
9993 Condition::BytesPrefix { column_id, prefix } => {
9994 if let Some(b) = self.bitmap.get(column_id) {
9995 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9996 let mut acc = roaring::RoaringBitmap::new();
9997 for key in b.keys() {
9998 if key.starts_with(&lookup_prefix) {
9999 acc |= b.get(&key);
10000 }
10001 }
10002 RowIdSet::from_roaring(acc)
10003 } else {
10004 RowIdSet::empty()
10005 }
10006 }
10007 Condition::FmContains { column_id, pattern } => self
10008 .fm
10009 .get(column_id)
10010 .map(|f| {
10011 RowIdSet::from_unsorted(
10012 f.locate(pattern).into_iter().map(|r| r.0).collect(),
10013 )
10014 })
10015 .unwrap_or_else(RowIdSet::empty),
10016 Condition::FmContainsAll {
10017 column_id,
10018 patterns,
10019 } => {
10020 if let Some(f) = self.fm.get(column_id) {
10021 let sets: Vec<RowIdSet> = patterns
10022 .iter()
10023 .map(|pat| {
10024 RowIdSet::from_unsorted(
10025 f.locate(pat).into_iter().map(|r| r.0).collect(),
10026 )
10027 })
10028 .collect();
10029 RowIdSet::intersect_many(sets)
10030 } else {
10031 RowIdSet::empty()
10032 }
10033 }
10034 Condition::Ann {
10035 column_id,
10036 query,
10037 k,
10038 } => RowIdSet::from_unsorted(
10039 self.retrieve_filtered(
10040 &crate::query::Retriever::Ann {
10041 column_id: *column_id,
10042 query: query.clone(),
10043 k: *k,
10044 },
10045 snapshot,
10046 None,
10047 None,
10048 None,
10049 None,
10050 )?
10051 .into_iter()
10052 .map(|hit| hit.row_id.0)
10053 .collect(),
10054 ),
10055 Condition::SparseMatch {
10056 column_id,
10057 query,
10058 k,
10059 } => RowIdSet::from_unsorted(
10060 self.retrieve_filtered(
10061 &crate::query::Retriever::Sparse {
10062 column_id: *column_id,
10063 query: query.clone(),
10064 k: *k,
10065 },
10066 snapshot,
10067 None,
10068 None,
10069 None,
10070 None,
10071 )?
10072 .into_iter()
10073 .map(|hit| hit.row_id.0)
10074 .collect(),
10075 ),
10076 Condition::MinHashSimilar {
10077 column_id,
10078 query,
10079 k,
10080 } => match self.minhash.get(column_id) {
10081 Some(index) => {
10082 let candidates = index.candidate_row_ids(query);
10083 let eligible =
10084 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
10085 RowIdSet::from_unsorted(
10086 index
10087 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
10088 .into_iter()
10089 .map(|(row_id, _)| row_id.0)
10090 .collect(),
10091 )
10092 }
10093 None => RowIdSet::empty(),
10094 },
10095 Condition::Range { column_id, lo, hi } => {
10096 if let Some(li) = self.learned_range.get(column_id) {
10097 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10098 } else {
10099 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
10100 }
10101 }
10102 Condition::RangeF64 {
10103 column_id,
10104 lo,
10105 lo_inclusive,
10106 hi,
10107 hi_inclusive,
10108 } => {
10109 if let Some(li) = self.learned_range.get(column_id) {
10110 RowIdSet::from_unsorted(
10111 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10112 .into_iter()
10113 .collect(),
10114 )
10115 } else {
10116 reader.range_row_id_set_f64(
10117 *column_id,
10118 *lo,
10119 *lo_inclusive,
10120 *hi,
10121 *hi_inclusive,
10122 )?
10123 }
10124 }
10125 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
10126 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
10127 };
10128 sets.push(s);
10129 }
10130 Ok(RowIdSet::intersect_many(sets))
10131 }
10132
10133 pub fn scan_cursor(
10154 &self,
10155 snapshot: Snapshot,
10156 projection: Vec<(u16, TypeId)>,
10157 conditions: &[crate::query::Condition],
10158 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
10159 if self.ttl.is_some() {
10160 return Ok(None);
10161 }
10162 if !conditions.is_empty() && !self.indexes_complete {
10168 return Ok(None);
10169 }
10170 if self.run_refs.len() == 1 {
10171 Ok(self
10172 .native_page_cursor(snapshot, projection, conditions)?
10173 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10174 } else {
10175 Ok(self
10176 .native_multi_run_cursor(snapshot, projection, conditions)?
10177 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10178 }
10179 }
10180
10181 pub fn aggregate_native(
10195 &self,
10196 snapshot: Snapshot,
10197 column: Option<u16>,
10198 conditions: &[crate::query::Condition],
10199 agg: NativeAgg,
10200 ) -> Result<Option<NativeAggResult>> {
10201 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
10202 }
10203
10204 pub fn aggregate_native_with_control(
10205 &self,
10206 snapshot: Snapshot,
10207 column: Option<u16>,
10208 conditions: &[crate::query::Condition],
10209 agg: NativeAgg,
10210 control: &crate::ExecutionControl,
10211 ) -> Result<Option<NativeAggResult>> {
10212 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
10213 }
10214
10215 fn aggregate_native_inner(
10216 &self,
10217 snapshot: Snapshot,
10218 column: Option<u16>,
10219 conditions: &[crate::query::Condition],
10220 agg: NativeAgg,
10221 control: Option<&crate::ExecutionControl>,
10222 ) -> Result<Option<NativeAggResult>> {
10223 execution_checkpoint(control, 0)?;
10224 if self.ttl.is_some() {
10225 return Ok(None);
10226 }
10227 if self.run_refs.len() == 1 && conditions.is_empty() {
10229 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
10230 return Ok(Some(res));
10231 }
10232 }
10233 if matches!(agg, NativeAgg::Count) && column.is_none() {
10237 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
10238 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
10239 }
10240 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10241 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
10242 }
10243 let cid = match column {
10246 Some(c) => c,
10247 None => return Ok(None),
10248 };
10249 let ty = self.column_type(cid);
10250 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
10251 execution_checkpoint(control, 0)?;
10252 return match ty {
10253 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10254 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
10255 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10256 }
10257 TypeId::Float64 => {
10258 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
10259 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10260 }
10261 _ => Ok(None),
10262 };
10263 }
10264 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10266 execution_checkpoint(control, 0)?;
10267 match ty {
10268 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10269 let mut count = 0u64;
10270 let mut sum = 0i128;
10271 let mut mn = i64::MAX;
10272 let mut mx = i64::MIN;
10273 for row in &rows {
10274 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
10275 count += 1;
10276 sum += i128::from(*v);
10277 mn = mn.min(*v);
10278 mx = mx.max(*v);
10279 }
10280 }
10281 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10282 }
10283 TypeId::Float64 => {
10284 let mut count = 0u64;
10285 let mut sum = 0.0f64;
10286 let mut mn = f64::INFINITY;
10287 let mut mx = f64::NEG_INFINITY;
10288 for row in &rows {
10289 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
10290 count += 1;
10291 sum += *v;
10292 mn = mn.min(*v);
10293 mx = mx.max(*v);
10294 }
10295 }
10296 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10297 }
10298 _ => Ok(None),
10299 }
10300 }
10301
10302 fn visible_rows_filtered(
10304 &self,
10305 snapshot: Snapshot,
10306 conditions: &[crate::query::Condition],
10307 control: Option<&crate::ExecutionControl>,
10308 ) -> Result<Vec<Row>> {
10309 let rows = if let Some(control) = control {
10310 self.visible_rows_controlled(snapshot, control)?
10311 } else {
10312 self.visible_rows(snapshot)?
10313 };
10314 if conditions.is_empty() {
10315 return Ok(rows);
10316 }
10317 Ok(rows
10318 .into_iter()
10319 .filter(|row| {
10320 conditions
10321 .iter()
10322 .all(|cond| condition_matches_row(cond, row, &self.schema))
10323 })
10324 .collect())
10325 }
10326
10327 fn aggregate_from_stats(
10335 &self,
10336 snapshot: Snapshot,
10337 column: Option<u16>,
10338 agg: NativeAgg,
10339 ) -> Result<Option<NativeAggResult>> {
10340 let cid = match (agg, column) {
10341 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
10342 _ => return Ok(None), };
10344 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
10345 return Ok(None);
10346 };
10347 let Some(cs) = stats.get(&cid) else {
10348 return Ok(None);
10349 };
10350 match agg {
10351 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
10353 self.live_count.saturating_sub(cs.null_count),
10354 ))),
10355 NativeAgg::Min | NativeAgg::Max => {
10356 let bound = if agg == NativeAgg::Min {
10357 &cs.min
10358 } else {
10359 &cs.max
10360 };
10361 match bound {
10362 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
10363 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
10364 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
10369 None => Ok(None),
10370 }
10371 }
10372 _ => Ok(None),
10373 }
10374 }
10375
10376 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
10385 if self.ttl.is_some() {
10386 return Ok(None);
10387 }
10388 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
10389 return Ok(None);
10390 }
10391 self.ensure_indexes_complete()?;
10394 let reader = self.open_reader(self.run_refs[0].run_id)?;
10395 if self.live_count != reader.row_count() as u64 {
10396 return Ok(None);
10397 }
10398 let Some(bm) = self.bitmap.get(&column_id) else {
10399 return Ok(None); };
10401 let mut distinct = bm.value_count() as u64;
10402 if !bm.get(&Value::Null.encode_key()).is_empty() {
10405 distinct = distinct.saturating_sub(1);
10406 }
10407 Ok(Some(distinct))
10408 }
10409
10410 pub fn aggregate_incremental(
10422 &mut self,
10423 cache_key: u64,
10424 conditions: &[crate::query::Condition],
10425 column: Option<u16>,
10426 agg: NativeAgg,
10427 ) -> Result<IncrementalAggResult> {
10428 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
10429 }
10430
10431 pub fn aggregate_incremental_with_control(
10432 &mut self,
10433 cache_key: u64,
10434 conditions: &[crate::query::Condition],
10435 column: Option<u16>,
10436 agg: NativeAgg,
10437 control: &crate::ExecutionControl,
10438 ) -> Result<IncrementalAggResult> {
10439 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
10440 }
10441
10442 fn aggregate_incremental_inner(
10443 &mut self,
10444 cache_key: u64,
10445 conditions: &[crate::query::Condition],
10446 column: Option<u16>,
10447 agg: NativeAgg,
10448 control: Option<&crate::ExecutionControl>,
10449 ) -> Result<IncrementalAggResult> {
10450 execution_checkpoint(control, 0)?;
10451 let snap = self.snapshot();
10452 let cur_wm = self.allocator.current().0;
10453 let cur_epoch = snap.epoch.0;
10454 let incremental_ok = self.ttl.is_none()
10461 && !self.had_deletes
10462 && self.memtable.is_empty()
10463 && self.mutable_run.is_empty();
10464
10465 if incremental_ok {
10468 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10469 if cached.epoch == cur_epoch {
10470 return Ok(IncrementalAggResult {
10471 state: cached.state,
10472 incremental: true,
10473 delta_rows: 0,
10474 });
10475 }
10476 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10477 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10478 let mut delta_rids = Vec::with_capacity(delta_len);
10479 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10480 execution_checkpoint(control, index)?;
10481 delta_rids.push(row_id);
10482 }
10483 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10484 execution_checkpoint(control, 0)?;
10485 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10486 let delta_state = agg_state_from_rows(
10487 &delta_rows,
10488 conditions,
10489 &index_sets,
10490 column,
10491 agg,
10492 &self.schema,
10493 control,
10494 )?;
10495 let merged = cached.state.merge(delta_state);
10496 let delta_n = delta_rids.len() as u64;
10497 Arc::make_mut(&mut self.agg_cache).insert(
10498 cache_key,
10499 CachedAgg {
10500 state: merged.clone(),
10501 watermark: cur_wm,
10502 epoch: cur_epoch,
10503 },
10504 );
10505 return Ok(IncrementalAggResult {
10506 state: merged,
10507 incremental: true,
10508 delta_rows: delta_n,
10509 });
10510 }
10511 }
10512 }
10513
10514 let cursor_ok =
10519 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10520 let state = if cursor_ok && agg != NativeAgg::Avg {
10521 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10522 Some(result) => {
10523 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10524 }
10525 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10526 }
10527 } else {
10528 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10529 };
10530 if incremental_ok {
10532 Arc::make_mut(&mut self.agg_cache).insert(
10533 cache_key,
10534 CachedAgg {
10535 state: state.clone(),
10536 watermark: cur_wm,
10537 epoch: cur_epoch,
10538 },
10539 );
10540 }
10541 Ok(IncrementalAggResult {
10542 state,
10543 incremental: false,
10544 delta_rows: 0,
10545 })
10546 }
10547
10548 fn agg_state_full_scan(
10551 &self,
10552 conditions: &[crate::query::Condition],
10553 column: Option<u16>,
10554 agg: NativeAgg,
10555 snap: Snapshot,
10556 control: Option<&crate::ExecutionControl>,
10557 ) -> Result<AggState> {
10558 execution_checkpoint(control, 0)?;
10559 let rows = self.visible_rows(snap)?;
10560 execution_checkpoint(control, 0)?;
10561 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10562 agg_state_from_rows(
10563 &rows,
10564 conditions,
10565 &index_sets,
10566 column,
10567 agg,
10568 &self.schema,
10569 control,
10570 )
10571 }
10572
10573 fn resolve_index_conditions(
10576 &self,
10577 conditions: &[crate::query::Condition],
10578 snapshot: Snapshot,
10579 ) -> Result<Vec<RowIdSet>> {
10580 use crate::query::Condition;
10581 let mut sets = Vec::new();
10582 for c in conditions {
10583 if matches!(
10584 c,
10585 Condition::Ann { .. }
10586 | Condition::SparseMatch { .. }
10587 | Condition::MinHashSimilar { .. }
10588 ) {
10589 sets.push(self.resolve_condition(c, snapshot)?);
10590 }
10591 }
10592 Ok(sets)
10593 }
10594
10595 fn column_type(&self, cid: u16) -> TypeId {
10596 self.schema
10597 .columns
10598 .iter()
10599 .find(|c| c.id == cid)
10600 .map(|c| c.ty.clone())
10601 .unwrap_or(TypeId::Bytes)
10602 }
10603
10604 pub fn approx_aggregate(
10613 &mut self,
10614 conditions: &[crate::query::Condition],
10615 column: Option<u16>,
10616 agg: ApproxAgg,
10617 z: f64,
10618 ) -> Result<Option<ApproxResult>> {
10619 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10620 }
10621
10622 pub fn approx_aggregate_with_candidate_authorization(
10625 &mut self,
10626 conditions: &[crate::query::Condition],
10627 column: Option<u16>,
10628 agg: ApproxAgg,
10629 z: f64,
10630 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10631 ) -> Result<Option<ApproxResult>> {
10632 use crate::query::Condition;
10633 self.ensure_reservoir_complete()?;
10634 let mut snapshot = self.snapshot();
10639 let n_pop = self.count();
10640 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10641 if sample_rids.is_empty() {
10642 return Ok(None);
10643 }
10644 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10646 if live_sample.is_empty() {
10647 snapshot = Snapshot::unbounded();
10648 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10649 }
10650 let s = live_sample.len();
10651 if s == 0 {
10652 return Ok(None);
10653 }
10654 let authorized = authorization
10655 .map(|authorization| {
10656 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10657 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10658 })
10659 .transpose()?;
10660
10661 let mut index_sets: Vec<RowIdSet> = Vec::new();
10664 for c in conditions {
10665 if matches!(
10666 c,
10667 Condition::Ann { .. }
10668 | Condition::SparseMatch { .. }
10669 | Condition::MinHashSimilar { .. }
10670 ) {
10671 index_sets.push(self.resolve_condition(c, snapshot)?);
10672 }
10673 }
10674
10675 let cid = match (agg, column) {
10677 (ApproxAgg::Count, _) => None,
10678 (_, Some(c)) => Some(c),
10679 _ => return Ok(None),
10680 };
10681 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10682 for r in &live_sample {
10683 if authorized
10684 .as_ref()
10685 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10686 {
10687 continue;
10688 }
10689 if !conditions
10691 .iter()
10692 .all(|c| condition_matches_row(c, r, &self.schema))
10693 {
10694 continue;
10695 }
10696 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10698 continue;
10699 }
10700 if let Some(cid) = cid {
10701 let mut cells = r
10702 .columns
10703 .get(&cid)
10704 .cloned()
10705 .map(|value| vec![(cid, value)])
10706 .unwrap_or_default();
10707 if let Some(authorization) = authorization {
10708 authorization.security.apply_masks_to_cells(
10709 authorization.table,
10710 &mut cells,
10711 authorization.principal,
10712 );
10713 }
10714 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10715 passing_vals.push(v);
10716 } } else {
10718 passing_vals.push(0.0); }
10720 }
10721 let m = passing_vals.len();
10722
10723 let (point, half) = match agg {
10724 ApproxAgg::Count => {
10725 let p = m as f64 / s as f64;
10727 let point = n_pop as f64 * p;
10728 let var = if s > 1 {
10729 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10730 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10731 } else {
10732 0.0
10733 };
10734 (point, z * var.sqrt())
10735 }
10736 ApproxAgg::Sum => {
10737 let y: Vec<f64> = live_sample
10739 .iter()
10740 .map(|r| {
10741 let passes_row = authorized
10742 .as_ref()
10743 .is_none_or(|authorized| authorized.contains(&r.row_id))
10744 && conditions
10745 .iter()
10746 .all(|c| condition_matches_row(c, r, &self.schema))
10747 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10748 if passes_row {
10749 cid.and_then(|cid| {
10750 let mut cells = r
10751 .columns
10752 .get(&cid)
10753 .cloned()
10754 .map(|value| vec![(cid, value)])
10755 .unwrap_or_default();
10756 if let Some(authorization) = authorization {
10757 authorization.security.apply_masks_to_cells(
10758 authorization.table,
10759 &mut cells,
10760 authorization.principal,
10761 );
10762 }
10763 as_f64(cells.first().map(|(_, value)| value))
10764 })
10765 .unwrap_or(0.0)
10766 } else {
10767 0.0
10768 }
10769 })
10770 .collect();
10771 let mean_y = y.iter().sum::<f64>() / s as f64;
10772 let point = n_pop as f64 * mean_y;
10773 let var = if s > 1 {
10774 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10775 let var_y = ss / (s - 1) as f64;
10776 n_pop as f64 * n_pop as f64 * var_y / s as f64
10777 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10778 } else {
10779 0.0
10780 };
10781 (point, z * var.sqrt())
10782 }
10783 ApproxAgg::Avg => {
10784 if m == 0 {
10785 return Ok(Some(ApproxResult {
10786 point: 0.0,
10787 ci_low: 0.0,
10788 ci_high: 0.0,
10789 n_population: n_pop,
10790 n_sample_live: s,
10791 n_passing: 0,
10792 }));
10793 }
10794 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10795 let half = if m > 1 {
10796 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10797 let sd = (ss / (m - 1) as f64).sqrt();
10798 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10799 z * sd / (m as f64).sqrt() * fpc.sqrt()
10800 } else {
10801 0.0
10802 };
10803 (mean, half)
10804 }
10805 };
10806
10807 Ok(Some(ApproxResult {
10808 point,
10809 ci_low: point - half,
10810 ci_high: point + half,
10811 n_population: n_pop,
10812 n_sample_live: s,
10813 n_passing: m,
10814 }))
10815 }
10816
10817 pub fn exact_column_stats(
10825 &self,
10826 _snapshot: Snapshot,
10827 projection: &[u16],
10828 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10829 if self.ttl.is_some()
10830 || !(self.memtable.is_empty()
10831 && self.mutable_run.is_empty()
10832 && self.run_refs.len() == 1)
10833 {
10834 return Ok(None);
10835 }
10836 let reader = self.open_reader(self.run_refs[0].run_id)?;
10837 if self.live_count != reader.row_count() as u64 {
10838 return Ok(None);
10839 }
10840 let mut out = HashMap::new();
10841 for &cid in projection {
10842 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10843 Some(c) => c,
10844 None => continue,
10845 };
10846 let Some(stats) = reader.column_page_stats(cid) else {
10848 out.insert(
10849 cid,
10850 ColumnStat {
10851 min: None,
10852 max: None,
10853 null_count: self.live_count,
10854 },
10855 );
10856 continue;
10857 };
10858 let stat = match cdef.ty {
10859 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10860 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10861 min: mn.map(Value::Int64),
10862 max: mx.map(Value::Int64),
10863 null_count: n,
10864 })
10865 }
10866 TypeId::Float64 => {
10867 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10868 min: mn.map(Value::Float64),
10869 max: mx.map(Value::Float64),
10870 null_count: n,
10871 })
10872 }
10873 _ => None,
10874 };
10875 if let Some(s) = stat {
10876 out.insert(cid, s);
10877 }
10878 }
10879 Ok(Some(out))
10880 }
10881
10882 pub fn dir(&self) -> &Path {
10883 &self.dir
10884 }
10885
10886 pub fn schema(&self) -> &Schema {
10887 &self.schema
10888 }
10889
10890 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
10891 self.ann.get(&column_id)
10892 }
10893
10894 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
10895 self.ann.get_mut(&column_id)
10896 }
10897
10898 pub(crate) fn set_catalog_name(&mut self, name: String) {
10899 self.name = name;
10900 }
10901
10902 pub(crate) fn prepare_alter_column(
10903 &mut self,
10904 column_name: &str,
10905 change: &AlterColumn,
10906 ) -> Result<(ColumnDef, Option<Schema>)> {
10907 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10908 return Err(MongrelError::InvalidArgument(
10909 "ALTER COLUMN requires committing staged writes first".into(),
10910 ));
10911 }
10912 let old = self
10913 .schema
10914 .columns
10915 .iter()
10916 .find(|c| c.name == column_name)
10917 .cloned()
10918 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10919 let mut next = old.clone();
10920
10921 if let Some(name) = &change.name {
10922 let trimmed = name.trim();
10923 if trimmed.is_empty() {
10924 return Err(MongrelError::InvalidArgument(
10925 "ALTER COLUMN name must not be empty".into(),
10926 ));
10927 }
10928 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10929 return Err(MongrelError::Schema(format!(
10930 "column {trimmed} already exists"
10931 )));
10932 }
10933 next.name = trimmed.to_string();
10934 }
10935
10936 if let Some(ty) = &change.ty {
10937 next.ty = ty.clone();
10938 }
10939 if let Some(flags) = change.flags {
10940 validate_alter_column_flags(old.flags, flags)?;
10941 next.flags = flags;
10942 }
10943
10944 if let Some(default_change) = &change.default_value {
10945 next.default_value = default_change.clone();
10946 }
10947 if let Some(source_change) = &change.embedding_source {
10948 next.embedding_source = source_change.clone();
10949 }
10950
10951 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10952 if old.flags.contains(ColumnFlags::NULLABLE)
10953 && !next.flags.contains(ColumnFlags::NULLABLE)
10954 && self.column_has_nulls(old.id)?
10955 {
10956 return Err(MongrelError::InvalidArgument(format!(
10957 "column '{}' contains NULL values",
10958 old.name
10959 )));
10960 }
10961 if next == old {
10962 return Ok((next, None));
10963 }
10964 let mut schema = self.schema.clone();
10965 let index = schema
10966 .columns
10967 .iter()
10968 .position(|column| column.id == next.id)
10969 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10970 schema.columns[index] = next.clone();
10971 schema.schema_id = schema
10972 .schema_id
10973 .checked_add(1)
10974 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10975 schema.validate_auto_increment()?;
10976 schema.validate_defaults()?;
10977 Ok((next, Some(schema)))
10978 }
10979
10980 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10981 self.schema = schema;
10982 self.auto_inc = resolve_auto_inc(&self.schema);
10983 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10984 self.clear_result_cache();
10985 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10986 }
10987
10988 pub(crate) fn publish_index_schema_change(
10992 &mut self,
10993 schema: Schema,
10994 artifact: SecondaryIndexArtifact,
10995 ) -> Result<()> {
10996 self.apply_altered_schema_prepared(schema);
10997 self.prune_index_maps_to_schema();
10998 match artifact {
10999 SecondaryIndexArtifact::Bitmap(column_id, index) => {
11000 self.bitmap.insert(column_id, index);
11001 }
11002 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
11003 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
11004 }
11005 SecondaryIndexArtifact::Fm(column_id, index) => {
11006 self.fm.insert(column_id, *index);
11007 }
11008 SecondaryIndexArtifact::Ann(column_id, index) => {
11009 self.ann.insert(column_id, *index);
11010 }
11011 SecondaryIndexArtifact::Sparse(column_id, index) => {
11012 self.sparse.insert(column_id, index);
11013 }
11014 SecondaryIndexArtifact::MinHash(column_id, index) => {
11015 self.minhash.insert(column_id, index);
11016 }
11017 }
11018 self.indexes_complete = true;
11019 self.seal_generations();
11020 let view = Arc::new(self.capture_read_generation());
11021 self.published.store(Arc::clone(&view));
11022 checkpoint_current_schema(self)?;
11023 self.invalidate_index_checkpoint();
11026 Ok(())
11027 }
11028
11029 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
11035 self.apply_altered_schema_prepared(schema);
11036 self.prune_index_maps_to_schema();
11037 self.indexes_complete = true;
11038 self.seal_generations();
11039 let view = Arc::new(self.capture_read_generation());
11040 self.published.store(Arc::clone(&view));
11041 checkpoint_current_schema(self)?;
11042 self.invalidate_index_checkpoint();
11044 Ok(())
11045 }
11046
11047 fn prune_index_maps_to_schema(&mut self) {
11048 let keep_bitmap: std::collections::HashSet<u16> = self
11049 .schema
11050 .indexes
11051 .iter()
11052 .filter(|index| index.kind == IndexKind::Bitmap)
11053 .map(|index| index.column_id)
11054 .collect();
11055 let keep_ann: std::collections::HashSet<u16> = self
11056 .schema
11057 .indexes
11058 .iter()
11059 .filter(|index| index.kind == IndexKind::Ann)
11060 .map(|index| index.column_id)
11061 .collect();
11062 let keep_fm: std::collections::HashSet<u16> = self
11063 .schema
11064 .indexes
11065 .iter()
11066 .filter(|index| index.kind == IndexKind::FmIndex)
11067 .map(|index| index.column_id)
11068 .collect();
11069 let keep_sparse: std::collections::HashSet<u16> = self
11070 .schema
11071 .indexes
11072 .iter()
11073 .filter(|index| index.kind == IndexKind::Sparse)
11074 .map(|index| index.column_id)
11075 .collect();
11076 let keep_minhash: std::collections::HashSet<u16> = self
11077 .schema
11078 .indexes
11079 .iter()
11080 .filter(|index| index.kind == IndexKind::MinHash)
11081 .map(|index| index.column_id)
11082 .collect();
11083 let keep_learned: std::collections::HashSet<u16> = self
11084 .schema
11085 .indexes
11086 .iter()
11087 .filter(|index| index.kind == IndexKind::LearnedRange)
11088 .map(|index| index.column_id)
11089 .collect();
11090 self.bitmap
11091 .retain(|column_id, _| keep_bitmap.contains(column_id));
11092 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
11093 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
11094 self.sparse
11095 .retain(|column_id, _| keep_sparse.contains(column_id));
11096 self.minhash
11097 .retain(|column_id, _| keep_minhash.contains(column_id));
11098 {
11099 let learned = Arc::make_mut(&mut self.learned_range);
11100 learned.retain(|column_id, _| keep_learned.contains(column_id));
11101 }
11102 }
11103
11104 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
11105 checkpoint_current_schema(self)
11106 }
11107
11108 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
11109 self.ensure_writable()?;
11110 let previous_schema = self.schema.clone();
11111 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
11112 if let Some(schema) = schema {
11113 self.apply_altered_schema_prepared(schema);
11114 self.checkpoint_standalone_schema_change(previous_schema)?;
11115 }
11116 Ok(column)
11117 }
11118
11119 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
11120 if self.live_count == 0 {
11121 return Ok(false);
11122 }
11123 let snap = self.snapshot();
11124 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
11125 Ok(columns
11126 .first()
11127 .map(|(_, col)| col.null_count(col.len()) != 0)
11128 .unwrap_or(true))
11129 }
11130
11131 fn has_stored_versions(&self) -> bool {
11132 !self.memtable.is_empty()
11133 || !self.mutable_run.is_empty()
11134 || self.run_refs.iter().any(|r| r.row_count > 0)
11135 || !self.retiring.is_empty()
11136 }
11137
11138 pub fn add_column(
11143 &mut self,
11144 name: &str,
11145 ty: TypeId,
11146 flags: ColumnFlags,
11147 default_value: Option<crate::schema::DefaultExpr>,
11148 ) -> Result<u16> {
11149 self.add_column_with_id(name, ty, flags, default_value, None)
11150 }
11151
11152 pub fn add_column_with_id(
11153 &mut self,
11154 name: &str,
11155 ty: TypeId,
11156 flags: ColumnFlags,
11157 default_value: Option<crate::schema::DefaultExpr>,
11158 requested_id: Option<u16>,
11159 ) -> Result<u16> {
11160 self.ensure_writable()?;
11161 let previous_schema = self.schema.clone();
11162 let (column, schema) =
11163 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
11164 self.apply_altered_schema_prepared(schema);
11165 self.checkpoint_standalone_schema_change(previous_schema)?;
11166 Ok(column.id)
11167 }
11168
11169 pub(crate) fn prepare_add_column(
11170 &mut self,
11171 name: &str,
11172 ty: TypeId,
11173 flags: ColumnFlags,
11174 default_value: Option<crate::schema::DefaultExpr>,
11175 requested_id: Option<u16>,
11176 ) -> Result<(ColumnDef, Schema)> {
11177 self.ensure_writable()?;
11178 if self.schema.columns.iter().any(|c| c.name == name) {
11179 return Err(MongrelError::Schema(format!(
11180 "column {name} already exists"
11181 )));
11182 }
11183 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
11184 if self.schema.columns.iter().any(|c| c.id == id) {
11185 return Err(MongrelError::Schema(format!(
11186 "column id {id} already exists"
11187 )));
11188 }
11189 id
11190 } else {
11191 self.schema
11192 .columns
11193 .iter()
11194 .map(|c| c.id)
11195 .max()
11196 .unwrap_or(0)
11197 .checked_add(1)
11198 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
11199 };
11200 let column = ColumnDef {
11201 id,
11202 name: name.to_string(),
11203 ty,
11204 flags,
11205 default_value,
11206 embedding_source: None,
11207 };
11208 let mut next_schema = self.schema.clone();
11209 next_schema.columns.push(column.clone());
11210 next_schema.schema_id = next_schema
11211 .schema_id
11212 .checked_add(1)
11213 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11214 next_schema.validate_auto_increment()?;
11215 next_schema.validate_defaults()?;
11216 Ok((column, next_schema))
11217 }
11218
11219 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
11228 self.ensure_writable()?;
11229 let cid = self
11230 .schema
11231 .columns
11232 .iter()
11233 .find(|c| c.name == column_name)
11234 .map(|c| c.id)
11235 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11236 let ty = self
11237 .schema
11238 .columns
11239 .iter()
11240 .find(|c| c.id == cid)
11241 .map(|c| c.ty.clone())
11242 .unwrap_or(TypeId::Int64);
11243 if !matches!(
11244 ty,
11245 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
11246 ) {
11247 return Err(MongrelError::Schema(format!(
11248 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
11249 )));
11250 }
11251 if self
11252 .schema
11253 .indexes
11254 .iter()
11255 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
11256 {
11257 return Ok(()); }
11259 let previous_schema = self.schema.clone();
11260 let previous_learned_range = Arc::clone(&self.learned_range);
11261 let mut next_schema = previous_schema.clone();
11262 next_schema.indexes.push(IndexDef {
11263 name: format!("{}_learned_range", column_name),
11264 column_id: cid,
11265 kind: IndexKind::LearnedRange,
11266 predicate: None,
11267 options: Default::default(),
11268 });
11269 next_schema.schema_id = next_schema
11270 .schema_id
11271 .checked_add(1)
11272 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11273 self.apply_altered_schema_prepared(next_schema);
11274 if let Err(error) = self.build_learned_ranges() {
11275 self.apply_altered_schema_prepared(previous_schema);
11276 self.learned_range = previous_learned_range;
11277 return Err(error);
11278 }
11279 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
11280 if !matches!(
11281 &error,
11282 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
11283 ) {
11284 self.learned_range = previous_learned_range;
11285 }
11286 return Err(error);
11287 }
11288 Ok(())
11289 }
11290
11291 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
11292 let mut schema_published = false;
11293 let schema_result = match self._root_guard.as_deref() {
11294 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
11295 schema_published = true;
11296 }),
11297 None => write_schema_with_after(&self.dir, &self.schema, || {
11298 schema_published = true;
11299 }),
11300 };
11301 if schema_result.is_err() && !schema_published {
11302 self.apply_altered_schema_prepared(previous_schema);
11303 return schema_result;
11304 }
11305
11306 let manifest_result = self.persist_manifest(self.current_epoch());
11307 match (schema_result, manifest_result) {
11308 (_, Ok(())) => Ok(()),
11309 (Ok(()), Err(error)) => {
11310 self.poison_after_maintenance_publish_failure();
11311 Err(MongrelError::DurableCommit {
11312 epoch: self.current_epoch().0,
11313 message: format!(
11314 "schema is durable but matching manifest publication failed: {error}"
11315 ),
11316 })
11317 }
11318 (Err(schema_error), Err(manifest_error)) => {
11319 self.poison_after_maintenance_publish_failure();
11320 Err(MongrelError::CommitOutcomeUnknown {
11321 epoch: self.current_epoch().0,
11322 message: format!(
11323 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11324 ),
11325 })
11326 }
11327 }
11328 }
11329
11330 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
11333 self.sync_byte_threshold = threshold;
11334 if let WalSink::Private(w) = &mut self.wal {
11335 w.set_sync_byte_threshold(threshold);
11336 }
11337 }
11338
11339 pub fn page_cache_flush(&self) {
11343 self.page_cache.flush_to_disk();
11344 }
11345
11346 pub fn page_cache_len(&self) -> usize {
11348 self.page_cache.len()
11349 }
11350
11351 pub fn decoded_cache_len(&self) -> usize {
11354 self.decoded_cache.len()
11355 }
11356
11357 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
11360 self.memtable.drain_sorted()
11361 }
11362
11363 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
11364 self.runs_dir().join(format!("r-{run_id}.sr"))
11365 }
11366
11367 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
11368 match self.runs_root.as_deref() {
11369 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
11370 None => Ok(None),
11371 }
11372 }
11373
11374 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
11375 match self.runs_root.as_deref() {
11376 Some(root) => Ok(Some(root.create_regular_new(name)?)),
11377 None => Ok(None),
11378 }
11379 }
11380
11381 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
11382 match self.runs_root.as_deref() {
11383 Some(root) => match root.remove_file(name) {
11384 Ok(()) => Ok(()),
11385 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11386 Err(error) => Err(error.into()),
11387 },
11388 None => match std::fs::remove_file(self.runs_dir().join(name)) {
11389 Ok(()) => Ok(()),
11390 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11391 Err(error) => Err(error.into()),
11392 },
11393 }
11394 }
11395
11396 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
11397 match self.runs_root.as_deref() {
11398 Some(root) => root
11399 .rename_file_new(source, destination)
11400 .map_err(Into::into),
11401 None => crate::durable_file::rename(
11402 &self.runs_dir().join(source),
11403 &self.runs_dir().join(destination),
11404 )
11405 .map_err(Into::into),
11406 }
11407 }
11408
11409 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11410 self.run_refs.iter().map(|run| run.run_id)
11411 }
11412
11413 pub(crate) fn table_dir(&self) -> &Path {
11414 &self.dir
11415 }
11416
11417 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
11418 &self.schema
11419 }
11420
11421 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
11422 let id = self.next_run_id;
11423 self.next_run_id = self
11424 .next_run_id
11425 .checked_add(1)
11426 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
11427 Ok(id)
11428 }
11429
11430 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
11431 self.run_refs.push(run_ref);
11432 }
11433
11434 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
11444 self.retiring.push(crate::manifest::RetiredRun {
11445 run_id,
11446 retire_epoch,
11447 });
11448 }
11449
11450 pub(crate) fn reap_retiring(
11454 &mut self,
11455 min_active: Epoch,
11456 backup_pinned: &std::collections::HashSet<u128>,
11457 ) -> Result<usize> {
11458 if self.retiring.is_empty() {
11459 return Ok(0);
11460 }
11461 let mut reaped = 0;
11462 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11463 for r in std::mem::take(&mut self.retiring) {
11469 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11470 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11471 reaped += 1;
11472 } else {
11473 kept.push(r);
11474 }
11475 }
11476 self.retiring = kept;
11477 if reaped > 0 {
11478 self.persist_manifest(self.current_epoch())?;
11479 }
11480 Ok(reaped)
11481 }
11482
11483 pub(crate) fn has_reapable_retiring(
11484 &self,
11485 min_active: Epoch,
11486 backup_pinned: &std::collections::HashSet<u128>,
11487 ) -> bool {
11488 self.retiring
11489 .iter()
11490 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11491 }
11492
11493 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11494 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11495 return false;
11496 }
11497 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11498 self.run_refs.push(run_ref);
11499 self.indexes_complete = false;
11500 true
11501 }
11502
11503 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11504 self.kek.as_ref()
11505 }
11506
11507 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11508 let mut reader = match self.runs_root.as_deref() {
11509 Some(root) => RunReader::open_file_with_cache(
11510 root.open_regular(format!("r-{run_id}.sr"))?,
11511 self.schema.clone(),
11512 self.kek.clone(),
11513 Some(self.page_cache.clone()),
11514 Some(self.decoded_cache.clone()),
11515 self.table_id,
11516 Some(&self.verified_runs),
11517 None,
11518 )?,
11519 None => RunReader::open_with_cache(
11520 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11521 self.schema.clone(),
11522 self.kek.clone(),
11523 Some(self.page_cache.clone()),
11524 Some(self.decoded_cache.clone()),
11525 self.table_id,
11526 Some(&self.verified_runs),
11527 )?,
11528 };
11529 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11533 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11534 }
11535 Ok(reader)
11536 }
11537
11538 pub(crate) fn run_refs(&self) -> &[RunRef] {
11539 &self.run_refs
11540 }
11541
11542 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11543 self.retiring.iter().map(|run| run.run_id)
11544 }
11545
11546 pub(crate) fn runs_dir(&self) -> PathBuf {
11547 self.runs_root
11548 .as_deref()
11549 .and_then(|root| root.io_path().ok())
11550 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11551 }
11552
11553 pub(crate) fn wal_dir(&self) -> PathBuf {
11554 self.dir.join(WAL_DIR)
11555 }
11556
11557 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11558 self.run_refs = refs;
11559 }
11560
11561 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11562 self.compaction_zstd_level
11563 }
11564
11565 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11566 self.kek.clone()
11567 }
11568
11569 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11573 self.kek.as_ref().map(|k| k.derive_idx_key())
11574 }
11575
11576 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11580 self.kek.as_ref().map(|k| *k.derive_meta_key())
11581 }
11582
11583 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11586 self.column_keys
11587 .iter()
11588 .map(|(&id, &(_, scheme))| (id, scheme))
11589 .collect()
11590 }
11591
11592 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11597 self.tokenize_value_enc(column_id, v)
11598 }
11599
11600 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11601 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11602 let (key, scheme) = self.column_keys.get(&column_id)?;
11603 let token: Vec<u8> = match (*scheme, v) {
11604 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11605 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11606 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11607 _ => hmac_token(key, &v.encode_key()).to_vec(),
11608 };
11609 Some(Value::Bytes(token))
11610 }
11611
11612 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11614 self.index_lookup_key_bytes(column_id, &v.encode_key())
11615 }
11616
11617 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11620 {
11621 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11622 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11623 if *scheme == SCHEME_HMAC_EQ {
11624 return hmac_token(key, encoded).to_vec();
11625 }
11626 }
11627 }
11628 let _ = column_id;
11629 encoded.to_vec()
11630 }
11631}
11632
11633fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11634 let columnar::NativeColumn::Int64 { data, validity } = col else {
11635 return false;
11636 };
11637 if data.len() < n || !columnar::all_non_null(validity, n) {
11638 return false;
11639 }
11640 data.iter()
11641 .take(n)
11642 .zip(data.iter().skip(1))
11643 .all(|(a, b)| a < b)
11644}
11645
11646#[derive(Debug, Clone)]
11650pub struct ColumnStat {
11651 pub min: Option<Value>,
11652 pub max: Option<Value>,
11653 pub null_count: u64,
11654}
11655
11656#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11658pub enum NativeAgg {
11659 Count,
11660 Sum,
11661 Min,
11662 Max,
11663 Avg,
11664}
11665
11666#[derive(Debug, Clone, PartialEq)]
11668pub enum NativeAggResult {
11669 Count(u64),
11670 Int(i64),
11671 Float(f64),
11672 Null,
11674}
11675
11676#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11678pub enum ApproxAgg {
11679 Count,
11680 Sum,
11681 Avg,
11682}
11683
11684#[derive(Debug, Clone)]
11688pub struct ApproxResult {
11689 pub point: f64,
11691 pub ci_low: f64,
11693 pub ci_high: f64,
11695 pub n_population: u64,
11697 pub n_sample_live: usize,
11699 pub n_passing: usize,
11701}
11702
11703#[derive(Debug, Clone, PartialEq)]
11708pub enum AggState {
11709 Count(u64),
11711 SumI {
11713 sum: i128,
11714 count: u64,
11715 },
11716 SumF {
11718 sum: f64,
11719 count: u64,
11720 },
11721 AvgI {
11723 sum: i128,
11724 count: u64,
11725 },
11726 AvgF {
11728 sum: f64,
11729 count: u64,
11730 },
11731 MinI(i64),
11733 MaxI(i64),
11734 MinF(f64),
11736 MaxF(f64),
11737 Empty,
11739}
11740
11741impl AggState {
11742 pub fn merge(self, other: AggState) -> AggState {
11744 use AggState::*;
11745 match (self, other) {
11746 (Empty, x) | (x, Empty) => x,
11747 (Count(a), Count(b)) => Count(a + b),
11748 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11749 sum: sa + sb,
11750 count: ca + cb,
11751 },
11752 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11753 sum: sa + sb,
11754 count: ca + cb,
11755 },
11756 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11757 sum: sa + sb,
11758 count: ca + cb,
11759 },
11760 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11761 sum: sa + sb,
11762 count: ca + cb,
11763 },
11764 (MinI(a), MinI(b)) => MinI(a.min(b)),
11765 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11766 (MinF(a), MinF(b)) => MinF(a.min(b)),
11767 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11768 _ => Empty, }
11770 }
11771
11772 pub fn point(&self) -> Option<f64> {
11774 match self {
11775 AggState::Count(n) => Some(*n as f64),
11776 AggState::SumI { sum, .. } => Some(*sum as f64),
11777 AggState::SumF { sum, .. } => Some(*sum),
11778 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11779 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11780 AggState::MinI(n) => Some(*n as f64),
11781 AggState::MaxI(n) => Some(*n as f64),
11782 AggState::MinF(n) => Some(*n),
11783 AggState::MaxF(n) => Some(*n),
11784 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11785 }
11786 }
11787
11788 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11792 let is_float = matches!(ty, Some(TypeId::Float64));
11793 match (agg, result) {
11794 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11795 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11796 sum: x as i128,
11797 count: 1, },
11799 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11800 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11801 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11802 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11803 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11804 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11805 (NativeAgg::Count, _) => AggState::Empty,
11806 (_, NativeAggResult::Null) => AggState::Empty,
11807 _ => {
11808 let _ = is_float;
11809 AggState::Empty
11810 }
11811 }
11812 }
11813}
11814
11815#[derive(Debug, Clone)]
11818pub struct CachedAgg {
11819 pub state: AggState,
11820 pub watermark: u64,
11821 pub epoch: u64,
11822}
11823
11824#[derive(Debug, Clone)]
11826pub struct IncrementalAggResult {
11827 pub state: AggState,
11829 pub incremental: bool,
11832 pub delta_rows: u64,
11834}
11835
11836fn agg_state_from_rows(
11840 rows: &[Row],
11841 conditions: &[crate::query::Condition],
11842 index_sets: &[RowIdSet],
11843 column: Option<u16>,
11844 agg: NativeAgg,
11845 schema: &Schema,
11846 control: Option<&crate::ExecutionControl>,
11847) -> Result<AggState> {
11848 let mut count: u64 = 0;
11849 let mut sum_i: i128 = 0;
11850 let mut sum_f: f64 = 0.0;
11851 let mut mn_i: i64 = i64::MAX;
11852 let mut mx_i: i64 = i64::MIN;
11853 let mut mn_f: f64 = f64::INFINITY;
11854 let mut mx_f: f64 = f64::NEG_INFINITY;
11855 let mut saw_int = false;
11856 let mut saw_float = false;
11857 for (index, r) in rows.iter().enumerate() {
11858 execution_checkpoint(control, index)?;
11859 if !conditions
11860 .iter()
11861 .all(|c| condition_matches_row(c, r, schema))
11862 {
11863 continue;
11864 }
11865 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11866 continue;
11867 }
11868 match agg {
11869 NativeAgg::Count => match column {
11870 None => count += 1,
11872 Some(cid) => match r.columns.get(&cid) {
11875 None | Some(Value::Null) => {}
11876 Some(_) => count += 1,
11877 },
11878 },
11879 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11880 Some(Value::Int64(n)) => {
11881 count += 1;
11882 sum_i += *n as i128;
11883 mn_i = mn_i.min(*n);
11884 mx_i = mx_i.max(*n);
11885 saw_int = true;
11886 }
11887 Some(Value::Float64(f)) => {
11888 count += 1;
11889 sum_f += f;
11890 mn_f = mn_f.min(*f);
11891 mx_f = mx_f.max(*f);
11892 saw_float = true;
11893 }
11894 _ => {}
11895 },
11896 }
11897 }
11898 Ok(match agg {
11899 NativeAgg::Count => {
11900 if count == 0 {
11901 AggState::Empty
11902 } else {
11903 AggState::Count(count)
11904 }
11905 }
11906 NativeAgg::Sum => {
11907 if count == 0 {
11908 AggState::Empty
11909 } else if saw_int {
11910 AggState::SumI { sum: sum_i, count }
11911 } else {
11912 AggState::SumF { sum: sum_f, count }
11913 }
11914 }
11915 NativeAgg::Avg => {
11916 if count == 0 {
11917 AggState::Empty
11918 } else if saw_int {
11919 AggState::AvgI { sum: sum_i, count }
11920 } else {
11921 AggState::AvgF { sum: sum_f, count }
11922 }
11923 }
11924 NativeAgg::Min => {
11925 if !saw_int && !saw_float {
11926 AggState::Empty
11927 } else if saw_int {
11928 AggState::MinI(mn_i)
11929 } else {
11930 AggState::MinF(mn_f)
11931 }
11932 }
11933 NativeAgg::Max => {
11934 if !saw_int && !saw_float {
11935 AggState::Empty
11936 } else if saw_int {
11937 AggState::MaxI(mx_i)
11938 } else {
11939 AggState::MaxF(mx_f)
11940 }
11941 }
11942 })
11943}
11944
11945fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11949 use crate::query::Condition;
11950 match c {
11951 Condition::Pk(key) => match schema.primary_key() {
11952 Some(pk) => row
11953 .columns
11954 .get(&pk.id)
11955 .map(|v| v.encode_key() == *key)
11956 .unwrap_or(false),
11957 None => false,
11958 },
11959 Condition::BitmapEq { column_id, value } => row
11960 .columns
11961 .get(column_id)
11962 .map(|v| v.encode_key() == *value)
11963 .unwrap_or(false),
11964 Condition::BitmapIn { column_id, values } => {
11965 let key = row.columns.get(column_id).map(|v| v.encode_key());
11966 match key {
11967 Some(k) => values.contains(&k),
11968 None => false,
11969 }
11970 }
11971 Condition::BytesPrefix { column_id, prefix } => row
11972 .columns
11973 .get(column_id)
11974 .map(|v| v.encode_key().starts_with(prefix))
11975 .unwrap_or(false),
11976 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11977 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11978 _ => false,
11979 },
11980 Condition::RangeF64 {
11981 column_id,
11982 lo,
11983 lo_inclusive,
11984 hi,
11985 hi_inclusive,
11986 } => match row.columns.get(column_id) {
11987 Some(Value::Float64(n)) => {
11988 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11989 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11990 lo_ok && hi_ok
11991 }
11992 _ => false,
11993 },
11994 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11995 Some(Value::Bytes(b)) => {
11996 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11997 }
11998 _ => false,
11999 },
12000 Condition::FmContainsAll {
12001 column_id,
12002 patterns,
12003 } => match row.columns.get(column_id) {
12004 Some(Value::Bytes(b)) => patterns
12005 .iter()
12006 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
12007 _ => false,
12008 },
12009 Condition::Ann { .. }
12010 | Condition::SparseMatch { .. }
12011 | Condition::MinHashSimilar { .. } => true,
12012 Condition::IsNull { column_id } => {
12013 matches!(row.columns.get(column_id), Some(Value::Null) | None)
12014 }
12015 Condition::IsNotNull { column_id } => {
12016 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
12017 }
12018 }
12019}
12020
12021fn as_f64(v: Option<&Value>) -> Option<f64> {
12023 match v {
12024 Some(Value::Int64(n)) => Some(*n as f64),
12025 Some(Value::Float64(f)) => Some(*f),
12026 _ => None,
12027 }
12028}
12029
12030fn accumulate_int(
12034 cursor: &mut dyn crate::cursor::Cursor,
12035 control: Option<&crate::ExecutionControl>,
12036) -> Result<(u64, i128, i64, i64)> {
12037 let mut count: u64 = 0;
12038 let mut sum: i128 = 0;
12039 let mut mn: i64 = i64::MAX;
12040 let mut mx: i64 = i64::MIN;
12041 while let Some(cols) = cursor.next_batch()? {
12042 execution_checkpoint(control, 0)?;
12043 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
12044 if crate::columnar::all_non_null(validity, data.len()) {
12045 count += data.len() as u64;
12047 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12048 execution_checkpoint(control, chunk_index * 1024)?;
12049 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
12050 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
12051 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
12052 }
12053 } else {
12054 for (i, &v) in data.iter().enumerate() {
12055 execution_checkpoint(control, i)?;
12056 if crate::columnar::validity_bit(validity, i) {
12057 count += 1;
12058 sum += v as i128;
12059 mn = mn.min(v);
12060 mx = mx.max(v);
12061 }
12062 }
12063 }
12064 }
12065 }
12066 Ok((count, sum, mn, mx))
12067}
12068
12069fn accumulate_float(
12071 cursor: &mut dyn crate::cursor::Cursor,
12072 control: Option<&crate::ExecutionControl>,
12073) -> Result<(u64, f64, f64, f64)> {
12074 let mut count: u64 = 0;
12075 let mut sum: f64 = 0.0;
12076 let mut mn: f64 = f64::INFINITY;
12077 let mut mx: f64 = f64::NEG_INFINITY;
12078 while let Some(cols) = cursor.next_batch()? {
12079 execution_checkpoint(control, 0)?;
12080 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
12081 if crate::columnar::all_non_null(validity, data.len()) {
12082 count += data.len() as u64;
12083 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12084 execution_checkpoint(control, chunk_index * 1024)?;
12085 sum += chunk.iter().sum::<f64>();
12086 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
12087 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
12088 }
12089 } else {
12090 for (i, &v) in data.iter().enumerate() {
12091 execution_checkpoint(control, i)?;
12092 if crate::columnar::validity_bit(validity, i) {
12093 count += 1;
12094 sum += v;
12095 mn = mn.min(v);
12096 mx = mx.max(v);
12097 }
12098 }
12099 }
12100 }
12101 }
12102 Ok((count, sum, mn, mx))
12103}
12104
12105#[inline]
12106fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
12107 if index.is_multiple_of(256) {
12108 control
12109 .map(crate::ExecutionControl::checkpoint)
12110 .transpose()?;
12111 }
12112 Ok(())
12113}
12114
12115fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
12116 if count == 0 && !matches!(agg, NativeAgg::Count) {
12117 return NativeAggResult::Null;
12118 }
12119 match agg {
12120 NativeAgg::Count => NativeAggResult::Count(count),
12121 NativeAgg::Sum => match sum.try_into() {
12124 Ok(v) => NativeAggResult::Int(v),
12125 Err(_) => NativeAggResult::Null,
12126 },
12127 NativeAgg::Min => NativeAggResult::Int(mn),
12128 NativeAgg::Max => NativeAggResult::Int(mx),
12129 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
12130 }
12131}
12132
12133fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
12134 if count == 0 && !matches!(agg, NativeAgg::Count) {
12135 return NativeAggResult::Null;
12136 }
12137 match agg {
12138 NativeAgg::Count => NativeAggResult::Count(count),
12139 NativeAgg::Sum => NativeAggResult::Float(sum),
12140 NativeAgg::Min => NativeAggResult::Float(mn),
12141 NativeAgg::Max => NativeAggResult::Float(mx),
12142 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
12143 }
12144}
12145
12146fn agg_int(
12149 stats: &[crate::page::PageStat],
12150 decode: fn(Option<&[u8]>) -> Option<i64>,
12151) -> Option<(Option<i64>, Option<i64>, u64)> {
12152 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
12153 let mut any = false;
12154 for s in stats {
12155 if let Some(v) = decode(s.min.as_deref()) {
12156 mn = mn.min(v);
12157 any = true;
12158 }
12159 if let Some(v) = decode(s.max.as_deref()) {
12160 mx = mx.max(v);
12161 any = true;
12162 }
12163 nulls += s.null_count;
12164 }
12165 any.then_some((Some(mn), Some(mx), nulls))
12166}
12167
12168fn agg_float(
12170 stats: &[crate::page::PageStat],
12171 decode: fn(Option<&[u8]>) -> Option<f64>,
12172) -> Option<(Option<f64>, Option<f64>, u64)> {
12173 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
12174 let mut any = false;
12175 for s in stats {
12176 if let Some(v) = decode(s.min.as_deref()) {
12177 mn = mn.min(v);
12178 any = true;
12179 }
12180 if let Some(v) = decode(s.max.as_deref()) {
12181 mx = mx.max(v);
12182 any = true;
12183 }
12184 nulls += s.null_count;
12185 }
12186 any.then_some((Some(mn), Some(mx), nulls))
12187}
12188
12189type SecondaryIndexes = (
12191 HashMap<u16, BitmapIndex>,
12192 HashMap<u16, AnnIndex>,
12193 HashMap<u16, FmIndex>,
12194 HashMap<u16, SparseIndex>,
12195 HashMap<u16, MinHashIndex>,
12196);
12197
12198fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
12199 let mut bitmap = HashMap::new();
12200 let mut ann = HashMap::new();
12201 let mut fm = HashMap::new();
12202 let mut sparse = HashMap::new();
12203 let mut minhash = HashMap::new();
12204 for idef in &schema.indexes {
12205 match idef.kind {
12206 IndexKind::Bitmap => {
12207 bitmap.insert(idef.column_id, BitmapIndex::new());
12208 }
12209 IndexKind::Ann => {
12210 let dim = schema
12211 .columns
12212 .iter()
12213 .find(|c| c.id == idef.column_id)
12214 .and_then(|c| match c.ty {
12215 TypeId::Embedding { dim } => Some(dim as usize),
12216 _ => None,
12217 })
12218 .unwrap_or(0);
12219 let options = idef.options.ann.clone().unwrap_or_default();
12220 ann.insert(
12221 idef.column_id,
12222 AnnIndex::with_full_options(
12223 dim,
12224 options.m,
12225 options.ef_construction,
12226 options.ef_search,
12227 &options,
12228 ),
12229 );
12230 }
12231 IndexKind::FmIndex => {
12232 fm.insert(idef.column_id, FmIndex::new());
12233 }
12234 IndexKind::Sparse => {
12235 sparse.insert(idef.column_id, SparseIndex::new());
12236 }
12237 IndexKind::MinHash => {
12238 let options = idef.options.minhash.clone().unwrap_or_default();
12239 minhash.insert(
12240 idef.column_id,
12241 MinHashIndex::with_options(options.permutations, options.bands),
12242 );
12243 }
12244 _ => {}
12245 }
12246 }
12247 (bitmap, ann, fm, sparse, minhash)
12248}
12249
12250const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
12251 | ColumnFlags::AUTO_INCREMENT
12252 | ColumnFlags::ENCRYPTED
12253 | ColumnFlags::ENCRYPTED_INDEXABLE
12254 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
12255
12256fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
12257 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
12258 return Err(MongrelError::Schema(
12259 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
12260 ));
12261 }
12262 Ok(())
12263}
12264
12265fn validate_alter_column_type(
12266 schema: &Schema,
12267 old: &ColumnDef,
12268 next: &ColumnDef,
12269 has_stored_versions: bool,
12270) -> Result<()> {
12271 if old.ty == next.ty {
12272 return Ok(());
12273 }
12274 if schema.indexes.iter().any(|i| i.column_id == old.id) {
12275 return Err(MongrelError::Schema(format!(
12276 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
12277 old.name
12278 )));
12279 }
12280 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
12281 return Ok(());
12282 }
12283 Err(MongrelError::Schema(format!(
12284 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
12285 old.ty, next.ty
12286 )))
12287}
12288
12289fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
12290 matches!(
12291 (old, new),
12292 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
12293 )
12294}
12295
12296fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
12302 let mut prev: Option<i64> = None;
12303 for r in rows {
12304 match r.columns.get(&pk_id) {
12305 Some(Value::Int64(v)) => {
12306 if prev.is_some_and(|p| p >= *v) {
12307 return false;
12308 }
12309 prev = Some(*v);
12310 }
12311 _ => return false,
12312 }
12313 }
12314 true
12315}
12316
12317#[allow(clippy::too_many_arguments)]
12318fn index_into(
12319 schema: &Schema,
12320 row: &Row,
12321 hot: &mut HotIndex,
12322 bitmap: &mut HashMap<u16, BitmapIndex>,
12323 ann: &mut HashMap<u16, AnnIndex>,
12324 fm: &mut HashMap<u16, FmIndex>,
12325 sparse: &mut HashMap<u16, SparseIndex>,
12326 minhash: &mut HashMap<u16, MinHashIndex>,
12327) {
12328 for idef in &schema.indexes {
12329 let Some(val) = row.columns.get(&idef.column_id) else {
12330 continue;
12331 };
12332 match idef.kind {
12333 IndexKind::Bitmap => {
12334 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12335 b.insert(val.encode_key(), row.row_id);
12336 }
12337 }
12338 IndexKind::Ann => {
12339 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12340 if let Some(meta) = val.generated_embedding_metadata() {
12341 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12343 continue;
12344 }
12345 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12346 .is_err()
12347 {
12348 continue;
12349 }
12350 }
12351 a.insert_validated(v, row.row_id);
12352 }
12353 }
12354 IndexKind::FmIndex => {
12355 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12356 f.insert(b.clone(), row.row_id);
12357 }
12358 }
12359 IndexKind::Sparse => {
12360 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12361 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12364 s.insert(&terms, row.row_id);
12365 }
12366 }
12367 }
12368 IndexKind::MinHash => {
12369 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12370 let tokens = crate::index::token_hashes_from_bytes(b);
12373 mh.insert(&tokens, row.row_id);
12374 }
12375 }
12376 _ => {}
12377 }
12378 }
12379 if let Some(pk_col) = schema.primary_key() {
12380 if let Some(pk_val) = row.columns.get(&pk_col.id) {
12381 hot.insert(pk_val.encode_key(), row.row_id);
12382 }
12383 }
12384}
12385
12386#[allow(clippy::too_many_arguments)]
12389fn index_into_single(
12390 idef: &IndexDef,
12391 _schema: &Schema,
12392 row: &Row,
12393 _hot: &mut HotIndex,
12394 bitmap: &mut HashMap<u16, BitmapIndex>,
12395 ann: &mut HashMap<u16, AnnIndex>,
12396 fm: &mut HashMap<u16, FmIndex>,
12397 sparse: &mut HashMap<u16, SparseIndex>,
12398 minhash: &mut HashMap<u16, MinHashIndex>,
12399) {
12400 let Some(val) = row.columns.get(&idef.column_id) else {
12401 return;
12402 };
12403 match idef.kind {
12404 IndexKind::Bitmap => {
12405 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12406 b.insert(val.encode_key(), row.row_id);
12407 }
12408 }
12409 IndexKind::Ann => {
12410 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12411 if let Some(meta) = val.generated_embedding_metadata() {
12412 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12414 return;
12415 }
12416 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12417 .is_err()
12418 {
12419 return;
12420 }
12421 }
12422 a.insert_validated(v, row.row_id);
12423 }
12424 }
12425 IndexKind::FmIndex => {
12426 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12427 f.insert(b.clone(), row.row_id);
12428 }
12429 }
12430 IndexKind::Sparse => {
12431 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12432 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12433 s.insert(&terms, row.row_id);
12434 }
12435 }
12436 }
12437 IndexKind::MinHash => {
12438 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12439 let tokens = crate::index::token_hashes_from_bytes(b);
12440 mh.insert(&tokens, row.row_id);
12441 }
12442 }
12443 _ => {}
12444 }
12445}
12446
12447fn eval_partial_predicate(
12453 pred: &str,
12454 columns_map: &HashMap<u16, &Value>,
12455 name_to_id: &HashMap<&str, u16>,
12456) -> bool {
12457 let lower = pred.trim().to_ascii_lowercase();
12458 if let Some(rest) = lower.strip_suffix(" is not null") {
12460 let col_name = rest.trim();
12461 if let Some(col_id) = name_to_id.get(col_name) {
12462 return columns_map
12463 .get(col_id)
12464 .is_some_and(|v| !matches!(v, Value::Null));
12465 }
12466 }
12467 if let Some(rest) = lower.strip_suffix(" is null") {
12469 let col_name = rest.trim();
12470 if let Some(col_id) = name_to_id.get(col_name) {
12471 return columns_map
12472 .get(col_id)
12473 .is_none_or(|v| matches!(v, Value::Null));
12474 }
12475 }
12476 true
12479}
12480
12481#[allow(dead_code)]
12487fn bulk_index_key(
12488 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12489 column_id: u16,
12490 ty: TypeId,
12491 col: &columnar::NativeColumn,
12492 i: usize,
12493) -> Option<Vec<u8>> {
12494 let encoded = columnar::encode_key_native(ty, col, i)?;
12495 {
12496 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12497 if let Some((key, scheme)) = column_keys.get(&column_id) {
12498 return Some(match (*scheme, col) {
12499 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12500 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12501 ope_token_i64(key, data[i]).to_vec()
12502 }
12503 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12504 ope_token_f64(key, data[i]).to_vec()
12505 }
12506 _ => hmac_token(key, &encoded).to_vec(),
12507 });
12508 }
12509 }
12510 Some(encoded)
12511}
12512
12513pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12514 write_schema_with_after(dir, schema, || {})
12515}
12516
12517pub(crate) fn write_schema_durable(
12518 root: &crate::durable_file::DurableRoot,
12519 schema: &Schema,
12520) -> Result<()> {
12521 write_schema_durable_with_after(root, schema, || {})
12522}
12523
12524fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12525where
12526 F: FnOnce(),
12527{
12528 let json = serde_json::to_string_pretty(schema)
12529 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12530 crate::durable_file::write_atomic_with_after(
12531 &dir.join(SCHEMA_FILENAME),
12532 json.as_bytes(),
12533 after_publish,
12534 )?;
12535 Ok(())
12536}
12537
12538fn write_schema_durable_with_after<F>(
12539 root: &crate::durable_file::DurableRoot,
12540 schema: &Schema,
12541 after_publish: F,
12542) -> Result<()>
12543where
12544 F: FnOnce(),
12545{
12546 let json = serde_json::to_string_pretty(schema)
12547 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12548 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12549 Ok(())
12550}
12551
12552fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12553 let mut schema_published = false;
12554 let schema_result = match table._root_guard.as_deref() {
12555 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12556 schema_published = true;
12557 }),
12558 None => write_schema_with_after(&table.dir, &table.schema, || {
12559 schema_published = true;
12560 }),
12561 };
12562 if schema_result.is_err() && !schema_published {
12563 return schema_result;
12564 }
12565 match table.persist_manifest(table.current_epoch()) {
12566 Ok(()) => Ok(()),
12567 Err(manifest_error) => Err(match schema_result {
12568 Ok(()) => manifest_error,
12569 Err(schema_error) => MongrelError::Other(format!(
12570 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12571 )),
12572 }),
12573 }
12574}
12575
12576fn read_schema(dir: &Path) -> Result<Schema> {
12577 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12578 read_schema_file(file)
12579}
12580
12581fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12582 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12583 use std::io::Read;
12584
12585 let length = file.metadata()?.len();
12586 if length > MAX_SCHEMA_BYTES {
12587 return Err(MongrelError::ResourceLimitExceeded {
12588 resource: "schema bytes",
12589 requested: usize::try_from(length).unwrap_or(usize::MAX),
12590 limit: MAX_SCHEMA_BYTES as usize,
12591 });
12592 }
12593 let mut bytes = Vec::with_capacity(length as usize);
12594 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12595 if bytes.len() as u64 != length {
12596 return Err(MongrelError::Schema(
12597 "schema length changed while reading".into(),
12598 ));
12599 }
12600 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12601}
12602
12603fn preflight_standalone_open(
12604 dir: &Path,
12605 runs_root: Option<&crate::durable_file::DurableRoot>,
12606 idx_root: Option<&crate::durable_file::DurableRoot>,
12607 manifest: &Manifest,
12608 schema: &Schema,
12609 records: &[crate::wal::Record],
12610 kek: Option<Arc<Kek>>,
12611) -> Result<()> {
12612 crate::wal::validate_shared_transaction_framing(records)?;
12613 if manifest.schema_id > schema.schema_id
12614 || manifest.flushed_epoch > manifest.current_epoch
12615 || manifest.global_idx_epoch > manifest.current_epoch
12616 || manifest.next_row_id == u64::MAX
12617 || manifest.auto_inc_next < 0
12618 || manifest.auto_inc_next == i64::MAX
12619 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12620 {
12621 return Err(MongrelError::InvalidArgument(
12622 "manifest counters or schema identity are invalid".into(),
12623 ));
12624 }
12625 let mut run_ids = HashSet::new();
12626 let mut maximum_row_id = None::<u64>;
12627 for run in &manifest.runs {
12628 if run.run_id >= u64::MAX as u128
12629 || !run_ids.insert(run.run_id)
12630 || run.epoch_created > manifest.current_epoch
12631 {
12632 return Err(MongrelError::InvalidArgument(
12633 "manifest contains an invalid or duplicate active run".into(),
12634 ));
12635 }
12636 let mut reader = match runs_root {
12637 Some(root) => RunReader::open_file(
12638 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12639 schema.clone(),
12640 kek.clone(),
12641 )?,
12642 None => RunReader::open(
12643 dir.join(RUNS_DIR)
12644 .join(format!("r-{}.sr", run.run_id as u64)),
12645 schema.clone(),
12646 kek.clone(),
12647 )?,
12648 };
12649 let header = reader.header();
12650 if header.run_id != run.run_id
12651 || header.level != run.level
12652 || header.row_count != run.row_count
12653 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12654 || header.is_uniform_epoch() && header.epoch_created != 0
12655 || header.schema_id > schema.schema_id
12656 {
12657 return Err(MongrelError::InvalidArgument(format!(
12658 "run {} differs from its manifest",
12659 run.run_id
12660 )));
12661 }
12662 if header.row_count != 0 {
12663 maximum_row_id = Some(
12664 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12665 );
12666 }
12667 reader.validate_all_pages()?;
12668 }
12669 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12670 return Err(MongrelError::InvalidArgument(
12671 "manifest next_row_id does not advance beyond persisted rows".into(),
12672 ));
12673 }
12674 for run in &manifest.retiring {
12675 if run.run_id >= u64::MAX as u128
12676 || run.retire_epoch > manifest.current_epoch
12677 || !run_ids.insert(run.run_id)
12678 {
12679 return Err(MongrelError::InvalidArgument(
12680 "manifest contains an invalid or duplicate retired run".into(),
12681 ));
12682 }
12683 }
12684 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12685 match idx_root {
12686 Some(root) => {
12687 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12688 }
12689 None => {
12690 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12691 }
12692 }
12693
12694 let committed = records
12695 .iter()
12696 .filter_map(|record| match record.op {
12697 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12698 _ => None,
12699 })
12700 .collect::<HashMap<_, _>>();
12701 for record in records {
12702 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12703 continue;
12704 };
12705 match &record.op {
12706 Op::Put { table_id, rows } => {
12707 if *table_id != manifest.table_id {
12708 return Err(MongrelError::CorruptWal {
12709 offset: record.seq.0,
12710 reason: format!(
12711 "private WAL record references table {table_id}, expected {}",
12712 manifest.table_id
12713 ),
12714 });
12715 }
12716 let rows: Vec<Row> =
12717 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12718 offset: record.seq.0,
12719 reason: format!("committed Put payload could not be decoded: {error}"),
12720 })?;
12721 for row in rows {
12722 if row.deleted || row.row_id.0 == u64::MAX {
12723 return Err(MongrelError::CorruptWal {
12724 offset: record.seq.0,
12725 reason: "committed Put contains an invalid row identity".into(),
12726 });
12727 }
12728 let cells = row.columns.into_iter().collect::<Vec<_>>();
12729 schema
12730 .validate_values(&cells)
12731 .map_err(|error| MongrelError::CorruptWal {
12732 offset: record.seq.0,
12733 reason: format!("committed Put violates table schema: {error}"),
12734 })?;
12735 if schema.auto_increment_column().is_some_and(|column| {
12736 matches!(
12737 cells.iter().find(|(id, _)| *id == column.id),
12738 Some((_, Value::Int64(value))) if *value == i64::MAX
12739 )
12740 }) {
12741 return Err(MongrelError::CorruptWal {
12742 offset: record.seq.0,
12743 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12744 });
12745 }
12746 }
12747 }
12748 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12749 if *table_id != manifest.table_id =>
12750 {
12751 return Err(MongrelError::CorruptWal {
12752 offset: record.seq.0,
12753 reason: format!(
12754 "private WAL record references table {table_id}, expected {}",
12755 manifest.table_id
12756 ),
12757 });
12758 }
12759 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12760 return Err(MongrelError::CorruptWal {
12761 offset: record.seq.0,
12762 reason: "private WAL contains shared spilled-run metadata".into(),
12763 });
12764 }
12765 _ => {}
12766 }
12767 }
12768 Ok(())
12769}
12770
12771fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12772 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12773}
12774
12775fn wal_segment_number(path: &Path) -> Option<u64> {
12776 path.file_stem()
12777 .and_then(|stem| stem.to_str())
12778 .and_then(|stem| stem.strip_prefix("seg-"))
12779 .and_then(|number| number.parse().ok())
12780}
12781
12782fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12783 let n = list_wal_numbers(wal_dir)?;
12784 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12785}
12786
12787fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12788 list_wal_numbers(wal_dir)?
12789 .map(|maximum| {
12790 maximum
12791 .checked_add(1)
12792 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12793 })
12794 .unwrap_or(Ok(0))
12795}
12796
12797fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12798 let mut max_n = None;
12799 let entries = match std::fs::read_dir(wal_dir) {
12800 Ok(entries) => entries,
12801 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12802 Err(error) => return Err(error.into()),
12803 };
12804 for entry in entries {
12805 let entry = entry?;
12806 let fname = entry.file_name();
12807 let Some(s) = fname.to_str() else {
12808 continue;
12809 };
12810 let Some(stripped) = s.strip_prefix("seg-") else {
12811 continue;
12812 };
12813 let Some(number) = stripped.strip_suffix(".wal") else {
12814 return Err(MongrelError::CorruptWal {
12815 offset: 0,
12816 reason: format!("malformed WAL segment name {s:?}"),
12817 });
12818 };
12819 let n = number
12820 .parse::<u32>()
12821 .map_err(|_| MongrelError::CorruptWal {
12822 offset: 0,
12823 reason: format!("malformed WAL segment name {s:?}"),
12824 })?;
12825 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12826 return Err(MongrelError::CorruptWal {
12827 offset: n as u64,
12828 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12829 });
12830 }
12831 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12832 }
12833 Ok(max_n)
12834}