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 std::fs::create_dir_all(&dir)?;
1907 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
1908 let pinned = root.io_path()?;
1909 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1910 ctx.root_guard = Some(root.clone());
1911 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
1912 root.finalize_deferred_sync_shallow()?;
1918 Ok(table)
1919 }
1920
1921 pub fn create_encrypted(
1932 dir: impl AsRef<Path>,
1933 schema: Schema,
1934 table_id: u64,
1935 passphrase: &str,
1936 ) -> Result<Self> {
1937 let dir = dir.as_ref().to_path_buf();
1938 std::fs::create_dir_all(&dir)?;
1939 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
1940 root.create_directory_all(META_DIR)?;
1941 let salt = crate::encryption::random_salt()?;
1942 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1943 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1944 let pinned = root.io_path()?;
1945 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1946 ctx.root_guard = Some(root.clone());
1947 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
1948 root.finalize_deferred_sync()?;
1949 Ok(table)
1950 }
1951
1952 pub fn create_with_key(
1957 dir: impl AsRef<Path>,
1958 schema: Schema,
1959 table_id: u64,
1960 key: &[u8],
1961 ) -> Result<Self> {
1962 let dir = dir.as_ref().to_path_buf();
1963 std::fs::create_dir_all(&dir)?;
1964 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
1965 root.create_directory_all(META_DIR)?;
1966 let salt = crate::encryption::random_salt()?;
1967 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1968 let kek: Arc<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.clone());
1972 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
1973 root.finalize_deferred_sync()?;
1974 Ok(table)
1975 }
1976
1977 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1979 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1980 let salt = read_table_encryption_salt_root(&root)?;
1981 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1982 let pinned = root.io_path()?;
1983 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1984 ctx.root_guard = Some(root);
1985 Self::open_in(&pinned, ctx)
1986 }
1987
1988 pub(crate) fn create_in(
1989 dir: impl AsRef<Path>,
1990 schema: Schema,
1991 table_id: u64,
1992 ctx: SharedCtx,
1993 ) -> Result<Self> {
1994 schema.validate_auto_increment()?;
1995 schema.validate_defaults()?;
1996 schema.validate_ai()?;
1997 for index in &schema.indexes {
1998 index.validate_options()?;
1999 }
2000 let dir = dir.as_ref().to_path_buf();
2001 let runs_root = match ctx.root_guard.as_ref() {
2002 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
2003 None => {
2004 crate::durable_file::create_directory_all(&dir)?;
2005 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
2006 None
2007 }
2008 };
2009 match ctx.root_guard.as_deref() {
2010 Some(root) => write_schema_durable(root, &schema)?,
2011 None => write_schema(&dir, &schema)?,
2012 }
2013 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
2014 let (wal, current_txn_id) = match ctx.shared.clone() {
2017 Some(s) => (WalSink::Shared(s), 0),
2018 None => {
2019 let pinned_wal_root = match ctx.root_guard.as_deref() {
2020 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
2021 None => None,
2022 };
2023 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
2024 root.io_path()?
2025 } else {
2026 let wal_dir = dir.join(WAL_DIR);
2027 std::fs::create_dir_all(&wal_dir)?;
2028 wal_dir
2029 };
2030 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
2031 (Some(root), Some(dk)) => {
2032 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
2033 }
2034 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
2035 (None, Some(dk)) => Wal::create_with_cipher(
2036 wal_dir.join("seg-000000.wal"),
2037 Epoch(0),
2038 Some(make_cipher(dk)),
2039 0,
2040 )?,
2041 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
2042 };
2043 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2044 (WalSink::Private(w), 1)
2045 }
2046 };
2047 let mut manifest = Manifest::new(table_id, schema.schema_id);
2048 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2053 match ctx.root_guard.as_deref() {
2054 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
2055 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
2056 }
2057 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
2058 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2059 let auto_inc = resolve_auto_inc(&schema);
2060 let rcache_dir = dir.join(RCACHE_DIR);
2061 let initial_view = ReadGeneration::empty(&schema);
2062 Ok(Self {
2063 dir,
2064 _root_guard: ctx.root_guard,
2065 runs_root,
2066 idx_root: None,
2067 table_id,
2068 name: ctx.table_name.unwrap_or_default(),
2069 auth: ctx.auth,
2070 read_only: ctx.read_only,
2071 durable_commit_failed: false,
2072 wal,
2073 memtable: Memtable::new(),
2074 mutable_run: MutableRun::new(),
2075 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2076 compaction_zstd_level: 3,
2077 allocator: RowIdAllocator::new(0),
2078 epoch: ctx.epoch,
2079 data_generation: 0,
2080 schema,
2081 hot: HotIndex::new(),
2082 kek: ctx.kek,
2083 column_keys,
2084 run_refs: Vec::new(),
2085 retiring: Vec::new(),
2086 next_run_id: 1,
2087 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2088 current_txn_id,
2089 pending_private_mutations: false,
2090 bitmap,
2091 ann,
2092 fm,
2093 sparse,
2094 minhash,
2095 learned_range: Arc::new(HashMap::new()),
2096 pk_by_row: ReversePkMap::new(),
2097 pinned: BTreeMap::new(),
2098 live_count: 0,
2099 reservoir: crate::reservoir::Reservoir::default(),
2100 reservoir_complete: true,
2101 had_deletes: false,
2102 agg_cache: Arc::new(HashMap::new()),
2103 global_idx_epoch: 0,
2104 indexes_complete: true,
2105 index_build_policy: IndexBuildPolicy::default(),
2106 pk_by_row_complete: false,
2107 flushed_epoch: 0,
2108 page_cache: ctx.page_cache,
2109 decoded_cache: ctx.decoded_cache,
2110 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2111 snapshots: ctx.snapshots,
2112 commit_lock: ctx.commit_lock,
2113 result_cache: Arc::new(parking_lot::Mutex::new(
2114 ResultCache::new()
2115 .with_dir(rcache_dir)
2116 .with_cache_dek(cache_dek.clone()),
2117 )),
2118 pending_delete_rids: roaring::RoaringBitmap::new(),
2119 pending_put_cols: std::collections::HashSet::new(),
2120 pending_rows: Vec::new(),
2121 pending_rows_auto_inc: Vec::new(),
2122 pending_dels: Vec::new(),
2123 pending_truncate: None,
2124 wal_dek,
2125 auto_inc,
2126 ttl: None,
2127 pins: Arc::new(crate::retention::PinRegistry::new()),
2128 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2129 read_generation_pin: None,
2130 })
2131 }
2132
2133 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
2137 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2138 let pinned = root.io_path()?;
2139 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2140 ctx.root_guard = Some(root);
2141 Self::open_in(&pinned, ctx)
2142 }
2143
2144 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
2147 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2148 let salt = read_table_encryption_salt_root(&root)?;
2149 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2150 let pinned = root.io_path()?;
2151 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2152 ctx.root_guard = Some(root);
2153 let t = Self::open_in(&pinned, ctx)?;
2154 Ok(t)
2155 }
2156
2157 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
2158 let dir = dir.as_ref().to_path_buf();
2159 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2160 let mut manifest = match ctx.root_guard.as_ref() {
2161 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
2162 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
2163 };
2164 let schema: Schema = match ctx.root_guard.as_ref() {
2165 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
2166 None => read_schema(&dir)?,
2167 };
2168 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
2174 let runs_root = match ctx.root_guard.as_ref() {
2175 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
2176 None => None,
2177 };
2178 let idx_root = match ctx.root_guard.as_ref() {
2179 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
2180 Ok(root) => Some(Arc::new(root)),
2181 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2182 Err(error) => return Err(error.into()),
2183 },
2184 None => None,
2185 };
2186 schema.validate_auto_increment()?;
2187 schema.validate_defaults()?;
2188 schema.validate_ai()?;
2189 for index in &schema.indexes {
2190 index.validate_options()?;
2191 }
2192 let replay_epoch = Epoch(manifest.current_epoch);
2193 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
2194 let private_replayed = if ctx.shared.is_none() {
2195 match latest_wal_segment(&dir.join(WAL_DIR))? {
2196 Some(path) => {
2197 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
2198 crate::wal::replay_with_cipher(path, cipher)?
2199 }
2200 None => Vec::new(),
2201 }
2202 } else {
2203 Vec::new()
2204 };
2205 if ctx.shared.is_none() {
2206 preflight_standalone_open(
2207 &dir,
2208 runs_root.as_deref(),
2209 idx_root.as_deref(),
2210 &manifest,
2211 &schema,
2212 &private_replayed,
2213 ctx.kek.clone(),
2214 )?;
2215 }
2216 let next_run_id = derive_next_run_id(
2217 &dir,
2218 runs_root.as_deref(),
2219 &manifest.runs,
2220 &manifest.retiring,
2221 )?;
2222 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2226 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2227 None => {
2228 let replayed = private_replayed;
2229 let wal_dir = dir.join(WAL_DIR);
2235 crate::durable_file::create_directory_all(&wal_dir)?;
2236 let segment = next_wal_segment(&wal_dir)?;
2237 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2238 let temporary = wal_dir.join(format!(
2239 ".recovery-{}-{}-{segment_no:06}.tmp",
2240 std::process::id(),
2241 std::time::SystemTime::now()
2242 .duration_since(std::time::UNIX_EPOCH)
2243 .unwrap_or_default()
2244 .as_nanos()
2245 ));
2246 let mut w = Wal::create_with_cipher(
2247 &temporary,
2248 replay_epoch,
2249 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2250 segment_no,
2251 )?;
2252 for record in &replayed {
2253 w.append_txn(record.txn_id, record.op.clone())?;
2254 }
2255 let mut w = w.publish_as(segment)?;
2256 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2257 let next_txn_id = replayed
2258 .iter()
2259 .map(|record| record.txn_id)
2260 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2261 .max()
2262 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2263 .unwrap_or(1);
2264 (WalSink::Private(w), replayed, next_txn_id)
2265 }
2266 };
2267
2268 let mut memtable = Memtable::new();
2269 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2270 let persisted_epoch = manifest.current_epoch;
2271 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2278 s.next = manifest.auto_inc_next;
2279 s.seeded = manifest.auto_inc_next > 0;
2280 s
2281 });
2282
2283 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2290 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2291 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2292 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2293 std::collections::BTreeMap::new();
2294 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2295 let mut recovered_epoch = manifest.current_epoch;
2296 let mut recovered_manifest_dirty = schema_manifest_repair;
2297 let mut saw_delete = false;
2298 for record in replayed {
2299 let txn_id = record.txn_id;
2300 match record.op {
2301 Op::Put { rows, .. } => {
2302 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2303 for row in &rows {
2304 allocator.advance_to(row.row_id)?;
2305 if let Some(ai) = auto_inc.as_mut() {
2306 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2307 let next = n.checked_add(1).ok_or_else(|| {
2308 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2309 })?;
2310 if next > ai.next {
2311 ai.next = next;
2312 }
2313 }
2314 }
2315 }
2316 staged_puts.entry(txn_id).or_default().extend(rows);
2317 }
2318 Op::Delete { row_ids, .. } => {
2319 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2320 }
2321 Op::TxnCommit { epoch, .. } => {
2322 let commit_epoch = Epoch(epoch);
2323 recovered_epoch = recovered_epoch.max(epoch);
2324 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2325 memtable = Memtable::new();
2326 replayed_puts.clear();
2327 replayed_deletes.clear();
2328 manifest.runs.clear();
2329 manifest.retiring.clear();
2330 manifest.live_count = 0;
2331 manifest.global_idx_epoch = 0;
2332 manifest.current_epoch = manifest.current_epoch.max(epoch);
2333 recovered_manifest_dirty = true;
2334 saw_delete = true;
2335 }
2336 if let Some(puts) = staged_puts.remove(&txn_id) {
2337 if commit_epoch.0 > manifest.flushed_epoch {
2338 for row in &puts {
2339 memtable.upsert(row.clone());
2340 }
2341 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2342 }
2343 }
2344 if let Some(dels) = staged_deletes.remove(&txn_id) {
2345 saw_delete = true;
2346 if commit_epoch.0 > manifest.flushed_epoch {
2347 for rid in dels {
2348 memtable.tombstone(rid, commit_epoch);
2349 replayed_deletes.push((rid, commit_epoch));
2350 }
2351 }
2352 }
2353 }
2354 Op::TxnAbort => {
2355 staged_puts.remove(&txn_id);
2356 staged_deletes.remove(&txn_id);
2357 staged_truncates.remove(&txn_id);
2358 }
2359 Op::TruncateTable { .. } => {
2360 staged_puts.remove(&txn_id);
2361 staged_deletes.remove(&txn_id);
2362 staged_truncates.insert(txn_id);
2363 }
2364 Op::ExternalTableState { .. }
2365 | Op::Flush { .. }
2366 | Op::Ddl(_)
2367 | Op::BeforeImage { .. }
2368 | Op::CommitTimestamp { .. }
2369 | Op::SpilledRows { .. } => {}
2370 }
2371 }
2372
2373 let rcache_dir = dir.join(RCACHE_DIR);
2374 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2375 let initial_view = ReadGeneration::empty(&schema);
2376 let mut db = Self {
2377 dir,
2378 _root_guard: ctx.root_guard,
2379 runs_root,
2380 idx_root,
2381 table_id: manifest.table_id,
2382 name: ctx.table_name.unwrap_or_default(),
2383 auth: ctx.auth,
2384 read_only: ctx.read_only,
2385 durable_commit_failed: false,
2386 wal,
2387 memtable,
2388 mutable_run: MutableRun::new(),
2389 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2390 compaction_zstd_level: 3,
2391 allocator,
2392 epoch: ctx.epoch,
2393 data_generation: persisted_epoch,
2394 schema,
2395 hot: HotIndex::new(),
2396 kek: ctx.kek,
2397 column_keys,
2398 run_refs: manifest.runs.clone(),
2399 retiring: manifest.retiring.clone(),
2400 next_run_id,
2401 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2402 current_txn_id,
2403 pending_private_mutations: false,
2404 bitmap: HashMap::new(),
2405 ann: HashMap::new(),
2406 fm: HashMap::new(),
2407 sparse: HashMap::new(),
2408 minhash: HashMap::new(),
2409 learned_range: Arc::new(HashMap::new()),
2410 pk_by_row: ReversePkMap::new(),
2411 pinned: BTreeMap::new(),
2412 live_count: manifest.live_count,
2413 reservoir: crate::reservoir::Reservoir::default(),
2414 reservoir_complete: false,
2415 had_deletes: saw_delete
2416 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2417 != manifest.live_count,
2418 agg_cache: Arc::new(HashMap::new()),
2419 global_idx_epoch: manifest.global_idx_epoch,
2420 indexes_complete: true,
2421 index_build_policy: IndexBuildPolicy::default(),
2422 pk_by_row_complete: false,
2423 flushed_epoch: manifest.flushed_epoch,
2424 page_cache: ctx.page_cache,
2425 decoded_cache: ctx.decoded_cache,
2426 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2427 snapshots: ctx.snapshots,
2428 commit_lock: ctx.commit_lock,
2429 result_cache: Arc::new(parking_lot::Mutex::new(
2430 ResultCache::new()
2431 .with_dir(rcache_dir)
2432 .with_cache_dek(cache_dek.clone()),
2433 )),
2434 pending_delete_rids: roaring::RoaringBitmap::new(),
2435 pending_put_cols: std::collections::HashSet::new(),
2436 pending_rows: Vec::new(),
2437 pending_rows_auto_inc: Vec::new(),
2438 pending_dels: Vec::new(),
2439 pending_truncate: None,
2440 wal_dek,
2441 auto_inc,
2442 ttl: manifest.ttl,
2443 pins: Arc::new(crate::retention::PinRegistry::new()),
2444 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2445 read_generation_pin: None,
2446 };
2447
2448 db.epoch.advance_recovered(Epoch(recovered_epoch));
2451
2452 let checkpoint = match db.idx_root.as_deref() {
2457 Some(root) => {
2458 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2459 }
2460 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2461 };
2462 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2463 c.epoch_built == manifest.global_idx_epoch
2464 && manifest.global_idx_epoch > 0
2465 && manifest
2466 .runs
2467 .iter()
2468 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2469 });
2470 if let Some(loaded) = checkpoint {
2471 if checkpoint_valid {
2472 db.hot = loaded.hot;
2473 db.bitmap = loaded.bitmap;
2474 db.ann = loaded.ann;
2475 db.fm = loaded.fm;
2476 db.sparse = loaded.sparse;
2477 db.minhash = loaded.minhash;
2478 db.learned_range = Arc::new(loaded.learned_range);
2479 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2484 for (cid, idx) in bitmap0 {
2485 db.bitmap.entry(cid).or_insert(idx);
2486 }
2487 for (cid, idx) in ann0 {
2488 db.ann.entry(cid).or_insert(idx);
2489 }
2490 for (cid, idx) in fm0 {
2491 db.fm.entry(cid).or_insert(idx);
2492 }
2493 for (cid, idx) in sparse0 {
2494 db.sparse.entry(cid).or_insert(idx);
2495 }
2496 for (cid, idx) in minhash0 {
2497 db.minhash.entry(cid).or_insert(idx);
2498 }
2499 }
2502 }
2503 if !checkpoint_valid {
2504 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2505 db.bitmap = bitmap;
2506 db.ann = ann;
2507 db.fm = fm;
2508 db.sparse = sparse;
2509 db.minhash = minhash;
2510 db.rebuild_indexes_from_runs()?;
2511 db.build_learned_ranges()?;
2512 }
2513
2514 for (epoch, group) in replayed_puts {
2519 let (losers, winner_pks) = db.partition_pk_winners(&group);
2520 for (key, &row_id) in &winner_pks {
2521 if let Some(old_rid) = db.hot.get(key) {
2522 if old_rid != row_id {
2523 db.tombstone_row(old_rid, epoch, None, false);
2524 }
2525 }
2526 }
2527 for &loser_rid in &losers {
2528 db.tombstone_row(loser_rid, epoch, None, false);
2529 }
2530 for (key, row_id) in winner_pks {
2531 db.insert_hot_pk(key, row_id);
2532 }
2533 if db.schema.primary_key().is_none() {
2534 for r in &group {
2535 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2536 }
2537 }
2538 for r in &group {
2539 if !losers.contains(&r.row_id) {
2540 db.index_row(r);
2541 }
2542 }
2543 }
2544 for (rid, epoch) in &replayed_deletes {
2548 db.remove_hot_for_row(*rid, *epoch);
2549 }
2550
2551 if recovered_manifest_dirty {
2552 let rows = db.visible_rows(Snapshot::unbounded())?;
2553 db.live_count = rows.len() as u64;
2554 db.persist_manifest(Epoch(recovered_epoch))?;
2555 }
2556
2557 db.result_cache.lock().load_persistent();
2564 Ok(db)
2565 }
2566
2567 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2573 if self.reservoir_complete {
2574 return Ok(());
2575 }
2576 self.rebuild_reservoir()?;
2577 self.reservoir_complete = true;
2578 Ok(())
2579 }
2580
2581 fn rebuild_reservoir(&mut self) -> Result<()> {
2584 let snap = self.snapshot();
2585 let rows = self.visible_rows(snap)?;
2586 self.reservoir.reset();
2587 for r in rows {
2588 self.reservoir.offer(r.row_id.0);
2589 }
2590 Ok(())
2591 }
2592
2593 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2594 self.rebuild_indexes_from_runs_inner(None)
2595 }
2596
2597 fn rebuild_indexes_from_runs_inner(
2598 &mut self,
2599 control: Option<&crate::ExecutionControl>,
2600 ) -> Result<()> {
2601 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2604 crate::retention::PinSource::OnlineIndexBuild,
2605 self.current_epoch(),
2606 );
2607 self.hot = HotIndex::new();
2608 self.pk_by_row.clear();
2609 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2610 self.bitmap = bitmap;
2611 self.ann = ann;
2612 self.fm = fm;
2613 self.sparse = sparse;
2614 self.minhash = minhash;
2615 let snapshot = Epoch(u64::MAX);
2616 let ttl_now = unix_nanos_now();
2617 let mut scanned = 0_usize;
2618 for rr in self.run_refs.clone() {
2619 if let Some(control) = control {
2620 control.checkpoint()?;
2621 }
2622 let mut reader = self.open_reader(rr.run_id)?;
2623 for row in reader.visible_rows(snapshot)? {
2624 if scanned.is_multiple_of(256) {
2625 if let Some(control) = control {
2626 control.checkpoint()?;
2627 }
2628 }
2629 scanned += 1;
2630 if self.row_expired_at(&row, ttl_now) {
2631 continue;
2632 }
2633 let tok_row = self.tokenized_for_indexes(&row);
2634 index_into(
2635 &self.schema,
2636 &tok_row,
2637 &mut self.hot,
2638 &mut self.bitmap,
2639 &mut self.ann,
2640 &mut self.fm,
2641 &mut self.sparse,
2642 &mut self.minhash,
2643 );
2644 }
2645 }
2646 for row in self.mutable_run.visible_versions(snapshot) {
2647 if scanned.is_multiple_of(256) {
2648 if let Some(control) = control {
2649 control.checkpoint()?;
2650 }
2651 }
2652 scanned += 1;
2653 if row.deleted {
2654 self.remove_hot_for_row(row.row_id, snapshot);
2655 } else if !self.row_expired_at(&row, ttl_now) {
2656 self.index_row(&row);
2657 }
2658 }
2659 for row in self.memtable.visible_versions(snapshot) {
2660 if scanned.is_multiple_of(256) {
2661 if let Some(control) = control {
2662 control.checkpoint()?;
2663 }
2664 }
2665 scanned += 1;
2666 if row.deleted {
2667 self.remove_hot_for_row(row.row_id, snapshot);
2668 } else if !self.row_expired_at(&row, ttl_now) {
2669 self.index_row(&row);
2670 }
2671 }
2672 self.refresh_pk_by_row_from_hot();
2673 Ok(())
2674 }
2675
2676 fn refresh_pk_by_row_from_hot(&mut self) {
2677 self.pk_by_row_complete = true;
2678 if self.schema.primary_key().is_none() {
2679 self.pk_by_row.clear();
2680 return;
2681 }
2682 self.pk_by_row = ReversePkMap::from_entries(
2688 self.hot
2689 .entries()
2690 .into_iter()
2691 .map(|(key, row_id)| (row_id, key)),
2692 );
2693 }
2694
2695 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2696 if self.schema.primary_key().is_some() {
2697 self.pk_by_row.insert(row_id, key.clone());
2698 }
2699 self.hot.insert(key, row_id);
2700 }
2701
2702 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2706 self.build_learned_ranges_inner(None)
2707 }
2708
2709 fn build_learned_ranges_inner(
2710 &mut self,
2711 control: Option<&crate::ExecutionControl>,
2712 ) -> Result<()> {
2713 self.learned_range = Arc::new(HashMap::new());
2714 if self.run_refs.len() != 1 {
2715 return Ok(());
2716 }
2717 let cols: Vec<(u16, usize)> = self
2718 .schema
2719 .indexes
2720 .iter()
2721 .filter(|i| i.kind == IndexKind::LearnedRange)
2722 .map(|i| {
2723 (
2724 i.column_id,
2725 i.options
2726 .learned_range
2727 .as_ref()
2728 .map(|options| options.epsilon)
2729 .unwrap_or(16),
2730 )
2731 })
2732 .collect();
2733 if cols.is_empty() {
2734 return Ok(());
2735 }
2736 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2737 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2738 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2739 _ => return Ok(()),
2740 };
2741 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2742 if column_index % 256 == 0 {
2743 if let Some(control) = control {
2744 control.checkpoint()?;
2745 }
2746 }
2747 let ty = self
2748 .schema
2749 .columns
2750 .iter()
2751 .find(|c| c.id == cid)
2752 .map(|c| c.ty.clone())
2753 .unwrap_or(TypeId::Int64);
2754 match ty {
2755 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2756 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2757 let pairs: Vec<(i64, u64)> = data
2758 .iter()
2759 .zip(row_ids.iter())
2760 .map(|(v, r)| (*v, *r))
2761 .collect();
2762 Arc::make_mut(&mut self.learned_range).insert(
2763 cid,
2764 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2765 );
2766 }
2767 }
2768 TypeId::Float64 => {
2769 if let columnar::NativeColumn::Float64 { data, .. } =
2770 reader.column_native(cid)?
2771 {
2772 let pairs: Vec<(f64, u64)> = data
2773 .iter()
2774 .zip(row_ids.iter())
2775 .map(|(v, r)| (*v, *r))
2776 .collect();
2777 Arc::make_mut(&mut self.learned_range).insert(
2778 cid,
2779 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2780 );
2781 }
2782 }
2783 _ => {}
2784 }
2785 }
2786 Ok(())
2787 }
2788
2789 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2796 if self.indexes_complete {
2797 crate::trace::QueryTrace::record(|t| {
2798 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2799 });
2800 return Ok(());
2801 }
2802 crate::trace::QueryTrace::record(|t| {
2803 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2804 });
2805 self.rebuild_indexes_from_runs()?;
2806 self.build_learned_ranges()?;
2807 self.indexes_complete = true;
2808 let epoch = self.current_epoch();
2809 self.checkpoint_indexes(epoch);
2810 Ok(())
2811 }
2812
2813 #[doc(hidden)]
2816 pub fn ensure_indexes_complete_controlled<F>(
2817 &mut self,
2818 control: &crate::ExecutionControl,
2819 before_publish: F,
2820 ) -> Result<bool>
2821 where
2822 F: FnOnce() -> bool,
2823 {
2824 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2825 .map(|(changed, _)| changed)
2826 }
2827
2828 #[doc(hidden)]
2831 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2832 &mut self,
2833 control: &crate::ExecutionControl,
2834 before_publish: F,
2835 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2836 where
2837 F: FnOnce() -> bool,
2838 {
2839 if self.indexes_complete {
2840 crate::trace::QueryTrace::record(|trace| {
2841 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2842 });
2843 return Ok((false, None));
2844 }
2845 crate::trace::QueryTrace::record(|trace| {
2846 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2847 });
2848 control.checkpoint()?;
2849 let maintenance_epoch = self.current_epoch();
2850 self.rebuild_indexes_from_runs_inner(Some(control))?;
2851 self.build_learned_ranges_inner(Some(control))?;
2852 control.checkpoint()?;
2853 if !before_publish() {
2854 return Err(MongrelError::Cancelled);
2855 }
2856 self.indexes_complete = true;
2857 self.checkpoint_indexes(maintenance_epoch);
2858 Ok((
2859 true,
2860 Some(MaintenanceReceipt {
2861 epoch: maintenance_epoch,
2862 }),
2863 ))
2864 }
2865
2866 fn pending_epoch(&self) -> Epoch {
2867 Epoch(self.epoch.visible().0 + 1)
2868 }
2869
2870 fn is_shared(&self) -> bool {
2873 matches!(self.wal, WalSink::Shared(_))
2874 }
2875
2876 fn ensure_txn_id(&mut self) -> Result<u64> {
2880 if self.current_txn_id == 0 {
2881 let id = match &self.wal {
2882 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2883 WalSink::Private(_) => {
2884 return Err(MongrelError::Full(
2885 "standalone transaction id namespace exhausted".into(),
2886 ))
2887 }
2888 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2889 };
2890 self.current_txn_id = id;
2891 }
2892 Ok(self.current_txn_id)
2893 }
2894
2895 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2898 self.ensure_writable()?;
2899 let txn_id = self.ensure_txn_id()?;
2900 let table_id = self.table_id;
2901 match &mut self.wal {
2902 WalSink::Private(w) => {
2903 w.append_txn(txn_id, op)?;
2904 self.pending_private_mutations = true;
2905 }
2906 WalSink::Shared(s) => {
2907 s.wal.lock().append(txn_id, table_id, op)?;
2908 }
2909 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2910 }
2911 Ok(())
2912 }
2913
2914 fn ensure_writable(&self) -> Result<()> {
2915 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2916 return Err(MongrelError::ReadOnlyReplica);
2917 }
2918 if self.durable_commit_failed {
2919 return Err(MongrelError::Other(
2920 "table poisoned by post-commit failure; reopen required".into(),
2921 ));
2922 }
2923 Ok(())
2924 }
2925
2926 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2937 match &self.auth {
2938 Some(checker) => checker.check(&self.name, perm),
2939 None => Ok(()),
2940 }
2941 }
2942 pub fn require_select(&self) -> Result<()> {
2947 self.require(crate::auth_state::RequiredPermission::Select)
2948 }
2949 fn require_insert(&self) -> Result<()> {
2950 self.require(crate::auth_state::RequiredPermission::Insert)
2951 }
2952 #[allow(dead_code)]
2956 fn require_update(&self) -> Result<()> {
2957 self.require(crate::auth_state::RequiredPermission::Update)
2958 }
2959 fn require_delete(&self) -> Result<()> {
2960 self.require(crate::auth_state::RequiredPermission::Delete)
2961 }
2962
2963 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2966 self.require_insert()?;
2967 Ok(self.put_returning(columns)?.0)
2968 }
2969
2970 pub fn put_returning(
2975 &mut self,
2976 mut columns: Vec<(u16, Value)>,
2977 ) -> Result<(RowId, Option<i64>)> {
2978 self.require_insert()?;
2979 let assigned = self.fill_auto_inc(&mut columns)?;
2980 self.apply_defaults(&mut columns)?;
2981 self.schema.validate_values(&columns)?;
2982 let row_id = if self.schema.clustered {
2987 self.derive_clustered_row_id(&columns)?
2988 } else {
2989 self.allocator.alloc()?
2990 };
2991 let epoch = self.pending_epoch();
2992 let mut row = Row::new(row_id, epoch);
2993 for (col_id, val) in columns {
2994 row.columns.insert(col_id, val);
2995 }
2996 self.commit_rows(vec![row], assigned.is_some())?;
2997 Ok((row_id, assigned))
2998 }
2999
3000 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
3003 self.require_insert()?;
3004 Ok(self
3005 .put_batch_returning(batch)?
3006 .into_iter()
3007 .map(|(r, _)| r)
3008 .collect())
3009 }
3010
3011 pub fn put_batch_returning(
3014 &mut self,
3015 batch: Vec<Vec<(u16, Value)>>,
3016 ) -> Result<Vec<(RowId, Option<i64>)>> {
3017 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
3018 for mut cols in batch {
3019 let assigned = self.fill_auto_inc(&mut cols)?;
3020 self.apply_defaults(&mut cols)?;
3021 filled.push((cols, assigned));
3022 }
3023 for (cols, _) in &filled {
3024 self.schema.validate_values(cols)?;
3025 }
3026 let epoch = self.pending_epoch();
3027 let mut rows = Vec::with_capacity(filled.len());
3028 let mut ids = Vec::with_capacity(filled.len());
3029 let first_row_id = if self.schema.clustered {
3030 None
3031 } else {
3032 let count = u64::try_from(filled.len())
3033 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
3034 Some(self.allocator.alloc_range(count)?.0)
3035 };
3036 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
3037 let row_id = match first_row_id {
3038 Some(first) => RowId(first + row_index as u64),
3039 None => self.derive_clustered_row_id(&cols)?,
3040 };
3041 let mut row = Row::new(row_id, epoch);
3042 for (c, v) in cols {
3043 row.columns.insert(c, v);
3044 }
3045 ids.push((row_id, assigned));
3046 rows.push(row);
3047 }
3048 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
3049 self.commit_rows(rows, all_auto_generated)?;
3050 Ok(ids)
3051 }
3052
3053 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
3059 self.ensure_writable()?;
3060 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3061 return Ok(None);
3062 };
3063 let pos = columns.iter().position(|(c, _)| *c == cid);
3064 let assigned = match pos {
3065 Some(i) => match &columns[i].1 {
3066 Value::Null => {
3067 let next = self.alloc_auto_inc_value()?;
3068 columns[i].1 = Value::Int64(next);
3069 Some(next)
3070 }
3071 Value::Int64(n) => {
3072 self.advance_auto_inc_past(*n)?;
3073 None
3074 }
3075 other => {
3076 return Err(MongrelError::InvalidArgument(format!(
3077 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
3078 other
3079 )))
3080 }
3081 },
3082 None => {
3083 let next = self.alloc_auto_inc_value()?;
3084 columns.push((cid, Value::Int64(next)));
3085 Some(next)
3086 }
3087 };
3088 Ok(assigned)
3089 }
3090
3091 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
3097 for col in &self.schema.columns {
3098 let Some(expr) = &col.default_value else {
3099 continue;
3100 };
3101 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
3103 continue;
3104 }
3105 let pos = columns.iter().position(|(c, _)| *c == col.id);
3106 let needs_default = match pos {
3107 None => true,
3108 Some(i) => matches!(columns[i].1, Value::Null),
3109 };
3110 if !needs_default {
3111 continue;
3112 }
3113 let v = match expr {
3114 crate::schema::DefaultExpr::Static(v) => v.clone(),
3115 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
3116 crate::schema::DefaultExpr::Uuid => {
3117 let mut buf = [0u8; 16];
3118 getrandom::getrandom(&mut buf)
3119 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
3120 Value::Uuid(buf)
3121 }
3122 };
3123 match pos {
3124 None => columns.push((col.id, v)),
3125 Some(i) => columns[i].1 = v,
3126 }
3127 }
3128 Ok(())
3129 }
3130
3131 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
3133 self.ensure_auto_inc_seeded()?;
3134 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3136 let v = ai.next;
3137 ai.next = ai
3138 .next
3139 .checked_add(1)
3140 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3141 Ok(v)
3142 }
3143
3144 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
3147 self.ensure_auto_inc_seeded()?;
3148 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3149 let floor = used
3150 .checked_add(1)
3151 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3152 .max(1);
3153 if ai.next < floor {
3154 ai.next = floor;
3155 }
3156 Ok(())
3157 }
3158
3159 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
3164 let needs_seed = match self.auto_inc {
3165 Some(ai) => !ai.seeded,
3166 None => return Ok(()),
3167 };
3168 if !needs_seed {
3169 return Ok(());
3170 }
3171 if self.seed_empty_auto_inc() {
3172 return Ok(());
3173 }
3174 let cid = self
3175 .auto_inc
3176 .as_ref()
3177 .expect("auto-inc column present")
3178 .column_id;
3179 let max = self.scan_max_int64(cid)?;
3180 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3181 let floor = max
3182 .checked_add(1)
3183 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3184 .max(1);
3185 if ai.next < floor {
3186 ai.next = floor;
3187 }
3188 ai.seeded = true;
3189 Ok(())
3190 }
3191
3192 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
3193 if n == 0 || self.auto_inc.is_none() {
3194 return Ok(None);
3195 }
3196 self.ensure_auto_inc_seeded()?;
3197 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3198 let start = ai.next;
3199 let count = i64::try_from(n)
3200 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
3201 ai.next = ai
3202 .next
3203 .checked_add(count)
3204 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3205 Ok(Some(start))
3206 }
3207
3208 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
3212 let mut max: i64 = 0;
3213 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
3214 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3215 if *n > max {
3216 max = *n;
3217 }
3218 }
3219 }
3220 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3221 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3222 if *n > max {
3223 max = *n;
3224 }
3225 }
3226 }
3227 for rr in self.run_refs.clone() {
3228 let reader = self.open_reader(rr.run_id)?;
3229 if let Some(stats) = reader.column_page_stats(column_id) {
3230 for s in stats {
3231 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3232 if n > max {
3233 max = n;
3234 }
3235 }
3236 }
3237 } else if reader.has_column(column_id) {
3238 if let columnar::NativeColumn::Int64 { data, validity } =
3239 reader.column_native_shared(column_id)?
3240 {
3241 for (i, n) in data.iter().enumerate() {
3242 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3243 {
3244 max = *n;
3245 }
3246 }
3247 }
3248 }
3249 }
3250 Ok(max)
3251 }
3252
3253 fn seed_empty_auto_inc(&mut self) -> bool {
3254 let Some(ai) = self.auto_inc.as_mut() else {
3255 return false;
3256 };
3257 if ai.seeded || self.live_count != 0 {
3258 return false;
3259 }
3260 if ai.next < 1 {
3261 ai.next = 1;
3262 }
3263 ai.seeded = true;
3264 true
3265 }
3266
3267 fn advance_auto_inc_from_native_columns(
3268 &mut self,
3269 columns: &[(u16, columnar::NativeColumn)],
3270 n: usize,
3271 live_before: u64,
3272 ) -> Result<()> {
3273 let Some(ai) = self.auto_inc.as_mut() else {
3274 return Ok(());
3275 };
3276 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3277 return Ok(());
3278 };
3279 let columnar::NativeColumn::Int64 { data, validity } = col else {
3280 return Err(MongrelError::InvalidArgument(format!(
3281 "AUTO_INCREMENT column {} must be Int64",
3282 ai.column_id
3283 )));
3284 };
3285 let max = if native_int64_strictly_increasing(col, n) {
3286 data.get(n.saturating_sub(1)).copied()
3287 } else {
3288 data.iter()
3289 .take(n)
3290 .enumerate()
3291 .filter_map(|(i, v)| {
3292 if validity.is_empty() || columnar::validity_bit(validity, i) {
3293 Some(*v)
3294 } else {
3295 None
3296 }
3297 })
3298 .max()
3299 };
3300 if let Some(max) = max {
3301 let floor = max
3302 .checked_add(1)
3303 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3304 .max(1);
3305 if ai.next < floor {
3306 ai.next = floor;
3307 }
3308 if ai.seeded || live_before == 0 {
3309 ai.seeded = true;
3310 }
3311 }
3312 Ok(())
3313 }
3314
3315 fn fill_auto_inc_native_columns(
3316 &mut self,
3317 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3318 n: usize,
3319 ) -> Result<()> {
3320 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3321 return Ok(());
3322 };
3323 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3324 if let Some(start) = self.alloc_auto_inc_range(n)? {
3325 columns.push((
3326 cid,
3327 columnar::NativeColumn::Int64 {
3328 data: (start..start.saturating_add(n as i64)).collect(),
3329 validity: vec![0xFF; n.div_ceil(8)],
3330 },
3331 ));
3332 }
3333 return Ok(());
3334 };
3335
3336 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3337 return Err(MongrelError::InvalidArgument(format!(
3338 "AUTO_INCREMENT column {cid} must be Int64"
3339 )));
3340 };
3341 if data.len() < n {
3342 return Err(MongrelError::InvalidArgument(format!(
3343 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3344 data.len()
3345 )));
3346 }
3347 if columnar::all_non_null(validity, n) {
3348 return Ok(());
3349 }
3350 if validity.iter().all(|b| *b == 0) {
3351 if let Some(start) = self.alloc_auto_inc_range(n)? {
3352 for (i, slot) in data.iter_mut().take(n).enumerate() {
3353 *slot = start.saturating_add(i as i64);
3354 }
3355 *validity = vec![0xFF; n.div_ceil(8)];
3356 }
3357 return Ok(());
3358 }
3359
3360 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3361 for (i, slot) in data.iter_mut().enumerate().take(n) {
3362 if columnar::validity_bit(validity, i) {
3363 self.advance_auto_inc_past(*slot)?;
3364 } else {
3365 *slot = self.alloc_auto_inc_value()?;
3366 }
3367 }
3368 *validity = new_validity;
3369 Ok(())
3370 }
3371
3372 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3386 self.ensure_writable()?;
3387 if self.auto_inc.is_none() {
3388 return Ok(None);
3389 }
3390 Ok(Some(self.alloc_auto_inc_value()?))
3391 }
3392
3393 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3399 let payload = bincode::serialize(&rows)?;
3400 self.wal_append_data(Op::Put {
3401 table_id: self.table_id,
3402 rows: payload,
3403 })?;
3404 if self.is_shared() {
3405 self.pending_rows_auto_inc
3406 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3407 self.pending_rows.extend(rows);
3408 } else {
3409 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3410 }
3411 Ok(())
3412 }
3413
3414 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3417 self.ensure_indexes_complete()
3418 }
3419
3420 pub(crate) fn prepare_durable_publish_controlled(
3421 &mut self,
3422 control: &crate::ExecutionControl,
3423 ) -> Result<()> {
3424 self.ensure_indexes_complete_controlled(control, || true)?;
3425 Ok(())
3426 }
3427
3428 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3429 self.apply_put_rows_inner_prepared(rows, true);
3430 }
3431
3432 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3433 if check_existing_pk {
3434 self.ensure_indexes_complete()?;
3435 }
3436 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3437 Ok(())
3438 }
3439
3440 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3444 if rows.len() == 1 {
3448 let row = rows.into_iter().next().expect("len checked");
3449 self.apply_put_row_single(row, check_existing_pk);
3450 return;
3451 }
3452 let pk_id = self.schema.primary_key().map(|c| c.id);
3469 let probe = match pk_id {
3470 Some(pid) => {
3471 check_existing_pk
3472 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3473 }
3474 None => false,
3475 };
3476 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3479 for r in rows {
3480 for &cid in r.columns.keys() {
3481 self.pending_put_cols.insert(cid);
3482 }
3483 let mut replaced_image: Option<Row> = None;
3484 match pk_id {
3485 Some(pid) if probe || maintain_pk_by_row => {
3486 if let Some(pk_val) = r.columns.get(&pid) {
3487 let key = self.index_lookup_key(pid, pk_val);
3488 if probe {
3489 if let Some(old_rid) = self.hot.get(&key) {
3490 if old_rid != r.row_id {
3491 replaced_image = self.get(old_rid, self.snapshot());
3492 self.tombstone_row(
3493 old_rid,
3494 r.committed_epoch,
3495 r.commit_ts,
3496 true,
3497 );
3498 }
3499 }
3500 }
3501 if maintain_pk_by_row {
3502 self.pk_by_row.insert(r.row_id, key);
3503 }
3504 }
3505 }
3506 Some(_) => {}
3507 None => {
3508 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3509 }
3510 }
3511 if let Some(old) = replaced_image {
3512 self.maintain_indexes_on_pk_replace(&old, &r);
3513 } else {
3514 self.index_row(&r);
3515 }
3516 self.reservoir.offer(r.row_id.0);
3517 self.memtable.upsert(r);
3518 self.live_count = self.live_count.saturating_add(1);
3521 }
3522 self.data_generation = self.data_generation.wrapping_add(1);
3523 }
3524
3525 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3535 for &cid in row.columns.keys() {
3536 self.pending_put_cols.insert(cid);
3537 }
3538 let epoch = row.committed_epoch;
3539 let mut replaced_image: Option<Row> = None;
3540 if let Some(pk_col) = self.schema.primary_key() {
3541 let pk_id = pk_col.id;
3542 if let Some(pk_val) = row.columns.get(&pk_id) {
3543 let maintain_pk_by_row = self.pk_by_row_complete;
3547 if check_existing_pk || maintain_pk_by_row {
3548 let key = self.index_lookup_key(pk_id, pk_val);
3549 if check_existing_pk {
3550 if let Some(old_rid) = self.hot.get(&key) {
3551 if old_rid != row.row_id {
3552 replaced_image = self.get(old_rid, self.snapshot());
3556 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
3557 }
3558 }
3559 }
3560 if maintain_pk_by_row {
3561 self.pk_by_row.insert(row.row_id, key);
3562 }
3563 }
3564 }
3565 } else {
3566 self.hot
3567 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3568 }
3569 if let Some(old) = replaced_image {
3570 self.maintain_indexes_on_pk_replace(&old, &row);
3571 } else {
3572 self.index_row(&row);
3573 }
3574 self.reservoir.offer(row.row_id.0);
3575 self.memtable.upsert(row);
3576 self.live_count = self.live_count.saturating_add(1);
3577 self.data_generation = self.data_generation.wrapping_add(1);
3578 }
3579
3580 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
3583 let has_partial = self
3584 .schema
3585 .indexes
3586 .iter()
3587 .any(|idx| idx.predicate.is_some());
3588 if has_partial {
3589 for idef in &self.schema.indexes {
3593 if idef.kind != crate::schema::IndexKind::Bitmap {
3594 continue;
3595 }
3596 if let Some(key) = crate::index::maintain::bitmap_key_for_column(old, idef.column_id)
3597 {
3598 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3599 b.remove(&key, old.row_id);
3600 }
3601 }
3602 }
3603 self.index_row(new);
3604 return;
3605 }
3606
3607 crate::index::maintain_bitmap_secondary_on_replace(
3608 &self.schema,
3609 &mut self.bitmap,
3610 old,
3611 new,
3612 );
3613 self.index_row_non_bitmap(new);
3614 }
3615
3616 fn index_row_non_bitmap(&mut self, row: &Row) {
3619 if row.deleted {
3620 return;
3621 }
3622 let effective = if self.column_keys.is_empty() {
3623 None
3624 } else {
3625 Some(self.tokenized_for_indexes(row))
3626 };
3627 let source = effective.as_ref().unwrap_or(row);
3628 for idef in &self.schema.indexes {
3629 if idef.kind == crate::schema::IndexKind::Bitmap {
3630 continue;
3631 }
3632 index_into_single(
3633 idef,
3634 &self.schema,
3635 source,
3636 &mut self.hot,
3637 &mut self.bitmap,
3638 &mut self.ann,
3639 &mut self.fm,
3640 &mut self.sparse,
3641 &mut self.minhash,
3642 );
3643 }
3644 if let Some(pk_col) = self.schema.primary_key() {
3645 if let Some(pk_val) = source.columns.get(&pk_col.id) {
3646 let key = if self.column_keys.is_empty() {
3647 pk_val.encode_key()
3648 } else {
3649 self.index_lookup_key(pk_col.id, pk_val)
3650 };
3651 self.hot.insert(key, source.row_id);
3652 }
3653 }
3654 }
3655
3656 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3659 self.allocator.alloc()
3660 }
3661
3662 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3668 let pk = self.schema.primary_key().ok_or_else(|| {
3669 MongrelError::Schema("clustered table requires a single-column primary key".into())
3670 })?;
3671 let pk_val = columns
3672 .iter()
3673 .find(|(id, _)| *id == pk.id)
3674 .map(|(_, v)| v)
3675 .ok_or_else(|| {
3676 MongrelError::Schema(format!(
3677 "clustered table missing primary key column {} ({})",
3678 pk.id, pk.name
3679 ))
3680 })?;
3681 Ok(clustered_row_id(pk_val))
3682 }
3683
3684 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3692 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3693 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3694 }
3695 let n = rows.len();
3696 for r in rows {
3697 for &cid in r.columns.keys() {
3698 self.pending_put_cols.insert(cid);
3699 }
3700 }
3701 let (losers, winner_pks) = self.partition_pk_winners(rows);
3702 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3703 let group_ts = rows.first().and_then(|r| r.commit_ts);
3705 for (key, &row_id) in &winner_pks {
3706 if let Some(old_rid) = self.hot.get(key) {
3707 if old_rid != row_id {
3708 self.tombstone_row(old_rid, epoch, group_ts, true);
3709 }
3710 }
3711 }
3712 for &loser_rid in &losers {
3715 self.tombstone_row(loser_rid, epoch, group_ts, false);
3716 }
3717 for (key, row_id) in winner_pks {
3719 self.insert_hot_pk(key, row_id);
3720 }
3721 if self.schema.primary_key().is_none() {
3722 for r in rows {
3723 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3724 }
3725 }
3726 for r in rows {
3727 self.allocator.advance_to(r.row_id)?;
3728 if !losers.contains(&r.row_id) {
3729 self.index_row(r);
3730 }
3731 }
3732 for r in rows {
3733 if !losers.contains(&r.row_id) {
3734 self.reservoir.offer(r.row_id.0);
3735 }
3736 }
3737 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3738 self.data_generation = self.data_generation.wrapping_add(1);
3739 Ok(())
3740 }
3741
3742 pub(crate) fn recover_apply(
3747 &mut self,
3748 rows: Vec<Row>,
3749 deletes: Vec<(RowId, Epoch)>,
3750 ) -> Result<()> {
3751 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3755 std::collections::BTreeMap::new();
3756 for row in rows {
3757 if row.row_id.0 >= u64::MAX - 1 {
3758 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3759 }
3760 self.allocator.advance_to(row.row_id)?;
3761 if let Some(ai) = self.auto_inc.as_mut() {
3766 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3767 let next = n.checked_add(1).ok_or_else(|| {
3768 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3769 })?;
3770 if next > ai.next {
3771 ai.next = next;
3772 }
3773 }
3774 }
3775 by_epoch.entry(row.committed_epoch).or_default().push(row);
3776 }
3777 for (epoch, group) in by_epoch {
3778 let (losers, winner_pks) = self.partition_pk_winners(&group);
3779 let group_ts = group.first().and_then(|r| r.commit_ts);
3781 for (key, &row_id) in &winner_pks {
3782 if let Some(old_rid) = self.hot.get(key) {
3783 if old_rid != row_id {
3784 self.tombstone_row(old_rid, epoch, group_ts, false);
3785 }
3786 }
3787 }
3788 for (key, row_id) in winner_pks {
3789 self.insert_hot_pk(key, row_id);
3790 }
3791 if self.schema.primary_key().is_none() {
3792 for r in &group {
3793 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3794 }
3795 }
3796 for r in &group {
3797 if !losers.contains(&r.row_id) {
3798 self.memtable.upsert(r.clone());
3799 self.index_row(r);
3800 }
3801 }
3802 }
3803 for (rid, epoch) in deletes {
3804 self.memtable.tombstone(rid, epoch);
3805 self.remove_hot_for_row(rid, epoch);
3806 }
3807 self.reservoir_complete = false;
3810 Ok(())
3811 }
3812
3813 pub(crate) fn flushed_epoch(&self) -> u64 {
3815 self.flushed_epoch
3816 }
3817
3818 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3819 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3820 }
3821
3822 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3824 self.schema.validate_values(cells)
3825 }
3826
3827 fn validate_columns_not_null(
3831 &self,
3832 columns: &[(u16, columnar::NativeColumn)],
3833 n: usize,
3834 ) -> Result<()> {
3835 let by_id: HashMap<u16, &columnar::NativeColumn> =
3836 columns.iter().map(|(id, c)| (*id, c)).collect();
3837 for col in &self.schema.columns {
3838 if !col.flags.contains(ColumnFlags::NULLABLE) {
3839 match by_id.get(&col.id) {
3840 None => {
3841 return Err(MongrelError::InvalidArgument(format!(
3842 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3843 col.name, col.id
3844 )));
3845 }
3846 Some(c) => {
3847 if c.null_count(n) != 0 {
3848 return Err(MongrelError::InvalidArgument(format!(
3849 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3850 col.name, col.id
3851 )));
3852 }
3853 }
3854 }
3855 }
3856 if let TypeId::Enum { variants } = &col.ty {
3857 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3858 if by_id.contains_key(&col.id) {
3859 return Err(MongrelError::InvalidArgument(format!(
3860 "column '{}' ({}) enum requires a bytes column",
3861 col.name, col.id
3862 )));
3863 }
3864 continue;
3865 };
3866 for index in 0..n {
3867 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3868 continue;
3869 };
3870 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3871 return Err(MongrelError::InvalidArgument(format!(
3872 "column '{}' ({}) enum value {:?} is not one of {:?}",
3873 col.name,
3874 col.id,
3875 String::from_utf8_lossy(value),
3876 variants
3877 )));
3878 }
3879 }
3880 }
3881 }
3882 Ok(())
3883 }
3884
3885 fn bulk_pk_winner_indices(
3890 &self,
3891 columns: &[(u16, columnar::NativeColumn)],
3892 n: usize,
3893 ) -> Option<Vec<usize>> {
3894 let pk_col = self.schema.primary_key()?;
3895 let pk_id = pk_col.id;
3896 let pk_ty = pk_col.ty.clone();
3897 let by_id: HashMap<u16, &columnar::NativeColumn> =
3898 columns.iter().map(|(id, c)| (*id, c)).collect();
3899 let pk_native = by_id.get(&pk_id)?;
3900 if native_int64_strictly_increasing(pk_native, n) {
3901 return None;
3902 }
3903 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3905 let mut null_pk_rows: Vec<usize> = Vec::new();
3906 for i in 0..n {
3907 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3908 Some(key) => {
3909 last.insert(key, i);
3910 }
3911 None => null_pk_rows.push(i),
3912 }
3913 }
3914 let mut winners: HashSet<usize> = last.values().copied().collect();
3915 for i in null_pk_rows {
3916 winners.insert(i);
3917 }
3918 Some((0..n).filter(|i| winners.contains(i)).collect())
3919 }
3920
3921 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3923 self.require_delete()?;
3924 let epoch = self.pending_epoch();
3925 self.wal_append_data(Op::Delete {
3926 table_id: self.table_id,
3927 row_ids: vec![row_id],
3928 })?;
3929 if self.is_shared() {
3930 self.pending_dels.push(row_id);
3931 } else {
3932 self.apply_delete(row_id, epoch);
3933 }
3934 Ok(())
3935 }
3936
3937 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3938 let pre = self.get(row_id, self.snapshot());
3939 self.delete(row_id)?;
3940 Ok(pre.map(|row| {
3941 let mut columns: Vec<_> = row.columns.into_iter().collect();
3942 columns.sort_by_key(|(id, _)| *id);
3943 OwnedRow { columns }
3944 }))
3945 }
3946
3947 pub fn truncate(&mut self) -> Result<()> {
3949 self.require_delete()?;
3950 let epoch = self.pending_epoch();
3951 self.wal_append_data(Op::TruncateTable {
3952 table_id: self.table_id,
3953 })?;
3954 self.pending_rows.clear();
3955 self.pending_rows_auto_inc.clear();
3956 self.pending_dels.clear();
3957 self.pending_truncate = Some(epoch);
3958 Ok(())
3959 }
3960
3961 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3963 self.run_refs.clear();
3968 self.retiring.clear();
3969 self.memtable = Memtable::new();
3970 self.mutable_run = MutableRun::new();
3971 self.hot = HotIndex::new();
3972 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3973 self.bitmap = bitmap;
3974 self.ann = ann;
3975 self.fm = fm;
3976 self.sparse = sparse;
3977 self.minhash = minhash;
3978 self.learned_range = Arc::new(HashMap::new());
3979 self.pk_by_row.clear();
3980 self.pk_by_row_complete = false;
3981 self.live_count = 0;
3982 self.reservoir = crate::reservoir::Reservoir::default();
3983 self.reservoir_complete = true;
3984 self.had_deletes = true;
3985 self.agg_cache = Arc::new(HashMap::new());
3986 self.global_idx_epoch = 0;
3987 self.indexes_complete = true;
3988 self.pending_delete_rids.clear();
3989 self.pending_put_cols.clear();
3990 self.pending_rows.clear();
3991 self.pending_rows_auto_inc.clear();
3992 self.pending_dels.clear();
3993 self.clear_result_cache();
3994 self.invalidate_index_checkpoint();
3995 self.data_generation = self.data_generation.wrapping_add(1);
3996 }
3997
3998 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
4001 self.apply_delete_at(row_id, epoch, None);
4002 }
4003
4004 pub(crate) fn apply_delete_at(
4006 &mut self,
4007 row_id: RowId,
4008 epoch: Epoch,
4009 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4010 ) {
4011 self.remove_hot_for_row(row_id, epoch);
4012 self.tombstone_row(row_id, epoch, commit_ts, true);
4013 self.data_generation = self.data_generation.wrapping_add(1);
4014 }
4015
4016 fn tombstone_row(
4020 &mut self,
4021 row_id: RowId,
4022 epoch: Epoch,
4023 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4024 adjust_live_count: bool,
4025 ) {
4026 let tombstone = Row {
4027 row_id,
4028 committed_epoch: epoch,
4029 columns: std::collections::HashMap::new(),
4030 deleted: true,
4031 commit_ts,
4032 };
4033 self.memtable.upsert(tombstone);
4034 self.pk_by_row.remove(&row_id);
4035 if adjust_live_count {
4036 self.live_count = self.live_count.saturating_sub(1);
4037 }
4038 self.pending_delete_rids.insert(row_id.0 as u32);
4040 self.had_deletes = true;
4043 self.agg_cache = Arc::new(HashMap::new());
4044 }
4045
4046 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
4050 let Some(pk_col) = self.schema.primary_key() else {
4051 return;
4052 };
4053 if self.pk_by_row_complete {
4056 if let Some(key) = self.pk_by_row.remove(&row_id) {
4057 if self.hot.get(&key) == Some(row_id) {
4058 self.hot.remove(&key);
4059 }
4060 }
4061 return;
4062 }
4063 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
4082 if self.indexes_complete {
4083 let pk_val = self
4084 .memtable
4085 .get_version(row_id, lookup_epoch)
4086 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4087 .or_else(|| {
4088 self.mutable_run
4089 .get_version(row_id, lookup_epoch)
4090 .filter(|(_, r)| !r.deleted)
4091 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4092 })
4093 .or_else(|| {
4094 self.run_refs.iter().find_map(|rr| {
4095 let mut reader = self.open_reader(rr.run_id).ok()?;
4096 let (_, deleted, val) = reader
4097 .get_version_column(row_id, lookup_epoch, pk_col.id)
4098 .ok()??;
4099 if deleted {
4100 return None;
4101 }
4102 val
4103 })
4104 });
4105 if let Some(pk_val) = pk_val {
4106 let key = self.index_lookup_key(pk_col.id, &pk_val);
4107 if self.hot.get(&key) == Some(row_id) {
4108 self.hot.remove(&key);
4109 }
4110 return;
4111 }
4112 }
4113 self.refresh_pk_by_row_from_hot();
4118 if let Some(key) = self.pk_by_row.remove(&row_id) {
4119 if self.hot.get(&key) == Some(row_id) {
4120 self.hot.remove(&key);
4121 }
4122 }
4123 }
4124
4125 fn partition_pk_winners(
4130 &self,
4131 rows: &[Row],
4132 ) -> (
4133 std::collections::HashSet<RowId>,
4134 std::collections::HashMap<Vec<u8>, RowId>,
4135 ) {
4136 let mut losers = std::collections::HashSet::new();
4137 let Some(pk_col) = self.schema.primary_key() else {
4138 return (losers, std::collections::HashMap::new());
4139 };
4140 let pk_id = pk_col.id;
4141 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
4142 std::collections::HashMap::new();
4143 for r in rows {
4144 let Some(pk_val) = r.columns.get(&pk_id) else {
4145 continue;
4146 };
4147 let key = self.index_lookup_key(pk_id, pk_val);
4148 if let Some(&old_rid) = winners.get(&key) {
4149 losers.insert(old_rid);
4150 }
4151 winners.insert(key, r.row_id);
4152 }
4153 (losers, winners)
4154 }
4155
4156 fn index_row(&mut self, row: &Row) {
4157 if row.deleted {
4158 return;
4159 }
4160 let any_predicate = self
4168 .schema
4169 .indexes
4170 .iter()
4171 .any(|idx| idx.predicate.is_some());
4172 if any_predicate {
4173 let columns_map: HashMap<u16, &Value> =
4174 row.columns.iter().map(|(k, v)| (*k, v)).collect();
4175 let name_to_id: HashMap<&str, u16> = self
4176 .schema
4177 .columns
4178 .iter()
4179 .map(|c| (c.name.as_str(), c.id))
4180 .collect();
4181 for idx in &self.schema.indexes {
4182 if let Some(pred) = &idx.predicate {
4183 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4184 continue; }
4186 }
4187 index_into_single(
4189 idx,
4190 &self.schema,
4191 row,
4192 &mut self.hot,
4193 &mut self.bitmap,
4194 &mut self.ann,
4195 &mut self.fm,
4196 &mut self.sparse,
4197 &mut self.minhash,
4198 );
4199 }
4200 return;
4201 }
4202 if self.column_keys.is_empty() {
4206 index_into(
4207 &self.schema,
4208 row,
4209 &mut self.hot,
4210 &mut self.bitmap,
4211 &mut self.ann,
4212 &mut self.fm,
4213 &mut self.sparse,
4214 &mut self.minhash,
4215 );
4216 return;
4217 }
4218 let effective_row = self.tokenized_for_indexes(row);
4219 index_into(
4220 &self.schema,
4221 &effective_row,
4222 &mut self.hot,
4223 &mut self.bitmap,
4224 &mut self.ann,
4225 &mut self.fm,
4226 &mut self.sparse,
4227 &mut self.minhash,
4228 );
4229 }
4230
4231 fn tokenized_for_indexes(&self, row: &Row) -> Row {
4237 if self.column_keys.is_empty() {
4238 return row.clone();
4239 }
4240 {
4241 use crate::encryption::SCHEME_HMAC_EQ;
4242 let mut tok = row.clone();
4243 for (&cid, &(_, scheme)) in &self.column_keys {
4244 if scheme != SCHEME_HMAC_EQ {
4245 continue;
4246 }
4247 if let Some(v) = tok.columns.get(&cid).cloned() {
4248 if let Some(t) = self.tokenize_value(cid, &v) {
4249 tok.columns.insert(cid, t);
4250 }
4251 }
4252 }
4253 tok
4254 }
4255 }
4256
4257 pub fn commit(&mut self) -> Result<Epoch> {
4262 self.commit_inner(None)
4263 }
4264
4265 #[doc(hidden)]
4268 pub fn commit_controlled<F>(
4269 &mut self,
4270 control: &crate::ExecutionControl,
4271 mut before_commit: F,
4272 ) -> Result<Epoch>
4273 where
4274 F: FnMut() -> Result<()>,
4275 {
4276 self.commit_inner(Some((control, &mut before_commit)))
4277 }
4278
4279 fn commit_inner(
4280 &mut self,
4281 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4282 ) -> Result<Epoch> {
4283 self.ensure_writable()?;
4284 if !self.has_pending_mutations() {
4285 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
4286 return Err(MongrelError::Full(
4287 "standalone transaction id namespace exhausted".into(),
4288 ));
4289 }
4290 return Ok(self.epoch.visible());
4291 }
4292 self.commit_new_epoch_inner(controlled)
4293 }
4294
4295 fn commit_new_epoch(&mut self) -> Result<Epoch> {
4299 self.commit_new_epoch_inner(None)
4300 }
4301
4302 fn commit_new_epoch_inner(
4303 &mut self,
4304 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4305 ) -> Result<Epoch> {
4306 self.ensure_writable()?;
4307 if self.is_shared() {
4308 self.commit_shared(controlled)
4309 } else {
4310 self.commit_private(controlled)
4311 }
4312 }
4313
4314 fn commit_private(
4316 &mut self,
4317 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4318 ) -> Result<Epoch> {
4319 let commit_lock = Arc::clone(&self.commit_lock);
4323 let _g = commit_lock.lock();
4324 let txn_id = self.ensure_txn_id()?;
4327 if let Some((control, before_commit)) = controlled {
4328 control.checkpoint()?;
4329 before_commit()?;
4330 }
4331 let new_epoch = self.epoch.bump_assigned();
4332 let epoch_authority = Arc::clone(&self.epoch);
4333 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4334 let wal_result = match &mut self.wal {
4338 WalSink::Private(w) => w
4339 .append_txn(
4340 txn_id,
4341 Op::TxnCommit {
4342 epoch: new_epoch.0,
4343 added_runs: Vec::new(),
4344 },
4345 )
4346 .and_then(|_| w.sync()),
4347 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4348 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4349 };
4350 if let Err(error) = wal_result {
4351 self.durable_commit_failed = true;
4352 return Err(MongrelError::CommitOutcomeUnknown {
4353 epoch: new_epoch.0,
4354 message: error.to_string(),
4355 });
4356 }
4357 if let Some(epoch) = self.pending_truncate.take() {
4360 self.apply_truncate(epoch);
4361 }
4362 self.invalidate_pending_cache();
4363 let publish_result = self.persist_manifest(new_epoch);
4364 self.epoch.publish_in_order(new_epoch);
4368 epoch_guard.disarm();
4369 if let Err(error) = publish_result {
4370 self.durable_commit_failed = true;
4371 return Err(MongrelError::DurableCommit {
4372 epoch: new_epoch.0,
4373 message: error.to_string(),
4374 });
4375 }
4376 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4377 self.pending_private_mutations = false;
4378 self.data_generation = self.data_generation.wrapping_add(1);
4379 Ok(new_epoch)
4380 }
4381
4382 fn commit_shared(
4388 &mut self,
4389 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4390 ) -> Result<Epoch> {
4391 use std::sync::atomic::Ordering;
4392 let s = match &self.wal {
4393 WalSink::Shared(s) => s.clone(),
4394 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4395 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4396 };
4397 if s.poisoned.load(Ordering::Relaxed) {
4398 return Err(MongrelError::Other(
4399 "database poisoned by fsync error".into(),
4400 ));
4401 }
4402 let commit_lock = Arc::clone(&self.commit_lock);
4409 let _g = commit_lock.lock();
4410 if !self.pending_rows.is_empty() {
4411 match controlled.as_ref() {
4412 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4413 None => self.prepare_durable_publish()?,
4414 }
4415 }
4416 let txn_id = self.ensure_txn_id()?;
4419 let mut wal = s.wal.lock();
4420 if let Some((control, before_commit)) = controlled {
4421 control.checkpoint()?;
4422 before_commit()?;
4423 }
4424 let new_epoch = self.epoch.bump_assigned();
4425 let epoch_authority = Arc::clone(&self.epoch);
4426 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4427 let commit_ts = s.hlc.now().map_err(|skew| {
4430 MongrelError::Other(format!(
4431 "clock skew rejected commit timestamp allocation: {skew}"
4432 ))
4433 })?;
4434 let commit_seq = match wal.append_commit_at(
4435 txn_id,
4436 new_epoch,
4437 &[],
4438 commit_ts.physical_micros.saturating_mul(1_000),
4439 ) {
4440 Ok(commit_seq) => commit_seq,
4441 Err(error) => {
4442 s.poisoned.store(true, Ordering::Relaxed);
4443 s.lifecycle.poison();
4444 return Err(MongrelError::CommitOutcomeUnknown {
4445 epoch: new_epoch.0,
4446 message: error.to_string(),
4447 });
4448 }
4449 };
4450 drop(wal);
4451 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4452 s.poisoned.store(true, Ordering::Relaxed);
4453 s.lifecycle.poison();
4454 return Err(MongrelError::CommitOutcomeUnknown {
4455 epoch: new_epoch.0,
4456 message: error.to_string(),
4457 });
4458 }
4459
4460 if self.pending_truncate.take().is_some() {
4463 self.apply_truncate(new_epoch);
4464 }
4465 let mut rows = std::mem::take(&mut self.pending_rows);
4466 if !rows.is_empty() {
4467 for r in &mut rows {
4468 r.committed_epoch = new_epoch;
4469 r.commit_ts = Some(commit_ts);
4470 }
4471 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4472 let all_auto_generated =
4473 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4474 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4475 } else {
4476 self.pending_rows_auto_inc.clear();
4477 }
4478 let dels = std::mem::take(&mut self.pending_dels);
4479 for rid in dels {
4480 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
4481 }
4482
4483 self.invalidate_pending_cache();
4484 let publish_result = self.persist_manifest(new_epoch);
4485 self.epoch.publish_in_order(new_epoch);
4486 epoch_guard.disarm();
4487 let _ = s.change_wake.send(());
4488 if let Err(error) = publish_result {
4489 self.durable_commit_failed = true;
4490 s.poisoned.store(true, Ordering::Relaxed);
4491 s.lifecycle.poison();
4492 return Err(MongrelError::DurableCommit {
4493 epoch: new_epoch.0,
4494 message: error.to_string(),
4495 });
4496 }
4497 self.current_txn_id = 0;
4499 self.data_generation = self.data_generation.wrapping_add(1);
4500 Ok(new_epoch)
4501 }
4502
4503 pub fn flush(&mut self) -> Result<Epoch> {
4511 self.flush_with_outcome().map(|(epoch, _)| epoch)
4512 }
4513
4514 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4516 self.flush_with_outcome_inner(None)
4517 }
4518
4519 #[doc(hidden)]
4522 pub fn flush_with_outcome_controlled<F>(
4523 &mut self,
4524 control: &crate::ExecutionControl,
4525 mut before_commit: F,
4526 ) -> Result<(Epoch, bool)>
4527 where
4528 F: FnMut() -> Result<()>,
4529 {
4530 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4531 }
4532
4533 fn flush_with_outcome_inner(
4534 &mut self,
4535 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4536 ) -> Result<(Epoch, bool)> {
4537 match controlled.as_ref() {
4538 Some((control, _)) => {
4539 self.ensure_indexes_complete_controlled(control, || true)?;
4540 }
4541 None => self.ensure_indexes_complete()?,
4542 }
4543 let committed = self.has_pending_mutations();
4544 let epoch = self.commit_inner(controlled)?;
4545 let finish: Result<(Epoch, bool)> = (|| {
4546 let rows = self.memtable.drain_sorted();
4547 if !rows.is_empty() {
4548 self.mutable_run.insert_many(rows);
4549 }
4550 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4551 self.spill_mutable_run(epoch)?;
4552 self.mark_flushed(epoch)?;
4556 self.persist_manifest(epoch)?;
4557 self.build_learned_ranges()?;
4558 self.checkpoint_indexes(epoch);
4561 }
4562 Ok((epoch, committed))
4565 })();
4566 let outcome = match finish {
4567 Err(error) if committed => Err(MongrelError::DurableCommit {
4568 epoch: epoch.0,
4569 message: error.to_string(),
4570 }),
4571 result => result,
4572 };
4573 if outcome.is_ok() {
4574 let _ = self.publish_read_generation();
4580 }
4581 outcome
4582 }
4583
4584 fn has_pending_mutations(&self) -> bool {
4585 self.pending_private_mutations
4586 || !self.pending_rows.is_empty()
4587 || !self.pending_dels.is_empty()
4588 || self.pending_truncate.is_some()
4589 }
4590
4591 pub fn has_pending_writes(&self) -> bool {
4592 self.has_pending_mutations()
4593 }
4594
4595 pub fn force_flush(&mut self) -> Result<Epoch> {
4604 let saved = self.mutable_run_spill_bytes;
4605 self.mutable_run_spill_bytes = 1;
4606 let result = self.flush();
4607 self.mutable_run_spill_bytes = saved;
4608 result
4609 }
4610
4611 pub fn close(&mut self) -> Result<()> {
4618 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4619 self.force_flush()?;
4620 }
4621 Ok(())
4622 }
4623
4624 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4631 let op = Op::Flush {
4632 table_id: self.table_id,
4633 flushed_epoch: epoch.0,
4634 };
4635 match &mut self.wal {
4636 WalSink::Private(w) => {
4637 w.append_system(op)?;
4638 w.sync()?;
4639 }
4640 WalSink::Shared(s) => {
4641 s.wal.lock().append_system(op)?;
4646 }
4647 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4648 }
4649 self.flushed_epoch = epoch.0;
4650 if matches!(self.wal, WalSink::Private(_)) {
4651 self.rotate_wal(epoch)?;
4652 }
4653 Ok(())
4654 }
4655
4656 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4660 if self.mutable_run.is_empty() {
4661 return Ok(());
4662 }
4663 let run_id = self.alloc_run_id()?;
4664 let rows = self.mutable_run.drain_sorted();
4665 if rows.is_empty() {
4666 return Ok(());
4667 }
4668 let path = self.run_path(run_id);
4669 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4670 if let Some(kek) = &self.kek {
4671 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4672 }
4673 let header = match self.create_run_file(run_id)? {
4674 Some(file) => writer.write_file(file, &rows)?,
4675 None => writer.write(&path, &rows)?,
4676 };
4677 self.run_refs.push(RunRef {
4678 run_id: run_id as u128,
4679 level: 0,
4680 epoch_created: epoch.0,
4681 row_count: header.row_count,
4682 });
4683 Ok(())
4684 }
4685
4686 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4690 self.mutable_run_spill_bytes = bytes.max(1);
4691 }
4692
4693 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4697 self.compaction_zstd_level = level;
4698 }
4699
4700 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4704 self.result_cache.lock().set_max_bytes(max_bytes);
4705 }
4706
4707 pub(crate) fn clear_result_cache(&mut self) {
4711 self.result_cache.lock().clear();
4712 }
4713
4714 pub fn mutable_run_len(&self) -> usize {
4716 self.mutable_run.len()
4717 }
4718
4719 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4722 self.mutable_run.drain_sorted()
4723 }
4724
4725 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4727 let mut snapshot = self.mutable_run.clone();
4728 snapshot.drain_sorted()
4729 }
4730
4731 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4736 self.ensure_writable()?;
4737 let n = batch.len();
4738 if n == 0 {
4739 return Ok(self.current_epoch());
4740 }
4741 for row in &batch {
4742 self.schema.validate_values(row)?;
4743 }
4744 let epoch = self.commit_new_epoch()?;
4745 let live_before = self.live_count;
4746 self.spill_mutable_run(epoch)?;
4750 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4751 && self.indexes_complete
4752 && self.run_refs.is_empty()
4753 && self.memtable.is_empty()
4754 && self.mutable_run.is_empty();
4755 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4761 use rayon::prelude::*;
4762 self.schema
4763 .columns
4764 .par_iter()
4765 .map(|cdef| {
4766 (
4767 cdef.id,
4768 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4769 )
4770 })
4771 .collect::<Vec<_>>()
4772 };
4773 drop(batch);
4774 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4779 self.validate_columns_not_null(&user_columns, n)?;
4780 let winner_idx = self
4781 .bulk_pk_winner_indices(&user_columns, n)
4782 .filter(|idx| idx.len() != n);
4783 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4784 match winner_idx.as_deref() {
4785 Some(idx) => {
4786 let compacted = user_columns
4787 .iter()
4788 .map(|(id, c)| (*id, c.gather(idx)))
4789 .collect();
4790 (compacted, idx.len())
4791 }
4792 None => (std::mem::take(&mut user_columns), n),
4793 };
4794 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4795 let first = self.allocator.alloc_range(write_n as u64)?.0;
4796 for rid in first..first + write_n as u64 {
4797 self.reservoir.offer(rid);
4798 }
4799 let run_id = self.alloc_run_id()?;
4800 let path = self.run_path(run_id);
4801 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4802 .clean(true)
4803 .with_lz4()
4804 .with_native_endian();
4805 if let Some(kek) = &self.kek {
4806 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4807 }
4808 let header = match self.create_run_file(run_id)? {
4809 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4810 None => writer.write_native(&path, &write_columns, write_n, first)?,
4811 };
4812 self.run_refs.push(RunRef {
4813 run_id: run_id as u128,
4814 level: 0,
4815 epoch_created: epoch.0,
4816 row_count: header.row_count,
4817 });
4818 self.live_count = self.live_count.saturating_add(write_n as u64);
4819 if eager_index_build {
4820 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4821 self.index_columns_bulk(&write_columns, &row_ids);
4822 self.indexes_complete = true;
4823 self.build_learned_ranges()?;
4824 } else {
4825 self.indexes_complete = false;
4826 }
4827 self.mark_flushed(epoch)?;
4828 self.persist_manifest(epoch)?;
4829 if eager_index_build {
4830 self.checkpoint_indexes(epoch);
4831 }
4832 self.clear_result_cache();
4833 Ok(epoch)
4834 }
4835
4836 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4839 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4840 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4841 let segment_no = segment
4844 .file_stem()
4845 .and_then(|s| s.to_str())
4846 .and_then(|s| s.strip_prefix("seg-"))
4847 .and_then(|s| s.parse::<u64>().ok())
4848 .unwrap_or(0);
4849 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4850 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4851 wal.sync()?;
4852 self.wal = WalSink::Private(wal);
4853 Ok(())
4854 }
4855
4856 pub(crate) fn invalidate_pending_cache(&mut self) {
4861 self.result_cache
4862 .lock()
4863 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4864 self.pending_delete_rids.clear();
4865 self.pending_put_cols.clear();
4866 }
4867
4868 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4869 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4870 m.current_epoch = epoch.0;
4871 m.next_row_id = self.allocator.current().0;
4872 m.runs = self.run_refs.clone();
4873 m.live_count = self.live_count;
4874 m.global_idx_epoch = self.global_idx_epoch;
4875 m.flushed_epoch = self.flushed_epoch;
4876 m.retiring = self.retiring.clone();
4877 m.auto_inc_next = match self.auto_inc {
4881 Some(ai) if ai.seeded => ai.next,
4882 _ => 0,
4883 };
4884 m.ttl = self.ttl;
4885 let meta_dek = self.manifest_meta_dek();
4886 match self._root_guard.as_deref() {
4887 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4888 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4889 }
4890 Ok(())
4891 }
4892
4893 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4894 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
4898 let live_count = u64::try_from(rows.len())
4899 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4900 let auto_inc = match self.auto_inc {
4901 Some(mut state) => {
4902 let maximum = self.scan_max_int64(state.column_id)?;
4903 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4904 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4905 })?;
4906 state.next = state.next.max(after_maximum).max(1);
4907 state.seeded = true;
4908 Some(state)
4909 }
4910 None => None,
4911 };
4912 Ok(RecoveryMetadataPlan {
4913 live_count,
4914 auto_inc,
4915 changed: live_count != self.live_count
4916 || auto_inc.is_some_and(|planned| {
4917 self.auto_inc.is_none_or(|current| {
4918 current.next != planned.next || current.seeded != planned.seeded
4919 })
4920 }),
4921 })
4922 }
4923
4924 pub(crate) fn apply_recovered_metadata(
4925 &mut self,
4926 plan: RecoveryMetadataPlan,
4927 epoch: Epoch,
4928 ) -> Result<()> {
4929 if !plan.changed {
4930 return Ok(());
4931 }
4932 self.live_count = plan.live_count;
4933 self.auto_inc = plan.auto_inc;
4934 self.persist_manifest(epoch)
4935 }
4936
4937 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4943 if !self.indexes_complete {
4946 return;
4947 }
4948 if crate::catalog::inject_hook("index.publish.before").is_err() {
4951 return;
4952 }
4953 if self.idx_root.is_none() {
4954 if let Some(root) = self._root_guard.as_ref() {
4955 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4956 return;
4957 };
4958 self.idx_root = Some(Arc::new(idx_root));
4959 }
4960 }
4961 let snap = global_idx::IndexSnapshot {
4962 hot: &self.hot,
4963 bitmap: &self.bitmap,
4964 ann: &self.ann,
4965 fm: &self.fm,
4966 sparse: &self.sparse,
4967 minhash: &self.minhash,
4968 learned_range: &self.learned_range,
4969 };
4970 let idx_dek = self.idx_dek();
4972 let written = match self.idx_root.as_deref() {
4973 Some(root) => global_idx::write_atomic_root(
4974 root,
4975 self.table_id,
4976 epoch.0,
4977 snap,
4978 idx_dek.as_deref(),
4979 ),
4980 None => global_idx::write_atomic(
4981 &self.dir,
4982 self.table_id,
4983 epoch.0,
4984 snap,
4985 idx_dek.as_deref(),
4986 ),
4987 };
4988 if written.is_ok() {
4989 self.global_idx_epoch = epoch.0;
4990 let _ = self.persist_manifest(epoch);
4991 let _ = crate::catalog::inject_hook("index.publish.after");
4993 }
4994 }
4995
4996 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4999 self.global_idx_epoch = 0;
5000 if let Some(root) = self.idx_root.as_deref() {
5001 let _ = root.remove_file(global_idx::IDX_FILENAME);
5002 } else {
5003 global_idx::remove(&self.dir);
5004 }
5005 let _ = self.persist_manifest(self.epoch.visible());
5006 }
5007
5008 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
5013 self.indexes_complete = false;
5014 self.global_idx_epoch = 0;
5015 if let Some(root) = self.idx_root.as_deref() {
5016 let _ = root.remove_file(global_idx::IDX_FILENAME);
5017 } else {
5018 global_idx::remove(&self.dir);
5019 }
5020 }
5021
5022 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
5023 self.indexes_complete = true;
5024 }
5025
5026 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
5032 self.durable_commit_failed = true;
5033 if let WalSink::Shared(shared) = &self.wal {
5034 shared
5035 .poisoned
5036 .store(true, std::sync::atomic::Ordering::Relaxed);
5037 }
5038 }
5039
5040 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
5044 self.durable_commit_failed = true;
5045 }
5046
5047 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
5056 let mut best: Option<Row> = None;
5057 let mut consider = |row: Row| {
5058 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5059 return;
5060 }
5061 if best.as_ref().is_none_or(|current| {
5062 Snapshot::version_is_newer(
5063 row.committed_epoch,
5064 row.commit_ts,
5065 current.committed_epoch,
5066 current.commit_ts,
5067 )
5068 }) {
5069 best = Some(row);
5070 }
5071 };
5072 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
5073 consider(row);
5074 }
5075 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
5076 consider(row);
5077 }
5078 for rr in &self.run_refs {
5079 let Ok(mut reader) = self.open_reader(rr.run_id) else {
5080 continue;
5081 };
5082 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
5085 continue;
5086 };
5087 consider(row);
5088 }
5089 let now_nanos = unix_nanos_now();
5090 match best {
5091 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
5092 Some(r) => Some(r),
5093 None => None,
5094 }
5095 }
5096
5097 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
5101 self.visible_rows_at_time(snapshot, unix_nanos_now())
5102 }
5103
5104 #[doc(hidden)]
5107 pub fn visible_rows_controlled(
5108 &self,
5109 snapshot: Snapshot,
5110 control: &crate::ExecutionControl,
5111 ) -> Result<Vec<Row>> {
5112 let mut out = Vec::new();
5113 self.for_each_visible_row_controlled(snapshot, control, |row| {
5114 out.push(row);
5115 Ok(())
5116 })?;
5117 Ok(out)
5118 }
5119
5120 #[doc(hidden)]
5123 pub fn for_each_visible_row_controlled<F>(
5124 &self,
5125 snapshot: Snapshot,
5126 control: &crate::ExecutionControl,
5127 mut visit: F,
5128 ) -> Result<()>
5129 where
5130 F: FnMut(Row) -> Result<()>,
5131 {
5132 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
5133 control.checkpoint()?;
5134 let memtable = self.memtable.visible_versions_at(snapshot);
5136 if !memtable.is_empty() {
5137 sources.push(ControlledVisibleSource::memory(memtable));
5138 }
5139 control.checkpoint()?;
5140 let mutable = self.mutable_run.visible_versions_at(snapshot);
5143 if !mutable.is_empty() {
5144 sources.push(ControlledVisibleSource::memory(mutable));
5145 }
5146 for run in &self.run_refs {
5147 control.checkpoint()?;
5148 let reader = self.open_reader(run.run_id)?;
5149 sources.push(ControlledVisibleSource::run(
5151 reader.into_visible_version_cursor(snapshot.epoch)?,
5152 ));
5153 }
5154 let now_nanos = unix_nanos_now();
5155 merge_controlled_visible_sources(
5156 &mut sources,
5157 control,
5158 |row| self.row_expired_at(row, now_nanos),
5159 |row| {
5160 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5163 return Ok(());
5164 }
5165 visit(row)
5166 },
5167 )
5168 }
5169
5170 #[doc(hidden)]
5171 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
5172 let mut best: HashMap<u64, Row> = HashMap::new();
5173 let mut fold = |row: Row| {
5174 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5175 return;
5176 }
5177 best.entry(row.row_id.0)
5178 .and_modify(|existing| {
5179 if Snapshot::version_is_newer(
5180 row.committed_epoch,
5181 row.commit_ts,
5182 existing.committed_epoch,
5183 existing.commit_ts,
5184 ) {
5185 *existing = row.clone();
5186 }
5187 })
5188 .or_insert(row);
5189 };
5190 for row in self.memtable.visible_versions_at(snapshot) {
5191 fold(row);
5192 }
5193 for row in self.mutable_run.visible_versions_at(snapshot) {
5194 fold(row);
5195 }
5196 for rr in &self.run_refs {
5197 let mut reader = self.open_reader(rr.run_id)?;
5198 for row in reader.visible_versions(snapshot.epoch)? {
5200 fold(row);
5201 }
5202 }
5203 let mut out: Vec<Row> = best
5204 .into_values()
5205 .filter_map(|r| {
5206 if r.deleted || self.row_expired_at(&r, now_nanos) {
5207 None
5208 } else {
5209 Some(r)
5210 }
5211 })
5212 .collect();
5213 out.sort_by_key(|r| r.row_id);
5214 Ok(out)
5215 }
5216
5217 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
5224 if self.ttl.is_none()
5225 && self.memtable.is_empty()
5226 && self.mutable_run.is_empty()
5227 && self.run_refs.len() == 1
5228 {
5229 let rr = self.run_refs[0].clone();
5230 let mut reader = self.open_reader(rr.run_id)?;
5231 let idxs = reader.visible_indices(snapshot.epoch)?;
5232 let mut cols = Vec::with_capacity(self.schema.columns.len());
5233 for cdef in &self.schema.columns {
5234 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
5235 }
5236 return Ok(cols);
5237 }
5238 let rows = self.visible_rows(snapshot)?;
5240 let mut cols: Vec<(u16, Vec<Value>)> = self
5241 .schema
5242 .columns
5243 .iter()
5244 .map(|c| (c.id, Vec::with_capacity(rows.len())))
5245 .collect();
5246 for r in &rows {
5247 for (cid, vec) in cols.iter_mut() {
5248 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
5249 }
5250 }
5251 Ok(cols)
5252 }
5253
5254 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
5256 let row_id = self.hot.get(key)?;
5257 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
5258 Some(row_id)
5259 } else {
5260 None
5261 }
5262 }
5263
5264 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
5269 self.query_at_with_allowed(q, self.snapshot(), None)
5270 }
5271
5272 pub fn query_controlled(
5275 &mut self,
5276 q: &crate::query::Query,
5277 control: &crate::ExecutionControl,
5278 ) -> Result<Vec<Row>> {
5279 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
5280 }
5281
5282 pub fn query_at_with_allowed(
5285 &mut self,
5286 q: &crate::query::Query,
5287 snapshot: Snapshot,
5288 allowed: Option<&std::collections::HashSet<RowId>>,
5289 ) -> Result<Vec<Row>> {
5290 self.query_at_with_allowed_after(q, snapshot, allowed, None)
5291 }
5292
5293 #[doc(hidden)]
5294 pub fn query_at_with_allowed_controlled(
5295 &mut self,
5296 q: &crate::query::Query,
5297 snapshot: Snapshot,
5298 allowed: Option<&std::collections::HashSet<RowId>>,
5299 control: &crate::ExecutionControl,
5300 ) -> Result<Vec<Row>> {
5301 self.require_select()?;
5302 self.ensure_indexes_complete_controlled(control, || true)?;
5303 self.validate_native_query(q)?;
5304 self.query_conditions_at(
5305 &q.conditions,
5306 snapshot,
5307 allowed,
5308 q.limit,
5309 q.offset,
5310 None,
5311 unix_nanos_now(),
5312 Some(control),
5313 )
5314 }
5315
5316 #[doc(hidden)]
5317 pub fn query_at_with_allowed_after(
5318 &mut self,
5319 q: &crate::query::Query,
5320 snapshot: Snapshot,
5321 allowed: Option<&std::collections::HashSet<RowId>>,
5322 after_row_id: Option<RowId>,
5323 ) -> Result<Vec<Row>> {
5324 self.query_at_with_allowed_after_at_time(
5325 q,
5326 snapshot,
5327 allowed,
5328 after_row_id,
5329 unix_nanos_now(),
5330 )
5331 }
5332
5333 #[doc(hidden)]
5334 pub fn query_at_with_allowed_after_at_time(
5335 &mut self,
5336 q: &crate::query::Query,
5337 snapshot: Snapshot,
5338 allowed: Option<&std::collections::HashSet<RowId>>,
5339 after_row_id: Option<RowId>,
5340 query_time_nanos: i64,
5341 ) -> Result<Vec<Row>> {
5342 self.require_select()?;
5343 self.ensure_indexes_complete()?;
5344 self.validate_native_query(q)?;
5345 self.query_conditions_at(
5346 &q.conditions,
5347 snapshot,
5348 allowed,
5349 q.limit,
5350 q.offset,
5351 after_row_id,
5352 query_time_nanos,
5353 None,
5354 )
5355 }
5356
5357 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
5358 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5359 return Err(MongrelError::InvalidArgument(format!(
5360 "query exceeds {} conditions",
5361 crate::query::MAX_HARD_CONDITIONS
5362 )));
5363 }
5364 if let Some(limit) = q.limit {
5365 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
5366 return Err(MongrelError::InvalidArgument(format!(
5367 "query limit must be between 1 and {}",
5368 crate::query::MAX_FINAL_LIMIT
5369 )));
5370 }
5371 }
5372 if q.offset > crate::query::MAX_QUERY_OFFSET {
5373 return Err(MongrelError::InvalidArgument(format!(
5374 "query offset exceeds {}",
5375 crate::query::MAX_QUERY_OFFSET
5376 )));
5377 }
5378 Ok(())
5379 }
5380
5381 #[doc(hidden)]
5384 pub fn query_all_at(
5385 &mut self,
5386 conditions: &[crate::query::Condition],
5387 snapshot: Snapshot,
5388 ) -> Result<Vec<Row>> {
5389 self.require_select()?;
5390 self.ensure_indexes_complete()?;
5391 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5392 return Err(MongrelError::InvalidArgument(format!(
5393 "query exceeds {} conditions",
5394 crate::query::MAX_HARD_CONDITIONS
5395 )));
5396 }
5397 self.query_conditions_at(
5398 conditions,
5399 snapshot,
5400 None,
5401 None,
5402 0,
5403 None,
5404 unix_nanos_now(),
5405 None,
5406 )
5407 }
5408
5409 #[allow(clippy::too_many_arguments)]
5410 fn query_conditions_at(
5411 &self,
5412 conditions: &[crate::query::Condition],
5413 snapshot: Snapshot,
5414 allowed: Option<&std::collections::HashSet<RowId>>,
5415 limit: Option<usize>,
5416 offset: usize,
5417 after_row_id: Option<RowId>,
5418 query_time_nanos: i64,
5419 control: Option<&crate::ExecutionControl>,
5420 ) -> Result<Vec<Row>> {
5421 control
5422 .map(crate::ExecutionControl::checkpoint)
5423 .transpose()?;
5424 crate::trace::QueryTrace::record(|t| {
5425 t.run_count = self.run_refs.len();
5426 t.memtable_rows = self.memtable.len();
5427 t.mutable_run_rows = self.mutable_run.len();
5428 });
5429 if conditions.is_empty() {
5433 crate::trace::QueryTrace::record(|t| {
5434 t.scan_mode = crate::trace::ScanMode::Materialized;
5435 t.row_materialized = true;
5436 });
5437 let mut rows = match control {
5438 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5439 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5440 };
5441 if let Some(allowed) = allowed {
5442 let mut filtered = Vec::with_capacity(rows.len());
5443 for (index, row) in rows.into_iter().enumerate() {
5444 if index & 255 == 0 {
5445 control
5446 .map(crate::ExecutionControl::checkpoint)
5447 .transpose()?;
5448 }
5449 if allowed.contains(&row.row_id) {
5450 filtered.push(row);
5451 }
5452 }
5453 rows = filtered;
5454 }
5455 if let Some(after_row_id) = after_row_id {
5456 rows.retain(|row| row.row_id > after_row_id);
5457 }
5458 rows.drain(..offset.min(rows.len()));
5459 if let Some(limit) = limit {
5460 rows.truncate(limit);
5461 }
5462 return Ok(rows);
5463 }
5464 crate::trace::QueryTrace::record(|t| {
5465 t.conditions_pushed = conditions.len();
5466 t.scan_mode = crate::trace::ScanMode::Materialized;
5467 t.row_materialized = true;
5468 });
5469 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5476 ordered.sort_by_key(|c| condition_cost_rank(c));
5477 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5478 for c in &ordered {
5479 control
5480 .map(crate::ExecutionControl::checkpoint)
5481 .transpose()?;
5482 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5483 let empty = s.is_empty();
5484 sets.push(s);
5485 if empty {
5486 break;
5487 }
5488 }
5489 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5490 if let Some(allowed) = allowed {
5491 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5492 }
5493 if let Some(after_row_id) = after_row_id {
5494 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5495 rids.drain(..first);
5496 }
5497 rids.drain(..offset.min(rids.len()));
5498 if let Some(limit) = limit {
5499 rids.truncate(limit);
5500 }
5501 control
5502 .map(crate::ExecutionControl::checkpoint)
5503 .transpose()?;
5504 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5505 }
5506
5507 pub fn retrieve(
5509 &mut self,
5510 retriever: &crate::query::Retriever,
5511 ) -> Result<Vec<crate::query::RetrieverHit>> {
5512 self.retrieve_with_allowed(retriever, None)
5513 }
5514
5515 pub fn retrieve_at(
5516 &mut self,
5517 retriever: &crate::query::Retriever,
5518 snapshot: Snapshot,
5519 allowed: Option<&std::collections::HashSet<RowId>>,
5520 ) -> Result<Vec<crate::query::RetrieverHit>> {
5521 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5522 }
5523
5524 pub fn retrieve_with_allowed(
5527 &mut self,
5528 retriever: &crate::query::Retriever,
5529 allowed: Option<&std::collections::HashSet<RowId>>,
5530 ) -> Result<Vec<crate::query::RetrieverHit>> {
5531 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5532 }
5533
5534 pub fn retrieve_at_with_allowed(
5535 &mut self,
5536 retriever: &crate::query::Retriever,
5537 snapshot: Snapshot,
5538 allowed: Option<&std::collections::HashSet<RowId>>,
5539 ) -> Result<Vec<crate::query::RetrieverHit>> {
5540 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5541 }
5542
5543 pub fn retrieve_at_with_allowed_and_context(
5544 &mut self,
5545 retriever: &crate::query::Retriever,
5546 snapshot: Snapshot,
5547 allowed: Option<&std::collections::HashSet<RowId>>,
5548 context: Option<&crate::query::AiExecutionContext>,
5549 ) -> Result<Vec<crate::query::RetrieverHit>> {
5550 self.require_select()?;
5551 self.ensure_indexes_complete()?;
5552 self.validate_retriever(retriever)?;
5553 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5554 }
5555
5556 pub fn retrieve_at_with_candidate_authorization_and_context(
5557 &mut self,
5558 retriever: &crate::query::Retriever,
5559 snapshot: Snapshot,
5560 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5561 context: Option<&crate::query::AiExecutionContext>,
5562 ) -> Result<Vec<crate::query::RetrieverHit>> {
5563 self.require_select()?;
5564 self.ensure_indexes_complete()?;
5565 self.retrieve_at_with_candidate_authorization_on_generation(
5566 retriever,
5567 snapshot,
5568 authorization,
5569 context,
5570 )
5571 }
5572
5573 #[doc(hidden)]
5574 pub fn retrieve_at_with_candidate_authorization_on_generation(
5575 &self,
5576 retriever: &crate::query::Retriever,
5577 snapshot: Snapshot,
5578 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5579 context: Option<&crate::query::AiExecutionContext>,
5580 ) -> Result<Vec<crate::query::RetrieverHit>> {
5581 self.require_select()?;
5582 self.validate_retriever(retriever)?;
5583 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5584 }
5585
5586 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5587 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5588 let (column_id, k) = match retriever {
5589 Retriever::Ann {
5590 column_id,
5591 query,
5592 k,
5593 } => {
5594 let index = self.ann.get(column_id).ok_or_else(|| {
5595 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5596 })?;
5597 if query.len() != index.dim() {
5598 return Err(MongrelError::InvalidArgument(format!(
5599 "ANN query dimension must be {}, got {}",
5600 index.dim(),
5601 query.len()
5602 )));
5603 }
5604 if query.iter().any(|value| !value.is_finite()) {
5605 return Err(MongrelError::InvalidArgument(
5606 "ANN query values must be finite".into(),
5607 ));
5608 }
5609 (*column_id, *k)
5610 }
5611 Retriever::Sparse {
5612 column_id,
5613 query,
5614 k,
5615 } => {
5616 if !self.sparse.contains_key(column_id) {
5617 return Err(MongrelError::InvalidArgument(format!(
5618 "column {column_id} has no Sparse index"
5619 )));
5620 }
5621 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5622 return Err(MongrelError::InvalidArgument(
5623 "Sparse query must be non-empty with finite weights".into(),
5624 ));
5625 }
5626 if query.len() > MAX_SPARSE_TERMS {
5627 return Err(MongrelError::InvalidArgument(format!(
5628 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5629 )));
5630 }
5631 (*column_id, *k)
5632 }
5633 Retriever::MinHash {
5634 column_id,
5635 members,
5636 k,
5637 } => {
5638 if !self.minhash.contains_key(column_id) {
5639 return Err(MongrelError::InvalidArgument(format!(
5640 "column {column_id} has no MinHash index"
5641 )));
5642 }
5643 if members.is_empty() {
5644 return Err(MongrelError::InvalidArgument(
5645 "MinHash members must not be empty".into(),
5646 ));
5647 }
5648 if members.len() > MAX_SET_MEMBERS {
5649 return Err(MongrelError::InvalidArgument(format!(
5650 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5651 )));
5652 }
5653 let mut total_bytes = 0usize;
5654 for member in members {
5655 let bytes = member.encoded_len();
5656 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5657 return Err(MongrelError::InvalidArgument(format!(
5658 "MinHash member exceeds {} bytes",
5659 crate::query::MAX_SET_MEMBER_BYTES
5660 )));
5661 }
5662 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5663 MongrelError::InvalidArgument("MinHash input size overflow".into())
5664 })?;
5665 }
5666 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5667 return Err(MongrelError::InvalidArgument(format!(
5668 "MinHash input exceeds {} bytes",
5669 crate::query::MAX_SET_INPUT_BYTES
5670 )));
5671 }
5672 (*column_id, *k)
5673 }
5674 };
5675 if k == 0 {
5676 return Err(MongrelError::InvalidArgument(
5677 "retriever k must be > 0".into(),
5678 ));
5679 }
5680 if k > MAX_RETRIEVER_K {
5681 return Err(MongrelError::InvalidArgument(format!(
5682 "retriever k exceeds {MAX_RETRIEVER_K}"
5683 )));
5684 }
5685 debug_assert!(self
5686 .schema
5687 .columns
5688 .iter()
5689 .any(|column| column.id == column_id));
5690 Ok(())
5691 }
5692
5693 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5694 use crate::query::Condition;
5695 match condition {
5696 Condition::Ann {
5697 column_id,
5698 query,
5699 k,
5700 } => self.validate_retriever(&crate::query::Retriever::Ann {
5701 column_id: *column_id,
5702 query: query.clone(),
5703 k: *k,
5704 }),
5705 Condition::SparseMatch {
5706 column_id,
5707 query,
5708 k,
5709 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5710 column_id: *column_id,
5711 query: query.clone(),
5712 k: *k,
5713 }),
5714 Condition::MinHashSimilar {
5715 column_id,
5716 query,
5717 k,
5718 } => {
5719 if !self.minhash.contains_key(column_id) {
5720 return Err(MongrelError::InvalidArgument(format!(
5721 "column {column_id} has no MinHash index"
5722 )));
5723 }
5724 if query.is_empty() || *k == 0 {
5725 return Err(MongrelError::InvalidArgument(
5726 "MinHash query must be non-empty and k must be > 0".into(),
5727 ));
5728 }
5729 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5730 {
5731 return Err(MongrelError::InvalidArgument(format!(
5732 "MinHash query must have <= {} members and k <= {}",
5733 crate::query::MAX_SET_MEMBERS,
5734 crate::query::MAX_RETRIEVER_K
5735 )));
5736 }
5737 Ok(())
5738 }
5739 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5740 Err(MongrelError::InvalidArgument(format!(
5741 "bitmap IN exceeds {} values",
5742 crate::query::MAX_SET_MEMBERS
5743 )))
5744 }
5745 Condition::FmContainsAll { patterns, .. }
5746 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5747 {
5748 Err(MongrelError::InvalidArgument(format!(
5749 "FM query exceeds {} patterns",
5750 crate::query::MAX_HARD_CONDITIONS
5751 )))
5752 }
5753 _ => Ok(()),
5754 }
5755 }
5756
5757 fn retrieve_filtered(
5758 &self,
5759 retriever: &crate::query::Retriever,
5760 snapshot: Snapshot,
5761 hard_filter: Option<&RowIdSet>,
5762 allowed: Option<&std::collections::HashSet<RowId>>,
5763 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5764 context: Option<&crate::query::AiExecutionContext>,
5765 ) -> Result<Vec<crate::query::RetrieverHit>> {
5766 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5767 let started = std::time::Instant::now();
5768 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5769 Retriever::Ann {
5770 column_id,
5771 query,
5772 k,
5773 } => {
5774 let Some(index) = self.ann.get(column_id) else {
5775 return Ok(Vec::new());
5776 };
5777 let cap = ann_candidate_cap(index.len(), context);
5778 if cap == 0 {
5779 return Ok(Vec::new());
5780 }
5781 let mut breadth = (*k).max(1).min(cap);
5782 let mut eligibility = std::collections::HashMap::new();
5783 let mut filtered = loop {
5784 let mut seen = std::collections::HashSet::new();
5785 if let Some(context) = context {
5786 context.checkpoint()?;
5787 }
5788 let raw = index.search_with_context(query, breadth, context)?;
5789 let unchecked: Vec<_> = raw
5790 .iter()
5791 .map(|(row_id, _)| *row_id)
5792 .filter(|row_id| !eligibility.contains_key(row_id))
5793 .filter(|row_id| {
5794 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5795 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5796 })
5797 .collect();
5798 let eligible = self.eligible_and_authorized_candidate_ids(
5799 &unchecked,
5800 *column_id,
5801 snapshot,
5802 candidate_authorization,
5803 context,
5804 )?;
5805 for row_id in unchecked {
5806 eligibility.insert(row_id, eligible.contains(&row_id));
5807 }
5808 let filtered: Vec<_> = raw
5809 .into_iter()
5810 .filter(|(row_id, _)| {
5811 seen.insert(*row_id)
5812 && eligibility.get(row_id).copied().unwrap_or(false)
5813 })
5814 .map(|(row_id, score)| {
5815 let score = match score {
5816 crate::index::AnnDistance::Hamming(d) => {
5817 RetrieverScore::AnnHammingDistance(d)
5818 }
5819 crate::index::AnnDistance::Cosine(d) => {
5820 RetrieverScore::AnnCosineDistance(d)
5821 }
5822 };
5823 (row_id, score)
5824 })
5825 .collect();
5826 if filtered.len() >= *k || breadth >= cap {
5827 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5828 crate::trace::QueryTrace::record(|trace| {
5829 trace.ann_candidate_cap_hit = true;
5830 });
5831 }
5832 break filtered;
5833 }
5834 breadth = breadth.saturating_mul(2).min(cap);
5835 };
5836 filtered.truncate(*k);
5837 filtered
5838 }
5839 Retriever::Sparse {
5840 column_id,
5841 query,
5842 k,
5843 } => self
5844 .sparse
5845 .get(column_id)
5846 .map(|index| -> Result<Vec<_>> {
5847 let mut breadth = (*k).max(1);
5848 let mut eligibility = std::collections::HashMap::new();
5849 loop {
5850 if let Some(context) = context {
5851 context.checkpoint()?;
5852 }
5853 let raw = index.search_with_context(query, breadth, context)?;
5854 let unchecked: Vec<_> = raw
5855 .iter()
5856 .map(|(row_id, _)| *row_id)
5857 .filter(|row_id| !eligibility.contains_key(row_id))
5858 .filter(|row_id| {
5859 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5860 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5861 })
5862 .collect();
5863 let eligible = self.eligible_and_authorized_candidate_ids(
5864 &unchecked,
5865 *column_id,
5866 snapshot,
5867 candidate_authorization,
5868 context,
5869 )?;
5870 for row_id in unchecked {
5871 eligibility.insert(row_id, eligible.contains(&row_id));
5872 }
5873 let filtered: Vec<_> = raw
5874 .iter()
5875 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5876 .take(*k)
5877 .map(|(row_id, score)| {
5878 (*row_id, RetrieverScore::SparseDotProduct(*score))
5879 })
5880 .collect();
5881 if filtered.len() >= *k || raw.len() < breadth {
5882 break Ok(filtered);
5883 }
5884 let next = breadth.saturating_mul(2);
5885 if next == breadth {
5886 break Ok(filtered);
5887 }
5888 breadth = next;
5889 }
5890 })
5891 .transpose()?
5892 .unwrap_or_default(),
5893 Retriever::MinHash {
5894 column_id,
5895 members,
5896 k,
5897 } => self
5898 .minhash
5899 .get(column_id)
5900 .map(|index| -> Result<Vec<_>> {
5901 let mut hashes = Vec::with_capacity(members.len());
5902 for member in members {
5903 if let Some(context) = context {
5904 context.consume(crate::query::work_units(
5905 member.encoded_len(),
5906 crate::query::PARSE_WORK_QUANTUM,
5907 ))?;
5908 }
5909 hashes.push(member.hash_v1());
5910 }
5911 let mut breadth = (*k).max(1);
5912 let mut eligibility = std::collections::HashMap::new();
5913 loop {
5914 if let Some(context) = context {
5915 context.checkpoint()?;
5916 }
5917 let raw = index.search_with_context(&hashes, breadth, context)?;
5918 let unchecked: Vec<_> = raw
5919 .iter()
5920 .map(|(row_id, _)| *row_id)
5921 .filter(|row_id| !eligibility.contains_key(row_id))
5922 .filter(|row_id| {
5923 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5924 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5925 })
5926 .collect();
5927 let eligible = self.eligible_and_authorized_candidate_ids(
5928 &unchecked,
5929 *column_id,
5930 snapshot,
5931 candidate_authorization,
5932 context,
5933 )?;
5934 for row_id in unchecked {
5935 eligibility.insert(row_id, eligible.contains(&row_id));
5936 }
5937 let filtered: Vec<_> = raw
5938 .iter()
5939 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5940 .take(*k)
5941 .map(|(row_id, score)| {
5942 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5943 })
5944 .collect();
5945 if filtered.len() >= *k || raw.len() < breadth {
5946 break Ok(filtered);
5947 }
5948 let next = breadth.saturating_mul(2);
5949 if next == breadth {
5950 break Ok(filtered);
5951 }
5952 breadth = next;
5953 }
5954 })
5955 .transpose()?
5956 .unwrap_or_default(),
5957 };
5958 let elapsed = started.elapsed().as_nanos() as u64;
5959 crate::trace::QueryTrace::record(|trace| {
5960 match retriever {
5961 Retriever::Ann { .. } => {
5962 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
5963 if let Retriever::Ann { column_id, .. } = retriever {
5964 if let Some(index) = self.ann.get(column_id) {
5965 trace.ann_algorithm = Some(index.algorithm());
5966 trace.ann_quantization = Some(index.quantization());
5967 trace.ann_backend = Some(index.backend_name());
5968 }
5969 }
5970 }
5971 Retriever::Sparse { .. } => {
5972 trace.sparse_candidate_nanos =
5973 trace.sparse_candidate_nanos.saturating_add(elapsed)
5974 }
5975 Retriever::MinHash { .. } => {
5976 trace.minhash_candidate_nanos =
5977 trace.minhash_candidate_nanos.saturating_add(elapsed)
5978 }
5979 }
5980 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5981 });
5982 Ok(scored
5983 .into_iter()
5984 .enumerate()
5985 .map(|(rank, (row_id, score))| RetrieverHit {
5986 row_id,
5987 rank: rank + 1,
5988 score,
5989 })
5990 .collect())
5991 }
5992
5993 fn eligible_candidate_ids(
5994 &self,
5995 candidates: &[RowId],
5996 _column_id: u16,
5997 snapshot: Snapshot,
5998 context: Option<&crate::query::AiExecutionContext>,
5999 ) -> Result<std::collections::HashSet<RowId>> {
6000 if !self.had_deletes
6001 && self.ttl.is_none()
6002 && self.pending_put_cols.is_empty()
6003 && snapshot.epoch == self.snapshot().epoch
6004 {
6005 return Ok(candidates.iter().copied().collect());
6006 }
6007 let mut readers: Vec<_> = self
6008 .run_refs
6009 .iter()
6010 .map(|run| self.open_reader(run.run_id))
6011 .collect::<Result<_>>()?;
6012 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6013 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
6014 for &row_id in candidates {
6015 if let Some(context) = context {
6016 context.consume(1)?;
6017 }
6018 let mem = self.memtable.get_version_at(row_id, snapshot);
6019 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
6020 let overlay = match (mem, mutable) {
6021 (Some(left), Some(right)) => Some(
6022 if Snapshot::version_is_newer(
6023 left.1.committed_epoch,
6024 left.1.commit_ts,
6025 right.1.committed_epoch,
6026 right.1.commit_ts,
6027 ) {
6028 left
6029 } else {
6030 right
6031 },
6032 ),
6033 (Some(value), None) | (None, Some(value)) => Some(value),
6034 (None, None) => None,
6035 };
6036 if let Some((_, row)) = overlay {
6037 if !row.deleted && !self.row_expired_at(&row, now) {
6038 eligible.insert(row_id);
6039 }
6040 continue;
6041 }
6042 let mut best: Option<(Epoch, bool, usize)> = None;
6043 for (index, reader) in readers.iter_mut().enumerate() {
6044 if let Some((epoch, deleted)) =
6045 reader.get_version_visibility(row_id, snapshot.epoch)?
6046 {
6047 if best
6048 .as_ref()
6049 .map(|(best_epoch, ..)| epoch > *best_epoch)
6050 .unwrap_or(true)
6051 {
6052 best = Some((epoch, deleted, index));
6053 }
6054 }
6055 }
6056 let Some((_, false, reader_index)) = best else {
6057 continue;
6058 };
6059 if let Some(ttl) = self.ttl {
6060 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6061 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6062 {
6063 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6064 continue;
6065 }
6066 }
6067 }
6068 eligible.insert(row_id);
6069 }
6070 Ok(eligible)
6071 }
6072
6073 fn eligible_and_authorized_candidate_ids(
6074 &self,
6075 candidates: &[RowId],
6076 column_id: u16,
6077 snapshot: Snapshot,
6078 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6079 context: Option<&crate::query::AiExecutionContext>,
6080 ) -> Result<std::collections::HashSet<RowId>> {
6081 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
6082 let Some(authorization) = authorization else {
6083 return Ok(eligible);
6084 };
6085 let candidates: Vec<_> = eligible.into_iter().collect();
6086 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
6087 }
6088
6089 fn policy_allowed_candidate_ids(
6090 &self,
6091 candidates: &[RowId],
6092 snapshot: Snapshot,
6093 authorization: &crate::security::CandidateAuthorization<'_>,
6094 context: Option<&crate::query::AiExecutionContext>,
6095 ) -> Result<std::collections::HashSet<RowId>> {
6096 let started = std::time::Instant::now();
6097 if candidates.is_empty()
6098 || authorization.principal.is_admin
6099 || !authorization.security.rls_enabled(authorization.table)
6100 {
6101 return Ok(candidates.iter().copied().collect());
6102 }
6103 if let Some(context) = context {
6104 context.checkpoint()?;
6105 }
6106 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
6107 let mut rows: std::collections::HashMap<RowId, Row> = candidates
6108 .iter()
6109 .map(|row_id| {
6110 (
6111 *row_id,
6112 Row {
6113 row_id: *row_id,
6114 committed_epoch: snapshot.epoch,
6115 columns: std::collections::HashMap::new(),
6116 deleted: false,
6117 commit_ts: None,
6118 },
6119 )
6120 })
6121 .collect();
6122 let columns = authorization
6123 .security
6124 .select_policy_columns(authorization.table, authorization.principal);
6125 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6126 let mut decoded = 0usize;
6127 for column_id in &columns {
6128 if let Some(context) = context {
6129 context.checkpoint()?;
6130 }
6131 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6132 &row_ids, *column_id, snapshot, query_now, context,
6133 )? {
6134 if let Some(row) = rows.get_mut(&row_id) {
6135 row.columns.insert(*column_id, value);
6136 decoded = decoded.saturating_add(1);
6137 }
6138 }
6139 }
6140 if let Some(context) = context {
6141 context.consume(candidates.len().saturating_add(decoded))?;
6142 }
6143 let allowed = rows
6144 .into_values()
6145 .filter_map(|row| {
6146 authorization
6147 .security
6148 .row_allowed(
6149 authorization.table,
6150 crate::security::PolicyCommand::Select,
6151 &row,
6152 authorization.principal,
6153 false,
6154 )
6155 .then_some(row.row_id)
6156 })
6157 .collect();
6158 crate::trace::QueryTrace::record(|trace| {
6159 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
6160 trace.rls_policy_columns_decoded =
6161 trace.rls_policy_columns_decoded.saturating_add(decoded);
6162 trace.authorization_nanos = trace
6163 .authorization_nanos
6164 .saturating_add(started.elapsed().as_nanos() as u64);
6165 });
6166 Ok(allowed)
6167 }
6168
6169 pub fn search(
6171 &mut self,
6172 request: &crate::query::SearchRequest,
6173 ) -> Result<Vec<crate::query::SearchHit>> {
6174 self.search_with_allowed(request, None)
6175 }
6176
6177 pub fn search_at(
6178 &mut self,
6179 request: &crate::query::SearchRequest,
6180 snapshot: Snapshot,
6181 authorized: Option<&std::collections::HashSet<RowId>>,
6182 ) -> Result<Vec<crate::query::SearchHit>> {
6183 self.search_at_with_allowed(request, snapshot, authorized)
6184 }
6185
6186 pub fn search_with_allowed(
6187 &mut self,
6188 request: &crate::query::SearchRequest,
6189 authorized: Option<&std::collections::HashSet<RowId>>,
6190 ) -> Result<Vec<crate::query::SearchHit>> {
6191 self.search_at_with_allowed(request, self.snapshot(), authorized)
6192 }
6193
6194 pub fn search_at_with_allowed(
6195 &mut self,
6196 request: &crate::query::SearchRequest,
6197 snapshot: Snapshot,
6198 authorized: Option<&std::collections::HashSet<RowId>>,
6199 ) -> Result<Vec<crate::query::SearchHit>> {
6200 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
6201 }
6202
6203 pub fn search_at_with_allowed_and_context(
6204 &mut self,
6205 request: &crate::query::SearchRequest,
6206 snapshot: Snapshot,
6207 authorized: Option<&std::collections::HashSet<RowId>>,
6208 context: Option<&crate::query::AiExecutionContext>,
6209 ) -> Result<Vec<crate::query::SearchHit>> {
6210 self.ensure_indexes_complete()?;
6211 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
6212 }
6213
6214 pub fn search_at_with_candidate_authorization_and_context(
6215 &mut self,
6216 request: &crate::query::SearchRequest,
6217 snapshot: Snapshot,
6218 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6219 context: Option<&crate::query::AiExecutionContext>,
6220 ) -> Result<Vec<crate::query::SearchHit>> {
6221 self.ensure_indexes_complete()?;
6222 self.search_at_with_filters_and_context(
6223 request,
6224 snapshot,
6225 None,
6226 authorization,
6227 context,
6228 None,
6229 )
6230 }
6231
6232 #[doc(hidden)]
6233 pub fn search_at_with_candidate_authorization_on_generation(
6234 &self,
6235 request: &crate::query::SearchRequest,
6236 snapshot: Snapshot,
6237 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6238 context: Option<&crate::query::AiExecutionContext>,
6239 ) -> Result<Vec<crate::query::SearchHit>> {
6240 self.search_at_with_filters_and_context(
6241 request,
6242 snapshot,
6243 None,
6244 authorization,
6245 context,
6246 None,
6247 )
6248 }
6249
6250 #[doc(hidden)]
6251 pub fn search_at_with_candidate_authorization_on_generation_after(
6252 &self,
6253 request: &crate::query::SearchRequest,
6254 snapshot: Snapshot,
6255 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6256 context: Option<&crate::query::AiExecutionContext>,
6257 after: Option<crate::query::SearchAfter>,
6258 ) -> Result<Vec<crate::query::SearchHit>> {
6259 self.search_at_with_filters_and_context(
6260 request,
6261 snapshot,
6262 None,
6263 authorization,
6264 context,
6265 after,
6266 )
6267 }
6268
6269 fn search_at_with_filters_and_context(
6270 &self,
6271 request: &crate::query::SearchRequest,
6272 snapshot: Snapshot,
6273 authorized: Option<&std::collections::HashSet<RowId>>,
6274 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6275 context: Option<&crate::query::AiExecutionContext>,
6276 after: Option<crate::query::SearchAfter>,
6277 ) -> Result<Vec<crate::query::SearchHit>> {
6278 use crate::query::{
6279 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
6280 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
6281 };
6282 let total_started = std::time::Instant::now();
6283 let rank_offset = after.map_or(0, |after| after.returned_count);
6284 self.require_select()?;
6285 if request.limit == 0 {
6286 return Err(MongrelError::InvalidArgument(
6287 "search limit must be > 0".into(),
6288 ));
6289 }
6290 if request.limit > MAX_FINAL_LIMIT {
6291 return Err(MongrelError::InvalidArgument(format!(
6292 "search limit exceeds {MAX_FINAL_LIMIT}"
6293 )));
6294 }
6295 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
6296 return Err(MongrelError::InvalidArgument(
6297 "search-after score must be finite".into(),
6298 ));
6299 }
6300 if request.retrievers.is_empty() {
6301 return Err(MongrelError::InvalidArgument(
6302 "search requires at least one retriever".into(),
6303 ));
6304 }
6305 if request.retrievers.len() > MAX_RETRIEVERS {
6306 return Err(MongrelError::InvalidArgument(format!(
6307 "search exceeds {MAX_RETRIEVERS} retrievers"
6308 )));
6309 }
6310 if request.must.len() > MAX_HARD_CONDITIONS {
6311 return Err(MongrelError::InvalidArgument(format!(
6312 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
6313 )));
6314 }
6315 for condition in &request.must {
6316 self.validate_condition(condition)?;
6317 }
6318 if request.must.iter().any(|condition| {
6319 matches!(
6320 condition,
6321 Condition::Ann { .. }
6322 | Condition::SparseMatch { .. }
6323 | Condition::MinHashSimilar { .. }
6324 )
6325 }) {
6326 return Err(MongrelError::InvalidArgument(
6327 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
6328 .into(),
6329 ));
6330 }
6331 let mut names = std::collections::HashSet::new();
6332 for named in &request.retrievers {
6333 if named.name.is_empty()
6334 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
6335 || !names.insert(named.name.as_str())
6336 {
6337 return Err(MongrelError::InvalidArgument(format!(
6338 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
6339 crate::query::MAX_RETRIEVER_NAME_BYTES
6340 )));
6341 }
6342 if !named.weight.is_finite()
6343 || named.weight < 0.0
6344 || named.weight > MAX_RETRIEVER_WEIGHT
6345 {
6346 return Err(MongrelError::InvalidArgument(format!(
6347 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6348 )));
6349 }
6350 self.validate_retriever(&named.retriever)?;
6351 }
6352 let projection = request
6353 .projection
6354 .clone()
6355 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
6356 if projection.len() > MAX_PROJECTION_COLUMNS {
6357 return Err(MongrelError::InvalidArgument(format!(
6358 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
6359 )));
6360 }
6361 for column_id in &projection {
6362 if !self
6363 .schema
6364 .columns
6365 .iter()
6366 .any(|column| column.id == *column_id)
6367 {
6368 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
6369 }
6370 }
6371 if let Some(crate::query::Rerank::ExactVector {
6372 embedding_column,
6373 query,
6374 candidate_limit,
6375 weight,
6376 ..
6377 }) = &request.rerank
6378 {
6379 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
6380 {
6381 return Err(MongrelError::InvalidArgument(format!(
6382 "rerank candidate_limit must be between search limit and {}",
6383 crate::query::MAX_RETRIEVER_K
6384 )));
6385 }
6386 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
6387 return Err(MongrelError::InvalidArgument(format!(
6388 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6389 )));
6390 }
6391 let column = self
6392 .schema
6393 .columns
6394 .iter()
6395 .find(|column| column.id == *embedding_column)
6396 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
6397 let crate::schema::TypeId::Embedding { dim } = column.ty else {
6398 return Err(MongrelError::InvalidArgument(format!(
6399 "rerank column {embedding_column} is not an embedding"
6400 )));
6401 };
6402 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
6403 return Err(MongrelError::InvalidArgument(format!(
6404 "rerank query must contain {dim} finite values"
6405 )));
6406 }
6407 }
6408
6409 let hard_filter_started = std::time::Instant::now();
6410 let hard_filter = if request.must.is_empty() {
6411 None
6412 } else {
6413 let mut sets = Vec::with_capacity(request.must.len());
6414 for condition in &request.must {
6415 if let Some(context) = context {
6416 context.checkpoint()?;
6417 }
6418 sets.push(self.resolve_condition(condition, snapshot)?);
6419 }
6420 Some(RowIdSet::intersect_many(sets))
6421 };
6422 crate::trace::QueryTrace::record(|trace| {
6423 trace.hard_filter_nanos = trace
6424 .hard_filter_nanos
6425 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6426 });
6427 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6428 return Ok(Vec::new());
6429 }
6430
6431 let constant = match request.fusion {
6432 Fusion::ReciprocalRank { constant } => constant,
6433 };
6434 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6435 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6436 let mut fusion_nanos = 0u64;
6437 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6438 std::collections::HashMap::new();
6439 for named in retrievers {
6440 if named.weight == 0.0 {
6441 continue;
6442 }
6443 if let Some(context) = context {
6444 context.checkpoint()?;
6445 }
6446 let hits = self.retrieve_filtered(
6447 &named.retriever,
6448 snapshot,
6449 hard_filter.as_ref(),
6450 authorized,
6451 candidate_authorization,
6452 context,
6453 )?;
6454 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6455 let fusion_started = std::time::Instant::now();
6456 for hit in hits {
6457 if let Some(context) = context {
6458 context.consume(1)?;
6459 }
6460 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6461 if !contribution.is_finite() {
6462 return Err(MongrelError::InvalidArgument(
6463 "retriever contribution must be finite".into(),
6464 ));
6465 }
6466 let max_fused_candidates = context.map_or(
6467 crate::query::MAX_FUSED_CANDIDATES,
6468 crate::query::AiExecutionContext::max_fused_candidates,
6469 );
6470 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6471 return Err(MongrelError::WorkBudgetExceeded);
6472 }
6473 let entry = fused.entry(hit.row_id).or_default();
6474 entry.0 += contribution;
6475 if !entry.0.is_finite() {
6476 return Err(MongrelError::InvalidArgument(
6477 "fused score must be finite".into(),
6478 ));
6479 }
6480 entry.1.push(ComponentScore {
6481 retriever_name: retriever_name.clone(),
6482 rank: hit.rank,
6483 raw_score: hit.score,
6484 contribution,
6485 });
6486 }
6487 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6488 }
6489 let union_size = fused.len();
6490 let mut ranked: Vec<_> = fused
6491 .into_iter()
6492 .map(|(row_id, (fused_score, components))| {
6493 (row_id, fused_score, components, None, fused_score)
6494 })
6495 .collect();
6496 let order = |(a_row, _, _, _, a_score): &(
6497 RowId,
6498 f64,
6499 Vec<ComponentScore>,
6500 Option<f32>,
6501 f64,
6502 ),
6503 (b_row, _, _, _, b_score): &(
6504 RowId,
6505 f64,
6506 Vec<ComponentScore>,
6507 Option<f32>,
6508 f64,
6509 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6510 if let Some(crate::query::Rerank::ExactVector {
6511 embedding_column,
6512 query,
6513 metric,
6514 candidate_limit,
6515 weight,
6516 }) = &request.rerank
6517 {
6518 let fused_order = |(a_row, a_score, ..): &(
6519 RowId,
6520 f64,
6521 Vec<ComponentScore>,
6522 Option<f32>,
6523 f64,
6524 ),
6525 (b_row, b_score, ..): &(
6526 RowId,
6527 f64,
6528 Vec<ComponentScore>,
6529 Option<f32>,
6530 f64,
6531 )| {
6532 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6533 };
6534 let selection_started = std::time::Instant::now();
6535 if let Some(context) = context {
6536 context.consume(ranked.len())?;
6537 }
6538 if ranked.len() > *candidate_limit {
6539 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6540 ranked.truncate(*candidate_limit);
6541 }
6542 ranked.sort_by(fused_order);
6543 fusion_nanos =
6544 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6545 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6546 if let Some(context) = context {
6547 context.consume(row_ids.len())?;
6548 }
6549 let query_now =
6550 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6551 let gather_started = std::time::Instant::now();
6552 let vectors = self.values_for_rids_batch_at_with_context(
6553 &row_ids,
6554 *embedding_column,
6555 snapshot,
6556 query_now,
6557 context,
6558 )?;
6559 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6560 let vector_work =
6561 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6562 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6563 if let Some(context) = context {
6564 context.consume(vector_work)?;
6565 }
6566 query
6567 .iter()
6568 .map(|value| f64::from(*value).powi(2))
6569 .sum::<f64>()
6570 .sqrt()
6571 } else {
6572 0.0
6573 };
6574 let score_started = std::time::Instant::now();
6575 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6576 for (row_id, value) in vectors {
6577 if let Some(meta) = value.generated_embedding_metadata() {
6578 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6579 continue;
6580 }
6581 }
6582 let Some(vector) = value.as_embedding() else {
6583 continue;
6584 };
6585 let score = match metric {
6586 crate::query::VectorMetric::DotProduct => {
6587 if let Some(context) = context {
6588 context.consume(vector_work)?;
6589 }
6590 query
6591 .iter()
6592 .zip(vector)
6593 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6594 .sum::<f64>()
6595 }
6596 crate::query::VectorMetric::Cosine => {
6597 if let Some(context) = context {
6598 context.consume(vector_work.saturating_mul(2))?;
6599 }
6600 let dot = query
6601 .iter()
6602 .zip(vector)
6603 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6604 .sum::<f64>();
6605 let norm = vector
6606 .iter()
6607 .map(|value| f64::from(*value).powi(2))
6608 .sum::<f64>()
6609 .sqrt();
6610 if query_norm == 0.0 || norm == 0.0 {
6611 0.0
6612 } else {
6613 dot / (query_norm * norm)
6614 }
6615 }
6616 crate::query::VectorMetric::Euclidean => {
6617 if let Some(context) = context {
6618 context.consume(vector_work)?;
6619 }
6620 query
6621 .iter()
6622 .zip(vector)
6623 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6624 .sum::<f64>()
6625 .sqrt()
6626 }
6627 };
6628 if !score.is_finite() {
6629 return Err(MongrelError::InvalidArgument(
6630 "exact rerank score must be finite".into(),
6631 ));
6632 }
6633 scores.insert(row_id, score as f32);
6634 }
6635 let mut reranked = Vec::with_capacity(ranked.len());
6636 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6637 let Some(score) = scores.get(&row_id).copied() else {
6638 continue;
6639 };
6640 let ordering_score = match metric {
6641 crate::query::VectorMetric::Euclidean => -f64::from(score),
6642 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6643 f64::from(score)
6644 }
6645 };
6646 let final_score = fused_score + *weight * ordering_score;
6647 if !final_score.is_finite() {
6648 return Err(MongrelError::InvalidArgument(
6649 "final rerank score must be finite".into(),
6650 ));
6651 }
6652 reranked.push((row_id, fused_score, components, Some(score), final_score));
6653 }
6654 ranked = reranked;
6655 ranked.sort_by(order);
6656 crate::trace::QueryTrace::record(|trace| {
6657 trace.exact_vector_gather_nanos =
6658 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6659 trace.exact_vector_score_nanos = trace
6660 .exact_vector_score_nanos
6661 .saturating_add(score_started.elapsed().as_nanos() as u64);
6662 });
6663 }
6664 if let Some(after) = after {
6665 ranked.retain(|(row_id, _, _, _, final_score)| {
6666 final_score.total_cmp(&after.final_score).is_lt()
6667 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6668 });
6669 }
6670 let projection_started = std::time::Instant::now();
6671 let sentinel = projection
6672 .first()
6673 .copied()
6674 .or_else(|| self.schema.columns.first().map(|column| column.id));
6675 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6676 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6677 let mut projection_rows = 0usize;
6678 let mut projection_cells = 0usize;
6679 while out.len() < request.limit && !ranked.is_empty() {
6680 if let Some(context) = context {
6681 context.checkpoint()?;
6682 context.consume(ranked.len())?;
6683 }
6684 let needed = request.limit - out.len();
6685 let window_size = ranked
6686 .len()
6687 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6688 let selection_started = std::time::Instant::now();
6689 let mut remainder = if ranked.len() > window_size {
6690 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6691 ranked.split_off(window_size)
6692 } else {
6693 Vec::new()
6694 };
6695 ranked.sort_by(order);
6696 fusion_nanos =
6697 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6698 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6699 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6700 if let Some(context) = context {
6701 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6702 }
6703 projection_rows = projection_rows.saturating_add(row_ids.len());
6704 projection_cells =
6705 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6706 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6707 std::collections::HashMap::new();
6708 if let Some(column_id) = sentinel {
6709 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6710 &row_ids, column_id, snapshot, query_now, context,
6711 )? {
6712 cells.entry(row_id).or_default().insert(column_id, value);
6713 }
6714 }
6715 for &column_id in &projection {
6716 if Some(column_id) == sentinel {
6717 continue;
6718 }
6719 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6720 &row_ids, column_id, snapshot, query_now, context,
6721 )? {
6722 cells.entry(row_id).or_default().insert(column_id, value);
6723 }
6724 }
6725 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6726 ranked.drain(..)
6727 {
6728 let Some(row_cells) = cells.remove(&row_id) else {
6729 continue;
6730 };
6731 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6732 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6733 out.push(SearchHit {
6734 row_id,
6735 cells: projection
6736 .iter()
6737 .filter_map(|column_id| {
6738 row_cells
6739 .get(column_id)
6740 .cloned()
6741 .map(|value| (*column_id, value))
6742 })
6743 .collect(),
6744 components,
6745 fused_score,
6746 exact_rerank_score,
6747 final_score,
6748 final_rank,
6749 });
6750 if out.len() == request.limit {
6751 break;
6752 }
6753 }
6754 ranked.append(&mut remainder);
6755 }
6756 crate::trace::QueryTrace::record(|trace| {
6757 trace.union_size = union_size;
6758 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6759 trace.projection_nanos = trace
6760 .projection_nanos
6761 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6762 trace.total_nanos = trace
6763 .total_nanos
6764 .saturating_add(total_started.elapsed().as_nanos() as u64);
6765 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6766 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6767 if let Some(context) = context {
6768 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6769 }
6770 });
6771 Ok(out)
6772 }
6773
6774 pub fn set_similarity(
6777 &mut self,
6778 request: &crate::query::SetSimilarityRequest,
6779 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6780 self.set_similarity_with_allowed(request, None)
6781 }
6782
6783 pub fn set_similarity_at(
6784 &mut self,
6785 request: &crate::query::SetSimilarityRequest,
6786 snapshot: Snapshot,
6787 allowed: Option<&std::collections::HashSet<RowId>>,
6788 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6789 self.set_similarity_explained_at(request, snapshot, allowed)
6790 .map(|(hits, _)| hits)
6791 }
6792
6793 pub fn ann_rerank(
6795 &mut self,
6796 request: &crate::query::AnnRerankRequest,
6797 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6798 self.ann_rerank_with_allowed(request, None)
6799 }
6800
6801 pub fn ann_rerank_with_allowed(
6802 &mut self,
6803 request: &crate::query::AnnRerankRequest,
6804 allowed: Option<&std::collections::HashSet<RowId>>,
6805 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6806 self.ann_rerank_at(request, self.snapshot(), allowed)
6807 }
6808
6809 pub fn ann_rerank_at(
6810 &mut self,
6811 request: &crate::query::AnnRerankRequest,
6812 snapshot: Snapshot,
6813 allowed: Option<&std::collections::HashSet<RowId>>,
6814 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6815 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6816 }
6817
6818 pub fn ann_rerank_at_with_context(
6819 &mut self,
6820 request: &crate::query::AnnRerankRequest,
6821 snapshot: Snapshot,
6822 allowed: Option<&std::collections::HashSet<RowId>>,
6823 context: Option<&crate::query::AiExecutionContext>,
6824 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6825 self.ensure_indexes_complete()?;
6826 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6827 }
6828
6829 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6830 &mut self,
6831 request: &crate::query::AnnRerankRequest,
6832 snapshot: Snapshot,
6833 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6834 context: Option<&crate::query::AiExecutionContext>,
6835 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6836 self.ensure_indexes_complete()?;
6837 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6838 }
6839
6840 #[doc(hidden)]
6841 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6842 &self,
6843 request: &crate::query::AnnRerankRequest,
6844 snapshot: Snapshot,
6845 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6846 context: Option<&crate::query::AiExecutionContext>,
6847 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6848 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6849 }
6850
6851 fn ann_rerank_at_with_filters_and_context(
6852 &self,
6853 request: &crate::query::AnnRerankRequest,
6854 snapshot: Snapshot,
6855 allowed: Option<&std::collections::HashSet<RowId>>,
6856 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6857 context: Option<&crate::query::AiExecutionContext>,
6858 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6859 use crate::query::{
6860 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6861 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6862 };
6863 if request.candidate_k == 0 || request.limit == 0 {
6864 return Err(MongrelError::InvalidArgument(
6865 "candidate_k and limit must be > 0".into(),
6866 ));
6867 }
6868 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6869 return Err(MongrelError::InvalidArgument(format!(
6870 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6871 )));
6872 }
6873 let retriever = Retriever::Ann {
6874 column_id: request.column_id,
6875 query: request.query.clone(),
6876 k: request.candidate_k,
6877 };
6878 self.require_select()?;
6879 self.validate_retriever(&retriever)?;
6880 let hits = self.retrieve_filtered(
6881 &retriever,
6882 snapshot,
6883 None,
6884 allowed,
6885 candidate_authorization,
6886 context,
6887 )?;
6888 let distances: std::collections::HashMap<_, _> = hits
6889 .iter()
6890 .filter_map(|hit| match hit.score {
6891 RetrieverScore::AnnHammingDistance(distance) => {
6892 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
6893 }
6894 RetrieverScore::AnnCosineDistance(distance) => {
6895 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
6896 }
6897 _ => None,
6898 })
6899 .collect();
6900 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6901 if let Some(context) = context {
6902 context.consume(row_ids.len())?;
6903 }
6904 let gather_started = std::time::Instant::now();
6905 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6906 let values = self.values_for_rids_batch_at_with_context(
6907 &row_ids,
6908 request.column_id,
6909 snapshot,
6910 query_now,
6911 context,
6912 )?;
6913 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6914 let score_started = std::time::Instant::now();
6915 let vector_work =
6916 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6917 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6918 if let Some(context) = context {
6919 context.consume(vector_work)?;
6920 }
6921 request
6922 .query
6923 .iter()
6924 .map(|value| f64::from(*value).powi(2))
6925 .sum::<f64>()
6926 .sqrt()
6927 } else {
6928 0.0
6929 };
6930 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6931 for (row_id, value) in values {
6932 if let Some(meta) = value.generated_embedding_metadata() {
6933 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6934 continue;
6935 }
6936 }
6937 let Some(vector) = value.as_embedding() else {
6938 continue;
6939 };
6940 let exact_score = match request.metric {
6941 VectorMetric::DotProduct => {
6942 if let Some(context) = context {
6943 context.consume(vector_work)?;
6944 }
6945 request
6946 .query
6947 .iter()
6948 .zip(vector)
6949 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6950 .sum::<f64>()
6951 }
6952 VectorMetric::Cosine => {
6953 if let Some(context) = context {
6954 context.consume(vector_work.saturating_mul(2))?;
6955 }
6956 let dot = request
6957 .query
6958 .iter()
6959 .zip(vector)
6960 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6961 .sum::<f64>();
6962 let norm = vector
6963 .iter()
6964 .map(|value| f64::from(*value).powi(2))
6965 .sum::<f64>()
6966 .sqrt();
6967 if query_norm == 0.0 || norm == 0.0 {
6968 0.0
6969 } else {
6970 dot / (query_norm * norm)
6971 }
6972 }
6973 VectorMetric::Euclidean => {
6974 if let Some(context) = context {
6975 context.consume(vector_work)?;
6976 }
6977 request
6978 .query
6979 .iter()
6980 .zip(vector)
6981 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6982 .sum::<f64>()
6983 .sqrt()
6984 }
6985 };
6986 let exact_score = exact_score as f32;
6987 if !exact_score.is_finite() {
6988 return Err(MongrelError::InvalidArgument(
6989 "exact ANN score must be finite".into(),
6990 ));
6991 }
6992 let Some(candidate_distance) = distances.get(&row_id).copied() else {
6993 continue;
6994 };
6995 reranked.push(AnnRerankHit {
6996 row_id,
6997 candidate_distance,
6998 exact_score,
6999 });
7000 }
7001 reranked.sort_by(|left, right| {
7002 let score = match request.metric {
7003 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
7004 VectorMetric::Cosine | VectorMetric::DotProduct => {
7005 right.exact_score.total_cmp(&left.exact_score)
7006 }
7007 };
7008 score.then_with(|| left.row_id.cmp(&right.row_id))
7009 });
7010 reranked.truncate(request.limit);
7011 crate::trace::QueryTrace::record(|trace| {
7012 trace.exact_vector_gather_nanos =
7013 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
7014 trace.exact_vector_score_nanos = trace
7015 .exact_vector_score_nanos
7016 .saturating_add(score_started.elapsed().as_nanos() as u64);
7017 });
7018 Ok(reranked)
7019 }
7020
7021 pub fn set_similarity_with_allowed(
7022 &mut self,
7023 request: &crate::query::SetSimilarityRequest,
7024 allowed: Option<&std::collections::HashSet<RowId>>,
7025 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7026 self.set_similarity_explained_at(request, self.snapshot(), allowed)
7027 .map(|(hits, _)| hits)
7028 }
7029
7030 pub fn set_similarity_explained(
7031 &mut self,
7032 request: &crate::query::SetSimilarityRequest,
7033 ) -> Result<(
7034 Vec<crate::query::SetSimilarityHit>,
7035 crate::query::SetSimilarityTrace,
7036 )> {
7037 self.set_similarity_explained_at(request, self.snapshot(), None)
7038 }
7039
7040 fn set_similarity_explained_at(
7041 &mut self,
7042 request: &crate::query::SetSimilarityRequest,
7043 snapshot: Snapshot,
7044 allowed: Option<&std::collections::HashSet<RowId>>,
7045 ) -> Result<(
7046 Vec<crate::query::SetSimilarityHit>,
7047 crate::query::SetSimilarityTrace,
7048 )> {
7049 self.ensure_indexes_complete()?;
7050 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
7051 }
7052
7053 pub fn set_similarity_at_with_context(
7054 &mut self,
7055 request: &crate::query::SetSimilarityRequest,
7056 snapshot: Snapshot,
7057 allowed: Option<&std::collections::HashSet<RowId>>,
7058 context: Option<&crate::query::AiExecutionContext>,
7059 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7060 self.ensure_indexes_complete()?;
7061 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
7062 .map(|(hits, _)| hits)
7063 }
7064
7065 pub fn set_similarity_at_with_candidate_authorization_and_context(
7066 &mut self,
7067 request: &crate::query::SetSimilarityRequest,
7068 snapshot: Snapshot,
7069 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7070 context: Option<&crate::query::AiExecutionContext>,
7071 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7072 self.ensure_indexes_complete()?;
7073 self.set_similarity_explained_at_with_context(
7074 request,
7075 snapshot,
7076 None,
7077 authorization,
7078 context,
7079 )
7080 .map(|(hits, _)| hits)
7081 }
7082
7083 #[doc(hidden)]
7084 pub fn set_similarity_at_with_candidate_authorization_on_generation(
7085 &self,
7086 request: &crate::query::SetSimilarityRequest,
7087 snapshot: Snapshot,
7088 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7089 context: Option<&crate::query::AiExecutionContext>,
7090 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7091 self.set_similarity_explained_at_with_context(
7092 request,
7093 snapshot,
7094 None,
7095 authorization,
7096 context,
7097 )
7098 .map(|(hits, _)| hits)
7099 }
7100
7101 fn set_similarity_explained_at_with_context(
7102 &self,
7103 request: &crate::query::SetSimilarityRequest,
7104 snapshot: Snapshot,
7105 allowed: Option<&std::collections::HashSet<RowId>>,
7106 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7107 context: Option<&crate::query::AiExecutionContext>,
7108 ) -> Result<(
7109 Vec<crate::query::SetSimilarityHit>,
7110 crate::query::SetSimilarityTrace,
7111 )> {
7112 use crate::query::{
7113 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
7114 MAX_SET_MEMBERS,
7115 };
7116 let mut trace = crate::query::SetSimilarityTrace::default();
7117 if request.members.is_empty() {
7118 return Ok((Vec::new(), trace));
7119 }
7120 if request.candidate_k == 0 || request.limit == 0 {
7121 return Err(MongrelError::InvalidArgument(
7122 "candidate_k and limit must be > 0".into(),
7123 ));
7124 }
7125 if request.candidate_k > MAX_RETRIEVER_K
7126 || request.limit > MAX_FINAL_LIMIT
7127 || request.members.len() > MAX_SET_MEMBERS
7128 {
7129 return Err(MongrelError::InvalidArgument(format!(
7130 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
7131 )));
7132 }
7133 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
7134 return Err(MongrelError::InvalidArgument(
7135 "min_jaccard must be finite and between 0 and 1".into(),
7136 ));
7137 }
7138 let started = std::time::Instant::now();
7139 let retriever = Retriever::MinHash {
7140 column_id: request.column_id,
7141 members: request.members.clone(),
7142 k: request.candidate_k,
7143 };
7144 self.require_select()?;
7145 self.validate_retriever(&retriever)?;
7146 let hits = self.retrieve_filtered(
7147 &retriever,
7148 snapshot,
7149 None,
7150 allowed,
7151 candidate_authorization,
7152 context,
7153 )?;
7154 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
7155 trace.candidate_count = hits.len();
7156 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7157 if let Some(context) = context {
7158 context.consume(row_ids.len())?;
7159 }
7160 let started = std::time::Instant::now();
7161 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7162 let values = self.values_for_rids_batch_at_with_context(
7163 &row_ids,
7164 request.column_id,
7165 snapshot,
7166 query_now,
7167 context,
7168 )?;
7169 trace.gather_us = started.elapsed().as_micros() as u64;
7170 if let Some(context) = context {
7171 context.consume(request.members.len())?;
7172 }
7173 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
7174 let estimates: std::collections::HashMap<_, _> = hits
7175 .into_iter()
7176 .filter_map(|hit| match hit.score {
7177 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
7178 _ => None,
7179 })
7180 .collect();
7181 let started = std::time::Instant::now();
7182 let mut parsed = Vec::with_capacity(values.len());
7183 for (row_id, value) in values {
7184 let Value::Bytes(bytes) = value else {
7185 continue;
7186 };
7187 if let Some(context) = context {
7188 context.consume(crate::query::work_units(
7189 bytes.len(),
7190 crate::query::PARSE_WORK_QUANTUM,
7191 ))?;
7192 }
7193 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
7194 continue;
7195 };
7196 if let Some(context) = context {
7197 context.consume(members.len())?;
7198 }
7199 let stored = members
7200 .into_iter()
7201 .filter_map(|member| match member {
7202 serde_json::Value::String(value) => {
7203 Some(crate::query::SetMember::String(value))
7204 }
7205 serde_json::Value::Number(value) => {
7206 Some(crate::query::SetMember::Number(value))
7207 }
7208 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
7209 _ => None,
7210 })
7211 .collect::<std::collections::HashSet<_>>();
7212 parsed.push((row_id, stored));
7213 }
7214 trace.parse_us = started.elapsed().as_micros() as u64;
7215 trace.verified_count = parsed.len();
7216 let started = std::time::Instant::now();
7217 let mut exact = Vec::new();
7218 for (row_id, stored) in parsed {
7219 if let Some(context) = context {
7220 context.consume(query.len().saturating_add(stored.len()))?;
7221 }
7222 let union = query.union(&stored).count();
7223 let score = if union == 0 {
7224 1.0
7225 } else {
7226 query.intersection(&stored).count() as f32 / union as f32
7227 };
7228 if score >= request.min_jaccard {
7229 exact.push(SetSimilarityHit {
7230 row_id,
7231 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
7232 exact_jaccard: score,
7233 });
7234 }
7235 }
7236 exact.sort_by(|a, b| {
7237 b.exact_jaccard
7238 .total_cmp(&a.exact_jaccard)
7239 .then_with(|| a.row_id.cmp(&b.row_id))
7240 });
7241 exact.truncate(request.limit);
7242 trace.score_us = started.elapsed().as_micros() as u64;
7243 crate::trace::QueryTrace::record(|query_trace| {
7244 query_trace.exact_set_gather_nanos = query_trace
7245 .exact_set_gather_nanos
7246 .saturating_add(trace.gather_us.saturating_mul(1_000));
7247 query_trace.exact_set_parse_nanos = query_trace
7248 .exact_set_parse_nanos
7249 .saturating_add(trace.parse_us.saturating_mul(1_000));
7250 query_trace.exact_set_score_nanos = query_trace
7251 .exact_set_score_nanos
7252 .saturating_add(trace.score_us.saturating_mul(1_000));
7253 });
7254 Ok((exact, trace))
7255 }
7256
7257 fn values_for_rids_batch_at(
7259 &self,
7260 row_ids: &[u64],
7261 column_id: u16,
7262 snapshot: Snapshot,
7263 now: i64,
7264 ) -> Result<Vec<(RowId, Value)>> {
7265 if self.ttl.is_none()
7266 && self.memtable.is_empty()
7267 && self.mutable_run.is_empty()
7268 && self.run_refs.len() == 1
7269 {
7270 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7271 if row_ids.len().saturating_mul(24) < reader.row_count() {
7276 let mut values = Vec::with_capacity(row_ids.len());
7277 for &raw_row_id in row_ids {
7278 let row_id = RowId(raw_row_id);
7279 if let Some((_, false, Some(value))) =
7280 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7281 {
7282 values.push((row_id, value));
7283 }
7284 }
7285 return Ok(values);
7286 }
7287 let (positions, visible_row_ids) =
7288 reader.visible_positions_with_rids(snapshot.epoch)?;
7289 let requested: Vec<(RowId, usize)> = row_ids
7290 .iter()
7291 .filter_map(|raw| {
7292 visible_row_ids
7293 .binary_search(&(*raw as i64))
7294 .ok()
7295 .map(|index| (RowId(*raw), positions[index]))
7296 })
7297 .collect();
7298 let values = reader.gather_column(
7299 column_id,
7300 &requested
7301 .iter()
7302 .map(|(_, position)| *position)
7303 .collect::<Vec<_>>(),
7304 )?;
7305 return Ok(requested
7306 .into_iter()
7307 .zip(values)
7308 .map(|((row_id, _), value)| (row_id, value))
7309 .collect());
7310 }
7311 self.values_for_rids_at(row_ids, column_id, snapshot, now)
7312 }
7313
7314 fn values_for_rids_batch_at_with_context(
7315 &self,
7316 row_ids: &[u64],
7317 column_id: u16,
7318 snapshot: Snapshot,
7319 now: i64,
7320 context: Option<&crate::query::AiExecutionContext>,
7321 ) -> Result<Vec<(RowId, Value)>> {
7322 let Some(context) = context else {
7323 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
7324 };
7325 let mut values = Vec::with_capacity(row_ids.len());
7326 for chunk in row_ids.chunks(256) {
7327 context.checkpoint()?;
7328 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
7329 }
7330 Ok(values)
7331 }
7332
7333 fn values_for_rids_at(
7335 &self,
7336 row_ids: &[u64],
7337 column_id: u16,
7338 snapshot: Snapshot,
7339 now: i64,
7340 ) -> Result<Vec<(RowId, Value)>> {
7341 let mut readers: Vec<_> = self
7342 .run_refs
7343 .iter()
7344 .map(|run| self.open_reader(run.run_id))
7345 .collect::<Result<_>>()?;
7346 let mut out = Vec::with_capacity(row_ids.len());
7347 for &raw_row_id in row_ids {
7348 let row_id = RowId(raw_row_id);
7349 let mem = self.memtable.get_version_at(row_id, snapshot);
7350 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
7351 let overlay = match (mem, mutable) {
7352 (Some((_, a)), Some((_, b))) => Some(
7353 if Snapshot::version_is_newer(
7354 a.committed_epoch,
7355 a.commit_ts,
7356 b.committed_epoch,
7357 b.commit_ts,
7358 ) {
7359 a
7360 } else {
7361 b
7362 },
7363 ),
7364 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
7365 (None, None) => None,
7366 };
7367 if let Some(row) = overlay {
7368 if !row.deleted && !self.row_expired_at(&row, now) {
7369 if let Some(value) = row.columns.get(&column_id) {
7370 out.push((row_id, value.clone()));
7371 }
7372 }
7373 continue;
7374 }
7375
7376 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
7377 for (index, reader) in readers.iter_mut().enumerate() {
7378 if let Some((epoch, deleted, value)) =
7379 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7380 {
7381 if best
7382 .as_ref()
7383 .map(|(best_epoch, ..)| epoch > *best_epoch)
7384 .unwrap_or(true)
7385 {
7386 best = Some((epoch, deleted, value, index));
7387 }
7388 }
7389 }
7390 let Some((_, false, Some(value), reader_index)) = best else {
7391 continue;
7392 };
7393 if let Some(ttl) = self.ttl {
7394 if ttl.column_id != column_id {
7395 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7396 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
7397 {
7398 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7399 continue;
7400 }
7401 }
7402 } else if let Value::Int64(timestamp) = value {
7403 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7404 continue;
7405 }
7406 }
7407 }
7408 out.push((row_id, value));
7409 }
7410 Ok(out)
7411 }
7412
7413 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
7418 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
7419 }
7420
7421 pub fn rows_for_rids_with_context(
7422 &self,
7423 rids: &[u64],
7424 snapshot: Snapshot,
7425 context: &crate::query::AiExecutionContext,
7426 ) -> Result<Vec<Row>> {
7427 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
7428 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
7429 }
7430
7431 fn rows_for_rids_at_time(
7432 &self,
7433 rids: &[u64],
7434 snapshot: Snapshot,
7435 ttl_now: i64,
7436 control: Option<&crate::ExecutionControl>,
7437 ) -> Result<Vec<Row>> {
7438 use std::collections::HashMap;
7439 let mut rows = Vec::with_capacity(rids.len());
7440 let tier_size = self.memtable.len() + self.mutable_run.len();
7455 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7456 if rids.len().saturating_mul(24) < tier_size {
7457 for &rid in rids {
7458 if overlay.len() & 255 == 0 {
7459 control
7460 .map(crate::ExecutionControl::checkpoint)
7461 .transpose()?;
7462 }
7463 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
7464 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
7465 let newest = match (mem, mrun) {
7466 (Some((_, mr)), Some((_, rr))) => Some(
7467 if Snapshot::version_is_newer(
7468 mr.committed_epoch,
7469 mr.commit_ts,
7470 rr.committed_epoch,
7471 rr.commit_ts,
7472 ) {
7473 mr
7474 } else {
7475 rr
7476 },
7477 ),
7478 (Some((_, mr)), None) => Some(mr),
7479 (None, Some((_, rr))) => Some(rr),
7480 (None, None) => None,
7481 };
7482 if let Some(row) = newest {
7483 overlay.insert(rid, row);
7484 }
7485 }
7486 } else {
7487 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7488 overlay
7489 .entry(row.row_id.0)
7490 .and_modify(|e| {
7491 if Snapshot::version_is_newer(
7492 row.committed_epoch,
7493 row.commit_ts,
7494 e.committed_epoch,
7495 e.commit_ts,
7496 ) {
7497 *e = row.clone();
7498 }
7499 })
7500 .or_insert(row);
7501 };
7502 for (index, row) in self
7503 .memtable
7504 .visible_versions_at(snapshot)
7505 .into_iter()
7506 .enumerate()
7507 {
7508 if index & 255 == 0 {
7509 control
7510 .map(crate::ExecutionControl::checkpoint)
7511 .transpose()?;
7512 }
7513 fold_newest(row, &mut overlay);
7514 }
7515 for (index, row) in self
7516 .mutable_run
7517 .visible_versions_at(snapshot)
7518 .into_iter()
7519 .enumerate()
7520 {
7521 if index & 255 == 0 {
7522 control
7523 .map(crate::ExecutionControl::checkpoint)
7524 .transpose()?;
7525 }
7526 fold_newest(row, &mut overlay);
7527 }
7528 }
7529 if self.run_refs.len() == 1 {
7530 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7531 if rids.len().saturating_mul(24) < reader.row_count() {
7539 for (index, &rid) in rids.iter().enumerate() {
7540 if index & 255 == 0 {
7541 control
7542 .map(crate::ExecutionControl::checkpoint)
7543 .transpose()?;
7544 }
7545 if let Some(r) = overlay.get(&rid) {
7546 if !r.deleted {
7547 rows.push(r.clone());
7548 }
7549 continue;
7550 }
7551 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7552 if !row.deleted {
7553 rows.push(row);
7554 }
7555 }
7556 }
7557 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7558 return Ok(rows);
7559 }
7560 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7569 enum Src {
7572 Overlay,
7573 Run,
7574 }
7575 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7576 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7577 for (index, rid) in rids.iter().enumerate() {
7578 if index & 255 == 0 {
7579 control
7580 .map(crate::ExecutionControl::checkpoint)
7581 .transpose()?;
7582 }
7583 if overlay.contains_key(rid) {
7584 plan.push(Src::Overlay);
7585 continue;
7586 }
7587 match vis_rids.binary_search(&(*rid as i64)) {
7588 Ok(i) => {
7589 plan.push(Src::Run);
7590 fetch.push(positions[i]);
7591 }
7592 Err(_) => { }
7593 }
7594 }
7595 let fetched = reader.materialize_batch(&fetch)?;
7596 let mut fetched_iter = fetched.into_iter();
7597 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7598 if index & 255 == 0 {
7599 control
7600 .map(crate::ExecutionControl::checkpoint)
7601 .transpose()?;
7602 }
7603 match src {
7604 Src::Overlay => {
7605 if let Some(r) = overlay.get(rid) {
7606 if !r.deleted {
7607 rows.push(r.clone());
7608 }
7609 }
7610 }
7611 Src::Run => {
7612 if let Some(row) = fetched_iter.next() {
7613 if !row.deleted {
7614 rows.push(row);
7615 }
7616 }
7617 }
7618 }
7619 }
7620 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7621 return Ok(rows);
7622 }
7623 let mut readers: Vec<_> = self
7627 .run_refs
7628 .iter()
7629 .map(|rr| self.open_reader(rr.run_id))
7630 .collect::<Result<Vec<_>>>()?;
7631 for (index, rid) in rids.iter().enumerate() {
7632 if index & 255 == 0 {
7633 control
7634 .map(crate::ExecutionControl::checkpoint)
7635 .transpose()?;
7636 }
7637 if let Some(r) = overlay.get(rid) {
7638 if !r.deleted {
7639 rows.push(r.clone());
7640 }
7641 continue;
7642 }
7643 let mut best: Option<(Epoch, Row)> = None;
7644 for reader in readers.iter_mut() {
7645 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7646 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7647 best = Some((epoch, row));
7648 }
7649 }
7650 }
7651 if let Some((_, r)) = best {
7652 if !r.deleted {
7653 rows.push(r);
7654 }
7655 }
7656 }
7657 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7658 Ok(rows)
7659 }
7660
7661 pub fn indexes_complete(&self) -> bool {
7671 self.indexes_complete
7672 }
7673
7674 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7676 self.index_build_policy
7677 }
7678
7679 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7683 self.index_build_policy = policy;
7684 }
7685
7686 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7691 if !self.indexes_complete {
7695 return None;
7696 }
7697 let b = self.bitmap.get(&column_id)?;
7698 let result: Vec<Vec<u8>> = b
7699 .keys()
7700 .into_iter()
7701 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7702 .collect();
7703 Some(result)
7704 }
7705
7706 pub fn fk_join_row_ids(
7707 &self,
7708 fk_column_id: u16,
7709 pk_values: &[Vec<u8>],
7710 fk_conditions: &[crate::query::Condition],
7711 snapshot: Snapshot,
7712 ) -> Result<Vec<u64>> {
7713 let Some(b) = self.bitmap.get(&fk_column_id) else {
7714 return Ok(Vec::new());
7715 };
7716 let mut join_set = {
7717 let mut acc = roaring::RoaringBitmap::new();
7718 for v in pk_values {
7719 acc |= b.get(v);
7720 }
7721 RowIdSet::from_roaring(acc)
7722 };
7723 if !fk_conditions.is_empty() {
7724 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7725 sets.push(join_set);
7726 for c in fk_conditions {
7727 sets.push(self.resolve_condition(c, snapshot)?);
7728 }
7729 join_set = RowIdSet::intersect_many(sets);
7730 }
7731 Ok(join_set.into_sorted_vec())
7732 }
7733
7734 pub fn fk_join_count(
7740 &self,
7741 fk_column_id: u16,
7742 pk_values: &[Vec<u8>],
7743 fk_conditions: &[crate::query::Condition],
7744 snapshot: Snapshot,
7745 ) -> Result<u64> {
7746 let Some(b) = self.bitmap.get(&fk_column_id) else {
7747 return Ok(0);
7748 };
7749 let mut acc = roaring::RoaringBitmap::new();
7750 for v in pk_values {
7751 acc |= b.get(v);
7752 }
7753 if fk_conditions.is_empty() {
7754 return Ok(acc.len());
7755 }
7756 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7757 sets.push(RowIdSet::from_roaring(acc));
7758 for c in fk_conditions {
7759 sets.push(self.resolve_condition(c, snapshot)?);
7760 }
7761 Ok(RowIdSet::intersect_many(sets).len() as u64)
7762 }
7763
7764 fn resolve_condition(
7769 &self,
7770 c: &crate::query::Condition,
7771 snapshot: Snapshot,
7772 ) -> Result<RowIdSet> {
7773 self.resolve_condition_with_allowed(c, snapshot, None)
7774 }
7775
7776 fn resolve_condition_with_allowed(
7777 &self,
7778 c: &crate::query::Condition,
7779 snapshot: Snapshot,
7780 allowed: Option<&std::collections::HashSet<RowId>>,
7781 ) -> Result<RowIdSet> {
7782 use crate::query::Condition;
7783 self.validate_condition(c)?;
7784 Ok(match c {
7785 Condition::Pk(key) => {
7786 let lookup = self
7787 .schema
7788 .primary_key()
7789 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7790 .unwrap_or_else(|| key.clone());
7791 self.hot
7792 .get(&lookup)
7793 .map(|r| RowIdSet::one(r.0))
7794 .unwrap_or_else(RowIdSet::empty)
7795 }
7796 Condition::BitmapEq { column_id, value } => {
7797 let lookup = self.index_lookup_key_bytes(*column_id, value);
7798 self.bitmap
7799 .get(column_id)
7800 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7801 .unwrap_or_else(RowIdSet::empty)
7802 }
7803 Condition::BitmapIn { column_id, values } => {
7804 let bm = self.bitmap.get(column_id);
7805 let mut acc = roaring::RoaringBitmap::new();
7806 if let Some(b) = bm {
7807 for v in values {
7808 let lookup = self.index_lookup_key_bytes(*column_id, v);
7809 acc |= b.get(&lookup);
7810 }
7811 }
7812 RowIdSet::from_roaring(acc)
7813 }
7814 Condition::BytesPrefix { column_id, prefix } => {
7815 if let Some(b) = self.bitmap.get(column_id) {
7820 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7821 let mut acc = roaring::RoaringBitmap::new();
7822 for key in b.keys() {
7823 if key.starts_with(&lookup_prefix) {
7824 acc |= b.get(&key);
7825 }
7826 }
7827 RowIdSet::from_roaring(acc)
7828 } else {
7829 RowIdSet::empty()
7830 }
7831 }
7832 Condition::FmContains { column_id, pattern } => self
7833 .fm
7834 .get(column_id)
7835 .map(|f| {
7836 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7837 })
7838 .unwrap_or_else(RowIdSet::empty),
7839 Condition::FmContainsAll {
7840 column_id,
7841 patterns,
7842 } => {
7843 if let Some(f) = self.fm.get(column_id) {
7846 let sets: Vec<RowIdSet> = patterns
7847 .iter()
7848 .map(|pat| {
7849 RowIdSet::from_unsorted(
7850 f.locate(pat).into_iter().map(|r| r.0).collect(),
7851 )
7852 })
7853 .collect();
7854 RowIdSet::intersect_many(sets)
7855 } else {
7856 RowIdSet::empty()
7857 }
7858 }
7859 Condition::Ann {
7860 column_id,
7861 query,
7862 k,
7863 } => RowIdSet::from_unsorted(
7864 self.retrieve_filtered(
7865 &crate::query::Retriever::Ann {
7866 column_id: *column_id,
7867 query: query.clone(),
7868 k: *k,
7869 },
7870 snapshot,
7871 None,
7872 allowed,
7873 None,
7874 None,
7875 )?
7876 .into_iter()
7877 .map(|hit| hit.row_id.0)
7878 .collect(),
7879 ),
7880 Condition::SparseMatch {
7881 column_id,
7882 query,
7883 k,
7884 } => RowIdSet::from_unsorted(
7885 self.retrieve_filtered(
7886 &crate::query::Retriever::Sparse {
7887 column_id: *column_id,
7888 query: query.clone(),
7889 k: *k,
7890 },
7891 snapshot,
7892 None,
7893 allowed,
7894 None,
7895 None,
7896 )?
7897 .into_iter()
7898 .map(|hit| hit.row_id.0)
7899 .collect(),
7900 ),
7901 Condition::MinHashSimilar {
7902 column_id,
7903 query,
7904 k,
7905 } => match self.minhash.get(column_id) {
7906 Some(index) => {
7907 let candidates = index.candidate_row_ids(query);
7908 let eligible =
7909 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7910 RowIdSet::from_unsorted(
7911 index
7912 .search_filtered(query, *k, |row_id| {
7913 eligible.contains(&row_id)
7914 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7915 })
7916 .into_iter()
7917 .map(|(row_id, _)| row_id.0)
7918 .collect(),
7919 )
7920 }
7921 None => RowIdSet::empty(),
7922 },
7923 Condition::Range { column_id, lo, hi } => {
7924 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7933 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7934 } else if self.run_refs.len() == 1 {
7935 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7936 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7937 } else {
7938 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7939 };
7940 set.remove_many(self.overlay_rid_set(snapshot));
7941 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7942 set
7943 }
7944 Condition::RangeF64 {
7945 column_id,
7946 lo,
7947 lo_inclusive,
7948 hi,
7949 hi_inclusive,
7950 } => {
7951 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7954 RowIdSet::from_unsorted(
7955 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7956 .into_iter()
7957 .collect(),
7958 )
7959 } else if self.run_refs.len() == 1 {
7960 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7961 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7962 } else {
7963 return self.range_scan_f64(
7964 *column_id,
7965 *lo,
7966 *lo_inclusive,
7967 *hi,
7968 *hi_inclusive,
7969 snapshot,
7970 );
7971 };
7972 set.remove_many(self.overlay_rid_set(snapshot));
7973 self.range_scan_overlay_f64(
7974 &mut set,
7975 *column_id,
7976 *lo,
7977 *lo_inclusive,
7978 *hi,
7979 *hi_inclusive,
7980 snapshot,
7981 );
7982 set
7983 }
7984 Condition::IsNull { column_id } => {
7985 let mut set = if self.run_refs.len() == 1 {
7986 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7987 r.null_row_id_set(*column_id, true)?
7988 } else {
7989 return self.null_scan(*column_id, true, snapshot);
7990 };
7991 set.remove_many(self.overlay_rid_set(snapshot));
7992 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7993 set
7994 }
7995 Condition::IsNotNull { column_id } => {
7996 let mut set = if self.run_refs.len() == 1 {
7997 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7998 r.null_row_id_set(*column_id, false)?
7999 } else {
8000 return self.null_scan(*column_id, false, snapshot);
8001 };
8002 set.remove_many(self.overlay_rid_set(snapshot));
8003 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
8004 set
8005 }
8006 })
8007 }
8008
8009 fn range_scan_i64(
8017 &self,
8018 column_id: u16,
8019 lo: i64,
8020 hi: i64,
8021 snapshot: Snapshot,
8022 ) -> Result<RowIdSet> {
8023 let mut row_ids = Vec::new();
8024 let overlay_rids = self.overlay_rid_set(snapshot);
8025 for rr in &self.run_refs {
8026 let mut reader = self.open_reader(rr.run_id)?;
8027 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
8028 for rid in matched {
8029 if !overlay_rids.contains(&rid) {
8030 row_ids.push(rid);
8031 }
8032 }
8033 }
8034 let mut s = RowIdSet::from_unsorted(row_ids);
8035 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
8036 Ok(s)
8037 }
8038
8039 fn range_scan_f64(
8042 &self,
8043 column_id: u16,
8044 lo: f64,
8045 lo_inclusive: bool,
8046 hi: f64,
8047 hi_inclusive: bool,
8048 snapshot: Snapshot,
8049 ) -> Result<RowIdSet> {
8050 let mut row_ids = Vec::new();
8051 let overlay_rids = self.overlay_rid_set(snapshot);
8052 for rr in &self.run_refs {
8053 let mut reader = self.open_reader(rr.run_id)?;
8054 let matched = reader.range_row_ids_visible_f64(
8055 column_id,
8056 lo,
8057 lo_inclusive,
8058 hi,
8059 hi_inclusive,
8060 snapshot.epoch,
8061 )?;
8062 for rid in matched {
8063 if !overlay_rids.contains(&rid) {
8064 row_ids.push(rid);
8065 }
8066 }
8067 }
8068 let mut s = RowIdSet::from_unsorted(row_ids);
8069 self.range_scan_overlay_f64(
8070 &mut s,
8071 column_id,
8072 lo,
8073 lo_inclusive,
8074 hi,
8075 hi_inclusive,
8076 snapshot,
8077 );
8078 Ok(s)
8079 }
8080
8081 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
8083 let mut s = HashSet::new();
8084 for row in self.memtable.visible_versions_at(snapshot) {
8085 s.insert(row.row_id.0);
8086 }
8087 for row in self.mutable_run.visible_versions_at(snapshot) {
8088 s.insert(row.row_id.0);
8089 }
8090 s
8091 }
8092
8093 fn range_scan_overlay_i64(
8094 &self,
8095 s: &mut RowIdSet,
8096 column_id: u16,
8097 lo: i64,
8098 hi: i64,
8099 snapshot: Snapshot,
8100 ) {
8101 let mut newest: HashMap<u64, Row> = HashMap::new();
8108 for r in self.mutable_run.visible_versions_at(snapshot) {
8109 newest.insert(r.row_id.0, r);
8110 }
8111 for r in self.memtable.visible_versions_at(snapshot) {
8112 newest
8113 .entry(r.row_id.0)
8114 .and_modify(|cur| {
8115 if Snapshot::version_is_newer(
8116 r.committed_epoch,
8117 r.commit_ts,
8118 cur.committed_epoch,
8119 cur.commit_ts,
8120 ) {
8121 *cur = r.clone();
8122 }
8123 })
8124 .or_insert(r);
8125 }
8126 for row in newest.values() {
8127 if !row.deleted {
8128 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
8129 if *v >= lo && *v <= hi {
8130 s.insert(row.row_id.0);
8131 }
8132 }
8133 }
8134 }
8135 }
8136
8137 #[allow(clippy::too_many_arguments)]
8138 fn range_scan_overlay_f64(
8139 &self,
8140 s: &mut RowIdSet,
8141 column_id: u16,
8142 lo: f64,
8143 lo_inclusive: bool,
8144 hi: f64,
8145 hi_inclusive: bool,
8146 snapshot: Snapshot,
8147 ) {
8148 let mut newest: HashMap<u64, Row> = HashMap::new();
8151 for r in self.mutable_run.visible_versions_at(snapshot) {
8152 newest.insert(r.row_id.0, r);
8153 }
8154 for r in self.memtable.visible_versions_at(snapshot) {
8155 newest
8156 .entry(r.row_id.0)
8157 .and_modify(|cur| {
8158 if Snapshot::version_is_newer(
8159 r.committed_epoch,
8160 r.commit_ts,
8161 cur.committed_epoch,
8162 cur.commit_ts,
8163 ) {
8164 *cur = r.clone();
8165 }
8166 })
8167 .or_insert(r);
8168 }
8169 for row in newest.values() {
8170 if !row.deleted {
8171 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
8172 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
8173 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
8174 if ok_lo && ok_hi {
8175 s.insert(row.row_id.0);
8176 }
8177 }
8178 }
8179 }
8180 }
8181
8182 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
8185 let mut row_ids = Vec::new();
8186 let overlay_rids = self.overlay_rid_set(snapshot);
8187 for rr in &self.run_refs {
8188 let mut reader = self.open_reader(rr.run_id)?;
8189 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
8190 for rid in matched {
8191 if !overlay_rids.contains(&rid) {
8192 row_ids.push(rid);
8193 }
8194 }
8195 }
8196 let mut s = RowIdSet::from_unsorted(row_ids);
8197 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
8198 Ok(s)
8199 }
8200
8201 fn null_scan_overlay(
8205 &self,
8206 s: &mut RowIdSet,
8207 column_id: u16,
8208 want_nulls: bool,
8209 snapshot: Snapshot,
8210 ) {
8211 let mut newest: HashMap<u64, Row> = HashMap::new();
8212 for r in self.mutable_run.visible_versions_at(snapshot) {
8213 newest.insert(r.row_id.0, r);
8214 }
8215 for r in self.memtable.visible_versions_at(snapshot) {
8216 newest
8217 .entry(r.row_id.0)
8218 .and_modify(|cur| {
8219 if Snapshot::version_is_newer(
8220 r.committed_epoch,
8221 r.commit_ts,
8222 cur.committed_epoch,
8223 cur.commit_ts,
8224 ) {
8225 *cur = r.clone();
8226 }
8227 })
8228 .or_insert(r);
8229 }
8230 for row in newest.values() {
8231 if row.deleted {
8232 continue;
8233 }
8234 let is_null = !row.columns.contains_key(&column_id)
8235 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
8236 if is_null == want_nulls {
8237 s.insert(row.row_id.0);
8238 }
8239 }
8240 }
8241
8242 pub fn snapshot(&self) -> Snapshot {
8243 let epoch = self.epoch.visible();
8244 match &self.wal {
8248 WalSink::Shared(shared) => match shared.hlc.now() {
8249 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
8250 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
8252 },
8253 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
8254 }
8255 }
8256
8257 pub fn data_generation(&self) -> u64 {
8259 self.data_generation
8260 }
8261
8262 pub(crate) fn bump_data_generation(&mut self) {
8263 self.data_generation = self.data_generation.wrapping_add(1);
8264 }
8265
8266 pub fn table_id(&self) -> u64 {
8268 self.table_id
8269 }
8270
8271 fn seal_generations(&mut self) {
8277 self.memtable.seal();
8278 self.mutable_run.seal();
8279 self.hot.seal();
8280 for index in self.bitmap.values_mut() {
8281 index.seal();
8282 }
8283 for index in self.ann.values_mut() {
8284 index.seal();
8285 }
8286 for index in self.fm.values_mut() {
8287 index.seal();
8288 }
8289 for index in self.sparse.values_mut() {
8290 index.seal();
8291 }
8292 for index in self.minhash.values_mut() {
8293 index.seal();
8294 }
8295 self.pk_by_row.seal();
8296 }
8297
8298 fn capture_read_generation(&self) -> ReadGeneration {
8303 let visible_through = self.current_epoch();
8304 ReadGeneration {
8305 schema: Arc::new(self.schema.clone()),
8306 base_runs: Arc::new(self.run_refs.clone()),
8307 deltas: TableDeltas {
8308 memtable: self.memtable.clone(),
8309 mutable_run: self.mutable_run.clone(),
8310 hot: self.hot.clone(),
8311 pk_by_row: self.pk_by_row.clone(),
8312 },
8313 indexes: Arc::new(IndexGeneration::capture(
8314 &self.bitmap,
8315 &self.learned_range,
8316 &self.fm,
8317 &self.ann,
8318 &self.sparse,
8319 &self.minhash,
8320 visible_through,
8321 match &self.wal {
8323 WalSink::Shared(shared) => shared
8324 .hlc
8325 .now()
8326 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
8327 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
8328 },
8329 )),
8330 visible_through,
8331 }
8332 }
8333
8334 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
8339 self.ensure_indexes_complete()?;
8340 self.seal_generations();
8341 let view = Arc::new(self.capture_read_generation());
8342 self.published.store(Arc::clone(&view));
8343 Ok(view)
8344 }
8345
8346 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
8352 self.published.load_full()
8353 }
8354
8355 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
8357 &self.pins
8358 }
8359
8360 pub fn version_gc_floor(&self) -> Epoch {
8365 self.min_active_snapshot()
8366 .unwrap_or_else(|| self.current_epoch())
8367 }
8368
8369 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
8376 let mut report = self.pins.report();
8377 let transaction_floor = [
8378 self.pinned.keys().next().copied(),
8379 self.snapshots.min_pinned(),
8380 ]
8381 .into_iter()
8382 .flatten()
8383 .min();
8384 if let Some(epoch) = transaction_floor {
8385 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
8386 }
8387 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
8388 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
8389 }
8390 report
8391 }
8392
8393 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
8394 self.publish_read_generation()?;
8395 let mut generation = self.clone();
8396 generation.read_only = true;
8397 generation.wal = WalSink::ReadOnly;
8398 generation.pending_delete_rids.clear();
8399 generation.pending_put_cols.clear();
8400 generation.pending_rows.clear();
8401 generation.pending_rows_auto_inc.clear();
8402 generation.pending_dels.clear();
8403 generation.pending_truncate = None;
8404 generation.agg_cache = Arc::new(HashMap::new());
8405 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
8408 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
8411 crate::retention::PinSource::ReadGeneration,
8412 self.current_epoch(),
8413 )));
8414 Ok(generation)
8415 }
8416
8417 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
8418 (std::mem::size_of::<Self>() as u64)
8419 .saturating_add(self.memtable.approx_bytes())
8420 .saturating_add(self.mutable_run.approx_bytes())
8421 .saturating_add(self.live_count.saturating_mul(64))
8422 }
8423
8424 pub fn pin_snapshot(&mut self) -> Snapshot {
8431 let snap = self.snapshot();
8432 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
8433 snap
8434 }
8435
8436 pub fn hlc_gc_floor(
8446 &self,
8447 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
8448 ) -> crate::epoch::GcFloor {
8449 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
8450 epoch
8451 .and_then(&mut project_epoch)
8452 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
8453 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
8454 };
8455 let transaction = [
8456 self.pinned.keys().next().copied(),
8457 self.snapshots.min_pinned(),
8458 ]
8459 .into_iter()
8460 .flatten()
8461 .min();
8462 let history = self.snapshots.history_floor(self.current_epoch());
8463 crate::epoch::GcFloor {
8464 transaction_snapshot: project(transaction),
8465 history_retention: project(history),
8466 backup_pitr: project(
8467 self.pins
8468 .oldest_for(crate::retention::PinSource::BackupPitr),
8469 ),
8470 replication: project(
8471 self.pins
8472 .oldest_for(crate::retention::PinSource::Replication),
8473 ),
8474 read_generation: project(
8475 self.pins
8476 .oldest_for(crate::retention::PinSource::ReadGeneration),
8477 ),
8478 online_index_build: project(
8479 self.pins
8480 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
8481 ),
8482 }
8483 }
8484
8485 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
8487 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
8488 *count -= 1;
8489 if *count == 0 {
8490 self.pinned.remove(&snap.epoch);
8491 }
8492 }
8493 }
8494
8495 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
8507 let local = self.pinned.keys().next().copied();
8508 let global = self.snapshots.min_pinned();
8509 let history = self.snapshots.history_floor(self.current_epoch());
8510 let pinned = self.pins.oldest_pinned();
8511 [local, global, history, pinned].into_iter().flatten().min()
8512 }
8513
8514 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
8518 self.ensure_writable()?;
8519 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
8520 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
8521 }
8522
8523 pub fn clear_ttl(&mut self) -> Result<()> {
8524 self.ensure_writable()?;
8525 self.apply_ttl_policy_at(None, self.current_epoch())
8526 }
8527
8528 pub fn ttl(&self) -> Option<TtlPolicy> {
8529 self.ttl
8530 }
8531
8532 pub(crate) fn prepare_ttl_policy(
8533 &self,
8534 column_name: &str,
8535 duration_nanos: u64,
8536 ) -> Result<TtlPolicy> {
8537 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
8538 return Err(MongrelError::InvalidArgument(
8539 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
8540 ));
8541 }
8542 let column = self
8543 .schema
8544 .columns
8545 .iter()
8546 .find(|column| column.name == column_name)
8547 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
8548 if column.ty != TypeId::TimestampNanos {
8549 return Err(MongrelError::Schema(format!(
8550 "TTL column {column_name} must be TimestampNanos, is {:?}",
8551 column.ty
8552 )));
8553 }
8554 Ok(TtlPolicy {
8555 column_id: column.id,
8556 duration_nanos,
8557 })
8558 }
8559
8560 pub(crate) fn apply_ttl_policy_at(
8561 &mut self,
8562 policy: Option<TtlPolicy>,
8563 epoch: Epoch,
8564 ) -> Result<()> {
8565 if let Some(policy) = policy {
8566 let column = self
8567 .schema
8568 .columns
8569 .iter()
8570 .find(|column| column.id == policy.column_id)
8571 .ok_or_else(|| {
8572 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8573 })?;
8574 if column.ty != TypeId::TimestampNanos
8575 || policy.duration_nanos == 0
8576 || policy.duration_nanos > i64::MAX as u64
8577 {
8578 return Err(MongrelError::Schema("invalid TTL policy".into()));
8579 }
8580 }
8581 self.ttl = policy;
8582 self.agg_cache = Arc::new(HashMap::new());
8583 self.clear_result_cache();
8584 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8585 self.persist_manifest(epoch)
8586 }
8587
8588 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8589 let Some(policy) = self.ttl else {
8590 return false;
8591 };
8592 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8593 return false;
8594 };
8595 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8596 }
8597
8598 pub fn current_epoch(&self) -> Epoch {
8599 self.epoch.visible()
8600 }
8601
8602 pub fn memtable_len(&self) -> usize {
8603 self.memtable.len()
8604 }
8605
8606 pub fn count(&self) -> u64 {
8609 if self.ttl.is_none()
8610 && self.pending_put_cols.is_empty()
8611 && self.pending_delete_rids.is_empty()
8612 && self.pending_rows.is_empty()
8613 && self.pending_dels.is_empty()
8614 && self.pending_truncate.is_none()
8615 {
8616 self.live_count
8617 } else {
8618 self.visible_rows(self.snapshot())
8619 .map(|rows| rows.len() as u64)
8620 .unwrap_or(self.live_count)
8621 }
8622 }
8623
8624 pub fn count_conditions(
8628 &mut self,
8629 conditions: &[crate::query::Condition],
8630 snapshot: Snapshot,
8631 ) -> Result<Option<u64>> {
8632 use crate::query::Condition;
8633 if self.ttl.is_some() {
8634 if conditions.is_empty() {
8635 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8636 }
8637 let mut sets = Vec::with_capacity(conditions.len());
8638 for condition in conditions {
8639 sets.push(self.resolve_condition(condition, snapshot)?);
8640 }
8641 let survivors = RowIdSet::intersect_many(sets);
8642 let rows = self.visible_rows(snapshot)?;
8643 return Ok(Some(
8644 rows.into_iter()
8645 .filter(|row| survivors.contains(row.row_id.0))
8646 .count() as u64,
8647 ));
8648 }
8649 if conditions.is_empty() {
8650 return Ok(Some(self.count()));
8651 }
8652 let served = |c: &Condition| {
8653 matches!(
8654 c,
8655 Condition::Pk(_)
8656 | Condition::BitmapEq { .. }
8657 | Condition::BitmapIn { .. }
8658 | Condition::BytesPrefix { .. }
8659 | Condition::FmContains { .. }
8660 | Condition::FmContainsAll { .. }
8661 | Condition::Ann { .. }
8662 | Condition::Range { .. }
8663 | Condition::RangeF64 { .. }
8664 | Condition::SparseMatch { .. }
8665 | Condition::MinHashSimilar { .. }
8666 | Condition::IsNull { .. }
8667 | Condition::IsNotNull { .. }
8668 )
8669 };
8670 if !conditions.iter().all(served) {
8671 return Ok(None);
8672 }
8673 self.ensure_indexes_complete()?;
8674 if !self.pending_put_cols.is_empty()
8675 || !self.pending_delete_rids.is_empty()
8676 || !self.pending_rows.is_empty()
8677 || !self.pending_dels.is_empty()
8678 || self.pending_truncate.is_some()
8679 {
8680 let mut sets = Vec::with_capacity(conditions.len());
8681 for condition in conditions {
8682 sets.push(self.resolve_condition(condition, snapshot)?);
8683 }
8684 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8685 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8686 }
8687 let mut sets = Vec::with_capacity(conditions.len());
8688 for condition in conditions {
8689 sets.push(self.resolve_condition(condition, snapshot)?);
8690 }
8691 let mut rids = RowIdSet::intersect_many(sets);
8692 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8702 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8703 }
8704 let count = rids.len() as u64;
8705 crate::trace::QueryTrace::record(|t| {
8706 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8707 t.survivor_count = Some(count as usize);
8708 t.conditions_pushed = conditions.len();
8709 });
8710 Ok(Some(count))
8711 }
8712
8713 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8718 let mut out = Vec::new();
8719 for row in self.memtable.visible_versions(snapshot.epoch) {
8720 if row.deleted {
8721 out.push(row.row_id.0);
8722 }
8723 }
8724 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8725 if row.deleted {
8726 out.push(row.row_id.0);
8727 }
8728 }
8729 out
8730 }
8731
8732 pub fn bulk_load_columns(
8741 &mut self,
8742 user_columns: Vec<(u16, columnar::NativeColumn)>,
8743 ) -> Result<Epoch> {
8744 self.bulk_load_columns_with(user_columns, 3, false, true)
8745 }
8746
8747 pub fn bulk_load_fast(
8754 &mut self,
8755 user_columns: Vec<(u16, columnar::NativeColumn)>,
8756 ) -> Result<Epoch> {
8757 self.bulk_load_columns_with(user_columns, -1, true, false)
8758 }
8759
8760 fn bulk_load_columns_with(
8761 &mut self,
8762 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8763 zstd_level: i32,
8764 force_plain: bool,
8765 lz4: bool,
8766 ) -> Result<Epoch> {
8767 self.ensure_writable()?;
8768 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8769 if n == 0 {
8770 return Ok(self.current_epoch());
8771 }
8772 let epoch = self.commit_new_epoch()?;
8773 let live_before = self.live_count;
8774 self.spill_mutable_run(epoch)?;
8776 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8777 && self.indexes_complete
8778 && self.run_refs.is_empty()
8779 && self.memtable.is_empty()
8780 && self.mutable_run.is_empty();
8781 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8784 self.validate_columns_not_null(&user_columns, n)?;
8785 let winner_idx = self
8786 .bulk_pk_winner_indices(&user_columns, n)
8787 .filter(|idx| idx.len() != n);
8788 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8789 match winner_idx.as_deref() {
8790 Some(idx) => {
8791 let compacted = user_columns
8792 .iter()
8793 .map(|(id, c)| (*id, c.gather(idx)))
8794 .collect();
8795 (compacted, idx.len())
8796 }
8797 None => (user_columns, n),
8798 };
8799 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8800 let first = self.allocator.alloc_range(write_n as u64)?.0;
8801 for rid in first..first + write_n as u64 {
8802 self.reservoir.offer(rid);
8803 }
8804 let run_id = self.alloc_run_id()?;
8805 let path = self.run_path(run_id);
8806 let mut writer =
8807 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8808 if force_plain {
8809 writer = writer.with_plain();
8810 } else if lz4 {
8811 writer = writer.with_lz4();
8814 } else {
8815 writer = writer.with_zstd_level(zstd_level);
8816 }
8817 if let Some(kek) = &self.kek {
8818 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8819 }
8820 let header = match self.create_run_file(run_id)? {
8821 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8822 None => writer.write_native(&path, &write_columns, write_n, first)?,
8823 };
8824 self.run_refs.push(RunRef {
8825 run_id: run_id as u128,
8826 level: 0,
8827 epoch_created: epoch.0,
8828 row_count: header.row_count,
8829 });
8830 self.live_count = self.live_count.saturating_add(write_n as u64);
8831 if eager_index_build {
8832 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8833 self.index_columns_bulk(&write_columns, &row_ids);
8834 self.indexes_complete = true;
8835 self.build_learned_ranges()?;
8836 } else {
8837 self.indexes_complete = false;
8841 }
8842 self.mark_flushed(epoch)?;
8843 self.persist_manifest(epoch)?;
8844 if eager_index_build {
8845 self.checkpoint_indexes(epoch);
8846 }
8847 self.clear_result_cache();
8848 self.data_generation = self.data_generation.wrapping_add(1);
8849 Ok(epoch)
8850 }
8851
8852 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8870 let n = row_ids.len();
8871 if n == 0 {
8872 return;
8873 }
8874 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8875 columns.iter().map(|(id, c)| (*id, c)).collect();
8876 let ty_of: std::collections::HashMap<u16, TypeId> = self
8877 .schema
8878 .columns
8879 .iter()
8880 .map(|c| (c.id, c.ty.clone()))
8881 .collect();
8882 let pk_id = self.schema.primary_key().map(|c| c.id);
8883
8884 for (i, &rid) in row_ids.iter().enumerate() {
8885 let row_id = RowId(rid);
8886 if let Some(pid) = pk_id {
8887 if let Some(col) = by_id.get(&pid) {
8888 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8889 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8890 self.insert_hot_pk(key, row_id);
8891 }
8892 }
8893 }
8894 for idef in &self.schema.indexes {
8895 let Some(col) = by_id.get(&idef.column_id) else {
8896 continue;
8897 };
8898 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8899 match idef.kind {
8900 IndexKind::Bitmap => {
8901 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8902 if let Some(key) =
8903 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8904 {
8905 b.insert(key, row_id);
8906 }
8907 }
8908 }
8909 IndexKind::FmIndex => {
8910 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8911 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8912 f.insert(bytes.to_vec(), row_id);
8913 }
8914 }
8915 }
8916 IndexKind::Sparse => {
8917 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8918 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8919 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8920 s.insert(&terms, row_id);
8921 }
8922 }
8923 }
8924 }
8925 IndexKind::MinHash => {
8926 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8927 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8928 let tokens = crate::index::token_hashes_from_bytes(bytes);
8929 mh.insert(&tokens, row_id);
8930 }
8931 }
8932 }
8933 _ => {}
8934 }
8935 }
8936 }
8937 }
8938
8939 pub fn visible_columns_native(
8944 &self,
8945 snapshot: Snapshot,
8946 projection: Option<&[u16]>,
8947 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8948 self.visible_columns_native_inner(snapshot, projection, None)
8949 }
8950
8951 pub fn visible_columns_native_with_control(
8952 &self,
8953 snapshot: Snapshot,
8954 projection: Option<&[u16]>,
8955 control: &crate::ExecutionControl,
8956 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8957 self.visible_columns_native_inner(snapshot, projection, Some(control))
8958 }
8959
8960 fn visible_columns_native_inner(
8961 &self,
8962 snapshot: Snapshot,
8963 projection: Option<&[u16]>,
8964 control: Option<&crate::ExecutionControl>,
8965 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8966 execution_checkpoint(control, 0)?;
8967 let wanted: Vec<u16> = match projection {
8968 Some(p) => p.to_vec(),
8969 None => self.schema.columns.iter().map(|c| c.id).collect(),
8970 };
8971 if self.ttl.is_none()
8972 && self.memtable.is_empty()
8973 && self.mutable_run.is_empty()
8974 && self.run_refs.len() == 1
8975 {
8976 let rr = self.run_refs[0].clone();
8977 let mut reader = self.open_reader(rr.run_id)?;
8978 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8979 execution_checkpoint(control, 0)?;
8980 let all_visible = idxs.len() == reader.row_count();
8981 if reader.has_mmap() && control.is_none() {
8987 use rayon::prelude::*;
8988 let valid: Vec<u16> = wanted
8991 .iter()
8992 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8993 .copied()
8994 .collect();
8995 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8997 .par_iter()
8998 .filter_map(|cid| {
8999 reader
9000 .column_native_shared(*cid)
9001 .ok()
9002 .map(|col| (*cid, col))
9003 })
9004 .collect();
9005 let cols = decoded
9006 .into_iter()
9007 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
9008 .collect();
9009 return Ok(cols);
9010 }
9011 let mut cols = Vec::with_capacity(wanted.len());
9012 for (index, cid) in wanted.iter().enumerate() {
9013 execution_checkpoint(control, index)?;
9014 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
9015 Some(c) => c,
9016 None => continue,
9017 };
9018 let col = reader.column_native(cdef.id)?;
9019 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
9020 }
9021 return Ok(cols);
9022 }
9023 let vcols = self.visible_columns(snapshot)?;
9024 execution_checkpoint(control, 0)?;
9025 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
9026 let out: Vec<(u16, columnar::NativeColumn)> = vcols
9027 .into_iter()
9028 .filter(|(id, _)| want_set.contains(id))
9029 .map(|(id, vals)| {
9030 let ty = self
9031 .schema
9032 .columns
9033 .iter()
9034 .find(|c| c.id == id)
9035 .map(|c| c.ty.clone())
9036 .unwrap_or(TypeId::Bytes);
9037 (id, columnar::values_to_native(ty, &vals))
9038 })
9039 .collect();
9040 Ok(out)
9041 }
9042
9043 pub fn run_count(&self) -> usize {
9044 self.run_refs.len()
9045 }
9046
9047 pub fn memtable_is_empty(&self) -> bool {
9049 self.memtable.is_empty()
9050 }
9051
9052 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
9056 self.page_cache.stats()
9057 }
9058
9059 pub fn reset_page_cache_stats(&self) {
9061 self.page_cache.reset_stats();
9062 }
9063
9064 pub fn run_ids(&self) -> Vec<u128> {
9067 self.run_refs.iter().map(|r| r.run_id).collect()
9068 }
9069
9070 pub fn single_run_is_clean(&self) -> bool {
9074 if self.ttl.is_some() || self.run_refs.len() != 1 {
9075 return false;
9076 }
9077 self.open_reader(self.run_refs[0].run_id)
9078 .map(|r| r.is_clean())
9079 .unwrap_or(false)
9080 }
9081
9082 fn resolve_footprint(
9089 &self,
9090 conditions: &[crate::query::Condition],
9091 snapshot: Snapshot,
9092 ) -> roaring::RoaringBitmap {
9093 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
9094 return roaring::RoaringBitmap::new();
9095 }
9096 if self.run_refs.is_empty() {
9097 return roaring::RoaringBitmap::new();
9098 }
9099 if self.run_refs.len() == 1 {
9101 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
9102 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
9103 return rids.to_roaring_lossy();
9104 }
9105 }
9106 }
9107 roaring::RoaringBitmap::new()
9108 }
9109
9110 pub fn query_columns_native_cached(
9121 &mut self,
9122 conditions: &[crate::query::Condition],
9123 projection: Option<&[u16]>,
9124 snapshot: Snapshot,
9125 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9126 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
9127 }
9128
9129 pub fn query_columns_native_cached_with_control(
9130 &mut self,
9131 conditions: &[crate::query::Condition],
9132 projection: Option<&[u16]>,
9133 snapshot: Snapshot,
9134 control: &crate::ExecutionControl,
9135 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9136 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
9137 }
9138
9139 fn query_columns_native_cached_inner(
9140 &mut self,
9141 conditions: &[crate::query::Condition],
9142 projection: Option<&[u16]>,
9143 snapshot: Snapshot,
9144 control: Option<&crate::ExecutionControl>,
9145 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9146 execution_checkpoint(control, 0)?;
9147 if self.ttl.is_some() {
9150 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9151 }
9152 if conditions.is_empty() {
9153 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9154 }
9155 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
9159 if let Some(hit) = self.result_cache.lock().get_columns(key) {
9160 crate::trace::QueryTrace::record(|t| {
9161 t.result_cache_hit = true;
9162 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9163 });
9164 return Ok(Some((*hit).clone()));
9165 }
9166 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
9167 execution_checkpoint(control, 0)?;
9168 if let Some(cols) = &res {
9169 let footprint = self.resolve_footprint(conditions, snapshot);
9170 let condition_cols = crate::query::condition_columns(conditions);
9171 execution_checkpoint(control, 0)?;
9172 self.result_cache.lock().insert(
9173 key,
9174 CachedEntry {
9175 data: CachedData::Columns(Arc::new(cols.clone())),
9176 footprint,
9177 condition_cols,
9178 },
9179 );
9180 }
9181 Ok(res)
9182 }
9183
9184 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
9189 if self.ttl.is_some() {
9190 return self.query(q);
9191 }
9192 if q.conditions.is_empty() {
9193 return self.query(q);
9194 }
9195 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
9196 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
9197 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
9198 if let Some(hit) = self.result_cache.lock().get_rows(key) {
9199 crate::trace::QueryTrace::record(|t| {
9200 t.result_cache_hit = true;
9201 t.scan_mode = crate::trace::ScanMode::Materialized;
9202 });
9203 return Ok((*hit).clone());
9204 }
9205 let rows = self.query(q)?;
9206 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
9207 let condition_cols = crate::query::condition_columns(&q.conditions);
9208 self.result_cache.lock().insert(
9209 key,
9210 CachedEntry {
9211 data: CachedData::Rows(Arc::new(rows.clone())),
9212 footprint,
9213 condition_cols,
9214 },
9215 );
9216 Ok(rows)
9217 }
9218
9219 #[allow(clippy::type_complexity)]
9234 pub fn query_columns_native_traced(
9235 &mut self,
9236 conditions: &[crate::query::Condition],
9237 projection: Option<&[u16]>,
9238 snapshot: Snapshot,
9239 ) -> Result<(
9240 Option<Vec<(u16, columnar::NativeColumn)>>,
9241 crate::trace::QueryTrace,
9242 )> {
9243 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9244 self.query_columns_native(conditions, projection, snapshot)
9245 });
9246 Ok((result?, trace))
9247 }
9248
9249 #[allow(clippy::type_complexity)]
9252 pub fn query_columns_native_cached_traced(
9253 &mut self,
9254 conditions: &[crate::query::Condition],
9255 projection: Option<&[u16]>,
9256 snapshot: Snapshot,
9257 ) -> Result<(
9258 Option<Vec<(u16, columnar::NativeColumn)>>,
9259 crate::trace::QueryTrace,
9260 )> {
9261 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9262 self.query_columns_native_cached(conditions, projection, snapshot)
9263 });
9264 Ok((result?, trace))
9265 }
9266
9267 pub fn native_page_cursor_traced(
9269 &self,
9270 snapshot: Snapshot,
9271 projection: Vec<(u16, TypeId)>,
9272 conditions: &[crate::query::Condition],
9273 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
9274 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9275 self.native_page_cursor(snapshot, projection, conditions)
9276 });
9277 Ok((result?, trace))
9278 }
9279
9280 pub fn native_multi_run_cursor_traced(
9282 &self,
9283 snapshot: Snapshot,
9284 projection: Vec<(u16, TypeId)>,
9285 conditions: &[crate::query::Condition],
9286 ) -> Result<(
9287 Option<crate::cursor::MultiRunCursor>,
9288 crate::trace::QueryTrace,
9289 )> {
9290 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9291 self.native_multi_run_cursor(snapshot, projection, conditions)
9292 });
9293 Ok((result?, trace))
9294 }
9295
9296 pub fn count_conditions_traced(
9298 &mut self,
9299 conditions: &[crate::query::Condition],
9300 snapshot: Snapshot,
9301 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
9302 let (result, trace) =
9303 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
9304 Ok((result?, trace))
9305 }
9306
9307 pub fn query_traced(
9309 &mut self,
9310 q: &crate::query::Query,
9311 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
9312 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
9313 Ok((result?, trace))
9314 }
9315
9316 pub fn query_columns_native(
9321 &mut self,
9322 conditions: &[crate::query::Condition],
9323 projection: Option<&[u16]>,
9324 snapshot: Snapshot,
9325 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9326 self.query_columns_native_inner(conditions, projection, snapshot, None)
9327 }
9328
9329 pub fn query_columns_native_with_control(
9330 &mut self,
9331 conditions: &[crate::query::Condition],
9332 projection: Option<&[u16]>,
9333 snapshot: Snapshot,
9334 control: &crate::ExecutionControl,
9335 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9336 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
9337 }
9338
9339 fn query_columns_native_inner(
9340 &mut self,
9341 conditions: &[crate::query::Condition],
9342 projection: Option<&[u16]>,
9343 snapshot: Snapshot,
9344 control: Option<&crate::ExecutionControl>,
9345 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9346 use crate::query::Condition;
9347 execution_checkpoint(control, 0)?;
9348 if self.ttl.is_some() {
9351 return Ok(None);
9352 }
9353 if conditions.is_empty() {
9354 return Ok(None);
9355 }
9356 self.ensure_indexes_complete()?;
9357
9358 let served = |c: &Condition| {
9363 matches!(
9364 c,
9365 Condition::Pk(_)
9366 | Condition::BitmapEq { .. }
9367 | Condition::BitmapIn { .. }
9368 | Condition::BytesPrefix { .. }
9369 | Condition::FmContains { .. }
9370 | Condition::FmContainsAll { .. }
9371 | Condition::Ann { .. }
9372 | Condition::Range { .. }
9373 | Condition::RangeF64 { .. }
9374 | Condition::SparseMatch { .. }
9375 | Condition::MinHashSimilar { .. }
9376 | Condition::IsNull { .. }
9377 | Condition::IsNotNull { .. }
9378 )
9379 };
9380 if !conditions.iter().all(served) {
9381 return Ok(None);
9382 }
9383 let fast_path =
9384 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9385 crate::trace::QueryTrace::record(|t| {
9386 t.run_count = self.run_refs.len();
9387 t.memtable_rows = self.memtable.len();
9388 t.mutable_run_rows = self.mutable_run.len();
9389 t.conditions_pushed = conditions.len();
9390 t.learned_range_used = conditions.iter().any(|c| match c {
9391 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
9392 self.learned_range.contains_key(column_id)
9393 }
9394 _ => false,
9395 });
9396 });
9397 let col_ids: Vec<u16> = projection
9399 .map(|p| p.to_vec())
9400 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
9401 let proj_pairs: Vec<(u16, TypeId)> = col_ids
9402 .iter()
9403 .map(|&cid| {
9404 let ty = self
9405 .schema
9406 .columns
9407 .iter()
9408 .find(|c| c.id == cid)
9409 .map(|c| c.ty.clone())
9410 .unwrap_or(TypeId::Bytes);
9411 (cid, ty)
9412 })
9413 .collect();
9414
9415 if fast_path {
9421 let needs_column = conditions.iter().any(|c| match c {
9424 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
9425 Condition::RangeF64 { column_id, .. } => {
9426 !self.learned_range.contains_key(column_id)
9427 }
9428 _ => false,
9429 });
9430 let mut reader_opt: Option<RunReader> = if needs_column {
9431 Some(self.open_reader(self.run_refs[0].run_id)?)
9432 } else {
9433 None
9434 };
9435 let mut sets: Vec<RowIdSet> = Vec::new();
9436 for (index, c) in conditions.iter().enumerate() {
9437 execution_checkpoint(control, index)?;
9438 let s = match c {
9439 Condition::Range { column_id, lo, hi }
9440 if !self.learned_range.contains_key(column_id) =>
9441 {
9442 if reader_opt.is_none() {
9443 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9444 }
9445 reader_opt
9446 .as_mut()
9447 .expect("reader opened for range")
9448 .range_row_id_set_i64(*column_id, *lo, *hi)?
9449 }
9450 Condition::RangeF64 {
9451 column_id,
9452 lo,
9453 lo_inclusive,
9454 hi,
9455 hi_inclusive,
9456 } if !self.learned_range.contains_key(column_id) => {
9457 if reader_opt.is_none() {
9458 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9459 }
9460 reader_opt
9461 .as_mut()
9462 .expect("reader opened for range")
9463 .range_row_id_set_f64(
9464 *column_id,
9465 *lo,
9466 *lo_inclusive,
9467 *hi,
9468 *hi_inclusive,
9469 )?
9470 }
9471 _ => self.resolve_condition(c, snapshot)?,
9472 };
9473 sets.push(s);
9474 }
9475 let candidates = RowIdSet::intersect_many(sets);
9476 crate::trace::QueryTrace::record(|t| {
9477 t.survivor_count = Some(candidates.len());
9478 });
9479 if candidates.is_empty() {
9480 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
9481 .iter()
9482 .map(|&id| {
9483 (
9484 id,
9485 columnar::null_native(
9486 proj_pairs
9487 .iter()
9488 .find(|(c, _)| c == &id)
9489 .map(|(_, t)| t.clone())
9490 .unwrap_or(TypeId::Bytes),
9491 0,
9492 ),
9493 )
9494 })
9495 .collect();
9496 return Ok(Some(cols));
9497 }
9498 let mut reader = match reader_opt.take() {
9499 Some(r) => r,
9500 None => self.open_reader(self.run_refs[0].run_id)?,
9501 };
9502 let candidate_ids = candidates.into_sorted_vec();
9503 let (positions, fast_rid) = if let Some(positions) =
9504 reader.positions_for_row_ids_fast(&candidate_ids)
9505 {
9506 (positions, true)
9507 } else {
9508 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
9509 match col {
9510 columnar::NativeColumn::Int64 { data, .. } => {
9511 let mut p = Vec::with_capacity(candidate_ids.len());
9512 for (index, rid) in candidate_ids.iter().enumerate() {
9513 execution_checkpoint(control, index)?;
9514 if let Ok(position) = data.binary_search(&(*rid as i64)) {
9515 p.push(position);
9516 }
9517 }
9518 p.sort_unstable();
9519 (p, false)
9520 }
9521 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
9522 }
9523 };
9524 crate::trace::QueryTrace::record(|t| {
9525 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9526 t.fast_row_id_map = fast_rid;
9527 });
9528 let mut cols = Vec::with_capacity(col_ids.len());
9529 for (index, cid) in col_ids.iter().enumerate() {
9530 execution_checkpoint(control, index)?;
9531 let col = reader.column_native(*cid)?;
9532 cols.push((*cid, col.gather(&positions)));
9533 }
9534 return Ok(Some(cols));
9535 }
9536
9537 if !self.run_refs.is_empty() {
9550 use crate::cursor::{
9551 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
9552 };
9553 let remaining: usize;
9554 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9555 let c = self
9556 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9557 .expect("single-run cursor should build when run_refs.len() == 1");
9558 remaining = c.remaining_rows();
9559 Box::new(c)
9560 } else {
9561 let c = self
9562 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9563 .expect("multi-run cursor should build when run_refs.len() >= 1");
9564 remaining = c.remaining_rows();
9565 Box::new(c)
9566 };
9567 crate::trace::QueryTrace::record(|t| {
9568 if t.survivor_count.is_none() {
9569 t.survivor_count = Some(remaining);
9570 }
9571 });
9572 let cols = match control {
9573 Some(control) => {
9574 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9575 }
9576 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9577 };
9578 return Ok(Some(cols));
9579 }
9580
9581 crate::trace::QueryTrace::record(|t| {
9586 t.scan_mode = crate::trace::ScanMode::Materialized;
9587 t.row_materialized = true;
9588 });
9589 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9590 for (index, c) in conditions.iter().enumerate() {
9591 execution_checkpoint(control, index)?;
9592 sets.push(self.resolve_condition(c, snapshot)?);
9593 }
9594 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9595 let rows = self.rows_for_rids(&rids, snapshot)?;
9596 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9597 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9598 execution_checkpoint(control, index)?;
9599 let vals: Vec<Value> = rows
9600 .iter()
9601 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9602 .collect();
9603 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9604 }
9605 Ok(Some(cols))
9606 }
9607
9608 pub fn native_page_cursor(
9623 &self,
9624 snapshot: Snapshot,
9625 projection: Vec<(u16, TypeId)>,
9626 conditions: &[crate::query::Condition],
9627 ) -> Result<Option<NativePageCursor>> {
9628 use crate::cursor::build_page_plans;
9629 if self.ttl.is_some() {
9630 return Ok(None);
9631 }
9632 if !conditions.is_empty() && !self.indexes_complete {
9635 return Ok(None);
9636 }
9637 if self.run_refs.len() != 1 {
9638 return Ok(None);
9639 }
9640 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9641 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9642
9643 let overlay_rids: HashSet<u64> = {
9646 let mut s = HashSet::new();
9647 for row in self.memtable.visible_versions(snapshot.epoch) {
9648 s.insert(row.row_id.0);
9649 }
9650 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9651 s.insert(row.row_id.0);
9652 }
9653 s
9654 };
9655
9656 let survivors = if conditions.is_empty() {
9660 None
9661 } else {
9662 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9663 };
9664
9665 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9672 survivors.clone()
9673 } else if let Some(s) = &survivors {
9674 let mut run_set = s.clone();
9675 run_set.remove_many(overlay_rids.iter().copied());
9676 Some(run_set)
9677 } else {
9678 Some(RowIdSet::from_unsorted(
9679 rids.iter()
9680 .map(|&r| r as u64)
9681 .filter(|r| !overlay_rids.contains(r))
9682 .collect(),
9683 ))
9684 };
9685
9686 let overlay_rows = if overlay_rids.is_empty() {
9687 Vec::new()
9688 } else {
9689 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9690 self.overlay_visible_rows(snapshot, bound)
9691 };
9692
9693 let plans = if positions.is_empty() {
9695 Vec::new()
9696 } else {
9697 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9698 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9699 };
9700
9701 let overlay = if overlay_rows.is_empty() {
9703 None
9704 } else {
9705 let filtered =
9706 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9707 if filtered.is_empty() {
9708 None
9709 } else {
9710 Some(self.materialize_overlay(&filtered, &projection))
9711 }
9712 };
9713
9714 let overlay_row_count = overlay
9715 .as_ref()
9716 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9717 .unwrap_or(0);
9718 crate::trace::QueryTrace::record(|t| {
9719 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9720 t.run_count = self.run_refs.len();
9721 t.memtable_rows = self.memtable.len();
9722 t.mutable_run_rows = self.mutable_run.len();
9723 t.overlay_rows = overlay_row_count;
9724 t.conditions_pushed = conditions.len();
9725 t.pages_decoded = plans
9726 .iter()
9727 .map(|p| p.positions.len())
9728 .sum::<usize>()
9729 .min(1);
9730 });
9731
9732 Ok(Some(NativePageCursor::new_with_overlay(
9733 reader, projection, plans, overlay,
9734 )))
9735 }
9736 #[allow(clippy::type_complexity)]
9746 pub fn native_multi_run_cursor(
9747 &self,
9748 snapshot: Snapshot,
9749 projection: Vec<(u16, TypeId)>,
9750 conditions: &[crate::query::Condition],
9751 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9752 use crate::cursor::{MultiRunCursor, RunStream};
9753 use crate::sorted_run::SYS_ROW_ID;
9754 use std::collections::{BinaryHeap, HashMap, HashSet};
9755 if self.ttl.is_some() {
9756 return Ok(None);
9757 }
9758 if !conditions.is_empty() && !self.indexes_complete {
9761 return Ok(None);
9762 }
9763 if self.run_refs.is_empty() {
9764 return Ok(None);
9765 }
9766
9767 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9769 Vec::with_capacity(self.run_refs.len());
9770 for rr in &self.run_refs {
9771 let mut reader = self.open_reader(rr.run_id)?;
9772 let (rids, eps, del) = reader.system_columns_native()?;
9773 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9774 run_meta.push((reader, rids, eps, del, page_rows));
9775 }
9776
9777 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9781 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9782 for i in 0..rids.len() {
9783 let rid = rids[i] as u64;
9784 let e = eps[i] as u64;
9785 if e > snapshot.epoch.0 {
9786 continue;
9787 }
9788 let is_del = del[i] != 0;
9789 best.entry(rid)
9790 .and_modify(|cur| {
9791 if e > cur.0 {
9792 *cur = (e, run_idx, i, is_del);
9793 }
9794 })
9795 .or_insert((e, run_idx, i, is_del));
9796 }
9797 }
9798
9799 let overlay_rids: HashSet<u64> = {
9801 let mut s = HashSet::new();
9802 for row in self.memtable.visible_versions(snapshot.epoch) {
9803 s.insert(row.row_id.0);
9804 }
9805 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9806 s.insert(row.row_id.0);
9807 }
9808 s
9809 };
9810
9811 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9813 None
9814 } else {
9815 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9816 for c in conditions {
9817 sets.push(self.resolve_condition(c, snapshot)?);
9818 }
9819 Some(RowIdSet::intersect_many(sets))
9820 };
9821
9822 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9826 for (rid, (_, run_idx, pos, deleted)) in &best {
9827 if *deleted {
9828 continue;
9829 }
9830 if overlay_rids.contains(rid) {
9831 continue;
9832 }
9833 if let Some(s) = &survivors {
9834 if !s.contains(*rid) {
9835 continue;
9836 }
9837 }
9838 per_run[*run_idx].push((*rid, *pos));
9839 }
9840 for v in per_run.iter_mut() {
9841 v.sort_unstable_by_key(|&(rid, _)| rid);
9842 }
9843
9844 let mut streams = Vec::with_capacity(run_meta.len());
9846 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9847 let mut total = 0usize;
9848 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9849 let mut starts = Vec::with_capacity(page_rows.len());
9850 let mut acc = 0usize;
9851 for &r in &page_rows {
9852 starts.push(acc);
9853 acc += r;
9854 }
9855 let mut survivors_vec: Vec<(u64, usize, usize)> =
9856 Vec::with_capacity(per_run[run_idx].len());
9857 for &(rid, pos) in &per_run[run_idx] {
9858 let page_seq = match starts.partition_point(|&s| s <= pos) {
9859 0 => continue,
9860 p => p - 1,
9861 };
9862 let within = pos - starts[page_seq];
9863 survivors_vec.push((rid, page_seq, within));
9864 }
9865 total += survivors_vec.len();
9866 if let Some(&(rid, _, _)) = survivors_vec.first() {
9867 heap.push(std::cmp::Reverse((rid, run_idx)));
9868 }
9869 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9870 }
9871
9872 let overlay_rows = if overlay_rids.is_empty() {
9874 Vec::new()
9875 } else {
9876 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9877 self.overlay_visible_rows(snapshot, bound)
9878 };
9879 let overlay = if overlay_rows.is_empty() {
9880 None
9881 } else {
9882 let filtered =
9883 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9884 if filtered.is_empty() {
9885 None
9886 } else {
9887 Some(self.materialize_overlay(&filtered, &projection))
9888 }
9889 };
9890
9891 let overlay_row_count = overlay
9892 .as_ref()
9893 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9894 .unwrap_or(0);
9895 crate::trace::QueryTrace::record(|t| {
9896 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9897 t.run_count = self.run_refs.len();
9898 t.memtable_rows = self.memtable.len();
9899 t.mutable_run_rows = self.mutable_run.len();
9900 t.overlay_rows = overlay_row_count;
9901 t.conditions_pushed = conditions.len();
9902 t.survivor_count = Some(total);
9903 });
9904
9905 Ok(Some(MultiRunCursor::new(
9906 streams, projection, heap, total, overlay,
9907 )))
9908 }
9909
9910 fn overlay_materialization_bound<'a>(
9922 conditions: &[crate::query::Condition],
9923 survivors: &'a Option<RowIdSet>,
9924 ) -> Option<&'a RowIdSet> {
9925 use crate::query::Condition;
9926 let has_range = conditions
9927 .iter()
9928 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9929 if has_range {
9930 None
9931 } else {
9932 survivors.as_ref()
9933 }
9934 }
9935
9936 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9948 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9949 let mut fold = |row: Row| {
9950 if let Some(b) = bound {
9951 if !b.contains(row.row_id.0) {
9952 return;
9953 }
9954 }
9955 best.entry(row.row_id.0)
9956 .and_modify(|(be, br)| {
9957 if row.committed_epoch > *be {
9958 *be = row.committed_epoch;
9959 *br = row.clone();
9960 }
9961 })
9962 .or_insert_with(|| (row.committed_epoch, row));
9963 };
9964 for row in self.memtable.visible_versions(snapshot.epoch) {
9965 fold(row);
9966 }
9967 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9968 fold(row);
9969 }
9970 let mut out: Vec<Row> = best
9971 .into_values()
9972 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9973 .collect();
9974 out.sort_by_key(|r| r.row_id);
9975 out
9976 }
9977
9978 fn filter_overlay_rows(
9986 &self,
9987 rows: Vec<Row>,
9988 conditions: &[crate::query::Condition],
9989 survivors: Option<&RowIdSet>,
9990 snapshot: Snapshot,
9991 ) -> Result<Vec<Row>> {
9992 if conditions.is_empty() {
9993 return Ok(rows);
9994 }
9995 use crate::query::Condition;
9996 let all_index_served = !conditions
10000 .iter()
10001 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
10002 if all_index_served {
10003 return Ok(rows
10004 .into_iter()
10005 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
10006 .collect());
10007 }
10008 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10011 for c in conditions {
10012 let s = match c {
10013 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
10014 _ => self.resolve_condition(c, snapshot)?,
10015 };
10016 per_cond_sets.push(s);
10017 }
10018 Ok(rows
10019 .into_iter()
10020 .filter(|row| {
10021 conditions.iter().enumerate().all(|(i, c)| match c {
10022 Condition::Range { column_id, lo, hi } => {
10023 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
10024 }
10025 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
10026 match row.columns.get(column_id) {
10027 Some(Value::Float64(v)) => {
10028 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
10029 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
10030 lo_ok && hi_ok
10031 }
10032 _ => false,
10033 }
10034 }
10035 _ => per_cond_sets[i].contains(row.row_id.0),
10036 })
10037 })
10038 .collect())
10039 }
10040
10041 fn materialize_overlay(
10044 &self,
10045 rows: &[Row],
10046 projection: &[(u16, TypeId)],
10047 ) -> Vec<columnar::NativeColumn> {
10048 if projection.is_empty() {
10049 return vec![columnar::null_native(TypeId::Int64, rows.len())];
10050 }
10051 let mut cols = Vec::with_capacity(projection.len());
10052 for (cid, ty) in projection {
10053 let vals: Vec<Value> = rows
10054 .iter()
10055 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
10056 .collect();
10057 cols.push(columnar::values_to_native(ty.clone(), &vals));
10058 }
10059 cols
10060 }
10061
10062 fn resolve_survivor_rids(
10067 &self,
10068 conditions: &[crate::query::Condition],
10069 reader: &mut RunReader,
10070 snapshot: Snapshot,
10071 ) -> Result<RowIdSet> {
10072 use crate::query::Condition;
10073 let mut sets: Vec<RowIdSet> = Vec::new();
10074 for c in conditions {
10075 self.validate_condition(c)?;
10076 let s: RowIdSet = match c {
10077 Condition::Pk(key) => {
10078 let lookup = self
10079 .schema
10080 .primary_key()
10081 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
10082 .unwrap_or_else(|| key.clone());
10083 self.hot
10084 .get(&lookup)
10085 .map(|r| RowIdSet::one(r.0))
10086 .unwrap_or_else(RowIdSet::empty)
10087 }
10088 Condition::BitmapEq { column_id, value } => {
10089 let lookup = self.index_lookup_key_bytes(*column_id, value);
10090 self.bitmap
10091 .get(column_id)
10092 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
10093 .unwrap_or_else(RowIdSet::empty)
10094 }
10095 Condition::BitmapIn { column_id, values } => {
10096 let bm = self.bitmap.get(column_id);
10097 let mut acc = roaring::RoaringBitmap::new();
10098 if let Some(b) = bm {
10099 for v in values {
10100 let lookup = self.index_lookup_key_bytes(*column_id, v);
10101 acc |= b.get(&lookup);
10102 }
10103 }
10104 RowIdSet::from_roaring(acc)
10105 }
10106 Condition::BytesPrefix { column_id, prefix } => {
10107 if let Some(b) = self.bitmap.get(column_id) {
10108 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
10109 let mut acc = roaring::RoaringBitmap::new();
10110 for key in b.keys() {
10111 if key.starts_with(&lookup_prefix) {
10112 acc |= b.get(&key);
10113 }
10114 }
10115 RowIdSet::from_roaring(acc)
10116 } else {
10117 RowIdSet::empty()
10118 }
10119 }
10120 Condition::FmContains { column_id, pattern } => self
10121 .fm
10122 .get(column_id)
10123 .map(|f| {
10124 RowIdSet::from_unsorted(
10125 f.locate(pattern).into_iter().map(|r| r.0).collect(),
10126 )
10127 })
10128 .unwrap_or_else(RowIdSet::empty),
10129 Condition::FmContainsAll {
10130 column_id,
10131 patterns,
10132 } => {
10133 if let Some(f) = self.fm.get(column_id) {
10134 let sets: Vec<RowIdSet> = patterns
10135 .iter()
10136 .map(|pat| {
10137 RowIdSet::from_unsorted(
10138 f.locate(pat).into_iter().map(|r| r.0).collect(),
10139 )
10140 })
10141 .collect();
10142 RowIdSet::intersect_many(sets)
10143 } else {
10144 RowIdSet::empty()
10145 }
10146 }
10147 Condition::Ann {
10148 column_id,
10149 query,
10150 k,
10151 } => RowIdSet::from_unsorted(
10152 self.retrieve_filtered(
10153 &crate::query::Retriever::Ann {
10154 column_id: *column_id,
10155 query: query.clone(),
10156 k: *k,
10157 },
10158 snapshot,
10159 None,
10160 None,
10161 None,
10162 None,
10163 )?
10164 .into_iter()
10165 .map(|hit| hit.row_id.0)
10166 .collect(),
10167 ),
10168 Condition::SparseMatch {
10169 column_id,
10170 query,
10171 k,
10172 } => RowIdSet::from_unsorted(
10173 self.retrieve_filtered(
10174 &crate::query::Retriever::Sparse {
10175 column_id: *column_id,
10176 query: query.clone(),
10177 k: *k,
10178 },
10179 snapshot,
10180 None,
10181 None,
10182 None,
10183 None,
10184 )?
10185 .into_iter()
10186 .map(|hit| hit.row_id.0)
10187 .collect(),
10188 ),
10189 Condition::MinHashSimilar {
10190 column_id,
10191 query,
10192 k,
10193 } => match self.minhash.get(column_id) {
10194 Some(index) => {
10195 let candidates = index.candidate_row_ids(query);
10196 let eligible =
10197 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
10198 RowIdSet::from_unsorted(
10199 index
10200 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
10201 .into_iter()
10202 .map(|(row_id, _)| row_id.0)
10203 .collect(),
10204 )
10205 }
10206 None => RowIdSet::empty(),
10207 },
10208 Condition::Range { column_id, lo, hi } => {
10209 if let Some(li) = self.learned_range.get(column_id) {
10210 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10211 } else {
10212 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
10213 }
10214 }
10215 Condition::RangeF64 {
10216 column_id,
10217 lo,
10218 lo_inclusive,
10219 hi,
10220 hi_inclusive,
10221 } => {
10222 if let Some(li) = self.learned_range.get(column_id) {
10223 RowIdSet::from_unsorted(
10224 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10225 .into_iter()
10226 .collect(),
10227 )
10228 } else {
10229 reader.range_row_id_set_f64(
10230 *column_id,
10231 *lo,
10232 *lo_inclusive,
10233 *hi,
10234 *hi_inclusive,
10235 )?
10236 }
10237 }
10238 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
10239 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
10240 };
10241 sets.push(s);
10242 }
10243 Ok(RowIdSet::intersect_many(sets))
10244 }
10245
10246 pub fn scan_cursor(
10267 &self,
10268 snapshot: Snapshot,
10269 projection: Vec<(u16, TypeId)>,
10270 conditions: &[crate::query::Condition],
10271 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
10272 if self.ttl.is_some() {
10273 return Ok(None);
10274 }
10275 if !conditions.is_empty() && !self.indexes_complete {
10281 return Ok(None);
10282 }
10283 if self.run_refs.len() == 1 {
10284 Ok(self
10285 .native_page_cursor(snapshot, projection, conditions)?
10286 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10287 } else {
10288 Ok(self
10289 .native_multi_run_cursor(snapshot, projection, conditions)?
10290 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10291 }
10292 }
10293
10294 pub fn aggregate_native(
10308 &self,
10309 snapshot: Snapshot,
10310 column: Option<u16>,
10311 conditions: &[crate::query::Condition],
10312 agg: NativeAgg,
10313 ) -> Result<Option<NativeAggResult>> {
10314 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
10315 }
10316
10317 pub fn aggregate_native_with_control(
10318 &self,
10319 snapshot: Snapshot,
10320 column: Option<u16>,
10321 conditions: &[crate::query::Condition],
10322 agg: NativeAgg,
10323 control: &crate::ExecutionControl,
10324 ) -> Result<Option<NativeAggResult>> {
10325 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
10326 }
10327
10328 fn aggregate_native_inner(
10329 &self,
10330 snapshot: Snapshot,
10331 column: Option<u16>,
10332 conditions: &[crate::query::Condition],
10333 agg: NativeAgg,
10334 control: Option<&crate::ExecutionControl>,
10335 ) -> Result<Option<NativeAggResult>> {
10336 execution_checkpoint(control, 0)?;
10337 if self.ttl.is_some() {
10338 return Ok(None);
10339 }
10340 if self.run_refs.len() == 1 && conditions.is_empty() {
10342 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
10343 return Ok(Some(res));
10344 }
10345 }
10346 if matches!(agg, NativeAgg::Count) && column.is_none() {
10350 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
10351 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
10352 }
10353 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10354 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
10355 }
10356 let cid = match column {
10359 Some(c) => c,
10360 None => return Ok(None),
10361 };
10362 let ty = self.column_type(cid);
10363 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
10364 execution_checkpoint(control, 0)?;
10365 return match ty {
10366 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10367 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
10368 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10369 }
10370 TypeId::Float64 => {
10371 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
10372 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10373 }
10374 _ => Ok(None),
10375 };
10376 }
10377 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10379 execution_checkpoint(control, 0)?;
10380 match ty {
10381 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10382 let mut count = 0u64;
10383 let mut sum = 0i128;
10384 let mut mn = i64::MAX;
10385 let mut mx = i64::MIN;
10386 for row in &rows {
10387 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
10388 count += 1;
10389 sum += i128::from(*v);
10390 mn = mn.min(*v);
10391 mx = mx.max(*v);
10392 }
10393 }
10394 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10395 }
10396 TypeId::Float64 => {
10397 let mut count = 0u64;
10398 let mut sum = 0.0f64;
10399 let mut mn = f64::INFINITY;
10400 let mut mx = f64::NEG_INFINITY;
10401 for row in &rows {
10402 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
10403 count += 1;
10404 sum += *v;
10405 mn = mn.min(*v);
10406 mx = mx.max(*v);
10407 }
10408 }
10409 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10410 }
10411 _ => Ok(None),
10412 }
10413 }
10414
10415 fn visible_rows_filtered(
10417 &self,
10418 snapshot: Snapshot,
10419 conditions: &[crate::query::Condition],
10420 control: Option<&crate::ExecutionControl>,
10421 ) -> Result<Vec<Row>> {
10422 let rows = if let Some(control) = control {
10423 self.visible_rows_controlled(snapshot, control)?
10424 } else {
10425 self.visible_rows(snapshot)?
10426 };
10427 if conditions.is_empty() {
10428 return Ok(rows);
10429 }
10430 Ok(rows
10431 .into_iter()
10432 .filter(|row| {
10433 conditions
10434 .iter()
10435 .all(|cond| condition_matches_row(cond, row, &self.schema))
10436 })
10437 .collect())
10438 }
10439
10440 fn aggregate_from_stats(
10448 &self,
10449 snapshot: Snapshot,
10450 column: Option<u16>,
10451 agg: NativeAgg,
10452 ) -> Result<Option<NativeAggResult>> {
10453 let cid = match (agg, column) {
10454 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
10455 _ => return Ok(None), };
10457 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
10458 return Ok(None);
10459 };
10460 let Some(cs) = stats.get(&cid) else {
10461 return Ok(None);
10462 };
10463 match agg {
10464 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
10466 self.live_count.saturating_sub(cs.null_count),
10467 ))),
10468 NativeAgg::Min | NativeAgg::Max => {
10469 let bound = if agg == NativeAgg::Min {
10470 &cs.min
10471 } else {
10472 &cs.max
10473 };
10474 match bound {
10475 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
10476 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
10477 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
10482 None => Ok(None),
10483 }
10484 }
10485 _ => Ok(None),
10486 }
10487 }
10488
10489 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
10498 if self.ttl.is_some() {
10499 return Ok(None);
10500 }
10501 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
10502 return Ok(None);
10503 }
10504 self.ensure_indexes_complete()?;
10507 let reader = self.open_reader(self.run_refs[0].run_id)?;
10508 if self.live_count != reader.row_count() as u64 {
10509 return Ok(None);
10510 }
10511 let Some(bm) = self.bitmap.get(&column_id) else {
10512 return Ok(None); };
10514 let mut distinct = bm.value_count() as u64;
10515 if !bm.get(&Value::Null.encode_key()).is_empty() {
10518 distinct = distinct.saturating_sub(1);
10519 }
10520 Ok(Some(distinct))
10521 }
10522
10523 pub fn aggregate_incremental(
10535 &mut self,
10536 cache_key: u64,
10537 conditions: &[crate::query::Condition],
10538 column: Option<u16>,
10539 agg: NativeAgg,
10540 ) -> Result<IncrementalAggResult> {
10541 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
10542 }
10543
10544 pub fn aggregate_incremental_with_control(
10545 &mut self,
10546 cache_key: u64,
10547 conditions: &[crate::query::Condition],
10548 column: Option<u16>,
10549 agg: NativeAgg,
10550 control: &crate::ExecutionControl,
10551 ) -> Result<IncrementalAggResult> {
10552 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
10553 }
10554
10555 fn aggregate_incremental_inner(
10556 &mut self,
10557 cache_key: u64,
10558 conditions: &[crate::query::Condition],
10559 column: Option<u16>,
10560 agg: NativeAgg,
10561 control: Option<&crate::ExecutionControl>,
10562 ) -> Result<IncrementalAggResult> {
10563 execution_checkpoint(control, 0)?;
10564 let snap = self.snapshot();
10565 let cur_wm = self.allocator.current().0;
10566 let cur_epoch = snap.epoch.0;
10567 let incremental_ok = self.ttl.is_none()
10574 && !self.had_deletes
10575 && self.memtable.is_empty()
10576 && self.mutable_run.is_empty();
10577
10578 if incremental_ok {
10581 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10582 if cached.epoch == cur_epoch {
10583 return Ok(IncrementalAggResult {
10584 state: cached.state,
10585 incremental: true,
10586 delta_rows: 0,
10587 });
10588 }
10589 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10590 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10591 let mut delta_rids = Vec::with_capacity(delta_len);
10592 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10593 execution_checkpoint(control, index)?;
10594 delta_rids.push(row_id);
10595 }
10596 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10597 execution_checkpoint(control, 0)?;
10598 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10599 let delta_state = agg_state_from_rows(
10600 &delta_rows,
10601 conditions,
10602 &index_sets,
10603 column,
10604 agg,
10605 &self.schema,
10606 control,
10607 )?;
10608 let merged = cached.state.merge(delta_state);
10609 let delta_n = delta_rids.len() as u64;
10610 Arc::make_mut(&mut self.agg_cache).insert(
10611 cache_key,
10612 CachedAgg {
10613 state: merged.clone(),
10614 watermark: cur_wm,
10615 epoch: cur_epoch,
10616 },
10617 );
10618 return Ok(IncrementalAggResult {
10619 state: merged,
10620 incremental: true,
10621 delta_rows: delta_n,
10622 });
10623 }
10624 }
10625 }
10626
10627 let cursor_ok =
10632 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10633 let state = if cursor_ok && agg != NativeAgg::Avg {
10634 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10635 Some(result) => {
10636 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10637 }
10638 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10639 }
10640 } else {
10641 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10642 };
10643 if incremental_ok {
10645 Arc::make_mut(&mut self.agg_cache).insert(
10646 cache_key,
10647 CachedAgg {
10648 state: state.clone(),
10649 watermark: cur_wm,
10650 epoch: cur_epoch,
10651 },
10652 );
10653 }
10654 Ok(IncrementalAggResult {
10655 state,
10656 incremental: false,
10657 delta_rows: 0,
10658 })
10659 }
10660
10661 fn agg_state_full_scan(
10664 &self,
10665 conditions: &[crate::query::Condition],
10666 column: Option<u16>,
10667 agg: NativeAgg,
10668 snap: Snapshot,
10669 control: Option<&crate::ExecutionControl>,
10670 ) -> Result<AggState> {
10671 execution_checkpoint(control, 0)?;
10672 let rows = self.visible_rows(snap)?;
10673 execution_checkpoint(control, 0)?;
10674 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10675 agg_state_from_rows(
10676 &rows,
10677 conditions,
10678 &index_sets,
10679 column,
10680 agg,
10681 &self.schema,
10682 control,
10683 )
10684 }
10685
10686 fn resolve_index_conditions(
10689 &self,
10690 conditions: &[crate::query::Condition],
10691 snapshot: Snapshot,
10692 ) -> Result<Vec<RowIdSet>> {
10693 use crate::query::Condition;
10694 let mut sets = Vec::new();
10695 for c in conditions {
10696 if matches!(
10697 c,
10698 Condition::Ann { .. }
10699 | Condition::SparseMatch { .. }
10700 | Condition::MinHashSimilar { .. }
10701 ) {
10702 sets.push(self.resolve_condition(c, snapshot)?);
10703 }
10704 }
10705 Ok(sets)
10706 }
10707
10708 fn column_type(&self, cid: u16) -> TypeId {
10709 self.schema
10710 .columns
10711 .iter()
10712 .find(|c| c.id == cid)
10713 .map(|c| c.ty.clone())
10714 .unwrap_or(TypeId::Bytes)
10715 }
10716
10717 pub fn approx_aggregate(
10726 &mut self,
10727 conditions: &[crate::query::Condition],
10728 column: Option<u16>,
10729 agg: ApproxAgg,
10730 z: f64,
10731 ) -> Result<Option<ApproxResult>> {
10732 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10733 }
10734
10735 pub fn approx_aggregate_with_candidate_authorization(
10738 &mut self,
10739 conditions: &[crate::query::Condition],
10740 column: Option<u16>,
10741 agg: ApproxAgg,
10742 z: f64,
10743 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10744 ) -> Result<Option<ApproxResult>> {
10745 use crate::query::Condition;
10746 self.ensure_reservoir_complete()?;
10747 let mut snapshot = self.snapshot();
10752 let n_pop = self.count();
10753 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10754 if sample_rids.is_empty() {
10755 return Ok(None);
10756 }
10757 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10759 if live_sample.is_empty() {
10760 snapshot = Snapshot::unbounded();
10761 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10762 }
10763 let s = live_sample.len();
10764 if s == 0 {
10765 return Ok(None);
10766 }
10767 let authorized = authorization
10768 .map(|authorization| {
10769 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10770 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10771 })
10772 .transpose()?;
10773
10774 let mut index_sets: Vec<RowIdSet> = Vec::new();
10777 for c in conditions {
10778 if matches!(
10779 c,
10780 Condition::Ann { .. }
10781 | Condition::SparseMatch { .. }
10782 | Condition::MinHashSimilar { .. }
10783 ) {
10784 index_sets.push(self.resolve_condition(c, snapshot)?);
10785 }
10786 }
10787
10788 let cid = match (agg, column) {
10790 (ApproxAgg::Count, _) => None,
10791 (_, Some(c)) => Some(c),
10792 _ => return Ok(None),
10793 };
10794 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10795 for r in &live_sample {
10796 if authorized
10797 .as_ref()
10798 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10799 {
10800 continue;
10801 }
10802 if !conditions
10804 .iter()
10805 .all(|c| condition_matches_row(c, r, &self.schema))
10806 {
10807 continue;
10808 }
10809 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10811 continue;
10812 }
10813 if let Some(cid) = cid {
10814 let mut cells = r
10815 .columns
10816 .get(&cid)
10817 .cloned()
10818 .map(|value| vec![(cid, value)])
10819 .unwrap_or_default();
10820 if let Some(authorization) = authorization {
10821 authorization.security.apply_masks_to_cells(
10822 authorization.table,
10823 &mut cells,
10824 authorization.principal,
10825 );
10826 }
10827 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10828 passing_vals.push(v);
10829 } } else {
10831 passing_vals.push(0.0); }
10833 }
10834 let m = passing_vals.len();
10835
10836 let (point, half) = match agg {
10837 ApproxAgg::Count => {
10838 let p = m as f64 / s as f64;
10840 let point = n_pop as f64 * p;
10841 let var = if s > 1 {
10842 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10843 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10844 } else {
10845 0.0
10846 };
10847 (point, z * var.sqrt())
10848 }
10849 ApproxAgg::Sum => {
10850 let y: Vec<f64> = live_sample
10852 .iter()
10853 .map(|r| {
10854 let passes_row = authorized
10855 .as_ref()
10856 .is_none_or(|authorized| authorized.contains(&r.row_id))
10857 && conditions
10858 .iter()
10859 .all(|c| condition_matches_row(c, r, &self.schema))
10860 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10861 if passes_row {
10862 cid.and_then(|cid| {
10863 let mut cells = r
10864 .columns
10865 .get(&cid)
10866 .cloned()
10867 .map(|value| vec![(cid, value)])
10868 .unwrap_or_default();
10869 if let Some(authorization) = authorization {
10870 authorization.security.apply_masks_to_cells(
10871 authorization.table,
10872 &mut cells,
10873 authorization.principal,
10874 );
10875 }
10876 as_f64(cells.first().map(|(_, value)| value))
10877 })
10878 .unwrap_or(0.0)
10879 } else {
10880 0.0
10881 }
10882 })
10883 .collect();
10884 let mean_y = y.iter().sum::<f64>() / s as f64;
10885 let point = n_pop as f64 * mean_y;
10886 let var = if s > 1 {
10887 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10888 let var_y = ss / (s - 1) as f64;
10889 n_pop as f64 * n_pop as f64 * var_y / s as f64
10890 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10891 } else {
10892 0.0
10893 };
10894 (point, z * var.sqrt())
10895 }
10896 ApproxAgg::Avg => {
10897 if m == 0 {
10898 return Ok(Some(ApproxResult {
10899 point: 0.0,
10900 ci_low: 0.0,
10901 ci_high: 0.0,
10902 n_population: n_pop,
10903 n_sample_live: s,
10904 n_passing: 0,
10905 }));
10906 }
10907 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10908 let half = if m > 1 {
10909 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10910 let sd = (ss / (m - 1) as f64).sqrt();
10911 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10912 z * sd / (m as f64).sqrt() * fpc.sqrt()
10913 } else {
10914 0.0
10915 };
10916 (mean, half)
10917 }
10918 };
10919
10920 Ok(Some(ApproxResult {
10921 point,
10922 ci_low: point - half,
10923 ci_high: point + half,
10924 n_population: n_pop,
10925 n_sample_live: s,
10926 n_passing: m,
10927 }))
10928 }
10929
10930 pub fn exact_column_stats(
10938 &self,
10939 _snapshot: Snapshot,
10940 projection: &[u16],
10941 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10942 if self.ttl.is_some()
10943 || !(self.memtable.is_empty()
10944 && self.mutable_run.is_empty()
10945 && self.run_refs.len() == 1)
10946 {
10947 return Ok(None);
10948 }
10949 let reader = self.open_reader(self.run_refs[0].run_id)?;
10950 if self.live_count != reader.row_count() as u64 {
10951 return Ok(None);
10952 }
10953 let mut out = HashMap::new();
10954 for &cid in projection {
10955 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10956 Some(c) => c,
10957 None => continue,
10958 };
10959 let Some(stats) = reader.column_page_stats(cid) else {
10961 out.insert(
10962 cid,
10963 ColumnStat {
10964 min: None,
10965 max: None,
10966 null_count: self.live_count,
10967 },
10968 );
10969 continue;
10970 };
10971 let stat = match cdef.ty {
10972 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10973 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10974 min: mn.map(Value::Int64),
10975 max: mx.map(Value::Int64),
10976 null_count: n,
10977 })
10978 }
10979 TypeId::Float64 => {
10980 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10981 min: mn.map(Value::Float64),
10982 max: mx.map(Value::Float64),
10983 null_count: n,
10984 })
10985 }
10986 _ => None,
10987 };
10988 if let Some(s) = stat {
10989 out.insert(cid, s);
10990 }
10991 }
10992 Ok(Some(out))
10993 }
10994
10995 pub fn dir(&self) -> &Path {
10996 &self.dir
10997 }
10998
10999 pub fn schema(&self) -> &Schema {
11000 &self.schema
11001 }
11002
11003 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
11004 self.ann.get(&column_id)
11005 }
11006
11007 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
11008 self.ann.get_mut(&column_id)
11009 }
11010
11011 pub(crate) fn set_catalog_name(&mut self, name: String) {
11012 self.name = name;
11013 }
11014
11015 pub(crate) fn prepare_alter_column(
11016 &mut self,
11017 column_name: &str,
11018 change: &AlterColumn,
11019 ) -> Result<(ColumnDef, Option<Schema>)> {
11020 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
11021 return Err(MongrelError::InvalidArgument(
11022 "ALTER COLUMN requires committing staged writes first".into(),
11023 ));
11024 }
11025 let old = self
11026 .schema
11027 .columns
11028 .iter()
11029 .find(|c| c.name == column_name)
11030 .cloned()
11031 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11032 let mut next = old.clone();
11033
11034 if let Some(name) = &change.name {
11035 let trimmed = name.trim();
11036 if trimmed.is_empty() {
11037 return Err(MongrelError::InvalidArgument(
11038 "ALTER COLUMN name must not be empty".into(),
11039 ));
11040 }
11041 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
11042 return Err(MongrelError::Schema(format!(
11043 "column {trimmed} already exists"
11044 )));
11045 }
11046 next.name = trimmed.to_string();
11047 }
11048
11049 if let Some(ty) = &change.ty {
11050 next.ty = ty.clone();
11051 }
11052 if let Some(flags) = change.flags {
11053 validate_alter_column_flags(old.flags, flags)?;
11054 next.flags = flags;
11055 }
11056
11057 if let Some(default_change) = &change.default_value {
11058 next.default_value = default_change.clone();
11059 }
11060 if let Some(source_change) = &change.embedding_source {
11061 next.embedding_source = source_change.clone();
11062 }
11063
11064 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
11065 if old.flags.contains(ColumnFlags::NULLABLE)
11066 && !next.flags.contains(ColumnFlags::NULLABLE)
11067 && self.column_has_nulls(old.id)?
11068 {
11069 return Err(MongrelError::InvalidArgument(format!(
11070 "column '{}' contains NULL values",
11071 old.name
11072 )));
11073 }
11074 if next == old {
11075 return Ok((next, None));
11076 }
11077 let mut schema = self.schema.clone();
11078 let index = schema
11079 .columns
11080 .iter()
11081 .position(|column| column.id == next.id)
11082 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
11083 schema.columns[index] = next.clone();
11084 schema.schema_id = schema
11085 .schema_id
11086 .checked_add(1)
11087 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11088 schema.validate_auto_increment()?;
11089 schema.validate_defaults()?;
11090 Ok((next, Some(schema)))
11091 }
11092
11093 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
11094 self.schema = schema;
11095 self.auto_inc = resolve_auto_inc(&self.schema);
11096 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
11097 self.clear_result_cache();
11098 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
11099 }
11100
11101 pub(crate) fn publish_index_schema_change(
11105 &mut self,
11106 schema: Schema,
11107 artifact: SecondaryIndexArtifact,
11108 ) -> Result<()> {
11109 self.apply_altered_schema_prepared(schema);
11110 self.prune_index_maps_to_schema();
11111 match artifact {
11112 SecondaryIndexArtifact::Bitmap(column_id, index) => {
11113 self.bitmap.insert(column_id, index);
11114 }
11115 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
11116 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
11117 }
11118 SecondaryIndexArtifact::Fm(column_id, index) => {
11119 self.fm.insert(column_id, *index);
11120 }
11121 SecondaryIndexArtifact::Ann(column_id, index) => {
11122 self.ann.insert(column_id, *index);
11123 }
11124 SecondaryIndexArtifact::Sparse(column_id, index) => {
11125 self.sparse.insert(column_id, index);
11126 }
11127 SecondaryIndexArtifact::MinHash(column_id, index) => {
11128 self.minhash.insert(column_id, index);
11129 }
11130 }
11131 self.indexes_complete = true;
11132 self.seal_generations();
11133 let view = Arc::new(self.capture_read_generation());
11134 self.published.store(Arc::clone(&view));
11135 checkpoint_current_schema(self)?;
11136 self.invalidate_index_checkpoint();
11139 Ok(())
11140 }
11141
11142 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
11148 self.apply_altered_schema_prepared(schema);
11149 self.prune_index_maps_to_schema();
11150 self.indexes_complete = true;
11151 self.seal_generations();
11152 let view = Arc::new(self.capture_read_generation());
11153 self.published.store(Arc::clone(&view));
11154 checkpoint_current_schema(self)?;
11155 self.invalidate_index_checkpoint();
11157 Ok(())
11158 }
11159
11160 fn prune_index_maps_to_schema(&mut self) {
11161 let keep_bitmap: std::collections::HashSet<u16> = self
11162 .schema
11163 .indexes
11164 .iter()
11165 .filter(|index| index.kind == IndexKind::Bitmap)
11166 .map(|index| index.column_id)
11167 .collect();
11168 let keep_ann: std::collections::HashSet<u16> = self
11169 .schema
11170 .indexes
11171 .iter()
11172 .filter(|index| index.kind == IndexKind::Ann)
11173 .map(|index| index.column_id)
11174 .collect();
11175 let keep_fm: std::collections::HashSet<u16> = self
11176 .schema
11177 .indexes
11178 .iter()
11179 .filter(|index| index.kind == IndexKind::FmIndex)
11180 .map(|index| index.column_id)
11181 .collect();
11182 let keep_sparse: std::collections::HashSet<u16> = self
11183 .schema
11184 .indexes
11185 .iter()
11186 .filter(|index| index.kind == IndexKind::Sparse)
11187 .map(|index| index.column_id)
11188 .collect();
11189 let keep_minhash: std::collections::HashSet<u16> = self
11190 .schema
11191 .indexes
11192 .iter()
11193 .filter(|index| index.kind == IndexKind::MinHash)
11194 .map(|index| index.column_id)
11195 .collect();
11196 let keep_learned: std::collections::HashSet<u16> = self
11197 .schema
11198 .indexes
11199 .iter()
11200 .filter(|index| index.kind == IndexKind::LearnedRange)
11201 .map(|index| index.column_id)
11202 .collect();
11203 self.bitmap
11204 .retain(|column_id, _| keep_bitmap.contains(column_id));
11205 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
11206 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
11207 self.sparse
11208 .retain(|column_id, _| keep_sparse.contains(column_id));
11209 self.minhash
11210 .retain(|column_id, _| keep_minhash.contains(column_id));
11211 {
11212 let learned = Arc::make_mut(&mut self.learned_range);
11213 learned.retain(|column_id, _| keep_learned.contains(column_id));
11214 }
11215 }
11216
11217 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
11218 checkpoint_current_schema(self)
11219 }
11220
11221 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
11222 self.ensure_writable()?;
11223 let previous_schema = self.schema.clone();
11224 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
11225 if let Some(schema) = schema {
11226 self.apply_altered_schema_prepared(schema);
11227 self.checkpoint_standalone_schema_change(previous_schema)?;
11228 }
11229 Ok(column)
11230 }
11231
11232 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
11233 if self.live_count == 0 {
11234 return Ok(false);
11235 }
11236 let snap = self.snapshot();
11237 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
11238 Ok(columns
11239 .first()
11240 .map(|(_, col)| col.null_count(col.len()) != 0)
11241 .unwrap_or(true))
11242 }
11243
11244 fn has_stored_versions(&self) -> bool {
11245 !self.memtable.is_empty()
11246 || !self.mutable_run.is_empty()
11247 || self.run_refs.iter().any(|r| r.row_count > 0)
11248 || !self.retiring.is_empty()
11249 }
11250
11251 pub fn add_column(
11256 &mut self,
11257 name: &str,
11258 ty: TypeId,
11259 flags: ColumnFlags,
11260 default_value: Option<crate::schema::DefaultExpr>,
11261 ) -> Result<u16> {
11262 self.add_column_with_id(name, ty, flags, default_value, None)
11263 }
11264
11265 pub fn add_column_with_id(
11266 &mut self,
11267 name: &str,
11268 ty: TypeId,
11269 flags: ColumnFlags,
11270 default_value: Option<crate::schema::DefaultExpr>,
11271 requested_id: Option<u16>,
11272 ) -> Result<u16> {
11273 self.ensure_writable()?;
11274 let previous_schema = self.schema.clone();
11275 let (column, schema) =
11276 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
11277 self.apply_altered_schema_prepared(schema);
11278 self.checkpoint_standalone_schema_change(previous_schema)?;
11279 Ok(column.id)
11280 }
11281
11282 pub(crate) fn prepare_add_column(
11283 &mut self,
11284 name: &str,
11285 ty: TypeId,
11286 flags: ColumnFlags,
11287 default_value: Option<crate::schema::DefaultExpr>,
11288 requested_id: Option<u16>,
11289 ) -> Result<(ColumnDef, Schema)> {
11290 self.ensure_writable()?;
11291 if self.schema.columns.iter().any(|c| c.name == name) {
11292 return Err(MongrelError::Schema(format!(
11293 "column {name} already exists"
11294 )));
11295 }
11296 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
11297 if self.schema.columns.iter().any(|c| c.id == id) {
11298 return Err(MongrelError::Schema(format!(
11299 "column id {id} already exists"
11300 )));
11301 }
11302 id
11303 } else {
11304 self.schema
11305 .columns
11306 .iter()
11307 .map(|c| c.id)
11308 .max()
11309 .unwrap_or(0)
11310 .checked_add(1)
11311 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
11312 };
11313 let column = ColumnDef {
11314 id,
11315 name: name.to_string(),
11316 ty,
11317 flags,
11318 default_value,
11319 embedding_source: None,
11320 };
11321 let mut next_schema = self.schema.clone();
11322 next_schema.columns.push(column.clone());
11323 next_schema.schema_id = next_schema
11324 .schema_id
11325 .checked_add(1)
11326 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11327 next_schema.validate_auto_increment()?;
11328 next_schema.validate_defaults()?;
11329 Ok((column, next_schema))
11330 }
11331
11332 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
11341 self.ensure_writable()?;
11342 let cid = self
11343 .schema
11344 .columns
11345 .iter()
11346 .find(|c| c.name == column_name)
11347 .map(|c| c.id)
11348 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11349 let ty = self
11350 .schema
11351 .columns
11352 .iter()
11353 .find(|c| c.id == cid)
11354 .map(|c| c.ty.clone())
11355 .unwrap_or(TypeId::Int64);
11356 if !matches!(
11357 ty,
11358 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
11359 ) {
11360 return Err(MongrelError::Schema(format!(
11361 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
11362 )));
11363 }
11364 if self
11365 .schema
11366 .indexes
11367 .iter()
11368 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
11369 {
11370 return Ok(()); }
11372 let previous_schema = self.schema.clone();
11373 let previous_learned_range = Arc::clone(&self.learned_range);
11374 let mut next_schema = previous_schema.clone();
11375 next_schema.indexes.push(IndexDef {
11376 name: format!("{}_learned_range", column_name),
11377 column_id: cid,
11378 kind: IndexKind::LearnedRange,
11379 predicate: None,
11380 options: Default::default(),
11381 });
11382 next_schema.schema_id = next_schema
11383 .schema_id
11384 .checked_add(1)
11385 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11386 self.apply_altered_schema_prepared(next_schema);
11387 if let Err(error) = self.build_learned_ranges() {
11388 self.apply_altered_schema_prepared(previous_schema);
11389 self.learned_range = previous_learned_range;
11390 return Err(error);
11391 }
11392 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
11393 if !matches!(
11394 &error,
11395 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
11396 ) {
11397 self.learned_range = previous_learned_range;
11398 }
11399 return Err(error);
11400 }
11401 Ok(())
11402 }
11403
11404 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
11405 let mut schema_published = false;
11406 let schema_result = match self._root_guard.as_deref() {
11407 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
11408 schema_published = true;
11409 }),
11410 None => write_schema_with_after(&self.dir, &self.schema, || {
11411 schema_published = true;
11412 }),
11413 };
11414 if schema_result.is_err() && !schema_published {
11415 self.apply_altered_schema_prepared(previous_schema);
11416 return schema_result;
11417 }
11418
11419 let manifest_result = self.persist_manifest(self.current_epoch());
11420 match (schema_result, manifest_result) {
11421 (_, Ok(())) => Ok(()),
11422 (Ok(()), Err(error)) => {
11423 self.poison_after_maintenance_publish_failure();
11424 Err(MongrelError::DurableCommit {
11425 epoch: self.current_epoch().0,
11426 message: format!(
11427 "schema is durable but matching manifest publication failed: {error}"
11428 ),
11429 })
11430 }
11431 (Err(schema_error), Err(manifest_error)) => {
11432 self.poison_after_maintenance_publish_failure();
11433 Err(MongrelError::CommitOutcomeUnknown {
11434 epoch: self.current_epoch().0,
11435 message: format!(
11436 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11437 ),
11438 })
11439 }
11440 }
11441 }
11442
11443 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
11446 self.sync_byte_threshold = threshold;
11447 if let WalSink::Private(w) = &mut self.wal {
11448 w.set_sync_byte_threshold(threshold);
11449 }
11450 }
11451
11452 pub fn page_cache_flush(&self) {
11456 self.page_cache.flush_to_disk();
11457 }
11458
11459 pub fn page_cache_len(&self) -> usize {
11461 self.page_cache.len()
11462 }
11463
11464 pub fn decoded_cache_len(&self) -> usize {
11467 self.decoded_cache.len()
11468 }
11469
11470 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
11473 self.memtable.drain_sorted()
11474 }
11475
11476 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
11477 self.runs_dir().join(format!("r-{run_id}.sr"))
11478 }
11479
11480 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
11481 match self.runs_root.as_deref() {
11482 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
11483 None => Ok(None),
11484 }
11485 }
11486
11487 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
11488 match self.runs_root.as_deref() {
11489 Some(root) => Ok(Some(root.create_regular_new(name)?)),
11490 None => Ok(None),
11491 }
11492 }
11493
11494 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
11495 match self.runs_root.as_deref() {
11496 Some(root) => match root.remove_file(name) {
11497 Ok(()) => Ok(()),
11498 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11499 Err(error) => Err(error.into()),
11500 },
11501 None => match std::fs::remove_file(self.runs_dir().join(name)) {
11502 Ok(()) => Ok(()),
11503 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11504 Err(error) => Err(error.into()),
11505 },
11506 }
11507 }
11508
11509 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
11510 match self.runs_root.as_deref() {
11511 Some(root) => root
11512 .rename_file_new(source, destination)
11513 .map_err(Into::into),
11514 None => crate::durable_file::rename(
11515 &self.runs_dir().join(source),
11516 &self.runs_dir().join(destination),
11517 )
11518 .map_err(Into::into),
11519 }
11520 }
11521
11522 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11523 self.run_refs.iter().map(|run| run.run_id)
11524 }
11525
11526 pub(crate) fn table_dir(&self) -> &Path {
11527 &self.dir
11528 }
11529
11530 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
11531 &self.schema
11532 }
11533
11534 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
11535 let id = self.next_run_id;
11536 self.next_run_id = self
11537 .next_run_id
11538 .checked_add(1)
11539 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
11540 Ok(id)
11541 }
11542
11543 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
11544 self.run_refs.push(run_ref);
11545 }
11546
11547 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
11557 self.retiring.push(crate::manifest::RetiredRun {
11558 run_id,
11559 retire_epoch,
11560 });
11561 }
11562
11563 pub(crate) fn reap_retiring(
11567 &mut self,
11568 min_active: Epoch,
11569 backup_pinned: &std::collections::HashSet<u128>,
11570 ) -> Result<usize> {
11571 if self.retiring.is_empty() {
11572 return Ok(0);
11573 }
11574 let mut reaped = 0;
11575 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11576 for r in std::mem::take(&mut self.retiring) {
11582 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11583 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11584 reaped += 1;
11585 } else {
11586 kept.push(r);
11587 }
11588 }
11589 self.retiring = kept;
11590 if reaped > 0 {
11591 self.persist_manifest(self.current_epoch())?;
11592 }
11593 Ok(reaped)
11594 }
11595
11596 pub(crate) fn has_reapable_retiring(
11597 &self,
11598 min_active: Epoch,
11599 backup_pinned: &std::collections::HashSet<u128>,
11600 ) -> bool {
11601 self.retiring
11602 .iter()
11603 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11604 }
11605
11606 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11607 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11608 return false;
11609 }
11610 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11611 self.run_refs.push(run_ref);
11612 self.indexes_complete = false;
11613 true
11614 }
11615
11616 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11617 self.kek.as_ref()
11618 }
11619
11620 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11621 let mut reader = match self.runs_root.as_deref() {
11622 Some(root) => RunReader::open_file_with_cache(
11623 root.open_regular(format!("r-{run_id}.sr"))?,
11624 self.schema.clone(),
11625 self.kek.clone(),
11626 Some(self.page_cache.clone()),
11627 Some(self.decoded_cache.clone()),
11628 self.table_id,
11629 Some(&self.verified_runs),
11630 None,
11631 )?,
11632 None => RunReader::open_with_cache(
11633 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11634 self.schema.clone(),
11635 self.kek.clone(),
11636 Some(self.page_cache.clone()),
11637 Some(self.decoded_cache.clone()),
11638 self.table_id,
11639 Some(&self.verified_runs),
11640 )?,
11641 };
11642 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11646 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11647 }
11648 Ok(reader)
11649 }
11650
11651 pub(crate) fn run_refs(&self) -> &[RunRef] {
11652 &self.run_refs
11653 }
11654
11655 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11656 self.retiring.iter().map(|run| run.run_id)
11657 }
11658
11659 pub(crate) fn runs_dir(&self) -> PathBuf {
11660 self.runs_root
11661 .as_deref()
11662 .and_then(|root| root.io_path().ok())
11663 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11664 }
11665
11666 pub(crate) fn wal_dir(&self) -> PathBuf {
11667 self.dir.join(WAL_DIR)
11668 }
11669
11670 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11671 self.run_refs = refs;
11672 }
11673
11674 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11675 self.compaction_zstd_level
11676 }
11677
11678 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11679 self.kek.clone()
11680 }
11681
11682 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11686 self.kek.as_ref().map(|k| k.derive_idx_key())
11687 }
11688
11689 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11693 self.kek.as_ref().map(|k| *k.derive_meta_key())
11694 }
11695
11696 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11699 self.column_keys
11700 .iter()
11701 .map(|(&id, &(_, scheme))| (id, scheme))
11702 .collect()
11703 }
11704
11705 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11710 self.tokenize_value_enc(column_id, v)
11711 }
11712
11713 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11714 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11715 let (key, scheme) = self.column_keys.get(&column_id)?;
11716 let token: Vec<u8> = match (*scheme, v) {
11717 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11718 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11719 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11720 _ => hmac_token(key, &v.encode_key()).to_vec(),
11721 };
11722 Some(Value::Bytes(token))
11723 }
11724
11725 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11727 self.index_lookup_key_bytes(column_id, &v.encode_key())
11728 }
11729
11730 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11733 {
11734 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11735 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11736 if *scheme == SCHEME_HMAC_EQ {
11737 return hmac_token(key, encoded).to_vec();
11738 }
11739 }
11740 }
11741 let _ = column_id;
11742 encoded.to_vec()
11743 }
11744}
11745
11746fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11747 let columnar::NativeColumn::Int64 { data, validity } = col else {
11748 return false;
11749 };
11750 if data.len() < n || !columnar::all_non_null(validity, n) {
11751 return false;
11752 }
11753 data.iter()
11754 .take(n)
11755 .zip(data.iter().skip(1))
11756 .all(|(a, b)| a < b)
11757}
11758
11759#[derive(Debug, Clone)]
11763pub struct ColumnStat {
11764 pub min: Option<Value>,
11765 pub max: Option<Value>,
11766 pub null_count: u64,
11767}
11768
11769#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11771pub enum NativeAgg {
11772 Count,
11773 Sum,
11774 Min,
11775 Max,
11776 Avg,
11777}
11778
11779#[derive(Debug, Clone, PartialEq)]
11781pub enum NativeAggResult {
11782 Count(u64),
11783 Int(i64),
11784 Float(f64),
11785 Null,
11787}
11788
11789#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11791pub enum ApproxAgg {
11792 Count,
11793 Sum,
11794 Avg,
11795}
11796
11797#[derive(Debug, Clone)]
11801pub struct ApproxResult {
11802 pub point: f64,
11804 pub ci_low: f64,
11806 pub ci_high: f64,
11808 pub n_population: u64,
11810 pub n_sample_live: usize,
11812 pub n_passing: usize,
11814}
11815
11816#[derive(Debug, Clone, PartialEq)]
11821pub enum AggState {
11822 Count(u64),
11824 SumI {
11826 sum: i128,
11827 count: u64,
11828 },
11829 SumF {
11831 sum: f64,
11832 count: u64,
11833 },
11834 AvgI {
11836 sum: i128,
11837 count: u64,
11838 },
11839 AvgF {
11841 sum: f64,
11842 count: u64,
11843 },
11844 MinI(i64),
11846 MaxI(i64),
11847 MinF(f64),
11849 MaxF(f64),
11850 Empty,
11852}
11853
11854impl AggState {
11855 pub fn merge(self, other: AggState) -> AggState {
11857 use AggState::*;
11858 match (self, other) {
11859 (Empty, x) | (x, Empty) => x,
11860 (Count(a), Count(b)) => Count(a + b),
11861 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11862 sum: sa + sb,
11863 count: ca + cb,
11864 },
11865 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11866 sum: sa + sb,
11867 count: ca + cb,
11868 },
11869 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11870 sum: sa + sb,
11871 count: ca + cb,
11872 },
11873 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11874 sum: sa + sb,
11875 count: ca + cb,
11876 },
11877 (MinI(a), MinI(b)) => MinI(a.min(b)),
11878 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11879 (MinF(a), MinF(b)) => MinF(a.min(b)),
11880 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11881 _ => Empty, }
11883 }
11884
11885 pub fn point(&self) -> Option<f64> {
11887 match self {
11888 AggState::Count(n) => Some(*n as f64),
11889 AggState::SumI { sum, .. } => Some(*sum as f64),
11890 AggState::SumF { sum, .. } => Some(*sum),
11891 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11892 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11893 AggState::MinI(n) => Some(*n as f64),
11894 AggState::MaxI(n) => Some(*n as f64),
11895 AggState::MinF(n) => Some(*n),
11896 AggState::MaxF(n) => Some(*n),
11897 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11898 }
11899 }
11900
11901 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11905 let is_float = matches!(ty, Some(TypeId::Float64));
11906 match (agg, result) {
11907 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11908 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11909 sum: x as i128,
11910 count: 1, },
11912 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11913 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11914 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11915 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11916 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11917 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11918 (NativeAgg::Count, _) => AggState::Empty,
11919 (_, NativeAggResult::Null) => AggState::Empty,
11920 _ => {
11921 let _ = is_float;
11922 AggState::Empty
11923 }
11924 }
11925 }
11926}
11927
11928#[derive(Debug, Clone)]
11931pub struct CachedAgg {
11932 pub state: AggState,
11933 pub watermark: u64,
11934 pub epoch: u64,
11935}
11936
11937#[derive(Debug, Clone)]
11939pub struct IncrementalAggResult {
11940 pub state: AggState,
11942 pub incremental: bool,
11945 pub delta_rows: u64,
11947}
11948
11949fn agg_state_from_rows(
11953 rows: &[Row],
11954 conditions: &[crate::query::Condition],
11955 index_sets: &[RowIdSet],
11956 column: Option<u16>,
11957 agg: NativeAgg,
11958 schema: &Schema,
11959 control: Option<&crate::ExecutionControl>,
11960) -> Result<AggState> {
11961 let mut count: u64 = 0;
11962 let mut sum_i: i128 = 0;
11963 let mut sum_f: f64 = 0.0;
11964 let mut mn_i: i64 = i64::MAX;
11965 let mut mx_i: i64 = i64::MIN;
11966 let mut mn_f: f64 = f64::INFINITY;
11967 let mut mx_f: f64 = f64::NEG_INFINITY;
11968 let mut saw_int = false;
11969 let mut saw_float = false;
11970 for (index, r) in rows.iter().enumerate() {
11971 execution_checkpoint(control, index)?;
11972 if !conditions
11973 .iter()
11974 .all(|c| condition_matches_row(c, r, schema))
11975 {
11976 continue;
11977 }
11978 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11979 continue;
11980 }
11981 match agg {
11982 NativeAgg::Count => match column {
11983 None => count += 1,
11985 Some(cid) => match r.columns.get(&cid) {
11988 None | Some(Value::Null) => {}
11989 Some(_) => count += 1,
11990 },
11991 },
11992 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11993 Some(Value::Int64(n)) => {
11994 count += 1;
11995 sum_i += *n as i128;
11996 mn_i = mn_i.min(*n);
11997 mx_i = mx_i.max(*n);
11998 saw_int = true;
11999 }
12000 Some(Value::Float64(f)) => {
12001 count += 1;
12002 sum_f += f;
12003 mn_f = mn_f.min(*f);
12004 mx_f = mx_f.max(*f);
12005 saw_float = true;
12006 }
12007 _ => {}
12008 },
12009 }
12010 }
12011 Ok(match agg {
12012 NativeAgg::Count => {
12013 if count == 0 {
12014 AggState::Empty
12015 } else {
12016 AggState::Count(count)
12017 }
12018 }
12019 NativeAgg::Sum => {
12020 if count == 0 {
12021 AggState::Empty
12022 } else if saw_int {
12023 AggState::SumI { sum: sum_i, count }
12024 } else {
12025 AggState::SumF { sum: sum_f, count }
12026 }
12027 }
12028 NativeAgg::Avg => {
12029 if count == 0 {
12030 AggState::Empty
12031 } else if saw_int {
12032 AggState::AvgI { sum: sum_i, count }
12033 } else {
12034 AggState::AvgF { sum: sum_f, count }
12035 }
12036 }
12037 NativeAgg::Min => {
12038 if !saw_int && !saw_float {
12039 AggState::Empty
12040 } else if saw_int {
12041 AggState::MinI(mn_i)
12042 } else {
12043 AggState::MinF(mn_f)
12044 }
12045 }
12046 NativeAgg::Max => {
12047 if !saw_int && !saw_float {
12048 AggState::Empty
12049 } else if saw_int {
12050 AggState::MaxI(mx_i)
12051 } else {
12052 AggState::MaxF(mx_f)
12053 }
12054 }
12055 })
12056}
12057
12058fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
12062 use crate::query::Condition;
12063 match c {
12064 Condition::Pk(key) => match schema.primary_key() {
12065 Some(pk) => row
12066 .columns
12067 .get(&pk.id)
12068 .map(|v| v.encode_key() == *key)
12069 .unwrap_or(false),
12070 None => false,
12071 },
12072 Condition::BitmapEq { column_id, value } => row
12073 .columns
12074 .get(column_id)
12075 .map(|v| v.encode_key() == *value)
12076 .unwrap_or(false),
12077 Condition::BitmapIn { column_id, values } => {
12078 let key = row.columns.get(column_id).map(|v| v.encode_key());
12079 match key {
12080 Some(k) => values.contains(&k),
12081 None => false,
12082 }
12083 }
12084 Condition::BytesPrefix { column_id, prefix } => row
12085 .columns
12086 .get(column_id)
12087 .map(|v| v.encode_key().starts_with(prefix))
12088 .unwrap_or(false),
12089 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
12090 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
12091 _ => false,
12092 },
12093 Condition::RangeF64 {
12094 column_id,
12095 lo,
12096 lo_inclusive,
12097 hi,
12098 hi_inclusive,
12099 } => match row.columns.get(column_id) {
12100 Some(Value::Float64(n)) => {
12101 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
12102 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
12103 lo_ok && hi_ok
12104 }
12105 _ => false,
12106 },
12107 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
12108 Some(Value::Bytes(b)) => {
12109 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
12110 }
12111 _ => false,
12112 },
12113 Condition::FmContainsAll {
12114 column_id,
12115 patterns,
12116 } => match row.columns.get(column_id) {
12117 Some(Value::Bytes(b)) => patterns
12118 .iter()
12119 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
12120 _ => false,
12121 },
12122 Condition::Ann { .. }
12123 | Condition::SparseMatch { .. }
12124 | Condition::MinHashSimilar { .. } => true,
12125 Condition::IsNull { column_id } => {
12126 matches!(row.columns.get(column_id), Some(Value::Null) | None)
12127 }
12128 Condition::IsNotNull { column_id } => {
12129 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
12130 }
12131 }
12132}
12133
12134fn as_f64(v: Option<&Value>) -> Option<f64> {
12136 match v {
12137 Some(Value::Int64(n)) => Some(*n as f64),
12138 Some(Value::Float64(f)) => Some(*f),
12139 _ => None,
12140 }
12141}
12142
12143fn accumulate_int(
12147 cursor: &mut dyn crate::cursor::Cursor,
12148 control: Option<&crate::ExecutionControl>,
12149) -> Result<(u64, i128, i64, i64)> {
12150 let mut count: u64 = 0;
12151 let mut sum: i128 = 0;
12152 let mut mn: i64 = i64::MAX;
12153 let mut mx: i64 = i64::MIN;
12154 while let Some(cols) = cursor.next_batch()? {
12155 execution_checkpoint(control, 0)?;
12156 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
12157 if crate::columnar::all_non_null(validity, data.len()) {
12158 count += data.len() as u64;
12160 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12161 execution_checkpoint(control, chunk_index * 1024)?;
12162 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
12163 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
12164 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
12165 }
12166 } else {
12167 for (i, &v) in data.iter().enumerate() {
12168 execution_checkpoint(control, i)?;
12169 if crate::columnar::validity_bit(validity, i) {
12170 count += 1;
12171 sum += v as i128;
12172 mn = mn.min(v);
12173 mx = mx.max(v);
12174 }
12175 }
12176 }
12177 }
12178 }
12179 Ok((count, sum, mn, mx))
12180}
12181
12182fn accumulate_float(
12184 cursor: &mut dyn crate::cursor::Cursor,
12185 control: Option<&crate::ExecutionControl>,
12186) -> Result<(u64, f64, f64, f64)> {
12187 let mut count: u64 = 0;
12188 let mut sum: f64 = 0.0;
12189 let mut mn: f64 = f64::INFINITY;
12190 let mut mx: f64 = f64::NEG_INFINITY;
12191 while let Some(cols) = cursor.next_batch()? {
12192 execution_checkpoint(control, 0)?;
12193 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
12194 if crate::columnar::all_non_null(validity, data.len()) {
12195 count += data.len() as u64;
12196 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12197 execution_checkpoint(control, chunk_index * 1024)?;
12198 sum += chunk.iter().sum::<f64>();
12199 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
12200 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
12201 }
12202 } else {
12203 for (i, &v) in data.iter().enumerate() {
12204 execution_checkpoint(control, i)?;
12205 if crate::columnar::validity_bit(validity, i) {
12206 count += 1;
12207 sum += v;
12208 mn = mn.min(v);
12209 mx = mx.max(v);
12210 }
12211 }
12212 }
12213 }
12214 }
12215 Ok((count, sum, mn, mx))
12216}
12217
12218#[inline]
12219fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
12220 if index.is_multiple_of(256) {
12221 control
12222 .map(crate::ExecutionControl::checkpoint)
12223 .transpose()?;
12224 }
12225 Ok(())
12226}
12227
12228fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
12229 if count == 0 && !matches!(agg, NativeAgg::Count) {
12230 return NativeAggResult::Null;
12231 }
12232 match agg {
12233 NativeAgg::Count => NativeAggResult::Count(count),
12234 NativeAgg::Sum => match sum.try_into() {
12237 Ok(v) => NativeAggResult::Int(v),
12238 Err(_) => NativeAggResult::Null,
12239 },
12240 NativeAgg::Min => NativeAggResult::Int(mn),
12241 NativeAgg::Max => NativeAggResult::Int(mx),
12242 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
12243 }
12244}
12245
12246fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
12247 if count == 0 && !matches!(agg, NativeAgg::Count) {
12248 return NativeAggResult::Null;
12249 }
12250 match agg {
12251 NativeAgg::Count => NativeAggResult::Count(count),
12252 NativeAgg::Sum => NativeAggResult::Float(sum),
12253 NativeAgg::Min => NativeAggResult::Float(mn),
12254 NativeAgg::Max => NativeAggResult::Float(mx),
12255 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
12256 }
12257}
12258
12259fn agg_int(
12262 stats: &[crate::page::PageStat],
12263 decode: fn(Option<&[u8]>) -> Option<i64>,
12264) -> Option<(Option<i64>, Option<i64>, u64)> {
12265 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
12266 let mut any = false;
12267 for s in stats {
12268 if let Some(v) = decode(s.min.as_deref()) {
12269 mn = mn.min(v);
12270 any = true;
12271 }
12272 if let Some(v) = decode(s.max.as_deref()) {
12273 mx = mx.max(v);
12274 any = true;
12275 }
12276 nulls += s.null_count;
12277 }
12278 any.then_some((Some(mn), Some(mx), nulls))
12279}
12280
12281fn agg_float(
12283 stats: &[crate::page::PageStat],
12284 decode: fn(Option<&[u8]>) -> Option<f64>,
12285) -> Option<(Option<f64>, Option<f64>, u64)> {
12286 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
12287 let mut any = false;
12288 for s in stats {
12289 if let Some(v) = decode(s.min.as_deref()) {
12290 mn = mn.min(v);
12291 any = true;
12292 }
12293 if let Some(v) = decode(s.max.as_deref()) {
12294 mx = mx.max(v);
12295 any = true;
12296 }
12297 nulls += s.null_count;
12298 }
12299 any.then_some((Some(mn), Some(mx), nulls))
12300}
12301
12302type SecondaryIndexes = (
12304 HashMap<u16, BitmapIndex>,
12305 HashMap<u16, AnnIndex>,
12306 HashMap<u16, FmIndex>,
12307 HashMap<u16, SparseIndex>,
12308 HashMap<u16, MinHashIndex>,
12309);
12310
12311fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
12312 let mut bitmap = HashMap::new();
12313 let mut ann = HashMap::new();
12314 let mut fm = HashMap::new();
12315 let mut sparse = HashMap::new();
12316 let mut minhash = HashMap::new();
12317 for idef in &schema.indexes {
12318 match idef.kind {
12319 IndexKind::Bitmap => {
12320 bitmap.insert(idef.column_id, BitmapIndex::new());
12321 }
12322 IndexKind::Ann => {
12323 let dim = schema
12324 .columns
12325 .iter()
12326 .find(|c| c.id == idef.column_id)
12327 .and_then(|c| match c.ty {
12328 TypeId::Embedding { dim } => Some(dim as usize),
12329 _ => None,
12330 })
12331 .unwrap_or(0);
12332 let options = idef.options.ann.clone().unwrap_or_default();
12333 ann.insert(
12334 idef.column_id,
12335 AnnIndex::with_full_options(
12336 dim,
12337 options.m,
12338 options.ef_construction,
12339 options.ef_search,
12340 &options,
12341 ),
12342 );
12343 }
12344 IndexKind::FmIndex => {
12345 fm.insert(idef.column_id, FmIndex::new());
12346 }
12347 IndexKind::Sparse => {
12348 sparse.insert(idef.column_id, SparseIndex::new());
12349 }
12350 IndexKind::MinHash => {
12351 let options = idef.options.minhash.clone().unwrap_or_default();
12352 minhash.insert(
12353 idef.column_id,
12354 MinHashIndex::with_options(options.permutations, options.bands),
12355 );
12356 }
12357 _ => {}
12358 }
12359 }
12360 (bitmap, ann, fm, sparse, minhash)
12361}
12362
12363const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
12364 | ColumnFlags::AUTO_INCREMENT
12365 | ColumnFlags::ENCRYPTED
12366 | ColumnFlags::ENCRYPTED_INDEXABLE
12367 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
12368
12369fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
12370 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
12371 return Err(MongrelError::Schema(
12372 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
12373 ));
12374 }
12375 Ok(())
12376}
12377
12378fn validate_alter_column_type(
12379 schema: &Schema,
12380 old: &ColumnDef,
12381 next: &ColumnDef,
12382 has_stored_versions: bool,
12383) -> Result<()> {
12384 if old.ty == next.ty {
12385 return Ok(());
12386 }
12387 if schema.indexes.iter().any(|i| i.column_id == old.id) {
12388 return Err(MongrelError::Schema(format!(
12389 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
12390 old.name
12391 )));
12392 }
12393 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
12394 return Ok(());
12395 }
12396 Err(MongrelError::Schema(format!(
12397 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
12398 old.ty, next.ty
12399 )))
12400}
12401
12402fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
12403 matches!(
12404 (old, new),
12405 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
12406 )
12407}
12408
12409fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
12415 let mut prev: Option<i64> = None;
12416 for r in rows {
12417 match r.columns.get(&pk_id) {
12418 Some(Value::Int64(v)) => {
12419 if prev.is_some_and(|p| p >= *v) {
12420 return false;
12421 }
12422 prev = Some(*v);
12423 }
12424 _ => return false,
12425 }
12426 }
12427 true
12428}
12429
12430#[allow(clippy::too_many_arguments)]
12431fn index_into(
12432 schema: &Schema,
12433 row: &Row,
12434 hot: &mut HotIndex,
12435 bitmap: &mut HashMap<u16, BitmapIndex>,
12436 ann: &mut HashMap<u16, AnnIndex>,
12437 fm: &mut HashMap<u16, FmIndex>,
12438 sparse: &mut HashMap<u16, SparseIndex>,
12439 minhash: &mut HashMap<u16, MinHashIndex>,
12440) {
12441 for idef in &schema.indexes {
12442 let Some(val) = row.columns.get(&idef.column_id) else {
12443 continue;
12444 };
12445 match idef.kind {
12446 IndexKind::Bitmap => {
12447 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12448 b.insert(val.encode_key(), row.row_id);
12449 }
12450 }
12451 IndexKind::Ann => {
12452 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12453 if let Some(meta) = val.generated_embedding_metadata() {
12454 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12456 continue;
12457 }
12458 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12459 .is_err()
12460 {
12461 continue;
12462 }
12463 }
12464 a.insert_validated(v, row.row_id);
12465 }
12466 }
12467 IndexKind::FmIndex => {
12468 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12469 f.insert(b.clone(), row.row_id);
12470 }
12471 }
12472 IndexKind::Sparse => {
12473 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12474 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12477 s.insert(&terms, row.row_id);
12478 }
12479 }
12480 }
12481 IndexKind::MinHash => {
12482 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12483 let tokens = crate::index::token_hashes_from_bytes(b);
12486 mh.insert(&tokens, row.row_id);
12487 }
12488 }
12489 _ => {}
12490 }
12491 }
12492 if let Some(pk_col) = schema.primary_key() {
12493 if let Some(pk_val) = row.columns.get(&pk_col.id) {
12494 hot.insert(pk_val.encode_key(), row.row_id);
12495 }
12496 }
12497}
12498
12499#[allow(clippy::too_many_arguments)]
12502fn index_into_single(
12503 idef: &IndexDef,
12504 _schema: &Schema,
12505 row: &Row,
12506 _hot: &mut HotIndex,
12507 bitmap: &mut HashMap<u16, BitmapIndex>,
12508 ann: &mut HashMap<u16, AnnIndex>,
12509 fm: &mut HashMap<u16, FmIndex>,
12510 sparse: &mut HashMap<u16, SparseIndex>,
12511 minhash: &mut HashMap<u16, MinHashIndex>,
12512) {
12513 let Some(val) = row.columns.get(&idef.column_id) else {
12514 return;
12515 };
12516 match idef.kind {
12517 IndexKind::Bitmap => {
12518 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12519 b.insert(val.encode_key(), row.row_id);
12520 }
12521 }
12522 IndexKind::Ann => {
12523 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12524 if let Some(meta) = val.generated_embedding_metadata() {
12525 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12527 return;
12528 }
12529 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12530 .is_err()
12531 {
12532 return;
12533 }
12534 }
12535 a.insert_validated(v, row.row_id);
12536 }
12537 }
12538 IndexKind::FmIndex => {
12539 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12540 f.insert(b.clone(), row.row_id);
12541 }
12542 }
12543 IndexKind::Sparse => {
12544 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12545 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12546 s.insert(&terms, row.row_id);
12547 }
12548 }
12549 }
12550 IndexKind::MinHash => {
12551 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12552 let tokens = crate::index::token_hashes_from_bytes(b);
12553 mh.insert(&tokens, row.row_id);
12554 }
12555 }
12556 _ => {}
12557 }
12558}
12559
12560fn eval_partial_predicate(
12566 pred: &str,
12567 columns_map: &HashMap<u16, &Value>,
12568 name_to_id: &HashMap<&str, u16>,
12569) -> bool {
12570 let lower = pred.trim().to_ascii_lowercase();
12571 if let Some(rest) = lower.strip_suffix(" is not null") {
12573 let col_name = rest.trim();
12574 if let Some(col_id) = name_to_id.get(col_name) {
12575 return columns_map
12576 .get(col_id)
12577 .is_some_and(|v| !matches!(v, Value::Null));
12578 }
12579 }
12580 if let Some(rest) = lower.strip_suffix(" is null") {
12582 let col_name = rest.trim();
12583 if let Some(col_id) = name_to_id.get(col_name) {
12584 return columns_map
12585 .get(col_id)
12586 .is_none_or(|v| matches!(v, Value::Null));
12587 }
12588 }
12589 true
12592}
12593
12594#[allow(dead_code)]
12600fn bulk_index_key(
12601 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12602 column_id: u16,
12603 ty: TypeId,
12604 col: &columnar::NativeColumn,
12605 i: usize,
12606) -> Option<Vec<u8>> {
12607 let encoded = columnar::encode_key_native(ty, col, i)?;
12608 {
12609 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12610 if let Some((key, scheme)) = column_keys.get(&column_id) {
12611 return Some(match (*scheme, col) {
12612 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12613 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12614 ope_token_i64(key, data[i]).to_vec()
12615 }
12616 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12617 ope_token_f64(key, data[i]).to_vec()
12618 }
12619 _ => hmac_token(key, &encoded).to_vec(),
12620 });
12621 }
12622 }
12623 Some(encoded)
12624}
12625
12626pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12627 write_schema_with_after(dir, schema, || {})
12628}
12629
12630pub(crate) fn write_schema_durable(
12631 root: &crate::durable_file::DurableRoot,
12632 schema: &Schema,
12633) -> Result<()> {
12634 write_schema_durable_with_after(root, schema, || {})
12635}
12636
12637fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12638where
12639 F: FnOnce(),
12640{
12641 let json = serde_json::to_string_pretty(schema)
12642 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12643 crate::durable_file::write_atomic_with_after(
12644 &dir.join(SCHEMA_FILENAME),
12645 json.as_bytes(),
12646 after_publish,
12647 )?;
12648 Ok(())
12649}
12650
12651fn write_schema_durable_with_after<F>(
12652 root: &crate::durable_file::DurableRoot,
12653 schema: &Schema,
12654 after_publish: F,
12655) -> Result<()>
12656where
12657 F: FnOnce(),
12658{
12659 let json = serde_json::to_string_pretty(schema)
12660 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12661 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12662 Ok(())
12663}
12664
12665fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12666 let mut schema_published = false;
12667 let schema_result = match table._root_guard.as_deref() {
12668 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12669 schema_published = true;
12670 }),
12671 None => write_schema_with_after(&table.dir, &table.schema, || {
12672 schema_published = true;
12673 }),
12674 };
12675 if schema_result.is_err() && !schema_published {
12676 return schema_result;
12677 }
12678 match table.persist_manifest(table.current_epoch()) {
12679 Ok(()) => Ok(()),
12680 Err(manifest_error) => Err(match schema_result {
12681 Ok(()) => manifest_error,
12682 Err(schema_error) => MongrelError::Other(format!(
12683 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12684 )),
12685 }),
12686 }
12687}
12688
12689fn read_schema(dir: &Path) -> Result<Schema> {
12690 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12691 read_schema_file(file)
12692}
12693
12694fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12695 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12696 use std::io::Read;
12697
12698 let length = file.metadata()?.len();
12699 if length > MAX_SCHEMA_BYTES {
12700 return Err(MongrelError::ResourceLimitExceeded {
12701 resource: "schema bytes",
12702 requested: usize::try_from(length).unwrap_or(usize::MAX),
12703 limit: MAX_SCHEMA_BYTES as usize,
12704 });
12705 }
12706 let mut bytes = Vec::with_capacity(length as usize);
12707 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12708 if bytes.len() as u64 != length {
12709 return Err(MongrelError::Schema(
12710 "schema length changed while reading".into(),
12711 ));
12712 }
12713 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12714}
12715
12716fn preflight_standalone_open(
12717 dir: &Path,
12718 runs_root: Option<&crate::durable_file::DurableRoot>,
12719 idx_root: Option<&crate::durable_file::DurableRoot>,
12720 manifest: &Manifest,
12721 schema: &Schema,
12722 records: &[crate::wal::Record],
12723 kek: Option<Arc<Kek>>,
12724) -> Result<()> {
12725 crate::wal::validate_shared_transaction_framing(records)?;
12726 if manifest.schema_id > schema.schema_id
12727 || manifest.flushed_epoch > manifest.current_epoch
12728 || manifest.global_idx_epoch > manifest.current_epoch
12729 || manifest.next_row_id == u64::MAX
12730 || manifest.auto_inc_next < 0
12731 || manifest.auto_inc_next == i64::MAX
12732 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12733 {
12734 return Err(MongrelError::InvalidArgument(
12735 "manifest counters or schema identity are invalid".into(),
12736 ));
12737 }
12738 let mut run_ids = HashSet::new();
12739 let mut maximum_row_id = None::<u64>;
12740 for run in &manifest.runs {
12741 if run.run_id >= u64::MAX as u128
12742 || !run_ids.insert(run.run_id)
12743 || run.epoch_created > manifest.current_epoch
12744 {
12745 return Err(MongrelError::InvalidArgument(
12746 "manifest contains an invalid or duplicate active run".into(),
12747 ));
12748 }
12749 let mut reader = match runs_root {
12750 Some(root) => RunReader::open_file(
12751 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12752 schema.clone(),
12753 kek.clone(),
12754 )?,
12755 None => RunReader::open(
12756 dir.join(RUNS_DIR)
12757 .join(format!("r-{}.sr", run.run_id as u64)),
12758 schema.clone(),
12759 kek.clone(),
12760 )?,
12761 };
12762 let header = reader.header();
12763 if header.run_id != run.run_id
12764 || header.level != run.level
12765 || header.row_count != run.row_count
12766 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12767 || header.is_uniform_epoch() && header.epoch_created != 0
12768 || header.schema_id > schema.schema_id
12769 {
12770 return Err(MongrelError::InvalidArgument(format!(
12771 "run {} differs from its manifest",
12772 run.run_id
12773 )));
12774 }
12775 if header.row_count != 0 {
12776 maximum_row_id = Some(
12777 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12778 );
12779 }
12780 reader.validate_all_pages()?;
12781 }
12782 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12783 return Err(MongrelError::InvalidArgument(
12784 "manifest next_row_id does not advance beyond persisted rows".into(),
12785 ));
12786 }
12787 for run in &manifest.retiring {
12788 if run.run_id >= u64::MAX as u128
12789 || run.retire_epoch > manifest.current_epoch
12790 || !run_ids.insert(run.run_id)
12791 {
12792 return Err(MongrelError::InvalidArgument(
12793 "manifest contains an invalid or duplicate retired run".into(),
12794 ));
12795 }
12796 }
12797 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12798 match idx_root {
12799 Some(root) => {
12800 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12801 }
12802 None => {
12803 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12804 }
12805 }
12806
12807 let committed = records
12808 .iter()
12809 .filter_map(|record| match record.op {
12810 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12811 _ => None,
12812 })
12813 .collect::<HashMap<_, _>>();
12814 for record in records {
12815 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12816 continue;
12817 };
12818 match &record.op {
12819 Op::Put { table_id, rows } => {
12820 if *table_id != manifest.table_id {
12821 return Err(MongrelError::CorruptWal {
12822 offset: record.seq.0,
12823 reason: format!(
12824 "private WAL record references table {table_id}, expected {}",
12825 manifest.table_id
12826 ),
12827 });
12828 }
12829 let rows: Vec<Row> =
12830 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12831 offset: record.seq.0,
12832 reason: format!("committed Put payload could not be decoded: {error}"),
12833 })?;
12834 for row in rows {
12835 if row.deleted || row.row_id.0 == u64::MAX {
12836 return Err(MongrelError::CorruptWal {
12837 offset: record.seq.0,
12838 reason: "committed Put contains an invalid row identity".into(),
12839 });
12840 }
12841 let cells = row.columns.into_iter().collect::<Vec<_>>();
12842 schema
12843 .validate_values(&cells)
12844 .map_err(|error| MongrelError::CorruptWal {
12845 offset: record.seq.0,
12846 reason: format!("committed Put violates table schema: {error}"),
12847 })?;
12848 if schema.auto_increment_column().is_some_and(|column| {
12849 matches!(
12850 cells.iter().find(|(id, _)| *id == column.id),
12851 Some((_, Value::Int64(value))) if *value == i64::MAX
12852 )
12853 }) {
12854 return Err(MongrelError::CorruptWal {
12855 offset: record.seq.0,
12856 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12857 });
12858 }
12859 }
12860 }
12861 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12862 if *table_id != manifest.table_id =>
12863 {
12864 return Err(MongrelError::CorruptWal {
12865 offset: record.seq.0,
12866 reason: format!(
12867 "private WAL record references table {table_id}, expected {}",
12868 manifest.table_id
12869 ),
12870 });
12871 }
12872 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12873 return Err(MongrelError::CorruptWal {
12874 offset: record.seq.0,
12875 reason: "private WAL contains shared spilled-run metadata".into(),
12876 });
12877 }
12878 _ => {}
12879 }
12880 }
12881 Ok(())
12882}
12883
12884fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12885 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12886}
12887
12888fn wal_segment_number(path: &Path) -> Option<u64> {
12889 path.file_stem()
12890 .and_then(|stem| stem.to_str())
12891 .and_then(|stem| stem.strip_prefix("seg-"))
12892 .and_then(|number| number.parse().ok())
12893}
12894
12895fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12896 let n = list_wal_numbers(wal_dir)?;
12897 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12898}
12899
12900fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12901 list_wal_numbers(wal_dir)?
12902 .map(|maximum| {
12903 maximum
12904 .checked_add(1)
12905 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12906 })
12907 .unwrap_or(Ok(0))
12908}
12909
12910fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12911 let mut max_n = None;
12912 let entries = match std::fs::read_dir(wal_dir) {
12913 Ok(entries) => entries,
12914 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12915 Err(error) => return Err(error.into()),
12916 };
12917 for entry in entries {
12918 let entry = entry?;
12919 let fname = entry.file_name();
12920 let Some(s) = fname.to_str() else {
12921 continue;
12922 };
12923 let Some(stripped) = s.strip_prefix("seg-") else {
12924 continue;
12925 };
12926 let Some(number) = stripped.strip_suffix(".wal") else {
12927 return Err(MongrelError::CorruptWal {
12928 offset: 0,
12929 reason: format!("malformed WAL segment name {s:?}"),
12930 });
12931 };
12932 let n = number
12933 .parse::<u32>()
12934 .map_err(|_| MongrelError::CorruptWal {
12935 offset: 0,
12936 reason: format!("malformed WAL segment name {s:?}"),
12937 })?;
12938 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12939 return Err(MongrelError::CorruptWal {
12940 offset: n as u64,
12941 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12942 });
12943 }
12944 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12945 }
12946 Ok(max_n)
12947}