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 fn rebuild_indexes(&mut self) -> Result<()> {
2597 self.rebuild_indexes_from_runs_inner(None)
2598 }
2599
2600 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2601 self.rebuild_indexes_from_runs_inner(None)
2602 }
2603
2604 fn pk_equality_fallback(
2607 &self,
2608 pk_column_id: u16,
2609 lookup: &[u8],
2610 snapshot: Snapshot,
2611 ) -> Result<RowIdSet> {
2612 for row in self.memtable.visible_versions_at(snapshot) {
2614 if row.deleted {
2615 continue;
2616 }
2617 if let Some(pk_val) = row.columns.get(&pk_column_id) {
2618 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
2619 return Ok(RowIdSet::one(row.row_id.0));
2620 }
2621 }
2622 }
2623 for row in self.mutable_run.visible_versions_at(snapshot) {
2624 if row.deleted {
2625 continue;
2626 }
2627 if let Some(pk_val) = row.columns.get(&pk_column_id) {
2628 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
2629 return Ok(RowIdSet::one(row.row_id.0));
2630 }
2631 }
2632 }
2633 if lookup.len() == 8 {
2636 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
2637 let n = i64::from_be_bytes(arr);
2638 return self.range_scan_i64(pk_column_id, n, n, snapshot);
2639 }
2640 }
2641 let mut found = Vec::new();
2644 let overlay = self.overlay_rid_set(snapshot);
2645 for rr in &self.run_refs {
2646 let mut reader = self.open_reader(rr.run_id)?;
2647 for row in reader.visible_rows(snapshot.epoch)? {
2648 if overlay.contains(&row.row_id.0) || row.deleted {
2649 continue;
2650 }
2651 if let Some(pk_val) = row.columns.get(&pk_column_id) {
2652 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
2653 found.push(row.row_id.0);
2654 }
2655 }
2656 }
2657 }
2658 Ok(RowIdSet::from_unsorted(found))
2659 }
2660
2661 fn rebuild_indexes_from_runs_inner(
2662 &mut self,
2663 control: Option<&crate::ExecutionControl>,
2664 ) -> Result<()> {
2665 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2668 crate::retention::PinSource::OnlineIndexBuild,
2669 self.current_epoch(),
2670 );
2671 self.hot = HotIndex::new();
2672 self.pk_by_row.clear();
2673 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2674 self.bitmap = bitmap;
2675 self.ann = ann;
2676 self.fm = fm;
2677 self.sparse = sparse;
2678 self.minhash = minhash;
2679 let snapshot = Epoch(u64::MAX);
2680 let ttl_now = unix_nanos_now();
2681 let mut scanned = 0_usize;
2682 for rr in self.run_refs.clone() {
2683 if let Some(control) = control {
2684 control.checkpoint()?;
2685 }
2686 let mut reader = self.open_reader(rr.run_id)?;
2687 for row in reader.visible_rows(snapshot)? {
2688 if scanned.is_multiple_of(256) {
2689 if let Some(control) = control {
2690 control.checkpoint()?;
2691 }
2692 }
2693 scanned += 1;
2694 if self.row_expired_at(&row, ttl_now) {
2695 continue;
2696 }
2697 let tok_row = self.tokenized_for_indexes(&row);
2698 index_into(
2699 &self.schema,
2700 &tok_row,
2701 &mut self.hot,
2702 &mut self.bitmap,
2703 &mut self.ann,
2704 &mut self.fm,
2705 &mut self.sparse,
2706 &mut self.minhash,
2707 );
2708 }
2709 }
2710 for row in self.mutable_run.visible_versions(snapshot) {
2711 if scanned.is_multiple_of(256) {
2712 if let Some(control) = control {
2713 control.checkpoint()?;
2714 }
2715 }
2716 scanned += 1;
2717 if row.deleted {
2718 self.remove_hot_for_row(row.row_id, snapshot);
2719 } else if !self.row_expired_at(&row, ttl_now) {
2720 self.index_row(&row);
2721 }
2722 }
2723 for row in self.memtable.visible_versions(snapshot) {
2724 if scanned.is_multiple_of(256) {
2725 if let Some(control) = control {
2726 control.checkpoint()?;
2727 }
2728 }
2729 scanned += 1;
2730 if row.deleted {
2731 self.remove_hot_for_row(row.row_id, snapshot);
2732 } else if !self.row_expired_at(&row, ttl_now) {
2733 self.index_row(&row);
2734 }
2735 }
2736 self.refresh_pk_by_row_from_hot();
2737 Ok(())
2738 }
2739
2740 fn refresh_pk_by_row_from_hot(&mut self) {
2741 self.pk_by_row_complete = true;
2742 if self.schema.primary_key().is_none() {
2743 self.pk_by_row.clear();
2744 return;
2745 }
2746 self.pk_by_row = ReversePkMap::from_entries(
2752 self.hot
2753 .entries()
2754 .into_iter()
2755 .map(|(key, row_id)| (row_id, key)),
2756 );
2757 }
2758
2759 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2760 if self.schema.primary_key().is_some() {
2761 self.pk_by_row.insert(row_id, key.clone());
2762 }
2763 self.hot.insert(key, row_id);
2764 }
2765
2766 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2770 self.build_learned_ranges_inner(None)
2771 }
2772
2773 fn build_learned_ranges_inner(
2774 &mut self,
2775 control: Option<&crate::ExecutionControl>,
2776 ) -> Result<()> {
2777 self.learned_range = Arc::new(HashMap::new());
2778 if self.run_refs.len() != 1 {
2779 return Ok(());
2780 }
2781 let cols: Vec<(u16, usize)> = self
2782 .schema
2783 .indexes
2784 .iter()
2785 .filter(|i| i.kind == IndexKind::LearnedRange)
2786 .map(|i| {
2787 (
2788 i.column_id,
2789 i.options
2790 .learned_range
2791 .as_ref()
2792 .map(|options| options.epsilon)
2793 .unwrap_or(16),
2794 )
2795 })
2796 .collect();
2797 if cols.is_empty() {
2798 return Ok(());
2799 }
2800 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2801 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2802 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2803 _ => return Ok(()),
2804 };
2805 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2806 if column_index % 256 == 0 {
2807 if let Some(control) = control {
2808 control.checkpoint()?;
2809 }
2810 }
2811 let ty = self
2812 .schema
2813 .columns
2814 .iter()
2815 .find(|c| c.id == cid)
2816 .map(|c| c.ty.clone())
2817 .unwrap_or(TypeId::Int64);
2818 match ty {
2819 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2820 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2821 let pairs: Vec<(i64, u64)> = data
2822 .iter()
2823 .zip(row_ids.iter())
2824 .map(|(v, r)| (*v, *r))
2825 .collect();
2826 Arc::make_mut(&mut self.learned_range).insert(
2827 cid,
2828 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2829 );
2830 }
2831 }
2832 TypeId::Float64 => {
2833 if let columnar::NativeColumn::Float64 { data, .. } =
2834 reader.column_native(cid)?
2835 {
2836 let pairs: Vec<(f64, u64)> = data
2837 .iter()
2838 .zip(row_ids.iter())
2839 .map(|(v, r)| (*v, *r))
2840 .collect();
2841 Arc::make_mut(&mut self.learned_range).insert(
2842 cid,
2843 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2844 );
2845 }
2846 }
2847 _ => {}
2848 }
2849 }
2850 Ok(())
2851 }
2852
2853 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2860 if self.indexes_complete {
2861 crate::trace::QueryTrace::record(|t| {
2862 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2863 });
2864 return Ok(());
2865 }
2866 crate::trace::QueryTrace::record(|t| {
2867 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2868 });
2869 self.rebuild_indexes_from_runs()?;
2870 self.build_learned_ranges()?;
2871 self.indexes_complete = true;
2872 let epoch = self.current_epoch();
2873 self.checkpoint_indexes(epoch);
2874 Ok(())
2875 }
2876
2877 #[doc(hidden)]
2880 pub fn ensure_indexes_complete_controlled<F>(
2881 &mut self,
2882 control: &crate::ExecutionControl,
2883 before_publish: F,
2884 ) -> Result<bool>
2885 where
2886 F: FnOnce() -> bool,
2887 {
2888 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2889 .map(|(changed, _)| changed)
2890 }
2891
2892 #[doc(hidden)]
2895 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2896 &mut self,
2897 control: &crate::ExecutionControl,
2898 before_publish: F,
2899 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2900 where
2901 F: FnOnce() -> bool,
2902 {
2903 if self.indexes_complete {
2904 crate::trace::QueryTrace::record(|trace| {
2905 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2906 });
2907 return Ok((false, None));
2908 }
2909 crate::trace::QueryTrace::record(|trace| {
2910 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2911 });
2912 control.checkpoint()?;
2913 let maintenance_epoch = self.current_epoch();
2914 self.rebuild_indexes_from_runs_inner(Some(control))?;
2915 self.build_learned_ranges_inner(Some(control))?;
2916 control.checkpoint()?;
2917 if !before_publish() {
2918 return Err(MongrelError::Cancelled);
2919 }
2920 self.indexes_complete = true;
2921 self.checkpoint_indexes(maintenance_epoch);
2922 Ok((
2923 true,
2924 Some(MaintenanceReceipt {
2925 epoch: maintenance_epoch,
2926 }),
2927 ))
2928 }
2929
2930 fn pending_epoch(&self) -> Epoch {
2931 Epoch(self.epoch.visible().0 + 1)
2932 }
2933
2934 fn is_shared(&self) -> bool {
2937 matches!(self.wal, WalSink::Shared(_))
2938 }
2939
2940 fn ensure_txn_id(&mut self) -> Result<u64> {
2944 if self.current_txn_id == 0 {
2945 let id = match &self.wal {
2946 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2947 WalSink::Private(_) => {
2948 return Err(MongrelError::Full(
2949 "standalone transaction id namespace exhausted".into(),
2950 ))
2951 }
2952 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2953 };
2954 self.current_txn_id = id;
2955 }
2956 Ok(self.current_txn_id)
2957 }
2958
2959 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2962 self.ensure_writable()?;
2963 let txn_id = self.ensure_txn_id()?;
2964 let table_id = self.table_id;
2965 match &mut self.wal {
2966 WalSink::Private(w) => {
2967 w.append_txn(txn_id, op)?;
2968 self.pending_private_mutations = true;
2969 }
2970 WalSink::Shared(s) => {
2971 s.wal.lock().append(txn_id, table_id, op)?;
2972 }
2973 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2974 }
2975 Ok(())
2976 }
2977
2978 fn ensure_writable(&self) -> Result<()> {
2979 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2980 return Err(MongrelError::ReadOnlyReplica);
2981 }
2982 if self.durable_commit_failed {
2983 return Err(MongrelError::Other(
2984 "table poisoned by post-commit failure; reopen required".into(),
2985 ));
2986 }
2987 Ok(())
2988 }
2989
2990 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
3001 match &self.auth {
3002 Some(checker) => checker.check(&self.name, perm),
3003 None => Ok(()),
3004 }
3005 }
3006 pub fn require_select(&self) -> Result<()> {
3011 self.require(crate::auth_state::RequiredPermission::Select)
3012 }
3013 fn require_insert(&self) -> Result<()> {
3014 self.require(crate::auth_state::RequiredPermission::Insert)
3015 }
3016 #[allow(dead_code)]
3020 fn require_update(&self) -> Result<()> {
3021 self.require(crate::auth_state::RequiredPermission::Update)
3022 }
3023 fn require_delete(&self) -> Result<()> {
3024 self.require(crate::auth_state::RequiredPermission::Delete)
3025 }
3026
3027 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
3030 self.require_insert()?;
3031 Ok(self.put_returning(columns)?.0)
3032 }
3033
3034 pub fn put_returning(
3039 &mut self,
3040 mut columns: Vec<(u16, Value)>,
3041 ) -> Result<(RowId, Option<i64>)> {
3042 self.require_insert()?;
3043 let assigned = self.fill_auto_inc(&mut columns)?;
3044 self.apply_defaults(&mut columns)?;
3045 self.schema.validate_values(&columns)?;
3046 let row_id = if self.schema.clustered {
3051 self.derive_clustered_row_id(&columns)?
3052 } else {
3053 self.allocator.alloc()?
3054 };
3055 let epoch = self.pending_epoch();
3056 let mut row = Row::new(row_id, epoch);
3057 for (col_id, val) in columns {
3058 row.columns.insert(col_id, val);
3059 }
3060 self.commit_rows(vec![row], assigned.is_some())?;
3061 Ok((row_id, assigned))
3062 }
3063
3064 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
3067 self.require_insert()?;
3068 Ok(self
3069 .put_batch_returning(batch)?
3070 .into_iter()
3071 .map(|(r, _)| r)
3072 .collect())
3073 }
3074
3075 pub fn put_batch_returning(
3078 &mut self,
3079 batch: Vec<Vec<(u16, Value)>>,
3080 ) -> Result<Vec<(RowId, Option<i64>)>> {
3081 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
3082 for mut cols in batch {
3083 let assigned = self.fill_auto_inc(&mut cols)?;
3084 self.apply_defaults(&mut cols)?;
3085 filled.push((cols, assigned));
3086 }
3087 for (cols, _) in &filled {
3088 self.schema.validate_values(cols)?;
3089 }
3090 let epoch = self.pending_epoch();
3091 let mut rows = Vec::with_capacity(filled.len());
3092 let mut ids = Vec::with_capacity(filled.len());
3093 let first_row_id = if self.schema.clustered {
3094 None
3095 } else {
3096 let count = u64::try_from(filled.len())
3097 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
3098 Some(self.allocator.alloc_range(count)?.0)
3099 };
3100 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
3101 let row_id = match first_row_id {
3102 Some(first) => RowId(first + row_index as u64),
3103 None => self.derive_clustered_row_id(&cols)?,
3104 };
3105 let mut row = Row::new(row_id, epoch);
3106 for (c, v) in cols {
3107 row.columns.insert(c, v);
3108 }
3109 ids.push((row_id, assigned));
3110 rows.push(row);
3111 }
3112 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
3113 self.commit_rows(rows, all_auto_generated)?;
3114 Ok(ids)
3115 }
3116
3117 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
3123 self.ensure_writable()?;
3124 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3125 return Ok(None);
3126 };
3127 let pos = columns.iter().position(|(c, _)| *c == cid);
3128 let assigned = match pos {
3129 Some(i) => match &columns[i].1 {
3130 Value::Null => {
3131 let next = self.alloc_auto_inc_value()?;
3132 columns[i].1 = Value::Int64(next);
3133 Some(next)
3134 }
3135 Value::Int64(n) => {
3136 self.advance_auto_inc_past(*n)?;
3137 None
3138 }
3139 other => {
3140 return Err(MongrelError::InvalidArgument(format!(
3141 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
3142 other
3143 )))
3144 }
3145 },
3146 None => {
3147 let next = self.alloc_auto_inc_value()?;
3148 columns.push((cid, Value::Int64(next)));
3149 Some(next)
3150 }
3151 };
3152 Ok(assigned)
3153 }
3154
3155 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
3161 for col in &self.schema.columns {
3162 let Some(expr) = &col.default_value else {
3163 continue;
3164 };
3165 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
3167 continue;
3168 }
3169 let pos = columns.iter().position(|(c, _)| *c == col.id);
3170 let needs_default = match pos {
3171 None => true,
3172 Some(i) => matches!(columns[i].1, Value::Null),
3173 };
3174 if !needs_default {
3175 continue;
3176 }
3177 let v = match expr {
3178 crate::schema::DefaultExpr::Static(v) => v.clone(),
3179 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
3180 crate::schema::DefaultExpr::Uuid => {
3181 let mut buf = [0u8; 16];
3182 getrandom::getrandom(&mut buf)
3183 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
3184 Value::Uuid(buf)
3185 }
3186 };
3187 match pos {
3188 None => columns.push((col.id, v)),
3189 Some(i) => columns[i].1 = v,
3190 }
3191 }
3192 Ok(())
3193 }
3194
3195 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
3197 self.ensure_auto_inc_seeded()?;
3198 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3200 let v = ai.next;
3201 ai.next = ai
3202 .next
3203 .checked_add(1)
3204 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3205 Ok(v)
3206 }
3207
3208 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
3211 self.ensure_auto_inc_seeded()?;
3212 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3213 let floor = used
3214 .checked_add(1)
3215 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3216 .max(1);
3217 if ai.next < floor {
3218 ai.next = floor;
3219 }
3220 Ok(())
3221 }
3222
3223 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
3228 let needs_seed = match self.auto_inc {
3229 Some(ai) => !ai.seeded,
3230 None => return Ok(()),
3231 };
3232 if !needs_seed {
3233 return Ok(());
3234 }
3235 if self.seed_empty_auto_inc() {
3236 return Ok(());
3237 }
3238 let cid = self
3239 .auto_inc
3240 .as_ref()
3241 .expect("auto-inc column present")
3242 .column_id;
3243 let max = self.scan_max_int64(cid)?;
3244 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3245 let floor = max
3246 .checked_add(1)
3247 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3248 .max(1);
3249 if ai.next < floor {
3250 ai.next = floor;
3251 }
3252 ai.seeded = true;
3253 Ok(())
3254 }
3255
3256 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
3257 if n == 0 || self.auto_inc.is_none() {
3258 return Ok(None);
3259 }
3260 self.ensure_auto_inc_seeded()?;
3261 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3262 let start = ai.next;
3263 let count = i64::try_from(n)
3264 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
3265 ai.next = ai
3266 .next
3267 .checked_add(count)
3268 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3269 Ok(Some(start))
3270 }
3271
3272 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
3276 let mut max: i64 = 0;
3277 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
3278 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3279 if *n > max {
3280 max = *n;
3281 }
3282 }
3283 }
3284 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3285 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3286 if *n > max {
3287 max = *n;
3288 }
3289 }
3290 }
3291 for rr in self.run_refs.clone() {
3292 let reader = self.open_reader(rr.run_id)?;
3293 if let Some(stats) = reader.column_page_stats(column_id) {
3294 for s in stats {
3295 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3296 if n > max {
3297 max = n;
3298 }
3299 }
3300 }
3301 } else if reader.has_column(column_id) {
3302 if let columnar::NativeColumn::Int64 { data, validity } =
3303 reader.column_native_shared(column_id)?
3304 {
3305 for (i, n) in data.iter().enumerate() {
3306 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3307 {
3308 max = *n;
3309 }
3310 }
3311 }
3312 }
3313 }
3314 Ok(max)
3315 }
3316
3317 fn seed_empty_auto_inc(&mut self) -> bool {
3318 let Some(ai) = self.auto_inc.as_mut() else {
3319 return false;
3320 };
3321 if ai.seeded || self.live_count != 0 {
3322 return false;
3323 }
3324 if ai.next < 1 {
3325 ai.next = 1;
3326 }
3327 ai.seeded = true;
3328 true
3329 }
3330
3331 fn advance_auto_inc_from_native_columns(
3332 &mut self,
3333 columns: &[(u16, columnar::NativeColumn)],
3334 n: usize,
3335 live_before: u64,
3336 ) -> Result<()> {
3337 let Some(ai) = self.auto_inc.as_mut() else {
3338 return Ok(());
3339 };
3340 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3341 return Ok(());
3342 };
3343 let columnar::NativeColumn::Int64 { data, validity } = col else {
3344 return Err(MongrelError::InvalidArgument(format!(
3345 "AUTO_INCREMENT column {} must be Int64",
3346 ai.column_id
3347 )));
3348 };
3349 let max = if native_int64_strictly_increasing(col, n) {
3350 data.get(n.saturating_sub(1)).copied()
3351 } else {
3352 data.iter()
3353 .take(n)
3354 .enumerate()
3355 .filter_map(|(i, v)| {
3356 if validity.is_empty() || columnar::validity_bit(validity, i) {
3357 Some(*v)
3358 } else {
3359 None
3360 }
3361 })
3362 .max()
3363 };
3364 if let Some(max) = max {
3365 let floor = max
3366 .checked_add(1)
3367 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3368 .max(1);
3369 if ai.next < floor {
3370 ai.next = floor;
3371 }
3372 if ai.seeded || live_before == 0 {
3373 ai.seeded = true;
3374 }
3375 }
3376 Ok(())
3377 }
3378
3379 fn fill_auto_inc_native_columns(
3380 &mut self,
3381 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3382 n: usize,
3383 ) -> Result<()> {
3384 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3385 return Ok(());
3386 };
3387 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3388 if let Some(start) = self.alloc_auto_inc_range(n)? {
3389 columns.push((
3390 cid,
3391 columnar::NativeColumn::Int64 {
3392 data: (start..start.saturating_add(n as i64)).collect(),
3393 validity: vec![0xFF; n.div_ceil(8)],
3394 },
3395 ));
3396 }
3397 return Ok(());
3398 };
3399
3400 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3401 return Err(MongrelError::InvalidArgument(format!(
3402 "AUTO_INCREMENT column {cid} must be Int64"
3403 )));
3404 };
3405 if data.len() < n {
3406 return Err(MongrelError::InvalidArgument(format!(
3407 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3408 data.len()
3409 )));
3410 }
3411 if columnar::all_non_null(validity, n) {
3412 return Ok(());
3413 }
3414 if validity.iter().all(|b| *b == 0) {
3415 if let Some(start) = self.alloc_auto_inc_range(n)? {
3416 for (i, slot) in data.iter_mut().take(n).enumerate() {
3417 *slot = start.saturating_add(i as i64);
3418 }
3419 *validity = vec![0xFF; n.div_ceil(8)];
3420 }
3421 return Ok(());
3422 }
3423
3424 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3425 for (i, slot) in data.iter_mut().enumerate().take(n) {
3426 if columnar::validity_bit(validity, i) {
3427 self.advance_auto_inc_past(*slot)?;
3428 } else {
3429 *slot = self.alloc_auto_inc_value()?;
3430 }
3431 }
3432 *validity = new_validity;
3433 Ok(())
3434 }
3435
3436 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3450 self.ensure_writable()?;
3451 if self.auto_inc.is_none() {
3452 return Ok(None);
3453 }
3454 Ok(Some(self.alloc_auto_inc_value()?))
3455 }
3456
3457 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3463 let payload = bincode::serialize(&rows)?;
3464 self.wal_append_data(Op::Put {
3465 table_id: self.table_id,
3466 rows: payload,
3467 })?;
3468 if self.is_shared() {
3469 self.pending_rows_auto_inc
3470 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3471 self.pending_rows.extend(rows);
3472 } else {
3473 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3474 }
3475 Ok(())
3476 }
3477
3478 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3481 self.ensure_indexes_complete()
3482 }
3483
3484 pub(crate) fn prepare_durable_publish_controlled(
3485 &mut self,
3486 control: &crate::ExecutionControl,
3487 ) -> Result<()> {
3488 self.ensure_indexes_complete_controlled(control, || true)?;
3489 Ok(())
3490 }
3491
3492 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3493 self.apply_put_rows_inner_prepared(rows, true);
3494 }
3495
3496 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3497 if check_existing_pk {
3498 self.ensure_indexes_complete()?;
3499 }
3500 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3501 Ok(())
3502 }
3503
3504 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3508 if rows.len() == 1 {
3512 let row = rows.into_iter().next().expect("len checked");
3513 self.apply_put_row_single(row, check_existing_pk);
3514 return;
3515 }
3516 let pk_id = self.schema.primary_key().map(|c| c.id);
3533 let probe = match pk_id {
3534 Some(pid) => {
3535 check_existing_pk
3536 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3537 }
3538 None => false,
3539 };
3540 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3543 for r in rows {
3544 for &cid in r.columns.keys() {
3545 self.pending_put_cols.insert(cid);
3546 }
3547 let mut replaced_image: Option<Row> = None;
3548 match pk_id {
3549 Some(pid) if probe || maintain_pk_by_row => {
3550 if let Some(pk_val) = r.columns.get(&pid) {
3551 let key = self.index_lookup_key(pid, pk_val);
3552 if probe {
3553 if let Some(old_rid) = self.hot.get(&key) {
3554 if old_rid != r.row_id {
3555 replaced_image = self.get(old_rid, self.snapshot());
3556 self.tombstone_row(
3557 old_rid,
3558 r.committed_epoch,
3559 r.commit_ts,
3560 true,
3561 );
3562 }
3563 }
3564 }
3565 if maintain_pk_by_row {
3566 self.pk_by_row.insert(r.row_id, key);
3567 }
3568 }
3569 }
3570 Some(_) => {}
3571 None => {
3572 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3573 }
3574 }
3575 if let Some(old) = replaced_image {
3576 self.maintain_indexes_on_pk_replace(&old, &r);
3577 } else {
3578 self.index_row(&r);
3579 }
3580 self.reservoir.offer(r.row_id.0);
3581 self.memtable.upsert(r);
3582 self.live_count = self.live_count.saturating_add(1);
3585 }
3586 self.data_generation = self.data_generation.wrapping_add(1);
3587 }
3588
3589 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3599 for &cid in row.columns.keys() {
3600 self.pending_put_cols.insert(cid);
3601 }
3602 let epoch = row.committed_epoch;
3603 let mut replaced_image: Option<Row> = None;
3604 if let Some(pk_col) = self.schema.primary_key() {
3605 let pk_id = pk_col.id;
3606 if let Some(pk_val) = row.columns.get(&pk_id) {
3607 let maintain_pk_by_row = self.pk_by_row_complete;
3611 if check_existing_pk || maintain_pk_by_row {
3612 let key = self.index_lookup_key(pk_id, pk_val);
3613 if check_existing_pk {
3614 if let Some(old_rid) = self.hot.get(&key) {
3615 if old_rid != row.row_id {
3616 replaced_image = self.get(old_rid, self.snapshot());
3620 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
3621 }
3622 }
3623 }
3624 if maintain_pk_by_row {
3625 self.pk_by_row.insert(row.row_id, key);
3626 }
3627 }
3628 }
3629 } else {
3630 self.hot
3631 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3632 }
3633 if let Some(old) = replaced_image {
3634 self.maintain_indexes_on_pk_replace(&old, &row);
3635 } else {
3636 self.index_row(&row);
3637 }
3638 self.reservoir.offer(row.row_id.0);
3639 self.memtable.upsert(row);
3640 self.live_count = self.live_count.saturating_add(1);
3641 self.data_generation = self.data_generation.wrapping_add(1);
3642 }
3643
3644 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
3647 let has_partial = self
3648 .schema
3649 .indexes
3650 .iter()
3651 .any(|idx| idx.predicate.is_some());
3652 if has_partial {
3653 for idef in &self.schema.indexes {
3657 if idef.kind != crate::schema::IndexKind::Bitmap {
3658 continue;
3659 }
3660 if let Some(key) = crate::index::maintain::bitmap_key_for_column(old, idef.column_id)
3661 {
3662 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3663 b.remove(&key, old.row_id);
3664 }
3665 }
3666 }
3667 self.index_row(new);
3668 return;
3669 }
3670
3671 crate::index::maintain_bitmap_secondary_on_replace(
3672 &self.schema,
3673 &mut self.bitmap,
3674 old,
3675 new,
3676 );
3677 self.index_row_non_bitmap(new);
3678 }
3679
3680 fn index_row_non_bitmap(&mut self, row: &Row) {
3683 if row.deleted {
3684 return;
3685 }
3686 let effective = if self.column_keys.is_empty() {
3687 None
3688 } else {
3689 Some(self.tokenized_for_indexes(row))
3690 };
3691 let source = effective.as_ref().unwrap_or(row);
3692 for idef in &self.schema.indexes {
3693 if idef.kind == crate::schema::IndexKind::Bitmap {
3694 continue;
3695 }
3696 index_into_single(
3697 idef,
3698 &self.schema,
3699 source,
3700 &mut self.hot,
3701 &mut self.bitmap,
3702 &mut self.ann,
3703 &mut self.fm,
3704 &mut self.sparse,
3705 &mut self.minhash,
3706 );
3707 }
3708 if let Some(pk_col) = self.schema.primary_key() {
3709 if let Some(pk_val) = source.columns.get(&pk_col.id) {
3710 let key = if self.column_keys.is_empty() {
3711 pk_val.encode_key()
3712 } else {
3713 self.index_lookup_key(pk_col.id, pk_val)
3714 };
3715 self.hot.insert(key, source.row_id);
3716 }
3717 }
3718 }
3719
3720 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3723 self.allocator.alloc()
3724 }
3725
3726 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3732 let pk = self.schema.primary_key().ok_or_else(|| {
3733 MongrelError::Schema("clustered table requires a single-column primary key".into())
3734 })?;
3735 let pk_val = columns
3736 .iter()
3737 .find(|(id, _)| *id == pk.id)
3738 .map(|(_, v)| v)
3739 .ok_or_else(|| {
3740 MongrelError::Schema(format!(
3741 "clustered table missing primary key column {} ({})",
3742 pk.id, pk.name
3743 ))
3744 })?;
3745 Ok(clustered_row_id(pk_val))
3746 }
3747
3748 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3756 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3757 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3758 }
3759 let n = rows.len();
3760 for r in rows {
3761 for &cid in r.columns.keys() {
3762 self.pending_put_cols.insert(cid);
3763 }
3764 }
3765 let (losers, winner_pks) = self.partition_pk_winners(rows);
3766 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3767 let group_ts = rows.first().and_then(|r| r.commit_ts);
3769 for (key, &row_id) in &winner_pks {
3770 if let Some(old_rid) = self.hot.get(key) {
3771 if old_rid != row_id {
3772 self.tombstone_row(old_rid, epoch, group_ts, true);
3773 }
3774 }
3775 }
3776 for &loser_rid in &losers {
3779 self.tombstone_row(loser_rid, epoch, group_ts, false);
3780 }
3781 for (key, row_id) in winner_pks {
3783 self.insert_hot_pk(key, row_id);
3784 }
3785 if self.schema.primary_key().is_none() {
3786 for r in rows {
3787 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3788 }
3789 }
3790 for r in rows {
3791 self.allocator.advance_to(r.row_id)?;
3792 if !losers.contains(&r.row_id) {
3793 self.index_row(r);
3794 }
3795 }
3796 for r in rows {
3797 if !losers.contains(&r.row_id) {
3798 self.reservoir.offer(r.row_id.0);
3799 }
3800 }
3801 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3802 self.data_generation = self.data_generation.wrapping_add(1);
3803 Ok(())
3804 }
3805
3806 pub(crate) fn recover_apply(
3811 &mut self,
3812 rows: Vec<Row>,
3813 deletes: Vec<(RowId, Epoch)>,
3814 ) -> Result<()> {
3815 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3819 std::collections::BTreeMap::new();
3820 for row in rows {
3821 if row.row_id.0 >= u64::MAX - 1 {
3822 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3823 }
3824 self.allocator.advance_to(row.row_id)?;
3825 if let Some(ai) = self.auto_inc.as_mut() {
3830 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3831 let next = n.checked_add(1).ok_or_else(|| {
3832 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3833 })?;
3834 if next > ai.next {
3835 ai.next = next;
3836 }
3837 }
3838 }
3839 by_epoch.entry(row.committed_epoch).or_default().push(row);
3840 }
3841 for (epoch, group) in by_epoch {
3842 let (losers, winner_pks) = self.partition_pk_winners(&group);
3843 let group_ts = group.first().and_then(|r| r.commit_ts);
3845 for (key, &row_id) in &winner_pks {
3846 if let Some(old_rid) = self.hot.get(key) {
3847 if old_rid != row_id {
3848 self.tombstone_row(old_rid, epoch, group_ts, false);
3849 }
3850 }
3851 }
3852 for (key, row_id) in winner_pks {
3853 self.insert_hot_pk(key, row_id);
3854 }
3855 if self.schema.primary_key().is_none() {
3856 for r in &group {
3857 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3858 }
3859 }
3860 for r in &group {
3861 if !losers.contains(&r.row_id) {
3862 self.memtable.upsert(r.clone());
3863 self.index_row(r);
3864 }
3865 }
3866 }
3867 for (rid, epoch) in deletes {
3868 self.memtable.tombstone(rid, epoch);
3869 self.remove_hot_for_row(rid, epoch);
3870 }
3871 self.reservoir_complete = false;
3874 Ok(())
3875 }
3876
3877 pub(crate) fn flushed_epoch(&self) -> u64 {
3879 self.flushed_epoch
3880 }
3881
3882 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3883 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3884 }
3885
3886 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3888 self.schema.validate_values(cells)
3889 }
3890
3891 fn validate_columns_not_null(
3895 &self,
3896 columns: &[(u16, columnar::NativeColumn)],
3897 n: usize,
3898 ) -> Result<()> {
3899 let by_id: HashMap<u16, &columnar::NativeColumn> =
3900 columns.iter().map(|(id, c)| (*id, c)).collect();
3901 for col in &self.schema.columns {
3902 if !col.flags.contains(ColumnFlags::NULLABLE) {
3903 match by_id.get(&col.id) {
3904 None => {
3905 return Err(MongrelError::InvalidArgument(format!(
3906 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3907 col.name, col.id
3908 )));
3909 }
3910 Some(c) => {
3911 if c.null_count(n) != 0 {
3912 return Err(MongrelError::InvalidArgument(format!(
3913 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3914 col.name, col.id
3915 )));
3916 }
3917 }
3918 }
3919 }
3920 if let TypeId::Enum { variants } = &col.ty {
3921 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3922 if by_id.contains_key(&col.id) {
3923 return Err(MongrelError::InvalidArgument(format!(
3924 "column '{}' ({}) enum requires a bytes column",
3925 col.name, col.id
3926 )));
3927 }
3928 continue;
3929 };
3930 for index in 0..n {
3931 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3932 continue;
3933 };
3934 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3935 return Err(MongrelError::InvalidArgument(format!(
3936 "column '{}' ({}) enum value {:?} is not one of {:?}",
3937 col.name,
3938 col.id,
3939 String::from_utf8_lossy(value),
3940 variants
3941 )));
3942 }
3943 }
3944 }
3945 }
3946 Ok(())
3947 }
3948
3949 fn bulk_pk_winner_indices(
3954 &self,
3955 columns: &[(u16, columnar::NativeColumn)],
3956 n: usize,
3957 ) -> Option<Vec<usize>> {
3958 let pk_col = self.schema.primary_key()?;
3959 let pk_id = pk_col.id;
3960 let pk_ty = pk_col.ty.clone();
3961 let by_id: HashMap<u16, &columnar::NativeColumn> =
3962 columns.iter().map(|(id, c)| (*id, c)).collect();
3963 let pk_native = by_id.get(&pk_id)?;
3964 if native_int64_strictly_increasing(pk_native, n) {
3965 return None;
3966 }
3967 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3969 let mut null_pk_rows: Vec<usize> = Vec::new();
3970 for i in 0..n {
3971 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3972 Some(key) => {
3973 last.insert(key, i);
3974 }
3975 None => null_pk_rows.push(i),
3976 }
3977 }
3978 let mut winners: HashSet<usize> = last.values().copied().collect();
3979 for i in null_pk_rows {
3980 winners.insert(i);
3981 }
3982 Some((0..n).filter(|i| winners.contains(i)).collect())
3983 }
3984
3985 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3987 self.require_delete()?;
3988 let epoch = self.pending_epoch();
3989 self.wal_append_data(Op::Delete {
3990 table_id: self.table_id,
3991 row_ids: vec![row_id],
3992 })?;
3993 if self.is_shared() {
3994 self.pending_dels.push(row_id);
3995 } else {
3996 self.apply_delete(row_id, epoch);
3997 }
3998 Ok(())
3999 }
4000
4001 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
4002 let pre = self.get(row_id, self.snapshot());
4003 self.delete(row_id)?;
4004 Ok(pre.map(|row| {
4005 let mut columns: Vec<_> = row.columns.into_iter().collect();
4006 columns.sort_by_key(|(id, _)| *id);
4007 OwnedRow { columns }
4008 }))
4009 }
4010
4011 pub fn truncate(&mut self) -> Result<()> {
4013 self.require_delete()?;
4014 let epoch = self.pending_epoch();
4015 self.wal_append_data(Op::TruncateTable {
4016 table_id: self.table_id,
4017 })?;
4018 self.pending_rows.clear();
4019 self.pending_rows_auto_inc.clear();
4020 self.pending_dels.clear();
4021 self.pending_truncate = Some(epoch);
4022 Ok(())
4023 }
4024
4025 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
4027 self.run_refs.clear();
4032 self.retiring.clear();
4033 self.memtable = Memtable::new();
4034 self.mutable_run = MutableRun::new();
4035 self.hot = HotIndex::new();
4036 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
4037 self.bitmap = bitmap;
4038 self.ann = ann;
4039 self.fm = fm;
4040 self.sparse = sparse;
4041 self.minhash = minhash;
4042 self.learned_range = Arc::new(HashMap::new());
4043 self.pk_by_row.clear();
4044 self.pk_by_row_complete = false;
4045 self.live_count = 0;
4046 self.reservoir = crate::reservoir::Reservoir::default();
4047 self.reservoir_complete = true;
4048 self.had_deletes = true;
4049 self.agg_cache = Arc::new(HashMap::new());
4050 self.global_idx_epoch = 0;
4051 self.indexes_complete = true;
4052 self.pending_delete_rids.clear();
4053 self.pending_put_cols.clear();
4054 self.pending_rows.clear();
4055 self.pending_rows_auto_inc.clear();
4056 self.pending_dels.clear();
4057 self.clear_result_cache();
4058 self.invalidate_index_checkpoint();
4059 self.data_generation = self.data_generation.wrapping_add(1);
4060 }
4061
4062 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
4065 self.apply_delete_at(row_id, epoch, None);
4066 }
4067
4068 pub(crate) fn apply_delete_at(
4070 &mut self,
4071 row_id: RowId,
4072 epoch: Epoch,
4073 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4074 ) {
4075 let preimage = self.get(row_id, self.snapshot());
4082 self.remove_hot_for_row(row_id, epoch);
4083 if let Some(row) = preimage.as_ref() {
4084 self.unindex_bitmap_membership(row);
4085 }
4086 self.tombstone_row(row_id, epoch, commit_ts, true);
4087 self.data_generation = self.data_generation.wrapping_add(1);
4088 }
4089
4090 fn unindex_bitmap_membership(&mut self, row: &Row) {
4094 if row.deleted {
4095 return;
4096 }
4097 for idef in &self.schema.indexes {
4098 if idef.kind != crate::schema::IndexKind::Bitmap {
4099 continue;
4100 }
4101 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
4102 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4103 b.remove(&key, row.row_id);
4104 }
4105 }
4106 }
4107 }
4108
4109 fn union_bitmap_point_i64(
4113 &self,
4114 set: &mut RowIdSet,
4115 column_id: u16,
4116 value: i64,
4117 snapshot: Snapshot,
4118 ) {
4119 let Some(b) = self.bitmap.get(&column_id) else {
4120 return;
4121 };
4122 let encoded = Value::Int64(value).encode_key();
4123 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
4124 for rid in b.get(&lookup).iter() {
4125 set.insert(u64::from(rid));
4126 }
4127 set.remove_many(self.overlay_tombstoned_rids(snapshot));
4131 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
4132 }
4133
4134 fn tombstone_row(
4138 &mut self,
4139 row_id: RowId,
4140 epoch: Epoch,
4141 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4142 adjust_live_count: bool,
4143 ) {
4144 let tombstone = Row {
4145 row_id,
4146 committed_epoch: epoch,
4147 columns: std::collections::HashMap::new(),
4148 deleted: true,
4149 commit_ts,
4150 };
4151 self.memtable.upsert(tombstone);
4152 self.pk_by_row.remove(&row_id);
4153 if adjust_live_count {
4154 self.live_count = self.live_count.saturating_sub(1);
4155 }
4156 self.pending_delete_rids.insert(row_id.0 as u32);
4158 self.had_deletes = true;
4161 self.agg_cache = Arc::new(HashMap::new());
4162 }
4163
4164 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
4168 let Some(pk_col) = self.schema.primary_key() else {
4169 return;
4170 };
4171 if self.pk_by_row_complete {
4174 if let Some(key) = self.pk_by_row.remove(&row_id) {
4175 if self.hot.get(&key) == Some(row_id) {
4176 self.hot.remove(&key);
4177 }
4178 }
4179 return;
4180 }
4181 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
4200 if self.indexes_complete {
4201 let pk_val = self
4202 .memtable
4203 .get_version(row_id, lookup_epoch)
4204 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4205 .or_else(|| {
4206 self.mutable_run
4207 .get_version(row_id, lookup_epoch)
4208 .filter(|(_, r)| !r.deleted)
4209 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4210 })
4211 .or_else(|| {
4212 self.run_refs.iter().find_map(|rr| {
4213 let mut reader = self.open_reader(rr.run_id).ok()?;
4214 let (_, deleted, val) = reader
4215 .get_version_column(row_id, lookup_epoch, pk_col.id)
4216 .ok()??;
4217 if deleted {
4218 return None;
4219 }
4220 val
4221 })
4222 });
4223 if let Some(pk_val) = pk_val {
4224 let key = self.index_lookup_key(pk_col.id, &pk_val);
4225 if self.hot.get(&key) == Some(row_id) {
4226 self.hot.remove(&key);
4227 }
4228 return;
4229 }
4230 }
4231 self.refresh_pk_by_row_from_hot();
4236 if let Some(key) = self.pk_by_row.remove(&row_id) {
4237 if self.hot.get(&key) == Some(row_id) {
4238 self.hot.remove(&key);
4239 }
4240 }
4241 }
4242
4243 fn partition_pk_winners(
4248 &self,
4249 rows: &[Row],
4250 ) -> (
4251 std::collections::HashSet<RowId>,
4252 std::collections::HashMap<Vec<u8>, RowId>,
4253 ) {
4254 let mut losers = std::collections::HashSet::new();
4255 let Some(pk_col) = self.schema.primary_key() else {
4256 return (losers, std::collections::HashMap::new());
4257 };
4258 let pk_id = pk_col.id;
4259 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
4260 std::collections::HashMap::new();
4261 for r in rows {
4262 let Some(pk_val) = r.columns.get(&pk_id) else {
4263 continue;
4264 };
4265 let key = self.index_lookup_key(pk_id, pk_val);
4266 if let Some(&old_rid) = winners.get(&key) {
4267 losers.insert(old_rid);
4268 }
4269 winners.insert(key, r.row_id);
4270 }
4271 (losers, winners)
4272 }
4273
4274 fn index_row(&mut self, row: &Row) {
4275 if row.deleted {
4276 return;
4277 }
4278 let any_predicate = self
4286 .schema
4287 .indexes
4288 .iter()
4289 .any(|idx| idx.predicate.is_some());
4290 if any_predicate {
4291 let columns_map: HashMap<u16, &Value> =
4292 row.columns.iter().map(|(k, v)| (*k, v)).collect();
4293 let name_to_id: HashMap<&str, u16> = self
4294 .schema
4295 .columns
4296 .iter()
4297 .map(|c| (c.name.as_str(), c.id))
4298 .collect();
4299 for idx in &self.schema.indexes {
4300 if let Some(pred) = &idx.predicate {
4301 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4302 continue; }
4304 }
4305 index_into_single(
4307 idx,
4308 &self.schema,
4309 row,
4310 &mut self.hot,
4311 &mut self.bitmap,
4312 &mut self.ann,
4313 &mut self.fm,
4314 &mut self.sparse,
4315 &mut self.minhash,
4316 );
4317 }
4318 return;
4319 }
4320 if self.column_keys.is_empty() {
4324 index_into(
4325 &self.schema,
4326 row,
4327 &mut self.hot,
4328 &mut self.bitmap,
4329 &mut self.ann,
4330 &mut self.fm,
4331 &mut self.sparse,
4332 &mut self.minhash,
4333 );
4334 return;
4335 }
4336 let effective_row = self.tokenized_for_indexes(row);
4337 index_into(
4338 &self.schema,
4339 &effective_row,
4340 &mut self.hot,
4341 &mut self.bitmap,
4342 &mut self.ann,
4343 &mut self.fm,
4344 &mut self.sparse,
4345 &mut self.minhash,
4346 );
4347 }
4348
4349 fn tokenized_for_indexes(&self, row: &Row) -> Row {
4355 if self.column_keys.is_empty() {
4356 return row.clone();
4357 }
4358 {
4359 use crate::encryption::SCHEME_HMAC_EQ;
4360 let mut tok = row.clone();
4361 for (&cid, &(_, scheme)) in &self.column_keys {
4362 if scheme != SCHEME_HMAC_EQ {
4363 continue;
4364 }
4365 if let Some(v) = tok.columns.get(&cid).cloned() {
4366 if let Some(t) = self.tokenize_value(cid, &v) {
4367 tok.columns.insert(cid, t);
4368 }
4369 }
4370 }
4371 tok
4372 }
4373 }
4374
4375 pub fn commit(&mut self) -> Result<Epoch> {
4380 self.commit_inner(None)
4381 }
4382
4383 #[doc(hidden)]
4386 pub fn commit_controlled<F>(
4387 &mut self,
4388 control: &crate::ExecutionControl,
4389 mut before_commit: F,
4390 ) -> Result<Epoch>
4391 where
4392 F: FnMut() -> Result<()>,
4393 {
4394 self.commit_inner(Some((control, &mut before_commit)))
4395 }
4396
4397 fn commit_inner(
4398 &mut self,
4399 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4400 ) -> Result<Epoch> {
4401 self.ensure_writable()?;
4402 if !self.has_pending_mutations() {
4403 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
4404 return Err(MongrelError::Full(
4405 "standalone transaction id namespace exhausted".into(),
4406 ));
4407 }
4408 return Ok(self.epoch.visible());
4409 }
4410 self.commit_new_epoch_inner(controlled)
4411 }
4412
4413 fn commit_new_epoch(&mut self) -> Result<Epoch> {
4417 self.commit_new_epoch_inner(None)
4418 }
4419
4420 fn commit_new_epoch_inner(
4421 &mut self,
4422 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4423 ) -> Result<Epoch> {
4424 self.ensure_writable()?;
4425 if self.is_shared() {
4426 self.commit_shared(controlled)
4427 } else {
4428 self.commit_private(controlled)
4429 }
4430 }
4431
4432 fn commit_private(
4434 &mut self,
4435 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4436 ) -> Result<Epoch> {
4437 let commit_lock = Arc::clone(&self.commit_lock);
4441 let _g = commit_lock.lock();
4442 let txn_id = self.ensure_txn_id()?;
4445 if let Some((control, before_commit)) = controlled {
4446 control.checkpoint()?;
4447 before_commit()?;
4448 }
4449 let new_epoch = self.epoch.bump_assigned();
4450 let epoch_authority = Arc::clone(&self.epoch);
4451 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4452 let wal_result = match &mut self.wal {
4456 WalSink::Private(w) => w
4457 .append_txn(
4458 txn_id,
4459 Op::TxnCommit {
4460 epoch: new_epoch.0,
4461 added_runs: Vec::new(),
4462 },
4463 )
4464 .and_then(|_| w.sync()),
4465 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4466 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4467 };
4468 if let Err(error) = wal_result {
4469 self.durable_commit_failed = true;
4470 return Err(MongrelError::CommitOutcomeUnknown {
4471 epoch: new_epoch.0,
4472 message: error.to_string(),
4473 });
4474 }
4475 if let Some(epoch) = self.pending_truncate.take() {
4478 self.apply_truncate(epoch);
4479 }
4480 self.invalidate_pending_cache();
4481 let publish_result = self.persist_manifest(new_epoch);
4482 self.epoch.publish_in_order(new_epoch);
4486 epoch_guard.disarm();
4487 if let Err(error) = publish_result {
4488 self.durable_commit_failed = true;
4489 return Err(MongrelError::DurableCommit {
4490 epoch: new_epoch.0,
4491 message: error.to_string(),
4492 });
4493 }
4494 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4495 self.pending_private_mutations = false;
4496 self.data_generation = self.data_generation.wrapping_add(1);
4497 Ok(new_epoch)
4498 }
4499
4500 fn commit_shared(
4506 &mut self,
4507 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4508 ) -> Result<Epoch> {
4509 use std::sync::atomic::Ordering;
4510 let s = match &self.wal {
4511 WalSink::Shared(s) => s.clone(),
4512 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4513 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4514 };
4515 if s.poisoned.load(Ordering::Relaxed) {
4516 return Err(MongrelError::Other(
4517 "database poisoned by fsync error".into(),
4518 ));
4519 }
4520 let commit_lock = Arc::clone(&self.commit_lock);
4527 let _g = commit_lock.lock();
4528 if !self.pending_rows.is_empty() {
4529 match controlled.as_ref() {
4530 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4531 None => self.prepare_durable_publish()?,
4532 }
4533 }
4534 let txn_id = self.ensure_txn_id()?;
4537 let mut wal = s.wal.lock();
4538 if let Some((control, before_commit)) = controlled {
4539 control.checkpoint()?;
4540 before_commit()?;
4541 }
4542 let new_epoch = self.epoch.bump_assigned();
4543 let epoch_authority = Arc::clone(&self.epoch);
4544 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4545 let commit_ts = s.hlc.now().map_err(|skew| {
4548 MongrelError::Other(format!(
4549 "clock skew rejected commit timestamp allocation: {skew}"
4550 ))
4551 })?;
4552 let commit_seq = match wal.append_commit_at(
4553 txn_id,
4554 new_epoch,
4555 &[],
4556 commit_ts.physical_micros.saturating_mul(1_000),
4557 ) {
4558 Ok(commit_seq) => commit_seq,
4559 Err(error) => {
4560 s.poisoned.store(true, Ordering::Relaxed);
4561 s.lifecycle.poison();
4562 return Err(MongrelError::CommitOutcomeUnknown {
4563 epoch: new_epoch.0,
4564 message: error.to_string(),
4565 });
4566 }
4567 };
4568 drop(wal);
4569 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4570 s.poisoned.store(true, Ordering::Relaxed);
4571 s.lifecycle.poison();
4572 return Err(MongrelError::CommitOutcomeUnknown {
4573 epoch: new_epoch.0,
4574 message: error.to_string(),
4575 });
4576 }
4577
4578 if self.pending_truncate.take().is_some() {
4581 self.apply_truncate(new_epoch);
4582 }
4583 let mut rows = std::mem::take(&mut self.pending_rows);
4584 if !rows.is_empty() {
4585 for r in &mut rows {
4586 r.committed_epoch = new_epoch;
4587 r.commit_ts = Some(commit_ts);
4588 }
4589 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4590 let all_auto_generated =
4591 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4592 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4593 } else {
4594 self.pending_rows_auto_inc.clear();
4595 }
4596 let dels = std::mem::take(&mut self.pending_dels);
4597 for rid in dels {
4598 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
4599 }
4600
4601 self.invalidate_pending_cache();
4602 let publish_result = self.persist_manifest(new_epoch);
4603 self.epoch.publish_in_order(new_epoch);
4604 epoch_guard.disarm();
4605 let _ = s.change_wake.send(());
4606 if let Err(error) = publish_result {
4607 self.durable_commit_failed = true;
4608 s.poisoned.store(true, Ordering::Relaxed);
4609 s.lifecycle.poison();
4610 return Err(MongrelError::DurableCommit {
4611 epoch: new_epoch.0,
4612 message: error.to_string(),
4613 });
4614 }
4615 self.current_txn_id = 0;
4617 self.data_generation = self.data_generation.wrapping_add(1);
4618 Ok(new_epoch)
4619 }
4620
4621 pub fn flush(&mut self) -> Result<Epoch> {
4629 self.flush_with_outcome().map(|(epoch, _)| epoch)
4630 }
4631
4632 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4634 self.flush_with_outcome_inner(None)
4635 }
4636
4637 #[doc(hidden)]
4640 pub fn flush_with_outcome_controlled<F>(
4641 &mut self,
4642 control: &crate::ExecutionControl,
4643 mut before_commit: F,
4644 ) -> Result<(Epoch, bool)>
4645 where
4646 F: FnMut() -> Result<()>,
4647 {
4648 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4649 }
4650
4651 fn flush_with_outcome_inner(
4652 &mut self,
4653 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4654 ) -> Result<(Epoch, bool)> {
4655 match controlled.as_ref() {
4656 Some((control, _)) => {
4657 self.ensure_indexes_complete_controlled(control, || true)?;
4658 }
4659 None => self.ensure_indexes_complete()?,
4660 }
4661 let committed = self.has_pending_mutations();
4662 let epoch = self.commit_inner(controlled)?;
4663 let finish: Result<(Epoch, bool)> = (|| {
4664 let rows = self.memtable.drain_sorted();
4665 if !rows.is_empty() {
4666 self.mutable_run.insert_many(rows);
4667 }
4668 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4669 self.spill_mutable_run(epoch)?;
4670 self.mark_flushed(epoch)?;
4674 self.persist_manifest(epoch)?;
4675 self.build_learned_ranges()?;
4676 self.checkpoint_indexes(epoch);
4679 }
4680 Ok((epoch, committed))
4683 })();
4684 let outcome = match finish {
4685 Err(error) if committed => Err(MongrelError::DurableCommit {
4686 epoch: epoch.0,
4687 message: error.to_string(),
4688 }),
4689 result => result,
4690 };
4691 if outcome.is_ok() {
4692 let _ = self.publish_read_generation();
4698 }
4699 outcome
4700 }
4701
4702 fn has_pending_mutations(&self) -> bool {
4703 self.pending_private_mutations
4704 || !self.pending_rows.is_empty()
4705 || !self.pending_dels.is_empty()
4706 || self.pending_truncate.is_some()
4707 }
4708
4709 pub fn has_pending_writes(&self) -> bool {
4710 self.has_pending_mutations()
4711 }
4712
4713 pub fn force_flush(&mut self) -> Result<Epoch> {
4722 let saved = self.mutable_run_spill_bytes;
4723 self.mutable_run_spill_bytes = 1;
4724 let result = self.flush();
4725 self.mutable_run_spill_bytes = saved;
4726 result
4727 }
4728
4729 pub fn close(&mut self) -> Result<()> {
4736 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4737 self.force_flush()?;
4738 }
4739 Ok(())
4740 }
4741
4742 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4749 let op = Op::Flush {
4750 table_id: self.table_id,
4751 flushed_epoch: epoch.0,
4752 };
4753 match &mut self.wal {
4754 WalSink::Private(w) => {
4755 w.append_system(op)?;
4756 w.sync()?;
4757 }
4758 WalSink::Shared(s) => {
4759 s.wal.lock().append_system(op)?;
4764 }
4765 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4766 }
4767 self.flushed_epoch = epoch.0;
4768 if matches!(self.wal, WalSink::Private(_)) {
4769 self.rotate_wal(epoch)?;
4770 }
4771 Ok(())
4772 }
4773
4774 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4778 if self.mutable_run.is_empty() {
4779 return Ok(());
4780 }
4781 let run_id = self.alloc_run_id()?;
4782 let rows = self.mutable_run.drain_sorted();
4783 if rows.is_empty() {
4784 return Ok(());
4785 }
4786 let path = self.run_path(run_id);
4787 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4788 if let Some(kek) = &self.kek {
4789 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4790 }
4791 let header = match self.create_run_file(run_id)? {
4792 Some(file) => writer.write_file(file, &rows)?,
4793 None => writer.write(&path, &rows)?,
4794 };
4795 self.run_refs.push(RunRef {
4796 run_id: run_id as u128,
4797 level: 0,
4798 epoch_created: epoch.0,
4799 row_count: header.row_count,
4800 });
4801 Ok(())
4802 }
4803
4804 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4808 self.mutable_run_spill_bytes = bytes.max(1);
4809 }
4810
4811 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4815 self.compaction_zstd_level = level;
4816 }
4817
4818 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4822 self.result_cache.lock().set_max_bytes(max_bytes);
4823 }
4824
4825 pub(crate) fn clear_result_cache(&mut self) {
4829 self.result_cache.lock().clear();
4830 }
4831
4832 pub fn mutable_run_len(&self) -> usize {
4834 self.mutable_run.len()
4835 }
4836
4837 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4840 self.mutable_run.drain_sorted()
4841 }
4842
4843 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4845 let mut snapshot = self.mutable_run.clone();
4846 snapshot.drain_sorted()
4847 }
4848
4849 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4854 self.ensure_writable()?;
4855 let n = batch.len();
4856 if n == 0 {
4857 return Ok(self.current_epoch());
4858 }
4859 for row in &batch {
4860 self.schema.validate_values(row)?;
4861 }
4862 let epoch = self.commit_new_epoch()?;
4863 let live_before = self.live_count;
4864 self.spill_mutable_run(epoch)?;
4868 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4869 && self.indexes_complete
4870 && self.run_refs.is_empty()
4871 && self.memtable.is_empty()
4872 && self.mutable_run.is_empty();
4873 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4879 use rayon::prelude::*;
4880 self.schema
4881 .columns
4882 .par_iter()
4883 .map(|cdef| {
4884 (
4885 cdef.id,
4886 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4887 )
4888 })
4889 .collect::<Vec<_>>()
4890 };
4891 drop(batch);
4892 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4897 self.validate_columns_not_null(&user_columns, n)?;
4898 let winner_idx = self
4899 .bulk_pk_winner_indices(&user_columns, n)
4900 .filter(|idx| idx.len() != n);
4901 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4902 match winner_idx.as_deref() {
4903 Some(idx) => {
4904 let compacted = user_columns
4905 .iter()
4906 .map(|(id, c)| (*id, c.gather(idx)))
4907 .collect();
4908 (compacted, idx.len())
4909 }
4910 None => (std::mem::take(&mut user_columns), n),
4911 };
4912 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4913 let first = self.allocator.alloc_range(write_n as u64)?.0;
4914 for rid in first..first + write_n as u64 {
4915 self.reservoir.offer(rid);
4916 }
4917 let run_id = self.alloc_run_id()?;
4918 let path = self.run_path(run_id);
4919 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4920 .clean(true)
4921 .with_lz4()
4922 .with_native_endian();
4923 if let Some(kek) = &self.kek {
4924 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4925 }
4926 let header = match self.create_run_file(run_id)? {
4927 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4928 None => writer.write_native(&path, &write_columns, write_n, first)?,
4929 };
4930 self.run_refs.push(RunRef {
4931 run_id: run_id as u128,
4932 level: 0,
4933 epoch_created: epoch.0,
4934 row_count: header.row_count,
4935 });
4936 self.live_count = self.live_count.saturating_add(write_n as u64);
4937 if eager_index_build {
4938 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4939 self.index_columns_bulk(&write_columns, &row_ids);
4940 self.indexes_complete = true;
4941 self.build_learned_ranges()?;
4942 } else {
4943 self.indexes_complete = false;
4944 }
4945 self.mark_flushed(epoch)?;
4946 self.persist_manifest(epoch)?;
4947 if eager_index_build {
4948 self.checkpoint_indexes(epoch);
4949 }
4950 self.clear_result_cache();
4951 Ok(epoch)
4952 }
4953
4954 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4957 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4958 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4959 let segment_no = segment
4962 .file_stem()
4963 .and_then(|s| s.to_str())
4964 .and_then(|s| s.strip_prefix("seg-"))
4965 .and_then(|s| s.parse::<u64>().ok())
4966 .unwrap_or(0);
4967 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4968 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4969 wal.sync()?;
4970 self.wal = WalSink::Private(wal);
4971 Ok(())
4972 }
4973
4974 pub(crate) fn invalidate_pending_cache(&mut self) {
4979 self.result_cache
4980 .lock()
4981 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4982 self.pending_delete_rids.clear();
4983 self.pending_put_cols.clear();
4984 }
4985
4986 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4987 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4988 m.current_epoch = epoch.0;
4989 m.next_row_id = self.allocator.current().0;
4990 m.runs = self.run_refs.clone();
4991 m.live_count = self.live_count;
4992 m.global_idx_epoch = self.global_idx_epoch;
4993 m.flushed_epoch = self.flushed_epoch;
4994 m.retiring = self.retiring.clone();
4995 m.auto_inc_next = match self.auto_inc {
4999 Some(ai) if ai.seeded => ai.next,
5000 _ => 0,
5001 };
5002 m.ttl = self.ttl;
5003 let meta_dek = self.manifest_meta_dek();
5004 match self._root_guard.as_deref() {
5005 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
5006 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
5007 }
5008 Ok(())
5009 }
5010
5011 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
5012 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
5016 let live_count = u64::try_from(rows.len())
5017 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
5018 let auto_inc = match self.auto_inc {
5019 Some(mut state) => {
5020 let maximum = self.scan_max_int64(state.column_id)?;
5021 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
5022 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
5023 })?;
5024 state.next = state.next.max(after_maximum).max(1);
5025 state.seeded = true;
5026 Some(state)
5027 }
5028 None => None,
5029 };
5030 Ok(RecoveryMetadataPlan {
5031 live_count,
5032 auto_inc,
5033 changed: live_count != self.live_count
5034 || auto_inc.is_some_and(|planned| {
5035 self.auto_inc.is_none_or(|current| {
5036 current.next != planned.next || current.seeded != planned.seeded
5037 })
5038 }),
5039 })
5040 }
5041
5042 pub(crate) fn apply_recovered_metadata(
5043 &mut self,
5044 plan: RecoveryMetadataPlan,
5045 epoch: Epoch,
5046 ) -> Result<()> {
5047 if !plan.changed {
5048 return Ok(());
5049 }
5050 self.live_count = plan.live_count;
5051 self.auto_inc = plan.auto_inc;
5052 self.persist_manifest(epoch)
5053 }
5054
5055 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
5061 if !self.indexes_complete {
5064 return;
5065 }
5066 if crate::catalog::inject_hook("index.publish.before").is_err() {
5069 return;
5070 }
5071 if self.idx_root.is_none() {
5072 if let Some(root) = self._root_guard.as_ref() {
5073 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
5074 return;
5075 };
5076 self.idx_root = Some(Arc::new(idx_root));
5077 }
5078 }
5079 let snap = global_idx::IndexSnapshot {
5080 hot: &self.hot,
5081 bitmap: &self.bitmap,
5082 ann: &self.ann,
5083 fm: &self.fm,
5084 sparse: &self.sparse,
5085 minhash: &self.minhash,
5086 learned_range: &self.learned_range,
5087 };
5088 let idx_dek = self.idx_dek();
5090 let written = match self.idx_root.as_deref() {
5091 Some(root) => global_idx::write_atomic_root(
5092 root,
5093 self.table_id,
5094 epoch.0,
5095 snap,
5096 idx_dek.as_deref(),
5097 ),
5098 None => global_idx::write_atomic(
5099 &self.dir,
5100 self.table_id,
5101 epoch.0,
5102 snap,
5103 idx_dek.as_deref(),
5104 ),
5105 };
5106 if written.is_ok() {
5107 self.global_idx_epoch = epoch.0;
5108 let _ = self.persist_manifest(epoch);
5109 let _ = crate::catalog::inject_hook("index.publish.after");
5111 }
5112 }
5113
5114 pub(crate) fn invalidate_index_checkpoint(&mut self) {
5117 self.global_idx_epoch = 0;
5118 if let Some(root) = self.idx_root.as_deref() {
5119 let _ = root.remove_file(global_idx::IDX_FILENAME);
5120 } else {
5121 global_idx::remove(&self.dir);
5122 }
5123 let _ = self.persist_manifest(self.epoch.visible());
5124 }
5125
5126 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
5131 self.indexes_complete = false;
5132 self.global_idx_epoch = 0;
5133 if let Some(root) = self.idx_root.as_deref() {
5134 let _ = root.remove_file(global_idx::IDX_FILENAME);
5135 } else {
5136 global_idx::remove(&self.dir);
5137 }
5138 }
5139
5140 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
5141 self.indexes_complete = true;
5142 }
5143
5144 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
5150 self.durable_commit_failed = true;
5151 if let WalSink::Shared(shared) = &self.wal {
5152 shared
5153 .poisoned
5154 .store(true, std::sync::atomic::Ordering::Relaxed);
5155 }
5156 }
5157
5158 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
5162 self.durable_commit_failed = true;
5163 }
5164
5165 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
5174 let mut best: Option<Row> = None;
5175 let mut consider = |row: Row| {
5176 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5177 return;
5178 }
5179 if best.as_ref().is_none_or(|current| {
5180 Snapshot::version_is_newer(
5181 row.committed_epoch,
5182 row.commit_ts,
5183 current.committed_epoch,
5184 current.commit_ts,
5185 )
5186 }) {
5187 best = Some(row);
5188 }
5189 };
5190 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
5191 consider(row);
5192 }
5193 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
5194 consider(row);
5195 }
5196 for rr in &self.run_refs {
5197 let Ok(mut reader) = self.open_reader(rr.run_id) else {
5198 continue;
5199 };
5200 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
5203 continue;
5204 };
5205 consider(row);
5206 }
5207 let now_nanos = unix_nanos_now();
5208 match best {
5209 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
5210 Some(r) => Some(r),
5211 None => None,
5212 }
5213 }
5214
5215 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
5219 self.visible_rows_at_time(snapshot, unix_nanos_now())
5220 }
5221
5222 #[doc(hidden)]
5225 pub fn visible_rows_controlled(
5226 &self,
5227 snapshot: Snapshot,
5228 control: &crate::ExecutionControl,
5229 ) -> Result<Vec<Row>> {
5230 let mut out = Vec::new();
5231 self.for_each_visible_row_controlled(snapshot, control, |row| {
5232 out.push(row);
5233 Ok(())
5234 })?;
5235 Ok(out)
5236 }
5237
5238 #[doc(hidden)]
5241 pub fn for_each_visible_row_controlled<F>(
5242 &self,
5243 snapshot: Snapshot,
5244 control: &crate::ExecutionControl,
5245 mut visit: F,
5246 ) -> Result<()>
5247 where
5248 F: FnMut(Row) -> Result<()>,
5249 {
5250 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
5251 control.checkpoint()?;
5252 let memtable = self.memtable.visible_versions_at(snapshot);
5254 if !memtable.is_empty() {
5255 sources.push(ControlledVisibleSource::memory(memtable));
5256 }
5257 control.checkpoint()?;
5258 let mutable = self.mutable_run.visible_versions_at(snapshot);
5261 if !mutable.is_empty() {
5262 sources.push(ControlledVisibleSource::memory(mutable));
5263 }
5264 for run in &self.run_refs {
5265 control.checkpoint()?;
5266 let reader = self.open_reader(run.run_id)?;
5267 sources.push(ControlledVisibleSource::run(
5269 reader.into_visible_version_cursor(snapshot.epoch)?,
5270 ));
5271 }
5272 let now_nanos = unix_nanos_now();
5273 merge_controlled_visible_sources(
5274 &mut sources,
5275 control,
5276 |row| self.row_expired_at(row, now_nanos),
5277 |row| {
5278 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5281 return Ok(());
5282 }
5283 visit(row)
5284 },
5285 )
5286 }
5287
5288 #[doc(hidden)]
5289 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
5290 let mut best: HashMap<u64, Row> = HashMap::new();
5291 let mut fold = |row: Row| {
5292 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5293 return;
5294 }
5295 best.entry(row.row_id.0)
5296 .and_modify(|existing| {
5297 if Snapshot::version_is_newer(
5298 row.committed_epoch,
5299 row.commit_ts,
5300 existing.committed_epoch,
5301 existing.commit_ts,
5302 ) {
5303 *existing = row.clone();
5304 }
5305 })
5306 .or_insert(row);
5307 };
5308 for row in self.memtable.visible_versions_at(snapshot) {
5309 fold(row);
5310 }
5311 for row in self.mutable_run.visible_versions_at(snapshot) {
5312 fold(row);
5313 }
5314 for rr in &self.run_refs {
5315 let mut reader = self.open_reader(rr.run_id)?;
5316 for row in reader.visible_versions(snapshot.epoch)? {
5318 fold(row);
5319 }
5320 }
5321 let mut out: Vec<Row> = best
5322 .into_values()
5323 .filter_map(|r| {
5324 if r.deleted || self.row_expired_at(&r, now_nanos) {
5325 None
5326 } else {
5327 Some(r)
5328 }
5329 })
5330 .collect();
5331 out.sort_by_key(|r| r.row_id);
5332 Ok(out)
5333 }
5334
5335 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
5342 if self.ttl.is_none()
5343 && self.memtable.is_empty()
5344 && self.mutable_run.is_empty()
5345 && self.run_refs.len() == 1
5346 {
5347 let rr = self.run_refs[0].clone();
5348 let mut reader = self.open_reader(rr.run_id)?;
5349 let idxs = reader.visible_indices(snapshot.epoch)?;
5350 let mut cols = Vec::with_capacity(self.schema.columns.len());
5351 for cdef in &self.schema.columns {
5352 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
5353 }
5354 return Ok(cols);
5355 }
5356 let rows = self.visible_rows(snapshot)?;
5358 let mut cols: Vec<(u16, Vec<Value>)> = self
5359 .schema
5360 .columns
5361 .iter()
5362 .map(|c| (c.id, Vec::with_capacity(rows.len())))
5363 .collect();
5364 for r in &rows {
5365 for (cid, vec) in cols.iter_mut() {
5366 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
5367 }
5368 }
5369 Ok(cols)
5370 }
5371
5372 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
5374 let row_id = self.hot.get(key)?;
5375 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
5376 Some(row_id)
5377 } else {
5378 None
5379 }
5380 }
5381
5382 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
5387 self.query_at_with_allowed(q, self.snapshot(), None)
5388 }
5389
5390 pub fn query_controlled(
5393 &mut self,
5394 q: &crate::query::Query,
5395 control: &crate::ExecutionControl,
5396 ) -> Result<Vec<Row>> {
5397 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
5398 }
5399
5400 pub fn query_at_with_allowed(
5403 &mut self,
5404 q: &crate::query::Query,
5405 snapshot: Snapshot,
5406 allowed: Option<&std::collections::HashSet<RowId>>,
5407 ) -> Result<Vec<Row>> {
5408 self.query_at_with_allowed_after(q, snapshot, allowed, None)
5409 }
5410
5411 #[doc(hidden)]
5412 pub fn query_at_with_allowed_controlled(
5413 &mut self,
5414 q: &crate::query::Query,
5415 snapshot: Snapshot,
5416 allowed: Option<&std::collections::HashSet<RowId>>,
5417 control: &crate::ExecutionControl,
5418 ) -> Result<Vec<Row>> {
5419 self.require_select()?;
5420 self.ensure_indexes_complete_controlled(control, || true)?;
5421 self.validate_native_query(q)?;
5422 self.query_conditions_at(
5423 &q.conditions,
5424 snapshot,
5425 allowed,
5426 q.limit,
5427 q.offset,
5428 None,
5429 unix_nanos_now(),
5430 Some(control),
5431 )
5432 }
5433
5434 #[doc(hidden)]
5435 pub fn query_at_with_allowed_after(
5436 &mut self,
5437 q: &crate::query::Query,
5438 snapshot: Snapshot,
5439 allowed: Option<&std::collections::HashSet<RowId>>,
5440 after_row_id: Option<RowId>,
5441 ) -> Result<Vec<Row>> {
5442 self.query_at_with_allowed_after_at_time(
5443 q,
5444 snapshot,
5445 allowed,
5446 after_row_id,
5447 unix_nanos_now(),
5448 )
5449 }
5450
5451 #[doc(hidden)]
5452 pub fn query_at_with_allowed_after_at_time(
5453 &mut self,
5454 q: &crate::query::Query,
5455 snapshot: Snapshot,
5456 allowed: Option<&std::collections::HashSet<RowId>>,
5457 after_row_id: Option<RowId>,
5458 query_time_nanos: i64,
5459 ) -> Result<Vec<Row>> {
5460 self.require_select()?;
5461 self.ensure_indexes_complete()?;
5462 self.validate_native_query(q)?;
5463 self.query_conditions_at(
5464 &q.conditions,
5465 snapshot,
5466 allowed,
5467 q.limit,
5468 q.offset,
5469 after_row_id,
5470 query_time_nanos,
5471 None,
5472 )
5473 }
5474
5475 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
5476 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5477 return Err(MongrelError::InvalidArgument(format!(
5478 "query exceeds {} conditions",
5479 crate::query::MAX_HARD_CONDITIONS
5480 )));
5481 }
5482 if let Some(limit) = q.limit {
5483 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
5484 return Err(MongrelError::InvalidArgument(format!(
5485 "query limit must be between 1 and {}",
5486 crate::query::MAX_FINAL_LIMIT
5487 )));
5488 }
5489 }
5490 if q.offset > crate::query::MAX_QUERY_OFFSET {
5491 return Err(MongrelError::InvalidArgument(format!(
5492 "query offset exceeds {}",
5493 crate::query::MAX_QUERY_OFFSET
5494 )));
5495 }
5496 Ok(())
5497 }
5498
5499 #[doc(hidden)]
5502 pub fn query_all_at(
5503 &mut self,
5504 conditions: &[crate::query::Condition],
5505 snapshot: Snapshot,
5506 ) -> Result<Vec<Row>> {
5507 self.require_select()?;
5508 self.ensure_indexes_complete()?;
5509 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5510 return Err(MongrelError::InvalidArgument(format!(
5511 "query exceeds {} conditions",
5512 crate::query::MAX_HARD_CONDITIONS
5513 )));
5514 }
5515 self.query_conditions_at(
5516 conditions,
5517 snapshot,
5518 None,
5519 None,
5520 0,
5521 None,
5522 unix_nanos_now(),
5523 None,
5524 )
5525 }
5526
5527 #[allow(clippy::too_many_arguments)]
5528 fn query_conditions_at(
5529 &self,
5530 conditions: &[crate::query::Condition],
5531 snapshot: Snapshot,
5532 allowed: Option<&std::collections::HashSet<RowId>>,
5533 limit: Option<usize>,
5534 offset: usize,
5535 after_row_id: Option<RowId>,
5536 query_time_nanos: i64,
5537 control: Option<&crate::ExecutionControl>,
5538 ) -> Result<Vec<Row>> {
5539 control
5540 .map(crate::ExecutionControl::checkpoint)
5541 .transpose()?;
5542 crate::trace::QueryTrace::record(|t| {
5543 t.run_count = self.run_refs.len();
5544 t.memtable_rows = self.memtable.len();
5545 t.mutable_run_rows = self.mutable_run.len();
5546 });
5547 if conditions.is_empty() {
5551 crate::trace::QueryTrace::record(|t| {
5552 t.scan_mode = crate::trace::ScanMode::Materialized;
5553 t.row_materialized = true;
5554 });
5555 let mut rows = match control {
5556 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5557 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5558 };
5559 if let Some(allowed) = allowed {
5560 let mut filtered = Vec::with_capacity(rows.len());
5561 for (index, row) in rows.into_iter().enumerate() {
5562 if index & 255 == 0 {
5563 control
5564 .map(crate::ExecutionControl::checkpoint)
5565 .transpose()?;
5566 }
5567 if allowed.contains(&row.row_id) {
5568 filtered.push(row);
5569 }
5570 }
5571 rows = filtered;
5572 }
5573 if let Some(after_row_id) = after_row_id {
5574 rows.retain(|row| row.row_id > after_row_id);
5575 }
5576 rows.drain(..offset.min(rows.len()));
5577 if let Some(limit) = limit {
5578 rows.truncate(limit);
5579 }
5580 return Ok(rows);
5581 }
5582 crate::trace::QueryTrace::record(|t| {
5583 t.conditions_pushed = conditions.len();
5584 t.scan_mode = crate::trace::ScanMode::Materialized;
5585 t.row_materialized = true;
5586 });
5587 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5594 ordered.sort_by_key(|c| condition_cost_rank(c));
5595 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5596 for c in &ordered {
5597 control
5598 .map(crate::ExecutionControl::checkpoint)
5599 .transpose()?;
5600 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5601 let empty = s.is_empty();
5602 sets.push(s);
5603 if empty {
5604 break;
5605 }
5606 }
5607 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5608 if let Some(allowed) = allowed {
5609 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5610 }
5611 if let Some(after_row_id) = after_row_id {
5612 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5613 rids.drain(..first);
5614 }
5615 rids.drain(..offset.min(rids.len()));
5616 if let Some(limit) = limit {
5617 rids.truncate(limit);
5618 }
5619 control
5620 .map(crate::ExecutionControl::checkpoint)
5621 .transpose()?;
5622 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5623 }
5624
5625 pub fn retrieve(
5627 &mut self,
5628 retriever: &crate::query::Retriever,
5629 ) -> Result<Vec<crate::query::RetrieverHit>> {
5630 self.retrieve_with_allowed(retriever, None)
5631 }
5632
5633 pub fn retrieve_at(
5634 &mut self,
5635 retriever: &crate::query::Retriever,
5636 snapshot: Snapshot,
5637 allowed: Option<&std::collections::HashSet<RowId>>,
5638 ) -> Result<Vec<crate::query::RetrieverHit>> {
5639 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5640 }
5641
5642 pub fn retrieve_with_allowed(
5645 &mut self,
5646 retriever: &crate::query::Retriever,
5647 allowed: Option<&std::collections::HashSet<RowId>>,
5648 ) -> Result<Vec<crate::query::RetrieverHit>> {
5649 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5650 }
5651
5652 pub fn retrieve_at_with_allowed(
5653 &mut self,
5654 retriever: &crate::query::Retriever,
5655 snapshot: Snapshot,
5656 allowed: Option<&std::collections::HashSet<RowId>>,
5657 ) -> Result<Vec<crate::query::RetrieverHit>> {
5658 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5659 }
5660
5661 pub fn retrieve_at_with_allowed_and_context(
5662 &mut self,
5663 retriever: &crate::query::Retriever,
5664 snapshot: Snapshot,
5665 allowed: Option<&std::collections::HashSet<RowId>>,
5666 context: Option<&crate::query::AiExecutionContext>,
5667 ) -> Result<Vec<crate::query::RetrieverHit>> {
5668 self.require_select()?;
5669 self.ensure_indexes_complete()?;
5670 self.validate_retriever(retriever)?;
5671 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5672 }
5673
5674 pub fn retrieve_at_with_candidate_authorization_and_context(
5675 &mut self,
5676 retriever: &crate::query::Retriever,
5677 snapshot: Snapshot,
5678 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5679 context: Option<&crate::query::AiExecutionContext>,
5680 ) -> Result<Vec<crate::query::RetrieverHit>> {
5681 self.require_select()?;
5682 self.ensure_indexes_complete()?;
5683 self.retrieve_at_with_candidate_authorization_on_generation(
5684 retriever,
5685 snapshot,
5686 authorization,
5687 context,
5688 )
5689 }
5690
5691 #[doc(hidden)]
5692 pub fn retrieve_at_with_candidate_authorization_on_generation(
5693 &self,
5694 retriever: &crate::query::Retriever,
5695 snapshot: Snapshot,
5696 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5697 context: Option<&crate::query::AiExecutionContext>,
5698 ) -> Result<Vec<crate::query::RetrieverHit>> {
5699 self.require_select()?;
5700 self.validate_retriever(retriever)?;
5701 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5702 }
5703
5704 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5705 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5706 let (column_id, k) = match retriever {
5707 Retriever::Ann {
5708 column_id,
5709 query,
5710 k,
5711 } => {
5712 let index = self.ann.get(column_id).ok_or_else(|| {
5713 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5714 })?;
5715 if query.len() != index.dim() {
5716 return Err(MongrelError::InvalidArgument(format!(
5717 "ANN query dimension must be {}, got {}",
5718 index.dim(),
5719 query.len()
5720 )));
5721 }
5722 if query.iter().any(|value| !value.is_finite()) {
5723 return Err(MongrelError::InvalidArgument(
5724 "ANN query values must be finite".into(),
5725 ));
5726 }
5727 (*column_id, *k)
5728 }
5729 Retriever::Sparse {
5730 column_id,
5731 query,
5732 k,
5733 } => {
5734 if !self.sparse.contains_key(column_id) {
5735 return Err(MongrelError::InvalidArgument(format!(
5736 "column {column_id} has no Sparse index"
5737 )));
5738 }
5739 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5740 return Err(MongrelError::InvalidArgument(
5741 "Sparse query must be non-empty with finite weights".into(),
5742 ));
5743 }
5744 if query.len() > MAX_SPARSE_TERMS {
5745 return Err(MongrelError::InvalidArgument(format!(
5746 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5747 )));
5748 }
5749 (*column_id, *k)
5750 }
5751 Retriever::MinHash {
5752 column_id,
5753 members,
5754 k,
5755 } => {
5756 if !self.minhash.contains_key(column_id) {
5757 return Err(MongrelError::InvalidArgument(format!(
5758 "column {column_id} has no MinHash index"
5759 )));
5760 }
5761 if members.is_empty() {
5762 return Err(MongrelError::InvalidArgument(
5763 "MinHash members must not be empty".into(),
5764 ));
5765 }
5766 if members.len() > MAX_SET_MEMBERS {
5767 return Err(MongrelError::InvalidArgument(format!(
5768 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5769 )));
5770 }
5771 let mut total_bytes = 0usize;
5772 for member in members {
5773 let bytes = member.encoded_len();
5774 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5775 return Err(MongrelError::InvalidArgument(format!(
5776 "MinHash member exceeds {} bytes",
5777 crate::query::MAX_SET_MEMBER_BYTES
5778 )));
5779 }
5780 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5781 MongrelError::InvalidArgument("MinHash input size overflow".into())
5782 })?;
5783 }
5784 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5785 return Err(MongrelError::InvalidArgument(format!(
5786 "MinHash input exceeds {} bytes",
5787 crate::query::MAX_SET_INPUT_BYTES
5788 )));
5789 }
5790 (*column_id, *k)
5791 }
5792 };
5793 if k == 0 {
5794 return Err(MongrelError::InvalidArgument(
5795 "retriever k must be > 0".into(),
5796 ));
5797 }
5798 if k > MAX_RETRIEVER_K {
5799 return Err(MongrelError::InvalidArgument(format!(
5800 "retriever k exceeds {MAX_RETRIEVER_K}"
5801 )));
5802 }
5803 debug_assert!(self
5804 .schema
5805 .columns
5806 .iter()
5807 .any(|column| column.id == column_id));
5808 Ok(())
5809 }
5810
5811 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5812 use crate::query::Condition;
5813 match condition {
5814 Condition::Ann {
5815 column_id,
5816 query,
5817 k,
5818 } => self.validate_retriever(&crate::query::Retriever::Ann {
5819 column_id: *column_id,
5820 query: query.clone(),
5821 k: *k,
5822 }),
5823 Condition::SparseMatch {
5824 column_id,
5825 query,
5826 k,
5827 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5828 column_id: *column_id,
5829 query: query.clone(),
5830 k: *k,
5831 }),
5832 Condition::MinHashSimilar {
5833 column_id,
5834 query,
5835 k,
5836 } => {
5837 if !self.minhash.contains_key(column_id) {
5838 return Err(MongrelError::InvalidArgument(format!(
5839 "column {column_id} has no MinHash index"
5840 )));
5841 }
5842 if query.is_empty() || *k == 0 {
5843 return Err(MongrelError::InvalidArgument(
5844 "MinHash query must be non-empty and k must be > 0".into(),
5845 ));
5846 }
5847 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5848 {
5849 return Err(MongrelError::InvalidArgument(format!(
5850 "MinHash query must have <= {} members and k <= {}",
5851 crate::query::MAX_SET_MEMBERS,
5852 crate::query::MAX_RETRIEVER_K
5853 )));
5854 }
5855 Ok(())
5856 }
5857 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5858 Err(MongrelError::InvalidArgument(format!(
5859 "bitmap IN exceeds {} values",
5860 crate::query::MAX_SET_MEMBERS
5861 )))
5862 }
5863 Condition::FmContainsAll { patterns, .. }
5864 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5865 {
5866 Err(MongrelError::InvalidArgument(format!(
5867 "FM query exceeds {} patterns",
5868 crate::query::MAX_HARD_CONDITIONS
5869 )))
5870 }
5871 _ => Ok(()),
5872 }
5873 }
5874
5875 fn retrieve_filtered(
5876 &self,
5877 retriever: &crate::query::Retriever,
5878 snapshot: Snapshot,
5879 hard_filter: Option<&RowIdSet>,
5880 allowed: Option<&std::collections::HashSet<RowId>>,
5881 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5882 context: Option<&crate::query::AiExecutionContext>,
5883 ) -> Result<Vec<crate::query::RetrieverHit>> {
5884 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5885 let started = std::time::Instant::now();
5886 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5887 Retriever::Ann {
5888 column_id,
5889 query,
5890 k,
5891 } => {
5892 let Some(index) = self.ann.get(column_id) else {
5893 return Ok(Vec::new());
5894 };
5895 let cap = ann_candidate_cap(index.len(), context);
5896 if cap == 0 {
5897 return Ok(Vec::new());
5898 }
5899 let mut breadth = (*k).max(1).min(cap);
5900 let mut eligibility = std::collections::HashMap::new();
5901 let mut filtered = loop {
5902 let mut seen = std::collections::HashSet::new();
5903 if let Some(context) = context {
5904 context.checkpoint()?;
5905 }
5906 let raw = index.search_with_context(query, breadth, context)?;
5907 let unchecked: Vec<_> = raw
5908 .iter()
5909 .map(|(row_id, _)| *row_id)
5910 .filter(|row_id| !eligibility.contains_key(row_id))
5911 .filter(|row_id| {
5912 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5913 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5914 })
5915 .collect();
5916 let eligible = self.eligible_and_authorized_candidate_ids(
5917 &unchecked,
5918 *column_id,
5919 snapshot,
5920 candidate_authorization,
5921 context,
5922 )?;
5923 for row_id in unchecked {
5924 eligibility.insert(row_id, eligible.contains(&row_id));
5925 }
5926 let filtered: Vec<_> = raw
5927 .into_iter()
5928 .filter(|(row_id, _)| {
5929 seen.insert(*row_id)
5930 && eligibility.get(row_id).copied().unwrap_or(false)
5931 })
5932 .map(|(row_id, score)| {
5933 let score = match score {
5934 crate::index::AnnDistance::Hamming(d) => {
5935 RetrieverScore::AnnHammingDistance(d)
5936 }
5937 crate::index::AnnDistance::Cosine(d) => {
5938 RetrieverScore::AnnCosineDistance(d)
5939 }
5940 };
5941 (row_id, score)
5942 })
5943 .collect();
5944 if filtered.len() >= *k || breadth >= cap {
5945 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5946 crate::trace::QueryTrace::record(|trace| {
5947 trace.ann_candidate_cap_hit = true;
5948 });
5949 }
5950 break filtered;
5951 }
5952 breadth = breadth.saturating_mul(2).min(cap);
5953 };
5954 filtered.truncate(*k);
5955 filtered
5956 }
5957 Retriever::Sparse {
5958 column_id,
5959 query,
5960 k,
5961 } => self
5962 .sparse
5963 .get(column_id)
5964 .map(|index| -> Result<Vec<_>> {
5965 let mut breadth = (*k).max(1);
5966 let mut eligibility = std::collections::HashMap::new();
5967 loop {
5968 if let Some(context) = context {
5969 context.checkpoint()?;
5970 }
5971 let raw = index.search_with_context(query, breadth, context)?;
5972 let unchecked: Vec<_> = raw
5973 .iter()
5974 .map(|(row_id, _)| *row_id)
5975 .filter(|row_id| !eligibility.contains_key(row_id))
5976 .filter(|row_id| {
5977 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5978 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5979 })
5980 .collect();
5981 let eligible = self.eligible_and_authorized_candidate_ids(
5982 &unchecked,
5983 *column_id,
5984 snapshot,
5985 candidate_authorization,
5986 context,
5987 )?;
5988 for row_id in unchecked {
5989 eligibility.insert(row_id, eligible.contains(&row_id));
5990 }
5991 let filtered: Vec<_> = raw
5992 .iter()
5993 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5994 .take(*k)
5995 .map(|(row_id, score)| {
5996 (*row_id, RetrieverScore::SparseDotProduct(*score))
5997 })
5998 .collect();
5999 if filtered.len() >= *k || raw.len() < breadth {
6000 break Ok(filtered);
6001 }
6002 let next = breadth.saturating_mul(2);
6003 if next == breadth {
6004 break Ok(filtered);
6005 }
6006 breadth = next;
6007 }
6008 })
6009 .transpose()?
6010 .unwrap_or_default(),
6011 Retriever::MinHash {
6012 column_id,
6013 members,
6014 k,
6015 } => self
6016 .minhash
6017 .get(column_id)
6018 .map(|index| -> Result<Vec<_>> {
6019 let mut hashes = Vec::with_capacity(members.len());
6020 for member in members {
6021 if let Some(context) = context {
6022 context.consume(crate::query::work_units(
6023 member.encoded_len(),
6024 crate::query::PARSE_WORK_QUANTUM,
6025 ))?;
6026 }
6027 hashes.push(member.hash_v1());
6028 }
6029 let mut breadth = (*k).max(1);
6030 let mut eligibility = std::collections::HashMap::new();
6031 loop {
6032 if let Some(context) = context {
6033 context.checkpoint()?;
6034 }
6035 let raw = index.search_with_context(&hashes, breadth, context)?;
6036 let unchecked: Vec<_> = raw
6037 .iter()
6038 .map(|(row_id, _)| *row_id)
6039 .filter(|row_id| !eligibility.contains_key(row_id))
6040 .filter(|row_id| {
6041 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
6042 && allowed.is_none_or(|allowed| allowed.contains(row_id))
6043 })
6044 .collect();
6045 let eligible = self.eligible_and_authorized_candidate_ids(
6046 &unchecked,
6047 *column_id,
6048 snapshot,
6049 candidate_authorization,
6050 context,
6051 )?;
6052 for row_id in unchecked {
6053 eligibility.insert(row_id, eligible.contains(&row_id));
6054 }
6055 let filtered: Vec<_> = raw
6056 .iter()
6057 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
6058 .take(*k)
6059 .map(|(row_id, score)| {
6060 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
6061 })
6062 .collect();
6063 if filtered.len() >= *k || raw.len() < breadth {
6064 break Ok(filtered);
6065 }
6066 let next = breadth.saturating_mul(2);
6067 if next == breadth {
6068 break Ok(filtered);
6069 }
6070 breadth = next;
6071 }
6072 })
6073 .transpose()?
6074 .unwrap_or_default(),
6075 };
6076 let elapsed = started.elapsed().as_nanos() as u64;
6077 crate::trace::QueryTrace::record(|trace| {
6078 match retriever {
6079 Retriever::Ann { .. } => {
6080 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
6081 if let Retriever::Ann { column_id, .. } = retriever {
6082 if let Some(index) = self.ann.get(column_id) {
6083 trace.ann_algorithm = Some(index.algorithm());
6084 trace.ann_quantization = Some(index.quantization());
6085 trace.ann_backend = Some(index.backend_name());
6086 }
6087 }
6088 }
6089 Retriever::Sparse { .. } => {
6090 trace.sparse_candidate_nanos =
6091 trace.sparse_candidate_nanos.saturating_add(elapsed)
6092 }
6093 Retriever::MinHash { .. } => {
6094 trace.minhash_candidate_nanos =
6095 trace.minhash_candidate_nanos.saturating_add(elapsed)
6096 }
6097 }
6098 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
6099 });
6100 Ok(scored
6101 .into_iter()
6102 .enumerate()
6103 .map(|(rank, (row_id, score))| RetrieverHit {
6104 row_id,
6105 rank: rank + 1,
6106 score,
6107 })
6108 .collect())
6109 }
6110
6111 fn eligible_candidate_ids(
6112 &self,
6113 candidates: &[RowId],
6114 _column_id: u16,
6115 snapshot: Snapshot,
6116 context: Option<&crate::query::AiExecutionContext>,
6117 ) -> Result<std::collections::HashSet<RowId>> {
6118 if !self.had_deletes
6119 && self.ttl.is_none()
6120 && self.pending_put_cols.is_empty()
6121 && snapshot.epoch == self.snapshot().epoch
6122 {
6123 return Ok(candidates.iter().copied().collect());
6124 }
6125 let mut readers: Vec<_> = self
6126 .run_refs
6127 .iter()
6128 .map(|run| self.open_reader(run.run_id))
6129 .collect::<Result<_>>()?;
6130 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6131 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
6132 for &row_id in candidates {
6133 if let Some(context) = context {
6134 context.consume(1)?;
6135 }
6136 let mem = self.memtable.get_version_at(row_id, snapshot);
6137 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
6138 let overlay = match (mem, mutable) {
6139 (Some(left), Some(right)) => Some(
6140 if Snapshot::version_is_newer(
6141 left.1.committed_epoch,
6142 left.1.commit_ts,
6143 right.1.committed_epoch,
6144 right.1.commit_ts,
6145 ) {
6146 left
6147 } else {
6148 right
6149 },
6150 ),
6151 (Some(value), None) | (None, Some(value)) => Some(value),
6152 (None, None) => None,
6153 };
6154 if let Some((_, row)) = overlay {
6155 if !row.deleted && !self.row_expired_at(&row, now) {
6156 eligible.insert(row_id);
6157 }
6158 continue;
6159 }
6160 let mut best: Option<(Epoch, bool, usize)> = None;
6161 for (index, reader) in readers.iter_mut().enumerate() {
6162 if let Some((epoch, deleted)) =
6163 reader.get_version_visibility(row_id, snapshot.epoch)?
6164 {
6165 if best
6166 .as_ref()
6167 .map(|(best_epoch, ..)| epoch > *best_epoch)
6168 .unwrap_or(true)
6169 {
6170 best = Some((epoch, deleted, index));
6171 }
6172 }
6173 }
6174 let Some((_, false, reader_index)) = best else {
6175 continue;
6176 };
6177 if let Some(ttl) = self.ttl {
6178 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6179 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6180 {
6181 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6182 continue;
6183 }
6184 }
6185 }
6186 eligible.insert(row_id);
6187 }
6188 Ok(eligible)
6189 }
6190
6191 fn eligible_and_authorized_candidate_ids(
6192 &self,
6193 candidates: &[RowId],
6194 column_id: u16,
6195 snapshot: Snapshot,
6196 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6197 context: Option<&crate::query::AiExecutionContext>,
6198 ) -> Result<std::collections::HashSet<RowId>> {
6199 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
6200 let Some(authorization) = authorization else {
6201 return Ok(eligible);
6202 };
6203 let candidates: Vec<_> = eligible.into_iter().collect();
6204 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
6205 }
6206
6207 fn policy_allowed_candidate_ids(
6208 &self,
6209 candidates: &[RowId],
6210 snapshot: Snapshot,
6211 authorization: &crate::security::CandidateAuthorization<'_>,
6212 context: Option<&crate::query::AiExecutionContext>,
6213 ) -> Result<std::collections::HashSet<RowId>> {
6214 let started = std::time::Instant::now();
6215 if candidates.is_empty()
6216 || authorization.principal.is_admin
6217 || !authorization.security.rls_enabled(authorization.table)
6218 {
6219 return Ok(candidates.iter().copied().collect());
6220 }
6221 if let Some(context) = context {
6222 context.checkpoint()?;
6223 }
6224 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
6225 let mut rows: std::collections::HashMap<RowId, Row> = candidates
6226 .iter()
6227 .map(|row_id| {
6228 (
6229 *row_id,
6230 Row {
6231 row_id: *row_id,
6232 committed_epoch: snapshot.epoch,
6233 columns: std::collections::HashMap::new(),
6234 deleted: false,
6235 commit_ts: None,
6236 },
6237 )
6238 })
6239 .collect();
6240 let columns = authorization
6241 .security
6242 .select_policy_columns(authorization.table, authorization.principal);
6243 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6244 let mut decoded = 0usize;
6245 for column_id in &columns {
6246 if let Some(context) = context {
6247 context.checkpoint()?;
6248 }
6249 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6250 &row_ids, *column_id, snapshot, query_now, context,
6251 )? {
6252 if let Some(row) = rows.get_mut(&row_id) {
6253 row.columns.insert(*column_id, value);
6254 decoded = decoded.saturating_add(1);
6255 }
6256 }
6257 }
6258 if let Some(context) = context {
6259 context.consume(candidates.len().saturating_add(decoded))?;
6260 }
6261 let allowed = rows
6262 .into_values()
6263 .filter_map(|row| {
6264 authorization
6265 .security
6266 .row_allowed(
6267 authorization.table,
6268 crate::security::PolicyCommand::Select,
6269 &row,
6270 authorization.principal,
6271 false,
6272 )
6273 .then_some(row.row_id)
6274 })
6275 .collect();
6276 crate::trace::QueryTrace::record(|trace| {
6277 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
6278 trace.rls_policy_columns_decoded =
6279 trace.rls_policy_columns_decoded.saturating_add(decoded);
6280 trace.authorization_nanos = trace
6281 .authorization_nanos
6282 .saturating_add(started.elapsed().as_nanos() as u64);
6283 });
6284 Ok(allowed)
6285 }
6286
6287 pub fn search(
6289 &mut self,
6290 request: &crate::query::SearchRequest,
6291 ) -> Result<Vec<crate::query::SearchHit>> {
6292 self.search_with_allowed(request, None)
6293 }
6294
6295 pub fn search_at(
6296 &mut self,
6297 request: &crate::query::SearchRequest,
6298 snapshot: Snapshot,
6299 authorized: Option<&std::collections::HashSet<RowId>>,
6300 ) -> Result<Vec<crate::query::SearchHit>> {
6301 self.search_at_with_allowed(request, snapshot, authorized)
6302 }
6303
6304 pub fn search_with_allowed(
6305 &mut self,
6306 request: &crate::query::SearchRequest,
6307 authorized: Option<&std::collections::HashSet<RowId>>,
6308 ) -> Result<Vec<crate::query::SearchHit>> {
6309 self.search_at_with_allowed(request, self.snapshot(), authorized)
6310 }
6311
6312 pub fn search_at_with_allowed(
6313 &mut self,
6314 request: &crate::query::SearchRequest,
6315 snapshot: Snapshot,
6316 authorized: Option<&std::collections::HashSet<RowId>>,
6317 ) -> Result<Vec<crate::query::SearchHit>> {
6318 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
6319 }
6320
6321 pub fn search_at_with_allowed_and_context(
6322 &mut self,
6323 request: &crate::query::SearchRequest,
6324 snapshot: Snapshot,
6325 authorized: Option<&std::collections::HashSet<RowId>>,
6326 context: Option<&crate::query::AiExecutionContext>,
6327 ) -> Result<Vec<crate::query::SearchHit>> {
6328 self.ensure_indexes_complete()?;
6329 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
6330 }
6331
6332 pub fn search_at_with_candidate_authorization_and_context(
6333 &mut self,
6334 request: &crate::query::SearchRequest,
6335 snapshot: Snapshot,
6336 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6337 context: Option<&crate::query::AiExecutionContext>,
6338 ) -> Result<Vec<crate::query::SearchHit>> {
6339 self.ensure_indexes_complete()?;
6340 self.search_at_with_filters_and_context(
6341 request,
6342 snapshot,
6343 None,
6344 authorization,
6345 context,
6346 None,
6347 )
6348 }
6349
6350 #[doc(hidden)]
6351 pub fn search_at_with_candidate_authorization_on_generation(
6352 &self,
6353 request: &crate::query::SearchRequest,
6354 snapshot: Snapshot,
6355 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6356 context: Option<&crate::query::AiExecutionContext>,
6357 ) -> Result<Vec<crate::query::SearchHit>> {
6358 self.search_at_with_filters_and_context(
6359 request,
6360 snapshot,
6361 None,
6362 authorization,
6363 context,
6364 None,
6365 )
6366 }
6367
6368 #[doc(hidden)]
6369 pub fn search_at_with_candidate_authorization_on_generation_after(
6370 &self,
6371 request: &crate::query::SearchRequest,
6372 snapshot: Snapshot,
6373 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6374 context: Option<&crate::query::AiExecutionContext>,
6375 after: Option<crate::query::SearchAfter>,
6376 ) -> Result<Vec<crate::query::SearchHit>> {
6377 self.search_at_with_filters_and_context(
6378 request,
6379 snapshot,
6380 None,
6381 authorization,
6382 context,
6383 after,
6384 )
6385 }
6386
6387 fn search_at_with_filters_and_context(
6388 &self,
6389 request: &crate::query::SearchRequest,
6390 snapshot: Snapshot,
6391 authorized: Option<&std::collections::HashSet<RowId>>,
6392 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6393 context: Option<&crate::query::AiExecutionContext>,
6394 after: Option<crate::query::SearchAfter>,
6395 ) -> Result<Vec<crate::query::SearchHit>> {
6396 use crate::query::{
6397 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
6398 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
6399 };
6400 let total_started = std::time::Instant::now();
6401 let rank_offset = after.map_or(0, |after| after.returned_count);
6402 self.require_select()?;
6403 if request.limit == 0 {
6404 return Err(MongrelError::InvalidArgument(
6405 "search limit must be > 0".into(),
6406 ));
6407 }
6408 if request.limit > MAX_FINAL_LIMIT {
6409 return Err(MongrelError::InvalidArgument(format!(
6410 "search limit exceeds {MAX_FINAL_LIMIT}"
6411 )));
6412 }
6413 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
6414 return Err(MongrelError::InvalidArgument(
6415 "search-after score must be finite".into(),
6416 ));
6417 }
6418 if request.retrievers.is_empty() {
6419 return Err(MongrelError::InvalidArgument(
6420 "search requires at least one retriever".into(),
6421 ));
6422 }
6423 if request.retrievers.len() > MAX_RETRIEVERS {
6424 return Err(MongrelError::InvalidArgument(format!(
6425 "search exceeds {MAX_RETRIEVERS} retrievers"
6426 )));
6427 }
6428 if request.must.len() > MAX_HARD_CONDITIONS {
6429 return Err(MongrelError::InvalidArgument(format!(
6430 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
6431 )));
6432 }
6433 for condition in &request.must {
6434 self.validate_condition(condition)?;
6435 }
6436 if request.must.iter().any(|condition| {
6437 matches!(
6438 condition,
6439 Condition::Ann { .. }
6440 | Condition::SparseMatch { .. }
6441 | Condition::MinHashSimilar { .. }
6442 )
6443 }) {
6444 return Err(MongrelError::InvalidArgument(
6445 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
6446 .into(),
6447 ));
6448 }
6449 let mut names = std::collections::HashSet::new();
6450 for named in &request.retrievers {
6451 if named.name.is_empty()
6452 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
6453 || !names.insert(named.name.as_str())
6454 {
6455 return Err(MongrelError::InvalidArgument(format!(
6456 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
6457 crate::query::MAX_RETRIEVER_NAME_BYTES
6458 )));
6459 }
6460 if !named.weight.is_finite()
6461 || named.weight < 0.0
6462 || named.weight > MAX_RETRIEVER_WEIGHT
6463 {
6464 return Err(MongrelError::InvalidArgument(format!(
6465 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6466 )));
6467 }
6468 self.validate_retriever(&named.retriever)?;
6469 }
6470 let projection = request
6471 .projection
6472 .clone()
6473 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
6474 if projection.len() > MAX_PROJECTION_COLUMNS {
6475 return Err(MongrelError::InvalidArgument(format!(
6476 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
6477 )));
6478 }
6479 for column_id in &projection {
6480 if !self
6481 .schema
6482 .columns
6483 .iter()
6484 .any(|column| column.id == *column_id)
6485 {
6486 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
6487 }
6488 }
6489 if let Some(crate::query::Rerank::ExactVector {
6490 embedding_column,
6491 query,
6492 candidate_limit,
6493 weight,
6494 ..
6495 }) = &request.rerank
6496 {
6497 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
6498 {
6499 return Err(MongrelError::InvalidArgument(format!(
6500 "rerank candidate_limit must be between search limit and {}",
6501 crate::query::MAX_RETRIEVER_K
6502 )));
6503 }
6504 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
6505 return Err(MongrelError::InvalidArgument(format!(
6506 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6507 )));
6508 }
6509 let column = self
6510 .schema
6511 .columns
6512 .iter()
6513 .find(|column| column.id == *embedding_column)
6514 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
6515 let crate::schema::TypeId::Embedding { dim } = column.ty else {
6516 return Err(MongrelError::InvalidArgument(format!(
6517 "rerank column {embedding_column} is not an embedding"
6518 )));
6519 };
6520 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
6521 return Err(MongrelError::InvalidArgument(format!(
6522 "rerank query must contain {dim} finite values"
6523 )));
6524 }
6525 }
6526
6527 let hard_filter_started = std::time::Instant::now();
6528 let hard_filter = if request.must.is_empty() {
6529 None
6530 } else {
6531 let mut sets = Vec::with_capacity(request.must.len());
6532 for condition in &request.must {
6533 if let Some(context) = context {
6534 context.checkpoint()?;
6535 }
6536 sets.push(self.resolve_condition(condition, snapshot)?);
6537 }
6538 Some(RowIdSet::intersect_many(sets))
6539 };
6540 crate::trace::QueryTrace::record(|trace| {
6541 trace.hard_filter_nanos = trace
6542 .hard_filter_nanos
6543 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6544 });
6545 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6546 return Ok(Vec::new());
6547 }
6548
6549 let constant = match request.fusion {
6550 Fusion::ReciprocalRank { constant } => constant,
6551 };
6552 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6553 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6554 let mut fusion_nanos = 0u64;
6555 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6556 std::collections::HashMap::new();
6557 for named in retrievers {
6558 if named.weight == 0.0 {
6559 continue;
6560 }
6561 if let Some(context) = context {
6562 context.checkpoint()?;
6563 }
6564 let hits = self.retrieve_filtered(
6565 &named.retriever,
6566 snapshot,
6567 hard_filter.as_ref(),
6568 authorized,
6569 candidate_authorization,
6570 context,
6571 )?;
6572 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6573 let fusion_started = std::time::Instant::now();
6574 for hit in hits {
6575 if let Some(context) = context {
6576 context.consume(1)?;
6577 }
6578 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6579 if !contribution.is_finite() {
6580 return Err(MongrelError::InvalidArgument(
6581 "retriever contribution must be finite".into(),
6582 ));
6583 }
6584 let max_fused_candidates = context.map_or(
6585 crate::query::MAX_FUSED_CANDIDATES,
6586 crate::query::AiExecutionContext::max_fused_candidates,
6587 );
6588 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6589 return Err(MongrelError::WorkBudgetExceeded);
6590 }
6591 let entry = fused.entry(hit.row_id).or_default();
6592 entry.0 += contribution;
6593 if !entry.0.is_finite() {
6594 return Err(MongrelError::InvalidArgument(
6595 "fused score must be finite".into(),
6596 ));
6597 }
6598 entry.1.push(ComponentScore {
6599 retriever_name: retriever_name.clone(),
6600 rank: hit.rank,
6601 raw_score: hit.score,
6602 contribution,
6603 });
6604 }
6605 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6606 }
6607 let union_size = fused.len();
6608 let mut ranked: Vec<_> = fused
6609 .into_iter()
6610 .map(|(row_id, (fused_score, components))| {
6611 (row_id, fused_score, components, None, fused_score)
6612 })
6613 .collect();
6614 let order = |(a_row, _, _, _, a_score): &(
6615 RowId,
6616 f64,
6617 Vec<ComponentScore>,
6618 Option<f32>,
6619 f64,
6620 ),
6621 (b_row, _, _, _, b_score): &(
6622 RowId,
6623 f64,
6624 Vec<ComponentScore>,
6625 Option<f32>,
6626 f64,
6627 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6628 if let Some(crate::query::Rerank::ExactVector {
6629 embedding_column,
6630 query,
6631 metric,
6632 candidate_limit,
6633 weight,
6634 }) = &request.rerank
6635 {
6636 let fused_order = |(a_row, a_score, ..): &(
6637 RowId,
6638 f64,
6639 Vec<ComponentScore>,
6640 Option<f32>,
6641 f64,
6642 ),
6643 (b_row, b_score, ..): &(
6644 RowId,
6645 f64,
6646 Vec<ComponentScore>,
6647 Option<f32>,
6648 f64,
6649 )| {
6650 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6651 };
6652 let selection_started = std::time::Instant::now();
6653 if let Some(context) = context {
6654 context.consume(ranked.len())?;
6655 }
6656 if ranked.len() > *candidate_limit {
6657 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6658 ranked.truncate(*candidate_limit);
6659 }
6660 ranked.sort_by(fused_order);
6661 fusion_nanos =
6662 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6663 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6664 if let Some(context) = context {
6665 context.consume(row_ids.len())?;
6666 }
6667 let query_now =
6668 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6669 let gather_started = std::time::Instant::now();
6670 let vectors = self.values_for_rids_batch_at_with_context(
6671 &row_ids,
6672 *embedding_column,
6673 snapshot,
6674 query_now,
6675 context,
6676 )?;
6677 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6678 let vector_work =
6679 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6680 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6681 if let Some(context) = context {
6682 context.consume(vector_work)?;
6683 }
6684 query
6685 .iter()
6686 .map(|value| f64::from(*value).powi(2))
6687 .sum::<f64>()
6688 .sqrt()
6689 } else {
6690 0.0
6691 };
6692 let score_started = std::time::Instant::now();
6693 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6694 for (row_id, value) in vectors {
6695 if let Some(meta) = value.generated_embedding_metadata() {
6696 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6697 continue;
6698 }
6699 }
6700 let Some(vector) = value.as_embedding() else {
6701 continue;
6702 };
6703 let score = match metric {
6704 crate::query::VectorMetric::DotProduct => {
6705 if let Some(context) = context {
6706 context.consume(vector_work)?;
6707 }
6708 query
6709 .iter()
6710 .zip(vector)
6711 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6712 .sum::<f64>()
6713 }
6714 crate::query::VectorMetric::Cosine => {
6715 if let Some(context) = context {
6716 context.consume(vector_work.saturating_mul(2))?;
6717 }
6718 let dot = query
6719 .iter()
6720 .zip(vector)
6721 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6722 .sum::<f64>();
6723 let norm = vector
6724 .iter()
6725 .map(|value| f64::from(*value).powi(2))
6726 .sum::<f64>()
6727 .sqrt();
6728 if query_norm == 0.0 || norm == 0.0 {
6729 0.0
6730 } else {
6731 dot / (query_norm * norm)
6732 }
6733 }
6734 crate::query::VectorMetric::Euclidean => {
6735 if let Some(context) = context {
6736 context.consume(vector_work)?;
6737 }
6738 query
6739 .iter()
6740 .zip(vector)
6741 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6742 .sum::<f64>()
6743 .sqrt()
6744 }
6745 };
6746 if !score.is_finite() {
6747 return Err(MongrelError::InvalidArgument(
6748 "exact rerank score must be finite".into(),
6749 ));
6750 }
6751 scores.insert(row_id, score as f32);
6752 }
6753 let mut reranked = Vec::with_capacity(ranked.len());
6754 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6755 let Some(score) = scores.get(&row_id).copied() else {
6756 continue;
6757 };
6758 let ordering_score = match metric {
6759 crate::query::VectorMetric::Euclidean => -f64::from(score),
6760 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6761 f64::from(score)
6762 }
6763 };
6764 let final_score = fused_score + *weight * ordering_score;
6765 if !final_score.is_finite() {
6766 return Err(MongrelError::InvalidArgument(
6767 "final rerank score must be finite".into(),
6768 ));
6769 }
6770 reranked.push((row_id, fused_score, components, Some(score), final_score));
6771 }
6772 ranked = reranked;
6773 ranked.sort_by(order);
6774 crate::trace::QueryTrace::record(|trace| {
6775 trace.exact_vector_gather_nanos =
6776 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6777 trace.exact_vector_score_nanos = trace
6778 .exact_vector_score_nanos
6779 .saturating_add(score_started.elapsed().as_nanos() as u64);
6780 });
6781 }
6782 if let Some(after) = after {
6783 ranked.retain(|(row_id, _, _, _, final_score)| {
6784 final_score.total_cmp(&after.final_score).is_lt()
6785 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6786 });
6787 }
6788 let projection_started = std::time::Instant::now();
6789 let sentinel = projection
6790 .first()
6791 .copied()
6792 .or_else(|| self.schema.columns.first().map(|column| column.id));
6793 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6794 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6795 let mut projection_rows = 0usize;
6796 let mut projection_cells = 0usize;
6797 while out.len() < request.limit && !ranked.is_empty() {
6798 if let Some(context) = context {
6799 context.checkpoint()?;
6800 context.consume(ranked.len())?;
6801 }
6802 let needed = request.limit - out.len();
6803 let window_size = ranked
6804 .len()
6805 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6806 let selection_started = std::time::Instant::now();
6807 let mut remainder = if ranked.len() > window_size {
6808 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6809 ranked.split_off(window_size)
6810 } else {
6811 Vec::new()
6812 };
6813 ranked.sort_by(order);
6814 fusion_nanos =
6815 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6816 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6817 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6818 if let Some(context) = context {
6819 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6820 }
6821 projection_rows = projection_rows.saturating_add(row_ids.len());
6822 projection_cells =
6823 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6824 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6825 std::collections::HashMap::new();
6826 if let Some(column_id) = sentinel {
6827 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6828 &row_ids, column_id, snapshot, query_now, context,
6829 )? {
6830 cells.entry(row_id).or_default().insert(column_id, value);
6831 }
6832 }
6833 for &column_id in &projection {
6834 if Some(column_id) == sentinel {
6835 continue;
6836 }
6837 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6838 &row_ids, column_id, snapshot, query_now, context,
6839 )? {
6840 cells.entry(row_id).or_default().insert(column_id, value);
6841 }
6842 }
6843 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6844 ranked.drain(..)
6845 {
6846 let Some(row_cells) = cells.remove(&row_id) else {
6847 continue;
6848 };
6849 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6850 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6851 out.push(SearchHit {
6852 row_id,
6853 cells: projection
6854 .iter()
6855 .filter_map(|column_id| {
6856 row_cells
6857 .get(column_id)
6858 .cloned()
6859 .map(|value| (*column_id, value))
6860 })
6861 .collect(),
6862 components,
6863 fused_score,
6864 exact_rerank_score,
6865 final_score,
6866 final_rank,
6867 });
6868 if out.len() == request.limit {
6869 break;
6870 }
6871 }
6872 ranked.append(&mut remainder);
6873 }
6874 crate::trace::QueryTrace::record(|trace| {
6875 trace.union_size = union_size;
6876 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6877 trace.projection_nanos = trace
6878 .projection_nanos
6879 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6880 trace.total_nanos = trace
6881 .total_nanos
6882 .saturating_add(total_started.elapsed().as_nanos() as u64);
6883 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6884 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6885 if let Some(context) = context {
6886 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6887 }
6888 });
6889 Ok(out)
6890 }
6891
6892 pub fn set_similarity(
6895 &mut self,
6896 request: &crate::query::SetSimilarityRequest,
6897 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6898 self.set_similarity_with_allowed(request, None)
6899 }
6900
6901 pub fn set_similarity_at(
6902 &mut self,
6903 request: &crate::query::SetSimilarityRequest,
6904 snapshot: Snapshot,
6905 allowed: Option<&std::collections::HashSet<RowId>>,
6906 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6907 self.set_similarity_explained_at(request, snapshot, allowed)
6908 .map(|(hits, _)| hits)
6909 }
6910
6911 pub fn ann_rerank(
6913 &mut self,
6914 request: &crate::query::AnnRerankRequest,
6915 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6916 self.ann_rerank_with_allowed(request, None)
6917 }
6918
6919 pub fn ann_rerank_with_allowed(
6920 &mut self,
6921 request: &crate::query::AnnRerankRequest,
6922 allowed: Option<&std::collections::HashSet<RowId>>,
6923 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6924 self.ann_rerank_at(request, self.snapshot(), allowed)
6925 }
6926
6927 pub fn ann_rerank_at(
6928 &mut self,
6929 request: &crate::query::AnnRerankRequest,
6930 snapshot: Snapshot,
6931 allowed: Option<&std::collections::HashSet<RowId>>,
6932 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6933 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6934 }
6935
6936 pub fn ann_rerank_at_with_context(
6937 &mut self,
6938 request: &crate::query::AnnRerankRequest,
6939 snapshot: Snapshot,
6940 allowed: Option<&std::collections::HashSet<RowId>>,
6941 context: Option<&crate::query::AiExecutionContext>,
6942 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6943 self.ensure_indexes_complete()?;
6944 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6945 }
6946
6947 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6948 &mut self,
6949 request: &crate::query::AnnRerankRequest,
6950 snapshot: Snapshot,
6951 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6952 context: Option<&crate::query::AiExecutionContext>,
6953 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6954 self.ensure_indexes_complete()?;
6955 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6956 }
6957
6958 #[doc(hidden)]
6959 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6960 &self,
6961 request: &crate::query::AnnRerankRequest,
6962 snapshot: Snapshot,
6963 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6964 context: Option<&crate::query::AiExecutionContext>,
6965 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6966 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6967 }
6968
6969 fn ann_rerank_at_with_filters_and_context(
6970 &self,
6971 request: &crate::query::AnnRerankRequest,
6972 snapshot: Snapshot,
6973 allowed: Option<&std::collections::HashSet<RowId>>,
6974 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6975 context: Option<&crate::query::AiExecutionContext>,
6976 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6977 use crate::query::{
6978 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6979 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6980 };
6981 if request.candidate_k == 0 || request.limit == 0 {
6982 return Err(MongrelError::InvalidArgument(
6983 "candidate_k and limit must be > 0".into(),
6984 ));
6985 }
6986 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6987 return Err(MongrelError::InvalidArgument(format!(
6988 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6989 )));
6990 }
6991 let retriever = Retriever::Ann {
6992 column_id: request.column_id,
6993 query: request.query.clone(),
6994 k: request.candidate_k,
6995 };
6996 self.require_select()?;
6997 self.validate_retriever(&retriever)?;
6998 let hits = self.retrieve_filtered(
6999 &retriever,
7000 snapshot,
7001 None,
7002 allowed,
7003 candidate_authorization,
7004 context,
7005 )?;
7006 let distances: std::collections::HashMap<_, _> = hits
7007 .iter()
7008 .filter_map(|hit| match hit.score {
7009 RetrieverScore::AnnHammingDistance(distance) => {
7010 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
7011 }
7012 RetrieverScore::AnnCosineDistance(distance) => {
7013 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
7014 }
7015 _ => None,
7016 })
7017 .collect();
7018 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7019 if let Some(context) = context {
7020 context.consume(row_ids.len())?;
7021 }
7022 let gather_started = std::time::Instant::now();
7023 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7024 let values = self.values_for_rids_batch_at_with_context(
7025 &row_ids,
7026 request.column_id,
7027 snapshot,
7028 query_now,
7029 context,
7030 )?;
7031 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
7032 let score_started = std::time::Instant::now();
7033 let vector_work =
7034 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
7035 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
7036 if let Some(context) = context {
7037 context.consume(vector_work)?;
7038 }
7039 request
7040 .query
7041 .iter()
7042 .map(|value| f64::from(*value).powi(2))
7043 .sum::<f64>()
7044 .sqrt()
7045 } else {
7046 0.0
7047 };
7048 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
7049 for (row_id, value) in values {
7050 if let Some(meta) = value.generated_embedding_metadata() {
7051 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
7052 continue;
7053 }
7054 }
7055 let Some(vector) = value.as_embedding() else {
7056 continue;
7057 };
7058 let exact_score = match request.metric {
7059 VectorMetric::DotProduct => {
7060 if let Some(context) = context {
7061 context.consume(vector_work)?;
7062 }
7063 request
7064 .query
7065 .iter()
7066 .zip(vector)
7067 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7068 .sum::<f64>()
7069 }
7070 VectorMetric::Cosine => {
7071 if let Some(context) = context {
7072 context.consume(vector_work.saturating_mul(2))?;
7073 }
7074 let dot = request
7075 .query
7076 .iter()
7077 .zip(vector)
7078 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7079 .sum::<f64>();
7080 let norm = vector
7081 .iter()
7082 .map(|value| f64::from(*value).powi(2))
7083 .sum::<f64>()
7084 .sqrt();
7085 if query_norm == 0.0 || norm == 0.0 {
7086 0.0
7087 } else {
7088 dot / (query_norm * norm)
7089 }
7090 }
7091 VectorMetric::Euclidean => {
7092 if let Some(context) = context {
7093 context.consume(vector_work)?;
7094 }
7095 request
7096 .query
7097 .iter()
7098 .zip(vector)
7099 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
7100 .sum::<f64>()
7101 .sqrt()
7102 }
7103 };
7104 let exact_score = exact_score as f32;
7105 if !exact_score.is_finite() {
7106 return Err(MongrelError::InvalidArgument(
7107 "exact ANN score must be finite".into(),
7108 ));
7109 }
7110 let Some(candidate_distance) = distances.get(&row_id).copied() else {
7111 continue;
7112 };
7113 reranked.push(AnnRerankHit {
7114 row_id,
7115 candidate_distance,
7116 exact_score,
7117 });
7118 }
7119 reranked.sort_by(|left, right| {
7120 let score = match request.metric {
7121 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
7122 VectorMetric::Cosine | VectorMetric::DotProduct => {
7123 right.exact_score.total_cmp(&left.exact_score)
7124 }
7125 };
7126 score.then_with(|| left.row_id.cmp(&right.row_id))
7127 });
7128 reranked.truncate(request.limit);
7129 crate::trace::QueryTrace::record(|trace| {
7130 trace.exact_vector_gather_nanos =
7131 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
7132 trace.exact_vector_score_nanos = trace
7133 .exact_vector_score_nanos
7134 .saturating_add(score_started.elapsed().as_nanos() as u64);
7135 });
7136 Ok(reranked)
7137 }
7138
7139 pub fn set_similarity_with_allowed(
7140 &mut self,
7141 request: &crate::query::SetSimilarityRequest,
7142 allowed: Option<&std::collections::HashSet<RowId>>,
7143 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7144 self.set_similarity_explained_at(request, self.snapshot(), allowed)
7145 .map(|(hits, _)| hits)
7146 }
7147
7148 pub fn set_similarity_explained(
7149 &mut self,
7150 request: &crate::query::SetSimilarityRequest,
7151 ) -> Result<(
7152 Vec<crate::query::SetSimilarityHit>,
7153 crate::query::SetSimilarityTrace,
7154 )> {
7155 self.set_similarity_explained_at(request, self.snapshot(), None)
7156 }
7157
7158 fn set_similarity_explained_at(
7159 &mut self,
7160 request: &crate::query::SetSimilarityRequest,
7161 snapshot: Snapshot,
7162 allowed: Option<&std::collections::HashSet<RowId>>,
7163 ) -> Result<(
7164 Vec<crate::query::SetSimilarityHit>,
7165 crate::query::SetSimilarityTrace,
7166 )> {
7167 self.ensure_indexes_complete()?;
7168 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
7169 }
7170
7171 pub fn set_similarity_at_with_context(
7172 &mut self,
7173 request: &crate::query::SetSimilarityRequest,
7174 snapshot: Snapshot,
7175 allowed: Option<&std::collections::HashSet<RowId>>,
7176 context: Option<&crate::query::AiExecutionContext>,
7177 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7178 self.ensure_indexes_complete()?;
7179 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
7180 .map(|(hits, _)| hits)
7181 }
7182
7183 pub fn set_similarity_at_with_candidate_authorization_and_context(
7184 &mut self,
7185 request: &crate::query::SetSimilarityRequest,
7186 snapshot: Snapshot,
7187 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7188 context: Option<&crate::query::AiExecutionContext>,
7189 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7190 self.ensure_indexes_complete()?;
7191 self.set_similarity_explained_at_with_context(
7192 request,
7193 snapshot,
7194 None,
7195 authorization,
7196 context,
7197 )
7198 .map(|(hits, _)| hits)
7199 }
7200
7201 #[doc(hidden)]
7202 pub fn set_similarity_at_with_candidate_authorization_on_generation(
7203 &self,
7204 request: &crate::query::SetSimilarityRequest,
7205 snapshot: Snapshot,
7206 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7207 context: Option<&crate::query::AiExecutionContext>,
7208 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7209 self.set_similarity_explained_at_with_context(
7210 request,
7211 snapshot,
7212 None,
7213 authorization,
7214 context,
7215 )
7216 .map(|(hits, _)| hits)
7217 }
7218
7219 fn set_similarity_explained_at_with_context(
7220 &self,
7221 request: &crate::query::SetSimilarityRequest,
7222 snapshot: Snapshot,
7223 allowed: Option<&std::collections::HashSet<RowId>>,
7224 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7225 context: Option<&crate::query::AiExecutionContext>,
7226 ) -> Result<(
7227 Vec<crate::query::SetSimilarityHit>,
7228 crate::query::SetSimilarityTrace,
7229 )> {
7230 use crate::query::{
7231 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
7232 MAX_SET_MEMBERS,
7233 };
7234 let mut trace = crate::query::SetSimilarityTrace::default();
7235 if request.members.is_empty() {
7236 return Ok((Vec::new(), trace));
7237 }
7238 if request.candidate_k == 0 || request.limit == 0 {
7239 return Err(MongrelError::InvalidArgument(
7240 "candidate_k and limit must be > 0".into(),
7241 ));
7242 }
7243 if request.candidate_k > MAX_RETRIEVER_K
7244 || request.limit > MAX_FINAL_LIMIT
7245 || request.members.len() > MAX_SET_MEMBERS
7246 {
7247 return Err(MongrelError::InvalidArgument(format!(
7248 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
7249 )));
7250 }
7251 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
7252 return Err(MongrelError::InvalidArgument(
7253 "min_jaccard must be finite and between 0 and 1".into(),
7254 ));
7255 }
7256 let started = std::time::Instant::now();
7257 let retriever = Retriever::MinHash {
7258 column_id: request.column_id,
7259 members: request.members.clone(),
7260 k: request.candidate_k,
7261 };
7262 self.require_select()?;
7263 self.validate_retriever(&retriever)?;
7264 let hits = self.retrieve_filtered(
7265 &retriever,
7266 snapshot,
7267 None,
7268 allowed,
7269 candidate_authorization,
7270 context,
7271 )?;
7272 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
7273 trace.candidate_count = hits.len();
7274 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7275 if let Some(context) = context {
7276 context.consume(row_ids.len())?;
7277 }
7278 let started = std::time::Instant::now();
7279 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7280 let values = self.values_for_rids_batch_at_with_context(
7281 &row_ids,
7282 request.column_id,
7283 snapshot,
7284 query_now,
7285 context,
7286 )?;
7287 trace.gather_us = started.elapsed().as_micros() as u64;
7288 if let Some(context) = context {
7289 context.consume(request.members.len())?;
7290 }
7291 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
7292 let estimates: std::collections::HashMap<_, _> = hits
7293 .into_iter()
7294 .filter_map(|hit| match hit.score {
7295 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
7296 _ => None,
7297 })
7298 .collect();
7299 let started = std::time::Instant::now();
7300 let mut parsed = Vec::with_capacity(values.len());
7301 for (row_id, value) in values {
7302 let Value::Bytes(bytes) = value else {
7303 continue;
7304 };
7305 if let Some(context) = context {
7306 context.consume(crate::query::work_units(
7307 bytes.len(),
7308 crate::query::PARSE_WORK_QUANTUM,
7309 ))?;
7310 }
7311 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
7312 continue;
7313 };
7314 if let Some(context) = context {
7315 context.consume(members.len())?;
7316 }
7317 let stored = members
7318 .into_iter()
7319 .filter_map(|member| match member {
7320 serde_json::Value::String(value) => {
7321 Some(crate::query::SetMember::String(value))
7322 }
7323 serde_json::Value::Number(value) => {
7324 Some(crate::query::SetMember::Number(value))
7325 }
7326 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
7327 _ => None,
7328 })
7329 .collect::<std::collections::HashSet<_>>();
7330 parsed.push((row_id, stored));
7331 }
7332 trace.parse_us = started.elapsed().as_micros() as u64;
7333 trace.verified_count = parsed.len();
7334 let started = std::time::Instant::now();
7335 let mut exact = Vec::new();
7336 for (row_id, stored) in parsed {
7337 if let Some(context) = context {
7338 context.consume(query.len().saturating_add(stored.len()))?;
7339 }
7340 let union = query.union(&stored).count();
7341 let score = if union == 0 {
7342 1.0
7343 } else {
7344 query.intersection(&stored).count() as f32 / union as f32
7345 };
7346 if score >= request.min_jaccard {
7347 exact.push(SetSimilarityHit {
7348 row_id,
7349 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
7350 exact_jaccard: score,
7351 });
7352 }
7353 }
7354 exact.sort_by(|a, b| {
7355 b.exact_jaccard
7356 .total_cmp(&a.exact_jaccard)
7357 .then_with(|| a.row_id.cmp(&b.row_id))
7358 });
7359 exact.truncate(request.limit);
7360 trace.score_us = started.elapsed().as_micros() as u64;
7361 crate::trace::QueryTrace::record(|query_trace| {
7362 query_trace.exact_set_gather_nanos = query_trace
7363 .exact_set_gather_nanos
7364 .saturating_add(trace.gather_us.saturating_mul(1_000));
7365 query_trace.exact_set_parse_nanos = query_trace
7366 .exact_set_parse_nanos
7367 .saturating_add(trace.parse_us.saturating_mul(1_000));
7368 query_trace.exact_set_score_nanos = query_trace
7369 .exact_set_score_nanos
7370 .saturating_add(trace.score_us.saturating_mul(1_000));
7371 });
7372 Ok((exact, trace))
7373 }
7374
7375 fn values_for_rids_batch_at(
7377 &self,
7378 row_ids: &[u64],
7379 column_id: u16,
7380 snapshot: Snapshot,
7381 now: i64,
7382 ) -> Result<Vec<(RowId, Value)>> {
7383 if self.ttl.is_none()
7384 && self.memtable.is_empty()
7385 && self.mutable_run.is_empty()
7386 && self.run_refs.len() == 1
7387 {
7388 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7389 if row_ids.len().saturating_mul(24) < reader.row_count() {
7394 let mut values = Vec::with_capacity(row_ids.len());
7395 for &raw_row_id in row_ids {
7396 let row_id = RowId(raw_row_id);
7397 if let Some((_, false, Some(value))) =
7398 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7399 {
7400 values.push((row_id, value));
7401 }
7402 }
7403 return Ok(values);
7404 }
7405 let (positions, visible_row_ids) =
7406 reader.visible_positions_with_rids(snapshot.epoch)?;
7407 let requested: Vec<(RowId, usize)> = row_ids
7408 .iter()
7409 .filter_map(|raw| {
7410 visible_row_ids
7411 .binary_search(&(*raw as i64))
7412 .ok()
7413 .map(|index| (RowId(*raw), positions[index]))
7414 })
7415 .collect();
7416 let values = reader.gather_column(
7417 column_id,
7418 &requested
7419 .iter()
7420 .map(|(_, position)| *position)
7421 .collect::<Vec<_>>(),
7422 )?;
7423 return Ok(requested
7424 .into_iter()
7425 .zip(values)
7426 .map(|((row_id, _), value)| (row_id, value))
7427 .collect());
7428 }
7429 self.values_for_rids_at(row_ids, column_id, snapshot, now)
7430 }
7431
7432 fn values_for_rids_batch_at_with_context(
7433 &self,
7434 row_ids: &[u64],
7435 column_id: u16,
7436 snapshot: Snapshot,
7437 now: i64,
7438 context: Option<&crate::query::AiExecutionContext>,
7439 ) -> Result<Vec<(RowId, Value)>> {
7440 let Some(context) = context else {
7441 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
7442 };
7443 let mut values = Vec::with_capacity(row_ids.len());
7444 for chunk in row_ids.chunks(256) {
7445 context.checkpoint()?;
7446 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
7447 }
7448 Ok(values)
7449 }
7450
7451 fn values_for_rids_at(
7453 &self,
7454 row_ids: &[u64],
7455 column_id: u16,
7456 snapshot: Snapshot,
7457 now: i64,
7458 ) -> Result<Vec<(RowId, Value)>> {
7459 let mut readers: Vec<_> = self
7460 .run_refs
7461 .iter()
7462 .map(|run| self.open_reader(run.run_id))
7463 .collect::<Result<_>>()?;
7464 let mut out = Vec::with_capacity(row_ids.len());
7465 for &raw_row_id in row_ids {
7466 let row_id = RowId(raw_row_id);
7467 let mem = self.memtable.get_version_at(row_id, snapshot);
7468 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
7469 let overlay = match (mem, mutable) {
7470 (Some((_, a)), Some((_, b))) => Some(
7471 if Snapshot::version_is_newer(
7472 a.committed_epoch,
7473 a.commit_ts,
7474 b.committed_epoch,
7475 b.commit_ts,
7476 ) {
7477 a
7478 } else {
7479 b
7480 },
7481 ),
7482 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
7483 (None, None) => None,
7484 };
7485 if let Some(row) = overlay {
7486 if !row.deleted && !self.row_expired_at(&row, now) {
7487 if let Some(value) = row.columns.get(&column_id) {
7488 out.push((row_id, value.clone()));
7489 }
7490 }
7491 continue;
7492 }
7493
7494 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
7495 for (index, reader) in readers.iter_mut().enumerate() {
7496 if let Some((epoch, deleted, value)) =
7497 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7498 {
7499 if best
7500 .as_ref()
7501 .map(|(best_epoch, ..)| epoch > *best_epoch)
7502 .unwrap_or(true)
7503 {
7504 best = Some((epoch, deleted, value, index));
7505 }
7506 }
7507 }
7508 let Some((_, false, Some(value), reader_index)) = best else {
7509 continue;
7510 };
7511 if let Some(ttl) = self.ttl {
7512 if ttl.column_id != column_id {
7513 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7514 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
7515 {
7516 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7517 continue;
7518 }
7519 }
7520 } else if let Value::Int64(timestamp) = value {
7521 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7522 continue;
7523 }
7524 }
7525 }
7526 out.push((row_id, value));
7527 }
7528 Ok(out)
7529 }
7530
7531 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
7536 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
7537 }
7538
7539 pub fn rows_for_rids_with_context(
7540 &self,
7541 rids: &[u64],
7542 snapshot: Snapshot,
7543 context: &crate::query::AiExecutionContext,
7544 ) -> Result<Vec<Row>> {
7545 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
7546 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
7547 }
7548
7549 fn rows_for_rids_at_time(
7550 &self,
7551 rids: &[u64],
7552 snapshot: Snapshot,
7553 ttl_now: i64,
7554 control: Option<&crate::ExecutionControl>,
7555 ) -> Result<Vec<Row>> {
7556 use std::collections::HashMap;
7557 let mut rows = Vec::with_capacity(rids.len());
7558 let tier_size = self.memtable.len() + self.mutable_run.len();
7573 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7574 if rids.len().saturating_mul(24) < tier_size {
7575 for &rid in rids {
7576 if overlay.len() & 255 == 0 {
7577 control
7578 .map(crate::ExecutionControl::checkpoint)
7579 .transpose()?;
7580 }
7581 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
7582 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
7583 let newest = match (mem, mrun) {
7584 (Some((_, mr)), Some((_, rr))) => Some(
7585 if Snapshot::version_is_newer(
7586 mr.committed_epoch,
7587 mr.commit_ts,
7588 rr.committed_epoch,
7589 rr.commit_ts,
7590 ) {
7591 mr
7592 } else {
7593 rr
7594 },
7595 ),
7596 (Some((_, mr)), None) => Some(mr),
7597 (None, Some((_, rr))) => Some(rr),
7598 (None, None) => None,
7599 };
7600 if let Some(row) = newest {
7601 overlay.insert(rid, row);
7602 }
7603 }
7604 } else {
7605 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7606 overlay
7607 .entry(row.row_id.0)
7608 .and_modify(|e| {
7609 if Snapshot::version_is_newer(
7610 row.committed_epoch,
7611 row.commit_ts,
7612 e.committed_epoch,
7613 e.commit_ts,
7614 ) {
7615 *e = row.clone();
7616 }
7617 })
7618 .or_insert(row);
7619 };
7620 for (index, row) in self
7621 .memtable
7622 .visible_versions_at(snapshot)
7623 .into_iter()
7624 .enumerate()
7625 {
7626 if index & 255 == 0 {
7627 control
7628 .map(crate::ExecutionControl::checkpoint)
7629 .transpose()?;
7630 }
7631 fold_newest(row, &mut overlay);
7632 }
7633 for (index, row) in self
7634 .mutable_run
7635 .visible_versions_at(snapshot)
7636 .into_iter()
7637 .enumerate()
7638 {
7639 if index & 255 == 0 {
7640 control
7641 .map(crate::ExecutionControl::checkpoint)
7642 .transpose()?;
7643 }
7644 fold_newest(row, &mut overlay);
7645 }
7646 }
7647 if self.run_refs.len() == 1 {
7648 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7649 if rids.len().saturating_mul(24) < reader.row_count() {
7657 for (index, &rid) in rids.iter().enumerate() {
7658 if index & 255 == 0 {
7659 control
7660 .map(crate::ExecutionControl::checkpoint)
7661 .transpose()?;
7662 }
7663 if let Some(r) = overlay.get(&rid) {
7664 if !r.deleted {
7665 rows.push(r.clone());
7666 }
7667 continue;
7668 }
7669 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7670 if !row.deleted {
7671 rows.push(row);
7672 }
7673 }
7674 }
7675 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7676 return Ok(rows);
7677 }
7678 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7687 enum Src {
7690 Overlay,
7691 Run,
7692 }
7693 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7694 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7695 for (index, rid) in rids.iter().enumerate() {
7696 if index & 255 == 0 {
7697 control
7698 .map(crate::ExecutionControl::checkpoint)
7699 .transpose()?;
7700 }
7701 if overlay.contains_key(rid) {
7702 plan.push(Src::Overlay);
7703 continue;
7704 }
7705 match vis_rids.binary_search(&(*rid as i64)) {
7706 Ok(i) => {
7707 plan.push(Src::Run);
7708 fetch.push(positions[i]);
7709 }
7710 Err(_) => { }
7711 }
7712 }
7713 let fetched = reader.materialize_batch(&fetch)?;
7714 let mut fetched_iter = fetched.into_iter();
7715 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7716 if index & 255 == 0 {
7717 control
7718 .map(crate::ExecutionControl::checkpoint)
7719 .transpose()?;
7720 }
7721 match src {
7722 Src::Overlay => {
7723 if let Some(r) = overlay.get(rid) {
7724 if !r.deleted {
7725 rows.push(r.clone());
7726 }
7727 }
7728 }
7729 Src::Run => {
7730 if let Some(row) = fetched_iter.next() {
7731 if !row.deleted {
7732 rows.push(row);
7733 }
7734 }
7735 }
7736 }
7737 }
7738 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7739 return Ok(rows);
7740 }
7741 let mut readers: Vec<_> = self
7745 .run_refs
7746 .iter()
7747 .map(|rr| self.open_reader(rr.run_id))
7748 .collect::<Result<Vec<_>>>()?;
7749 for (index, rid) in rids.iter().enumerate() {
7750 if index & 255 == 0 {
7751 control
7752 .map(crate::ExecutionControl::checkpoint)
7753 .transpose()?;
7754 }
7755 if let Some(r) = overlay.get(rid) {
7756 if !r.deleted {
7757 rows.push(r.clone());
7758 }
7759 continue;
7760 }
7761 let mut best: Option<(Epoch, Row)> = None;
7762 for reader in readers.iter_mut() {
7763 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7764 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7765 best = Some((epoch, row));
7766 }
7767 }
7768 }
7769 if let Some((_, r)) = best {
7770 if !r.deleted {
7771 rows.push(r);
7772 }
7773 }
7774 }
7775 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7776 Ok(rows)
7777 }
7778
7779 pub fn indexes_complete(&self) -> bool {
7789 self.indexes_complete
7790 }
7791
7792 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7794 self.index_build_policy
7795 }
7796
7797 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7801 self.index_build_policy = policy;
7802 }
7803
7804 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7809 if !self.indexes_complete {
7813 return None;
7814 }
7815 let b = self.bitmap.get(&column_id)?;
7816 let result: Vec<Vec<u8>> = b
7817 .keys()
7818 .into_iter()
7819 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7820 .collect();
7821 Some(result)
7822 }
7823
7824 pub fn fk_join_row_ids(
7825 &self,
7826 fk_column_id: u16,
7827 pk_values: &[Vec<u8>],
7828 fk_conditions: &[crate::query::Condition],
7829 snapshot: Snapshot,
7830 ) -> Result<Vec<u64>> {
7831 let Some(b) = self.bitmap.get(&fk_column_id) else {
7832 return Ok(Vec::new());
7833 };
7834 let mut join_set = {
7835 let mut acc = roaring::RoaringBitmap::new();
7836 for v in pk_values {
7837 acc |= b.get(v);
7838 }
7839 RowIdSet::from_roaring(acc)
7840 };
7841 if !fk_conditions.is_empty() {
7842 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7843 sets.push(join_set);
7844 for c in fk_conditions {
7845 sets.push(self.resolve_condition(c, snapshot)?);
7846 }
7847 join_set = RowIdSet::intersect_many(sets);
7848 }
7849 Ok(join_set.into_sorted_vec())
7850 }
7851
7852 pub fn fk_join_count(
7858 &self,
7859 fk_column_id: u16,
7860 pk_values: &[Vec<u8>],
7861 fk_conditions: &[crate::query::Condition],
7862 snapshot: Snapshot,
7863 ) -> Result<u64> {
7864 let Some(b) = self.bitmap.get(&fk_column_id) else {
7865 return Ok(0);
7866 };
7867 let mut acc = roaring::RoaringBitmap::new();
7868 for v in pk_values {
7869 acc |= b.get(v);
7870 }
7871 if fk_conditions.is_empty() {
7872 return Ok(acc.len());
7873 }
7874 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7875 sets.push(RowIdSet::from_roaring(acc));
7876 for c in fk_conditions {
7877 sets.push(self.resolve_condition(c, snapshot)?);
7878 }
7879 Ok(RowIdSet::intersect_many(sets).len() as u64)
7880 }
7881
7882 fn resolve_condition(
7887 &self,
7888 c: &crate::query::Condition,
7889 snapshot: Snapshot,
7890 ) -> Result<RowIdSet> {
7891 self.resolve_condition_with_allowed(c, snapshot, None)
7892 }
7893
7894 fn resolve_condition_with_allowed(
7895 &self,
7896 c: &crate::query::Condition,
7897 snapshot: Snapshot,
7898 allowed: Option<&std::collections::HashSet<RowId>>,
7899 ) -> Result<RowIdSet> {
7900 use crate::query::Condition;
7901 self.validate_condition(c)?;
7902 Ok(match c {
7903 Condition::Pk(key) => {
7904 let lookup = self
7905 .schema
7906 .primary_key()
7907 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7908 .unwrap_or_else(|| key.clone());
7909 if let Some(r) = self.hot.get(&lookup) {
7910 RowIdSet::one(r.0)
7911 } else if let Some(pk_col) = self.schema.primary_key() {
7912 self.pk_equality_fallback(pk_col.id, &lookup, snapshot)?
7918 } else {
7919 RowIdSet::empty()
7920 }
7921 }
7922 Condition::BitmapEq { column_id, value } => {
7923 let lookup = self.index_lookup_key_bytes(*column_id, value);
7924 self.bitmap
7925 .get(column_id)
7926 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7927 .unwrap_or_else(RowIdSet::empty)
7928 }
7929 Condition::BitmapIn { column_id, values } => {
7930 let bm = self.bitmap.get(column_id);
7931 let mut acc = roaring::RoaringBitmap::new();
7932 if let Some(b) = bm {
7933 for v in values {
7934 let lookup = self.index_lookup_key_bytes(*column_id, v);
7935 acc |= b.get(&lookup);
7936 }
7937 }
7938 RowIdSet::from_roaring(acc)
7939 }
7940 Condition::BytesPrefix { column_id, prefix } => {
7941 if let Some(b) = self.bitmap.get(column_id) {
7946 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7947 let mut acc = roaring::RoaringBitmap::new();
7948 for key in b.keys() {
7949 if key.starts_with(&lookup_prefix) {
7950 acc |= b.get(&key);
7951 }
7952 }
7953 RowIdSet::from_roaring(acc)
7954 } else {
7955 RowIdSet::empty()
7956 }
7957 }
7958 Condition::FmContains { column_id, pattern } => self
7959 .fm
7960 .get(column_id)
7961 .map(|f| {
7962 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7963 })
7964 .unwrap_or_else(RowIdSet::empty),
7965 Condition::FmContainsAll {
7966 column_id,
7967 patterns,
7968 } => {
7969 if let Some(f) = self.fm.get(column_id) {
7972 let sets: Vec<RowIdSet> = patterns
7973 .iter()
7974 .map(|pat| {
7975 RowIdSet::from_unsorted(
7976 f.locate(pat).into_iter().map(|r| r.0).collect(),
7977 )
7978 })
7979 .collect();
7980 RowIdSet::intersect_many(sets)
7981 } else {
7982 RowIdSet::empty()
7983 }
7984 }
7985 Condition::Ann {
7986 column_id,
7987 query,
7988 k,
7989 } => RowIdSet::from_unsorted(
7990 self.retrieve_filtered(
7991 &crate::query::Retriever::Ann {
7992 column_id: *column_id,
7993 query: query.clone(),
7994 k: *k,
7995 },
7996 snapshot,
7997 None,
7998 allowed,
7999 None,
8000 None,
8001 )?
8002 .into_iter()
8003 .map(|hit| hit.row_id.0)
8004 .collect(),
8005 ),
8006 Condition::SparseMatch {
8007 column_id,
8008 query,
8009 k,
8010 } => RowIdSet::from_unsorted(
8011 self.retrieve_filtered(
8012 &crate::query::Retriever::Sparse {
8013 column_id: *column_id,
8014 query: query.clone(),
8015 k: *k,
8016 },
8017 snapshot,
8018 None,
8019 allowed,
8020 None,
8021 None,
8022 )?
8023 .into_iter()
8024 .map(|hit| hit.row_id.0)
8025 .collect(),
8026 ),
8027 Condition::MinHashSimilar {
8028 column_id,
8029 query,
8030 k,
8031 } => match self.minhash.get(column_id) {
8032 Some(index) => {
8033 let candidates = index.candidate_row_ids(query);
8034 let eligible =
8035 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
8036 RowIdSet::from_unsorted(
8037 index
8038 .search_filtered(query, *k, |row_id| {
8039 eligible.contains(&row_id)
8040 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
8041 })
8042 .into_iter()
8043 .map(|(row_id, _)| row_id.0)
8044 .collect(),
8045 )
8046 }
8047 None => RowIdSet::empty(),
8048 },
8049 Condition::Range { column_id, lo, hi } => {
8050 let mut set = if let Some(li) = self.learned_range.get(column_id) {
8069 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
8070 } else if self.run_refs.len() == 1 {
8071 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8072 r.range_row_id_set_i64(*column_id, *lo, *hi)?
8073 } else {
8074 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
8077 if lo == hi {
8078 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
8079 }
8080 return Ok(multi);
8081 };
8082 set.remove_many(self.overlay_rid_set(snapshot));
8083 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
8084 if lo == hi {
8085 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
8086 }
8087 set
8088 }
8089 Condition::RangeF64 {
8090 column_id,
8091 lo,
8092 lo_inclusive,
8093 hi,
8094 hi_inclusive,
8095 } => {
8096 let mut set = if let Some(li) = self.learned_range.get(column_id) {
8099 RowIdSet::from_unsorted(
8100 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
8101 .into_iter()
8102 .collect(),
8103 )
8104 } else if self.run_refs.len() == 1 {
8105 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8106 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
8107 } else {
8108 return self.range_scan_f64(
8109 *column_id,
8110 *lo,
8111 *lo_inclusive,
8112 *hi,
8113 *hi_inclusive,
8114 snapshot,
8115 );
8116 };
8117 set.remove_many(self.overlay_rid_set(snapshot));
8118 self.range_scan_overlay_f64(
8119 &mut set,
8120 *column_id,
8121 *lo,
8122 *lo_inclusive,
8123 *hi,
8124 *hi_inclusive,
8125 snapshot,
8126 );
8127 set
8128 }
8129 Condition::IsNull { column_id } => {
8130 let mut set = if self.run_refs.len() == 1 {
8131 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8132 r.null_row_id_set(*column_id, true)?
8133 } else {
8134 return self.null_scan(*column_id, true, snapshot);
8135 };
8136 set.remove_many(self.overlay_rid_set(snapshot));
8137 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
8138 set
8139 }
8140 Condition::IsNotNull { column_id } => {
8141 let mut set = if self.run_refs.len() == 1 {
8142 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8143 r.null_row_id_set(*column_id, false)?
8144 } else {
8145 return self.null_scan(*column_id, false, snapshot);
8146 };
8147 set.remove_many(self.overlay_rid_set(snapshot));
8148 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
8149 set
8150 }
8151 })
8152 }
8153
8154 fn range_scan_i64(
8162 &self,
8163 column_id: u16,
8164 lo: i64,
8165 hi: i64,
8166 snapshot: Snapshot,
8167 ) -> Result<RowIdSet> {
8168 let mut row_ids = Vec::new();
8169 let overlay_rids = self.overlay_rid_set(snapshot);
8170 for rr in &self.run_refs {
8171 let mut reader = self.open_reader(rr.run_id)?;
8172 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
8173 for rid in matched {
8174 if !overlay_rids.contains(&rid) {
8175 row_ids.push(rid);
8176 }
8177 }
8178 }
8179 let mut s = RowIdSet::from_unsorted(row_ids);
8180 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
8181 Ok(s)
8182 }
8183
8184 fn range_scan_f64(
8187 &self,
8188 column_id: u16,
8189 lo: f64,
8190 lo_inclusive: bool,
8191 hi: f64,
8192 hi_inclusive: bool,
8193 snapshot: Snapshot,
8194 ) -> Result<RowIdSet> {
8195 let mut row_ids = Vec::new();
8196 let overlay_rids = self.overlay_rid_set(snapshot);
8197 for rr in &self.run_refs {
8198 let mut reader = self.open_reader(rr.run_id)?;
8199 let matched = reader.range_row_ids_visible_f64(
8200 column_id,
8201 lo,
8202 lo_inclusive,
8203 hi,
8204 hi_inclusive,
8205 snapshot.epoch,
8206 )?;
8207 for rid in matched {
8208 if !overlay_rids.contains(&rid) {
8209 row_ids.push(rid);
8210 }
8211 }
8212 }
8213 let mut s = RowIdSet::from_unsorted(row_ids);
8214 self.range_scan_overlay_f64(
8215 &mut s,
8216 column_id,
8217 lo,
8218 lo_inclusive,
8219 hi,
8220 hi_inclusive,
8221 snapshot,
8222 );
8223 Ok(s)
8224 }
8225
8226 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
8228 let mut s = HashSet::new();
8229 for row in self.memtable.visible_versions_at(snapshot) {
8230 s.insert(row.row_id.0);
8231 }
8232 for row in self.mutable_run.visible_versions_at(snapshot) {
8233 s.insert(row.row_id.0);
8234 }
8235 s
8236 }
8237
8238 fn range_scan_overlay_i64(
8239 &self,
8240 s: &mut RowIdSet,
8241 column_id: u16,
8242 lo: i64,
8243 hi: i64,
8244 snapshot: Snapshot,
8245 ) {
8246 let mut newest: HashMap<u64, Row> = HashMap::new();
8253 for r in self.mutable_run.visible_versions_at(snapshot) {
8254 newest.insert(r.row_id.0, r);
8255 }
8256 for r in self.memtable.visible_versions_at(snapshot) {
8257 newest
8258 .entry(r.row_id.0)
8259 .and_modify(|cur| {
8260 if Snapshot::version_is_newer(
8261 r.committed_epoch,
8262 r.commit_ts,
8263 cur.committed_epoch,
8264 cur.commit_ts,
8265 ) {
8266 *cur = r.clone();
8267 }
8268 })
8269 .or_insert(r);
8270 }
8271 for row in newest.values() {
8272 if !row.deleted {
8273 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
8274 if *v >= lo && *v <= hi {
8275 s.insert(row.row_id.0);
8276 }
8277 }
8278 }
8279 }
8280 }
8281
8282 #[allow(clippy::too_many_arguments)]
8283 fn range_scan_overlay_f64(
8284 &self,
8285 s: &mut RowIdSet,
8286 column_id: u16,
8287 lo: f64,
8288 lo_inclusive: bool,
8289 hi: f64,
8290 hi_inclusive: bool,
8291 snapshot: Snapshot,
8292 ) {
8293 let mut newest: HashMap<u64, Row> = HashMap::new();
8296 for r in self.mutable_run.visible_versions_at(snapshot) {
8297 newest.insert(r.row_id.0, r);
8298 }
8299 for r in self.memtable.visible_versions_at(snapshot) {
8300 newest
8301 .entry(r.row_id.0)
8302 .and_modify(|cur| {
8303 if Snapshot::version_is_newer(
8304 r.committed_epoch,
8305 r.commit_ts,
8306 cur.committed_epoch,
8307 cur.commit_ts,
8308 ) {
8309 *cur = r.clone();
8310 }
8311 })
8312 .or_insert(r);
8313 }
8314 for row in newest.values() {
8315 if !row.deleted {
8316 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
8317 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
8318 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
8319 if ok_lo && ok_hi {
8320 s.insert(row.row_id.0);
8321 }
8322 }
8323 }
8324 }
8325 }
8326
8327 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
8330 let mut row_ids = Vec::new();
8331 let overlay_rids = self.overlay_rid_set(snapshot);
8332 for rr in &self.run_refs {
8333 let mut reader = self.open_reader(rr.run_id)?;
8334 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
8335 for rid in matched {
8336 if !overlay_rids.contains(&rid) {
8337 row_ids.push(rid);
8338 }
8339 }
8340 }
8341 let mut s = RowIdSet::from_unsorted(row_ids);
8342 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
8343 Ok(s)
8344 }
8345
8346 fn null_scan_overlay(
8350 &self,
8351 s: &mut RowIdSet,
8352 column_id: u16,
8353 want_nulls: bool,
8354 snapshot: Snapshot,
8355 ) {
8356 let mut newest: HashMap<u64, Row> = HashMap::new();
8357 for r in self.mutable_run.visible_versions_at(snapshot) {
8358 newest.insert(r.row_id.0, r);
8359 }
8360 for r in self.memtable.visible_versions_at(snapshot) {
8361 newest
8362 .entry(r.row_id.0)
8363 .and_modify(|cur| {
8364 if Snapshot::version_is_newer(
8365 r.committed_epoch,
8366 r.commit_ts,
8367 cur.committed_epoch,
8368 cur.commit_ts,
8369 ) {
8370 *cur = r.clone();
8371 }
8372 })
8373 .or_insert(r);
8374 }
8375 for row in newest.values() {
8376 if row.deleted {
8377 continue;
8378 }
8379 let is_null = !row.columns.contains_key(&column_id)
8380 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
8381 if is_null == want_nulls {
8382 s.insert(row.row_id.0);
8383 }
8384 }
8385 }
8386
8387 pub fn snapshot(&self) -> Snapshot {
8388 let epoch = self.epoch.visible();
8389 match &self.wal {
8393 WalSink::Shared(shared) => match shared.hlc.now() {
8394 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
8395 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
8397 },
8398 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
8399 }
8400 }
8401
8402 pub fn data_generation(&self) -> u64 {
8404 self.data_generation
8405 }
8406
8407 pub(crate) fn bump_data_generation(&mut self) {
8408 self.data_generation = self.data_generation.wrapping_add(1);
8409 }
8410
8411 pub fn table_id(&self) -> u64 {
8413 self.table_id
8414 }
8415
8416 fn seal_generations(&mut self) {
8422 self.memtable.seal();
8423 self.mutable_run.seal();
8424 self.hot.seal();
8425 for index in self.bitmap.values_mut() {
8426 index.seal();
8427 }
8428 for index in self.ann.values_mut() {
8429 index.seal();
8430 }
8431 for index in self.fm.values_mut() {
8432 index.seal();
8433 }
8434 for index in self.sparse.values_mut() {
8435 index.seal();
8436 }
8437 for index in self.minhash.values_mut() {
8438 index.seal();
8439 }
8440 self.pk_by_row.seal();
8441 }
8442
8443 fn capture_read_generation(&self) -> ReadGeneration {
8448 let visible_through = self.current_epoch();
8449 ReadGeneration {
8450 schema: Arc::new(self.schema.clone()),
8451 base_runs: Arc::new(self.run_refs.clone()),
8452 deltas: TableDeltas {
8453 memtable: self.memtable.clone(),
8454 mutable_run: self.mutable_run.clone(),
8455 hot: self.hot.clone(),
8456 pk_by_row: self.pk_by_row.clone(),
8457 },
8458 indexes: Arc::new(IndexGeneration::capture(
8459 &self.bitmap,
8460 &self.learned_range,
8461 &self.fm,
8462 &self.ann,
8463 &self.sparse,
8464 &self.minhash,
8465 visible_through,
8466 match &self.wal {
8468 WalSink::Shared(shared) => shared
8469 .hlc
8470 .now()
8471 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
8472 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
8473 },
8474 )),
8475 visible_through,
8476 }
8477 }
8478
8479 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
8484 self.ensure_indexes_complete()?;
8485 self.seal_generations();
8486 let view = Arc::new(self.capture_read_generation());
8487 self.published.store(Arc::clone(&view));
8488 Ok(view)
8489 }
8490
8491 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
8497 self.published.load_full()
8498 }
8499
8500 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
8502 &self.pins
8503 }
8504
8505 pub fn version_gc_floor(&self) -> Epoch {
8510 self.min_active_snapshot()
8511 .unwrap_or_else(|| self.current_epoch())
8512 }
8513
8514 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
8521 let mut report = self.pins.report();
8522 let transaction_floor = [
8523 self.pinned.keys().next().copied(),
8524 self.snapshots.min_pinned(),
8525 ]
8526 .into_iter()
8527 .flatten()
8528 .min();
8529 if let Some(epoch) = transaction_floor {
8530 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
8531 }
8532 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
8533 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
8534 }
8535 report
8536 }
8537
8538 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
8539 self.publish_read_generation()?;
8540 let mut generation = self.clone();
8541 generation.read_only = true;
8542 generation.wal = WalSink::ReadOnly;
8543 generation.pending_delete_rids.clear();
8544 generation.pending_put_cols.clear();
8545 generation.pending_rows.clear();
8546 generation.pending_rows_auto_inc.clear();
8547 generation.pending_dels.clear();
8548 generation.pending_truncate = None;
8549 generation.agg_cache = Arc::new(HashMap::new());
8550 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
8553 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
8556 crate::retention::PinSource::ReadGeneration,
8557 self.current_epoch(),
8558 )));
8559 Ok(generation)
8560 }
8561
8562 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
8563 (std::mem::size_of::<Self>() as u64)
8564 .saturating_add(self.memtable.approx_bytes())
8565 .saturating_add(self.mutable_run.approx_bytes())
8566 .saturating_add(self.live_count.saturating_mul(64))
8567 }
8568
8569 pub fn pin_snapshot(&mut self) -> Snapshot {
8576 let snap = self.snapshot();
8577 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
8578 snap
8579 }
8580
8581 pub fn hlc_gc_floor(
8591 &self,
8592 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
8593 ) -> crate::epoch::GcFloor {
8594 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
8595 epoch
8596 .and_then(&mut project_epoch)
8597 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
8598 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
8599 };
8600 let transaction = [
8601 self.pinned.keys().next().copied(),
8602 self.snapshots.min_pinned(),
8603 ]
8604 .into_iter()
8605 .flatten()
8606 .min();
8607 let history = self.snapshots.history_floor(self.current_epoch());
8608 crate::epoch::GcFloor {
8609 transaction_snapshot: project(transaction),
8610 history_retention: project(history),
8611 backup_pitr: project(
8612 self.pins
8613 .oldest_for(crate::retention::PinSource::BackupPitr),
8614 ),
8615 replication: project(
8616 self.pins
8617 .oldest_for(crate::retention::PinSource::Replication),
8618 ),
8619 read_generation: project(
8620 self.pins
8621 .oldest_for(crate::retention::PinSource::ReadGeneration),
8622 ),
8623 online_index_build: project(
8624 self.pins
8625 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
8626 ),
8627 }
8628 }
8629
8630 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
8632 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
8633 *count -= 1;
8634 if *count == 0 {
8635 self.pinned.remove(&snap.epoch);
8636 }
8637 }
8638 }
8639
8640 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
8652 let local = self.pinned.keys().next().copied();
8653 let global = self.snapshots.min_pinned();
8654 let history = self.snapshots.history_floor(self.current_epoch());
8655 let pinned = self.pins.oldest_pinned();
8656 [local, global, history, pinned].into_iter().flatten().min()
8657 }
8658
8659 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
8663 self.ensure_writable()?;
8664 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
8665 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
8666 }
8667
8668 pub fn clear_ttl(&mut self) -> Result<()> {
8669 self.ensure_writable()?;
8670 self.apply_ttl_policy_at(None, self.current_epoch())
8671 }
8672
8673 pub fn ttl(&self) -> Option<TtlPolicy> {
8674 self.ttl
8675 }
8676
8677 pub(crate) fn prepare_ttl_policy(
8678 &self,
8679 column_name: &str,
8680 duration_nanos: u64,
8681 ) -> Result<TtlPolicy> {
8682 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
8683 return Err(MongrelError::InvalidArgument(
8684 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
8685 ));
8686 }
8687 let column = self
8688 .schema
8689 .columns
8690 .iter()
8691 .find(|column| column.name == column_name)
8692 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
8693 if column.ty != TypeId::TimestampNanos {
8694 return Err(MongrelError::Schema(format!(
8695 "TTL column {column_name} must be TimestampNanos, is {:?}",
8696 column.ty
8697 )));
8698 }
8699 Ok(TtlPolicy {
8700 column_id: column.id,
8701 duration_nanos,
8702 })
8703 }
8704
8705 pub(crate) fn apply_ttl_policy_at(
8706 &mut self,
8707 policy: Option<TtlPolicy>,
8708 epoch: Epoch,
8709 ) -> Result<()> {
8710 if let Some(policy) = policy {
8711 let column = self
8712 .schema
8713 .columns
8714 .iter()
8715 .find(|column| column.id == policy.column_id)
8716 .ok_or_else(|| {
8717 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8718 })?;
8719 if column.ty != TypeId::TimestampNanos
8720 || policy.duration_nanos == 0
8721 || policy.duration_nanos > i64::MAX as u64
8722 {
8723 return Err(MongrelError::Schema("invalid TTL policy".into()));
8724 }
8725 }
8726 self.ttl = policy;
8727 self.agg_cache = Arc::new(HashMap::new());
8728 self.clear_result_cache();
8729 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8730 self.persist_manifest(epoch)
8731 }
8732
8733 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8734 let Some(policy) = self.ttl else {
8735 return false;
8736 };
8737 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8738 return false;
8739 };
8740 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8741 }
8742
8743 pub fn current_epoch(&self) -> Epoch {
8744 self.epoch.visible()
8745 }
8746
8747 pub fn memtable_len(&self) -> usize {
8748 self.memtable.len()
8749 }
8750
8751 pub fn count(&self) -> u64 {
8754 if self.ttl.is_none()
8755 && self.pending_put_cols.is_empty()
8756 && self.pending_delete_rids.is_empty()
8757 && self.pending_rows.is_empty()
8758 && self.pending_dels.is_empty()
8759 && self.pending_truncate.is_none()
8760 {
8761 self.live_count
8762 } else {
8763 self.visible_rows(self.snapshot())
8764 .map(|rows| rows.len() as u64)
8765 .unwrap_or(self.live_count)
8766 }
8767 }
8768
8769 pub fn count_conditions(
8773 &mut self,
8774 conditions: &[crate::query::Condition],
8775 snapshot: Snapshot,
8776 ) -> Result<Option<u64>> {
8777 use crate::query::Condition;
8778 if self.ttl.is_some() {
8779 if conditions.is_empty() {
8780 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8781 }
8782 let mut sets = Vec::with_capacity(conditions.len());
8783 for condition in conditions {
8784 sets.push(self.resolve_condition(condition, snapshot)?);
8785 }
8786 let survivors = RowIdSet::intersect_many(sets);
8787 let rows = self.visible_rows(snapshot)?;
8788 return Ok(Some(
8789 rows.into_iter()
8790 .filter(|row| survivors.contains(row.row_id.0))
8791 .count() as u64,
8792 ));
8793 }
8794 if conditions.is_empty() {
8795 return Ok(Some(self.count()));
8796 }
8797 let served = |c: &Condition| {
8798 matches!(
8799 c,
8800 Condition::Pk(_)
8801 | Condition::BitmapEq { .. }
8802 | Condition::BitmapIn { .. }
8803 | Condition::BytesPrefix { .. }
8804 | Condition::FmContains { .. }
8805 | Condition::FmContainsAll { .. }
8806 | Condition::Ann { .. }
8807 | Condition::Range { .. }
8808 | Condition::RangeF64 { .. }
8809 | Condition::SparseMatch { .. }
8810 | Condition::MinHashSimilar { .. }
8811 | Condition::IsNull { .. }
8812 | Condition::IsNotNull { .. }
8813 )
8814 };
8815 if !conditions.iter().all(served) {
8816 return Ok(None);
8817 }
8818 self.ensure_indexes_complete()?;
8819 if !self.pending_put_cols.is_empty()
8820 || !self.pending_delete_rids.is_empty()
8821 || !self.pending_rows.is_empty()
8822 || !self.pending_dels.is_empty()
8823 || self.pending_truncate.is_some()
8824 {
8825 let mut sets = Vec::with_capacity(conditions.len());
8826 for condition in conditions {
8827 sets.push(self.resolve_condition(condition, snapshot)?);
8828 }
8829 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8830 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8831 }
8832 let mut sets = Vec::with_capacity(conditions.len());
8833 for condition in conditions {
8834 sets.push(self.resolve_condition(condition, snapshot)?);
8835 }
8836 let mut rids = RowIdSet::intersect_many(sets);
8837 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8847 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8848 }
8849 let count = rids.len() as u64;
8850 crate::trace::QueryTrace::record(|t| {
8851 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8852 t.survivor_count = Some(count as usize);
8853 t.conditions_pushed = conditions.len();
8854 });
8855 Ok(Some(count))
8856 }
8857
8858 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8863 let mut out = Vec::new();
8864 for row in self.memtable.visible_versions(snapshot.epoch) {
8865 if row.deleted {
8866 out.push(row.row_id.0);
8867 }
8868 }
8869 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8870 if row.deleted {
8871 out.push(row.row_id.0);
8872 }
8873 }
8874 out
8875 }
8876
8877 pub fn bulk_load_columns(
8886 &mut self,
8887 user_columns: Vec<(u16, columnar::NativeColumn)>,
8888 ) -> Result<Epoch> {
8889 self.bulk_load_columns_with(user_columns, 3, false, true)
8890 }
8891
8892 pub fn bulk_load_fast(
8899 &mut self,
8900 user_columns: Vec<(u16, columnar::NativeColumn)>,
8901 ) -> Result<Epoch> {
8902 self.bulk_load_columns_with(user_columns, -1, true, false)
8903 }
8904
8905 fn bulk_load_columns_with(
8906 &mut self,
8907 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8908 zstd_level: i32,
8909 force_plain: bool,
8910 lz4: bool,
8911 ) -> Result<Epoch> {
8912 self.ensure_writable()?;
8913 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8914 if n == 0 {
8915 return Ok(self.current_epoch());
8916 }
8917 let epoch = self.commit_new_epoch()?;
8918 let live_before = self.live_count;
8919 self.spill_mutable_run(epoch)?;
8921 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8922 && self.indexes_complete
8923 && self.run_refs.is_empty()
8924 && self.memtable.is_empty()
8925 && self.mutable_run.is_empty();
8926 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8929 self.validate_columns_not_null(&user_columns, n)?;
8930 let winner_idx = self
8931 .bulk_pk_winner_indices(&user_columns, n)
8932 .filter(|idx| idx.len() != n);
8933 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8934 match winner_idx.as_deref() {
8935 Some(idx) => {
8936 let compacted = user_columns
8937 .iter()
8938 .map(|(id, c)| (*id, c.gather(idx)))
8939 .collect();
8940 (compacted, idx.len())
8941 }
8942 None => (user_columns, n),
8943 };
8944 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8945 let first = self.allocator.alloc_range(write_n as u64)?.0;
8946 for rid in first..first + write_n as u64 {
8947 self.reservoir.offer(rid);
8948 }
8949 let run_id = self.alloc_run_id()?;
8950 let path = self.run_path(run_id);
8951 let mut writer =
8952 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8953 if force_plain {
8954 writer = writer.with_plain();
8955 } else if lz4 {
8956 writer = writer.with_lz4();
8959 } else {
8960 writer = writer.with_zstd_level(zstd_level);
8961 }
8962 if let Some(kek) = &self.kek {
8963 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8964 }
8965 let header = match self.create_run_file(run_id)? {
8966 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8967 None => writer.write_native(&path, &write_columns, write_n, first)?,
8968 };
8969 self.run_refs.push(RunRef {
8970 run_id: run_id as u128,
8971 level: 0,
8972 epoch_created: epoch.0,
8973 row_count: header.row_count,
8974 });
8975 self.live_count = self.live_count.saturating_add(write_n as u64);
8976 if eager_index_build {
8977 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8978 self.index_columns_bulk(&write_columns, &row_ids);
8979 self.indexes_complete = true;
8980 self.build_learned_ranges()?;
8981 } else {
8982 self.indexes_complete = false;
8986 }
8987 self.mark_flushed(epoch)?;
8988 self.persist_manifest(epoch)?;
8989 if eager_index_build {
8990 self.checkpoint_indexes(epoch);
8991 }
8992 self.clear_result_cache();
8993 self.data_generation = self.data_generation.wrapping_add(1);
8994 Ok(epoch)
8995 }
8996
8997 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
9015 let n = row_ids.len();
9016 if n == 0 {
9017 return;
9018 }
9019 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
9020 columns.iter().map(|(id, c)| (*id, c)).collect();
9021 let ty_of: std::collections::HashMap<u16, TypeId> = self
9022 .schema
9023 .columns
9024 .iter()
9025 .map(|c| (c.id, c.ty.clone()))
9026 .collect();
9027 let pk_id = self.schema.primary_key().map(|c| c.id);
9028
9029 for (i, &rid) in row_ids.iter().enumerate() {
9030 let row_id = RowId(rid);
9031 if let Some(pid) = pk_id {
9032 if let Some(col) = by_id.get(&pid) {
9033 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
9034 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
9035 self.insert_hot_pk(key, row_id);
9036 }
9037 }
9038 }
9039 for idef in &self.schema.indexes {
9040 let Some(col) = by_id.get(&idef.column_id) else {
9041 continue;
9042 };
9043 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
9044 match idef.kind {
9045 IndexKind::Bitmap => {
9046 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
9047 if let Some(key) =
9048 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
9049 {
9050 b.insert(key, row_id);
9051 }
9052 }
9053 }
9054 IndexKind::FmIndex => {
9055 if let Some(f) = self.fm.get_mut(&idef.column_id) {
9056 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9057 f.insert(bytes.to_vec(), row_id);
9058 }
9059 }
9060 }
9061 IndexKind::Sparse => {
9062 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
9063 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9064 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
9065 s.insert(&terms, row_id);
9066 }
9067 }
9068 }
9069 }
9070 IndexKind::MinHash => {
9071 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
9072 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9073 let tokens = crate::index::token_hashes_from_bytes(bytes);
9074 mh.insert(&tokens, row_id);
9075 }
9076 }
9077 }
9078 _ => {}
9079 }
9080 }
9081 }
9082 }
9083
9084 pub fn visible_columns_native(
9089 &self,
9090 snapshot: Snapshot,
9091 projection: Option<&[u16]>,
9092 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9093 self.visible_columns_native_inner(snapshot, projection, None)
9094 }
9095
9096 pub fn visible_columns_native_with_control(
9097 &self,
9098 snapshot: Snapshot,
9099 projection: Option<&[u16]>,
9100 control: &crate::ExecutionControl,
9101 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9102 self.visible_columns_native_inner(snapshot, projection, Some(control))
9103 }
9104
9105 fn visible_columns_native_inner(
9106 &self,
9107 snapshot: Snapshot,
9108 projection: Option<&[u16]>,
9109 control: Option<&crate::ExecutionControl>,
9110 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9111 execution_checkpoint(control, 0)?;
9112 let wanted: Vec<u16> = match projection {
9113 Some(p) => p.to_vec(),
9114 None => self.schema.columns.iter().map(|c| c.id).collect(),
9115 };
9116 if self.ttl.is_none()
9117 && self.memtable.is_empty()
9118 && self.mutable_run.is_empty()
9119 && self.run_refs.len() == 1
9120 {
9121 let rr = self.run_refs[0].clone();
9122 let mut reader = self.open_reader(rr.run_id)?;
9123 let idxs = reader.visible_indices_native(snapshot.epoch)?;
9124 execution_checkpoint(control, 0)?;
9125 let all_visible = idxs.len() == reader.row_count();
9126 if reader.has_mmap() && control.is_none() {
9132 use rayon::prelude::*;
9133 let valid: Vec<u16> = wanted
9136 .iter()
9137 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
9138 .copied()
9139 .collect();
9140 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
9142 .par_iter()
9143 .filter_map(|cid| {
9144 reader
9145 .column_native_shared(*cid)
9146 .ok()
9147 .map(|col| (*cid, col))
9148 })
9149 .collect();
9150 let cols = decoded
9151 .into_iter()
9152 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
9153 .collect();
9154 return Ok(cols);
9155 }
9156 let mut cols = Vec::with_capacity(wanted.len());
9157 for (index, cid) in wanted.iter().enumerate() {
9158 execution_checkpoint(control, index)?;
9159 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
9160 Some(c) => c,
9161 None => continue,
9162 };
9163 let col = reader.column_native(cdef.id)?;
9164 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
9165 }
9166 return Ok(cols);
9167 }
9168 let vcols = self.visible_columns(snapshot)?;
9169 execution_checkpoint(control, 0)?;
9170 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
9171 let out: Vec<(u16, columnar::NativeColumn)> = vcols
9172 .into_iter()
9173 .filter(|(id, _)| want_set.contains(id))
9174 .map(|(id, vals)| {
9175 let ty = self
9176 .schema
9177 .columns
9178 .iter()
9179 .find(|c| c.id == id)
9180 .map(|c| c.ty.clone())
9181 .unwrap_or(TypeId::Bytes);
9182 (id, columnar::values_to_native(ty, &vals))
9183 })
9184 .collect();
9185 Ok(out)
9186 }
9187
9188 pub fn run_count(&self) -> usize {
9189 self.run_refs.len()
9190 }
9191
9192 pub fn memtable_is_empty(&self) -> bool {
9194 self.memtable.is_empty()
9195 }
9196
9197 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
9201 self.page_cache.stats()
9202 }
9203
9204 pub fn reset_page_cache_stats(&self) {
9206 self.page_cache.reset_stats();
9207 }
9208
9209 pub fn run_ids(&self) -> Vec<u128> {
9212 self.run_refs.iter().map(|r| r.run_id).collect()
9213 }
9214
9215 pub fn single_run_is_clean(&self) -> bool {
9219 if self.ttl.is_some() || self.run_refs.len() != 1 {
9220 return false;
9221 }
9222 self.open_reader(self.run_refs[0].run_id)
9223 .map(|r| r.is_clean())
9224 .unwrap_or(false)
9225 }
9226
9227 fn resolve_footprint(
9234 &self,
9235 conditions: &[crate::query::Condition],
9236 snapshot: Snapshot,
9237 ) -> roaring::RoaringBitmap {
9238 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
9239 return roaring::RoaringBitmap::new();
9240 }
9241 if self.run_refs.is_empty() {
9242 return roaring::RoaringBitmap::new();
9243 }
9244 if self.run_refs.len() == 1 {
9246 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
9247 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
9248 return rids.to_roaring_lossy();
9249 }
9250 }
9251 }
9252 roaring::RoaringBitmap::new()
9253 }
9254
9255 pub fn query_columns_native_cached(
9266 &mut self,
9267 conditions: &[crate::query::Condition],
9268 projection: Option<&[u16]>,
9269 snapshot: Snapshot,
9270 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9271 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
9272 }
9273
9274 pub fn query_columns_native_cached_with_control(
9275 &mut self,
9276 conditions: &[crate::query::Condition],
9277 projection: Option<&[u16]>,
9278 snapshot: Snapshot,
9279 control: &crate::ExecutionControl,
9280 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9281 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
9282 }
9283
9284 fn query_columns_native_cached_inner(
9285 &mut self,
9286 conditions: &[crate::query::Condition],
9287 projection: Option<&[u16]>,
9288 snapshot: Snapshot,
9289 control: Option<&crate::ExecutionControl>,
9290 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9291 execution_checkpoint(control, 0)?;
9292 if self.ttl.is_some() {
9295 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9296 }
9297 if conditions.is_empty() {
9298 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9299 }
9300 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
9304 if let Some(hit) = self.result_cache.lock().get_columns(key) {
9305 crate::trace::QueryTrace::record(|t| {
9306 t.result_cache_hit = true;
9307 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9308 });
9309 return Ok(Some((*hit).clone()));
9310 }
9311 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
9312 execution_checkpoint(control, 0)?;
9313 if let Some(cols) = &res {
9314 let footprint = self.resolve_footprint(conditions, snapshot);
9315 let condition_cols = crate::query::condition_columns(conditions);
9316 execution_checkpoint(control, 0)?;
9317 self.result_cache.lock().insert(
9318 key,
9319 CachedEntry {
9320 data: CachedData::Columns(Arc::new(cols.clone())),
9321 footprint,
9322 condition_cols,
9323 },
9324 );
9325 }
9326 Ok(res)
9327 }
9328
9329 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
9334 if self.ttl.is_some() {
9335 return self.query(q);
9336 }
9337 if q.conditions.is_empty() {
9338 return self.query(q);
9339 }
9340 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
9341 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
9342 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
9343 if let Some(hit) = self.result_cache.lock().get_rows(key) {
9344 crate::trace::QueryTrace::record(|t| {
9345 t.result_cache_hit = true;
9346 t.scan_mode = crate::trace::ScanMode::Materialized;
9347 });
9348 return Ok((*hit).clone());
9349 }
9350 let rows = self.query(q)?;
9351 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
9352 let condition_cols = crate::query::condition_columns(&q.conditions);
9353 self.result_cache.lock().insert(
9354 key,
9355 CachedEntry {
9356 data: CachedData::Rows(Arc::new(rows.clone())),
9357 footprint,
9358 condition_cols,
9359 },
9360 );
9361 Ok(rows)
9362 }
9363
9364 #[allow(clippy::type_complexity)]
9379 pub fn query_columns_native_traced(
9380 &mut self,
9381 conditions: &[crate::query::Condition],
9382 projection: Option<&[u16]>,
9383 snapshot: Snapshot,
9384 ) -> Result<(
9385 Option<Vec<(u16, columnar::NativeColumn)>>,
9386 crate::trace::QueryTrace,
9387 )> {
9388 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9389 self.query_columns_native(conditions, projection, snapshot)
9390 });
9391 Ok((result?, trace))
9392 }
9393
9394 #[allow(clippy::type_complexity)]
9397 pub fn query_columns_native_cached_traced(
9398 &mut self,
9399 conditions: &[crate::query::Condition],
9400 projection: Option<&[u16]>,
9401 snapshot: Snapshot,
9402 ) -> Result<(
9403 Option<Vec<(u16, columnar::NativeColumn)>>,
9404 crate::trace::QueryTrace,
9405 )> {
9406 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9407 self.query_columns_native_cached(conditions, projection, snapshot)
9408 });
9409 Ok((result?, trace))
9410 }
9411
9412 pub fn native_page_cursor_traced(
9414 &self,
9415 snapshot: Snapshot,
9416 projection: Vec<(u16, TypeId)>,
9417 conditions: &[crate::query::Condition],
9418 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
9419 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9420 self.native_page_cursor(snapshot, projection, conditions)
9421 });
9422 Ok((result?, trace))
9423 }
9424
9425 pub fn native_multi_run_cursor_traced(
9427 &self,
9428 snapshot: Snapshot,
9429 projection: Vec<(u16, TypeId)>,
9430 conditions: &[crate::query::Condition],
9431 ) -> Result<(
9432 Option<crate::cursor::MultiRunCursor>,
9433 crate::trace::QueryTrace,
9434 )> {
9435 let (result, trace) = crate::trace::QueryTrace::capture(|| {
9436 self.native_multi_run_cursor(snapshot, projection, conditions)
9437 });
9438 Ok((result?, trace))
9439 }
9440
9441 pub fn count_conditions_traced(
9443 &mut self,
9444 conditions: &[crate::query::Condition],
9445 snapshot: Snapshot,
9446 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
9447 let (result, trace) =
9448 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
9449 Ok((result?, trace))
9450 }
9451
9452 pub fn query_traced(
9454 &mut self,
9455 q: &crate::query::Query,
9456 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
9457 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
9458 Ok((result?, trace))
9459 }
9460
9461 pub fn query_columns_native(
9466 &mut self,
9467 conditions: &[crate::query::Condition],
9468 projection: Option<&[u16]>,
9469 snapshot: Snapshot,
9470 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9471 self.query_columns_native_inner(conditions, projection, snapshot, None)
9472 }
9473
9474 pub fn query_columns_native_with_control(
9475 &mut self,
9476 conditions: &[crate::query::Condition],
9477 projection: Option<&[u16]>,
9478 snapshot: Snapshot,
9479 control: &crate::ExecutionControl,
9480 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9481 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
9482 }
9483
9484 fn query_columns_native_inner(
9485 &mut self,
9486 conditions: &[crate::query::Condition],
9487 projection: Option<&[u16]>,
9488 snapshot: Snapshot,
9489 control: Option<&crate::ExecutionControl>,
9490 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9491 use crate::query::Condition;
9492 execution_checkpoint(control, 0)?;
9493 if self.ttl.is_some() {
9496 return Ok(None);
9497 }
9498 if conditions.is_empty() {
9499 return Ok(None);
9500 }
9501 self.ensure_indexes_complete()?;
9502
9503 let served = |c: &Condition| {
9508 matches!(
9509 c,
9510 Condition::Pk(_)
9511 | Condition::BitmapEq { .. }
9512 | Condition::BitmapIn { .. }
9513 | Condition::BytesPrefix { .. }
9514 | Condition::FmContains { .. }
9515 | Condition::FmContainsAll { .. }
9516 | Condition::Ann { .. }
9517 | Condition::Range { .. }
9518 | Condition::RangeF64 { .. }
9519 | Condition::SparseMatch { .. }
9520 | Condition::MinHashSimilar { .. }
9521 | Condition::IsNull { .. }
9522 | Condition::IsNotNull { .. }
9523 )
9524 };
9525 if !conditions.iter().all(served) {
9526 return Ok(None);
9527 }
9528 let fast_path =
9529 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9530 crate::trace::QueryTrace::record(|t| {
9531 t.run_count = self.run_refs.len();
9532 t.memtable_rows = self.memtable.len();
9533 t.mutable_run_rows = self.mutable_run.len();
9534 t.conditions_pushed = conditions.len();
9535 t.learned_range_used = conditions.iter().any(|c| match c {
9536 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
9537 self.learned_range.contains_key(column_id)
9538 }
9539 _ => false,
9540 });
9541 });
9542 let col_ids: Vec<u16> = projection
9544 .map(|p| p.to_vec())
9545 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
9546 let proj_pairs: Vec<(u16, TypeId)> = col_ids
9547 .iter()
9548 .map(|&cid| {
9549 let ty = self
9550 .schema
9551 .columns
9552 .iter()
9553 .find(|c| c.id == cid)
9554 .map(|c| c.ty.clone())
9555 .unwrap_or(TypeId::Bytes);
9556 (cid, ty)
9557 })
9558 .collect();
9559
9560 if fast_path {
9566 let needs_column = conditions.iter().any(|c| match c {
9569 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
9570 Condition::RangeF64 { column_id, .. } => {
9571 !self.learned_range.contains_key(column_id)
9572 }
9573 _ => false,
9574 });
9575 let mut reader_opt: Option<RunReader> = if needs_column {
9576 Some(self.open_reader(self.run_refs[0].run_id)?)
9577 } else {
9578 None
9579 };
9580 let mut sets: Vec<RowIdSet> = Vec::new();
9581 for (index, c) in conditions.iter().enumerate() {
9582 execution_checkpoint(control, index)?;
9583 let s = match c {
9584 Condition::Range { column_id, lo, hi }
9585 if !self.learned_range.contains_key(column_id) =>
9586 {
9587 if reader_opt.is_none() {
9588 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9589 }
9590 reader_opt
9591 .as_mut()
9592 .expect("reader opened for range")
9593 .range_row_id_set_i64(*column_id, *lo, *hi)?
9594 }
9595 Condition::RangeF64 {
9596 column_id,
9597 lo,
9598 lo_inclusive,
9599 hi,
9600 hi_inclusive,
9601 } if !self.learned_range.contains_key(column_id) => {
9602 if reader_opt.is_none() {
9603 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9604 }
9605 reader_opt
9606 .as_mut()
9607 .expect("reader opened for range")
9608 .range_row_id_set_f64(
9609 *column_id,
9610 *lo,
9611 *lo_inclusive,
9612 *hi,
9613 *hi_inclusive,
9614 )?
9615 }
9616 _ => self.resolve_condition(c, snapshot)?,
9617 };
9618 sets.push(s);
9619 }
9620 let candidates = RowIdSet::intersect_many(sets);
9621 crate::trace::QueryTrace::record(|t| {
9622 t.survivor_count = Some(candidates.len());
9623 });
9624 if candidates.is_empty() {
9625 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
9626 .iter()
9627 .map(|&id| {
9628 (
9629 id,
9630 columnar::null_native(
9631 proj_pairs
9632 .iter()
9633 .find(|(c, _)| c == &id)
9634 .map(|(_, t)| t.clone())
9635 .unwrap_or(TypeId::Bytes),
9636 0,
9637 ),
9638 )
9639 })
9640 .collect();
9641 return Ok(Some(cols));
9642 }
9643 let mut reader = match reader_opt.take() {
9644 Some(r) => r,
9645 None => self.open_reader(self.run_refs[0].run_id)?,
9646 };
9647 let candidate_ids = candidates.into_sorted_vec();
9648 let (positions, fast_rid) = if let Some(positions) =
9649 reader.positions_for_row_ids_fast(&candidate_ids)
9650 {
9651 (positions, true)
9652 } else {
9653 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
9654 match col {
9655 columnar::NativeColumn::Int64 { data, .. } => {
9656 let mut p = Vec::with_capacity(candidate_ids.len());
9657 for (index, rid) in candidate_ids.iter().enumerate() {
9658 execution_checkpoint(control, index)?;
9659 if let Ok(position) = data.binary_search(&(*rid as i64)) {
9660 p.push(position);
9661 }
9662 }
9663 p.sort_unstable();
9664 (p, false)
9665 }
9666 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
9667 }
9668 };
9669 crate::trace::QueryTrace::record(|t| {
9670 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9671 t.fast_row_id_map = fast_rid;
9672 });
9673 let mut cols = Vec::with_capacity(col_ids.len());
9674 for (index, cid) in col_ids.iter().enumerate() {
9675 execution_checkpoint(control, index)?;
9676 let col = reader.column_native(*cid)?;
9677 cols.push((*cid, col.gather(&positions)));
9678 }
9679 return Ok(Some(cols));
9680 }
9681
9682 if !self.run_refs.is_empty() {
9695 use crate::cursor::{
9696 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
9697 };
9698 let remaining: usize;
9699 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9700 let c = self
9701 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9702 .expect("single-run cursor should build when run_refs.len() == 1");
9703 remaining = c.remaining_rows();
9704 Box::new(c)
9705 } else {
9706 let c = self
9707 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9708 .expect("multi-run cursor should build when run_refs.len() >= 1");
9709 remaining = c.remaining_rows();
9710 Box::new(c)
9711 };
9712 crate::trace::QueryTrace::record(|t| {
9713 if t.survivor_count.is_none() {
9714 t.survivor_count = Some(remaining);
9715 }
9716 });
9717 let cols = match control {
9718 Some(control) => {
9719 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9720 }
9721 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9722 };
9723 return Ok(Some(cols));
9724 }
9725
9726 crate::trace::QueryTrace::record(|t| {
9731 t.scan_mode = crate::trace::ScanMode::Materialized;
9732 t.row_materialized = true;
9733 });
9734 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9735 for (index, c) in conditions.iter().enumerate() {
9736 execution_checkpoint(control, index)?;
9737 sets.push(self.resolve_condition(c, snapshot)?);
9738 }
9739 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9740 let rows = self.rows_for_rids(&rids, snapshot)?;
9741 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9742 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9743 execution_checkpoint(control, index)?;
9744 let vals: Vec<Value> = rows
9745 .iter()
9746 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9747 .collect();
9748 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9749 }
9750 Ok(Some(cols))
9751 }
9752
9753 pub fn native_page_cursor(
9768 &self,
9769 snapshot: Snapshot,
9770 projection: Vec<(u16, TypeId)>,
9771 conditions: &[crate::query::Condition],
9772 ) -> Result<Option<NativePageCursor>> {
9773 use crate::cursor::build_page_plans;
9774 if self.ttl.is_some() {
9775 return Ok(None);
9776 }
9777 if !conditions.is_empty() && !self.indexes_complete {
9780 return Ok(None);
9781 }
9782 if self.run_refs.len() != 1 {
9783 return Ok(None);
9784 }
9785 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9786 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9787
9788 let overlay_rids: HashSet<u64> = {
9791 let mut s = HashSet::new();
9792 for row in self.memtable.visible_versions(snapshot.epoch) {
9793 s.insert(row.row_id.0);
9794 }
9795 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9796 s.insert(row.row_id.0);
9797 }
9798 s
9799 };
9800
9801 let survivors = if conditions.is_empty() {
9805 None
9806 } else {
9807 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9808 };
9809
9810 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9817 survivors.clone()
9818 } else if let Some(s) = &survivors {
9819 let mut run_set = s.clone();
9820 run_set.remove_many(overlay_rids.iter().copied());
9821 Some(run_set)
9822 } else {
9823 Some(RowIdSet::from_unsorted(
9824 rids.iter()
9825 .map(|&r| r as u64)
9826 .filter(|r| !overlay_rids.contains(r))
9827 .collect(),
9828 ))
9829 };
9830
9831 let overlay_rows = if overlay_rids.is_empty() {
9832 Vec::new()
9833 } else {
9834 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9835 self.overlay_visible_rows(snapshot, bound)
9836 };
9837
9838 let plans = if positions.is_empty() {
9840 Vec::new()
9841 } else {
9842 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9843 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9844 };
9845
9846 let overlay = if overlay_rows.is_empty() {
9848 None
9849 } else {
9850 let filtered =
9851 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9852 if filtered.is_empty() {
9853 None
9854 } else {
9855 Some(self.materialize_overlay(&filtered, &projection))
9856 }
9857 };
9858
9859 let overlay_row_count = overlay
9860 .as_ref()
9861 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9862 .unwrap_or(0);
9863 crate::trace::QueryTrace::record(|t| {
9864 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9865 t.run_count = self.run_refs.len();
9866 t.memtable_rows = self.memtable.len();
9867 t.mutable_run_rows = self.mutable_run.len();
9868 t.overlay_rows = overlay_row_count;
9869 t.conditions_pushed = conditions.len();
9870 t.pages_decoded = plans
9871 .iter()
9872 .map(|p| p.positions.len())
9873 .sum::<usize>()
9874 .min(1);
9875 });
9876
9877 Ok(Some(NativePageCursor::new_with_overlay(
9878 reader, projection, plans, overlay,
9879 )))
9880 }
9881 #[allow(clippy::type_complexity)]
9891 pub fn native_multi_run_cursor(
9892 &self,
9893 snapshot: Snapshot,
9894 projection: Vec<(u16, TypeId)>,
9895 conditions: &[crate::query::Condition],
9896 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9897 use crate::cursor::{MultiRunCursor, RunStream};
9898 use crate::sorted_run::SYS_ROW_ID;
9899 use std::collections::{BinaryHeap, HashMap, HashSet};
9900 if self.ttl.is_some() {
9901 return Ok(None);
9902 }
9903 if !conditions.is_empty() && !self.indexes_complete {
9906 return Ok(None);
9907 }
9908 if self.run_refs.is_empty() {
9909 return Ok(None);
9910 }
9911
9912 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9914 Vec::with_capacity(self.run_refs.len());
9915 for rr in &self.run_refs {
9916 let mut reader = self.open_reader(rr.run_id)?;
9917 let (rids, eps, del) = reader.system_columns_native()?;
9918 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9919 run_meta.push((reader, rids, eps, del, page_rows));
9920 }
9921
9922 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9926 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9927 for i in 0..rids.len() {
9928 let rid = rids[i] as u64;
9929 let e = eps[i] as u64;
9930 if e > snapshot.epoch.0 {
9931 continue;
9932 }
9933 let is_del = del[i] != 0;
9934 best.entry(rid)
9935 .and_modify(|cur| {
9936 if e > cur.0 {
9937 *cur = (e, run_idx, i, is_del);
9938 }
9939 })
9940 .or_insert((e, run_idx, i, is_del));
9941 }
9942 }
9943
9944 let overlay_rids: HashSet<u64> = {
9946 let mut s = HashSet::new();
9947 for row in self.memtable.visible_versions(snapshot.epoch) {
9948 s.insert(row.row_id.0);
9949 }
9950 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9951 s.insert(row.row_id.0);
9952 }
9953 s
9954 };
9955
9956 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9958 None
9959 } else {
9960 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9961 for c in conditions {
9962 sets.push(self.resolve_condition(c, snapshot)?);
9963 }
9964 Some(RowIdSet::intersect_many(sets))
9965 };
9966
9967 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9971 for (rid, (_, run_idx, pos, deleted)) in &best {
9972 if *deleted {
9973 continue;
9974 }
9975 if overlay_rids.contains(rid) {
9976 continue;
9977 }
9978 if let Some(s) = &survivors {
9979 if !s.contains(*rid) {
9980 continue;
9981 }
9982 }
9983 per_run[*run_idx].push((*rid, *pos));
9984 }
9985 for v in per_run.iter_mut() {
9986 v.sort_unstable_by_key(|&(rid, _)| rid);
9987 }
9988
9989 let mut streams = Vec::with_capacity(run_meta.len());
9991 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9992 let mut total = 0usize;
9993 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9994 let mut starts = Vec::with_capacity(page_rows.len());
9995 let mut acc = 0usize;
9996 for &r in &page_rows {
9997 starts.push(acc);
9998 acc += r;
9999 }
10000 let mut survivors_vec: Vec<(u64, usize, usize)> =
10001 Vec::with_capacity(per_run[run_idx].len());
10002 for &(rid, pos) in &per_run[run_idx] {
10003 let page_seq = match starts.partition_point(|&s| s <= pos) {
10004 0 => continue,
10005 p => p - 1,
10006 };
10007 let within = pos - starts[page_seq];
10008 survivors_vec.push((rid, page_seq, within));
10009 }
10010 total += survivors_vec.len();
10011 if let Some(&(rid, _, _)) = survivors_vec.first() {
10012 heap.push(std::cmp::Reverse((rid, run_idx)));
10013 }
10014 streams.push(RunStream::new(reader, survivors_vec, page_rows));
10015 }
10016
10017 let overlay_rows = if overlay_rids.is_empty() {
10019 Vec::new()
10020 } else {
10021 let bound = Self::overlay_materialization_bound(conditions, &survivors);
10022 self.overlay_visible_rows(snapshot, bound)
10023 };
10024 let overlay = if overlay_rows.is_empty() {
10025 None
10026 } else {
10027 let filtered =
10028 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
10029 if filtered.is_empty() {
10030 None
10031 } else {
10032 Some(self.materialize_overlay(&filtered, &projection))
10033 }
10034 };
10035
10036 let overlay_row_count = overlay
10037 .as_ref()
10038 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
10039 .unwrap_or(0);
10040 crate::trace::QueryTrace::record(|t| {
10041 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
10042 t.run_count = self.run_refs.len();
10043 t.memtable_rows = self.memtable.len();
10044 t.mutable_run_rows = self.mutable_run.len();
10045 t.overlay_rows = overlay_row_count;
10046 t.conditions_pushed = conditions.len();
10047 t.survivor_count = Some(total);
10048 });
10049
10050 Ok(Some(MultiRunCursor::new(
10051 streams, projection, heap, total, overlay,
10052 )))
10053 }
10054
10055 fn overlay_materialization_bound<'a>(
10067 conditions: &[crate::query::Condition],
10068 survivors: &'a Option<RowIdSet>,
10069 ) -> Option<&'a RowIdSet> {
10070 use crate::query::Condition;
10071 let has_range = conditions
10072 .iter()
10073 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
10074 if has_range {
10075 None
10076 } else {
10077 survivors.as_ref()
10078 }
10079 }
10080
10081 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
10093 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
10094 let mut fold = |row: Row| {
10095 if let Some(b) = bound {
10096 if !b.contains(row.row_id.0) {
10097 return;
10098 }
10099 }
10100 best.entry(row.row_id.0)
10101 .and_modify(|(be, br)| {
10102 if row.committed_epoch > *be {
10103 *be = row.committed_epoch;
10104 *br = row.clone();
10105 }
10106 })
10107 .or_insert_with(|| (row.committed_epoch, row));
10108 };
10109 for row in self.memtable.visible_versions(snapshot.epoch) {
10110 fold(row);
10111 }
10112 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10113 fold(row);
10114 }
10115 let mut out: Vec<Row> = best
10116 .into_values()
10117 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
10118 .collect();
10119 out.sort_by_key(|r| r.row_id);
10120 out
10121 }
10122
10123 fn filter_overlay_rows(
10131 &self,
10132 rows: Vec<Row>,
10133 conditions: &[crate::query::Condition],
10134 survivors: Option<&RowIdSet>,
10135 snapshot: Snapshot,
10136 ) -> Result<Vec<Row>> {
10137 if conditions.is_empty() {
10138 return Ok(rows);
10139 }
10140 use crate::query::Condition;
10141 let all_index_served = !conditions
10145 .iter()
10146 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
10147 if all_index_served {
10148 return Ok(rows
10149 .into_iter()
10150 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
10151 .collect());
10152 }
10153 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10156 for c in conditions {
10157 let s = match c {
10158 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
10159 _ => self.resolve_condition(c, snapshot)?,
10160 };
10161 per_cond_sets.push(s);
10162 }
10163 Ok(rows
10164 .into_iter()
10165 .filter(|row| {
10166 conditions.iter().enumerate().all(|(i, c)| match c {
10167 Condition::Range { column_id, lo, hi } => {
10168 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
10169 }
10170 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
10171 match row.columns.get(column_id) {
10172 Some(Value::Float64(v)) => {
10173 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
10174 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
10175 lo_ok && hi_ok
10176 }
10177 _ => false,
10178 }
10179 }
10180 _ => per_cond_sets[i].contains(row.row_id.0),
10181 })
10182 })
10183 .collect())
10184 }
10185
10186 fn materialize_overlay(
10189 &self,
10190 rows: &[Row],
10191 projection: &[(u16, TypeId)],
10192 ) -> Vec<columnar::NativeColumn> {
10193 if projection.is_empty() {
10194 return vec![columnar::null_native(TypeId::Int64, rows.len())];
10195 }
10196 let mut cols = Vec::with_capacity(projection.len());
10197 for (cid, ty) in projection {
10198 let vals: Vec<Value> = rows
10199 .iter()
10200 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
10201 .collect();
10202 cols.push(columnar::values_to_native(ty.clone(), &vals));
10203 }
10204 cols
10205 }
10206
10207 fn resolve_survivor_rids(
10212 &self,
10213 conditions: &[crate::query::Condition],
10214 reader: &mut RunReader,
10215 snapshot: Snapshot,
10216 ) -> Result<RowIdSet> {
10217 use crate::query::Condition;
10218 let mut sets: Vec<RowIdSet> = Vec::new();
10219 for c in conditions {
10220 self.validate_condition(c)?;
10221 let s: RowIdSet = match c {
10222 Condition::Pk(key) => {
10223 let lookup = self
10224 .schema
10225 .primary_key()
10226 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
10227 .unwrap_or_else(|| key.clone());
10228 self.hot
10229 .get(&lookup)
10230 .map(|r| RowIdSet::one(r.0))
10231 .unwrap_or_else(RowIdSet::empty)
10232 }
10233 Condition::BitmapEq { column_id, value } => {
10234 let lookup = self.index_lookup_key_bytes(*column_id, value);
10235 self.bitmap
10236 .get(column_id)
10237 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
10238 .unwrap_or_else(RowIdSet::empty)
10239 }
10240 Condition::BitmapIn { column_id, values } => {
10241 let bm = self.bitmap.get(column_id);
10242 let mut acc = roaring::RoaringBitmap::new();
10243 if let Some(b) = bm {
10244 for v in values {
10245 let lookup = self.index_lookup_key_bytes(*column_id, v);
10246 acc |= b.get(&lookup);
10247 }
10248 }
10249 RowIdSet::from_roaring(acc)
10250 }
10251 Condition::BytesPrefix { column_id, prefix } => {
10252 if let Some(b) = self.bitmap.get(column_id) {
10253 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
10254 let mut acc = roaring::RoaringBitmap::new();
10255 for key in b.keys() {
10256 if key.starts_with(&lookup_prefix) {
10257 acc |= b.get(&key);
10258 }
10259 }
10260 RowIdSet::from_roaring(acc)
10261 } else {
10262 RowIdSet::empty()
10263 }
10264 }
10265 Condition::FmContains { column_id, pattern } => self
10266 .fm
10267 .get(column_id)
10268 .map(|f| {
10269 RowIdSet::from_unsorted(
10270 f.locate(pattern).into_iter().map(|r| r.0).collect(),
10271 )
10272 })
10273 .unwrap_or_else(RowIdSet::empty),
10274 Condition::FmContainsAll {
10275 column_id,
10276 patterns,
10277 } => {
10278 if let Some(f) = self.fm.get(column_id) {
10279 let sets: Vec<RowIdSet> = patterns
10280 .iter()
10281 .map(|pat| {
10282 RowIdSet::from_unsorted(
10283 f.locate(pat).into_iter().map(|r| r.0).collect(),
10284 )
10285 })
10286 .collect();
10287 RowIdSet::intersect_many(sets)
10288 } else {
10289 RowIdSet::empty()
10290 }
10291 }
10292 Condition::Ann {
10293 column_id,
10294 query,
10295 k,
10296 } => RowIdSet::from_unsorted(
10297 self.retrieve_filtered(
10298 &crate::query::Retriever::Ann {
10299 column_id: *column_id,
10300 query: query.clone(),
10301 k: *k,
10302 },
10303 snapshot,
10304 None,
10305 None,
10306 None,
10307 None,
10308 )?
10309 .into_iter()
10310 .map(|hit| hit.row_id.0)
10311 .collect(),
10312 ),
10313 Condition::SparseMatch {
10314 column_id,
10315 query,
10316 k,
10317 } => RowIdSet::from_unsorted(
10318 self.retrieve_filtered(
10319 &crate::query::Retriever::Sparse {
10320 column_id: *column_id,
10321 query: query.clone(),
10322 k: *k,
10323 },
10324 snapshot,
10325 None,
10326 None,
10327 None,
10328 None,
10329 )?
10330 .into_iter()
10331 .map(|hit| hit.row_id.0)
10332 .collect(),
10333 ),
10334 Condition::MinHashSimilar {
10335 column_id,
10336 query,
10337 k,
10338 } => match self.minhash.get(column_id) {
10339 Some(index) => {
10340 let candidates = index.candidate_row_ids(query);
10341 let eligible =
10342 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
10343 RowIdSet::from_unsorted(
10344 index
10345 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
10346 .into_iter()
10347 .map(|(row_id, _)| row_id.0)
10348 .collect(),
10349 )
10350 }
10351 None => RowIdSet::empty(),
10352 },
10353 Condition::Range { column_id, lo, hi } => {
10354 if let Some(li) = self.learned_range.get(column_id) {
10355 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10356 } else {
10357 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
10358 }
10359 }
10360 Condition::RangeF64 {
10361 column_id,
10362 lo,
10363 lo_inclusive,
10364 hi,
10365 hi_inclusive,
10366 } => {
10367 if let Some(li) = self.learned_range.get(column_id) {
10368 RowIdSet::from_unsorted(
10369 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10370 .into_iter()
10371 .collect(),
10372 )
10373 } else {
10374 reader.range_row_id_set_f64(
10375 *column_id,
10376 *lo,
10377 *lo_inclusive,
10378 *hi,
10379 *hi_inclusive,
10380 )?
10381 }
10382 }
10383 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
10384 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
10385 };
10386 sets.push(s);
10387 }
10388 Ok(RowIdSet::intersect_many(sets))
10389 }
10390
10391 pub fn scan_cursor(
10412 &self,
10413 snapshot: Snapshot,
10414 projection: Vec<(u16, TypeId)>,
10415 conditions: &[crate::query::Condition],
10416 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
10417 if self.ttl.is_some() {
10418 return Ok(None);
10419 }
10420 if !conditions.is_empty() && !self.indexes_complete {
10426 return Ok(None);
10427 }
10428 if self.run_refs.len() == 1 {
10429 Ok(self
10430 .native_page_cursor(snapshot, projection, conditions)?
10431 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10432 } else {
10433 Ok(self
10434 .native_multi_run_cursor(snapshot, projection, conditions)?
10435 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
10436 }
10437 }
10438
10439 pub fn aggregate_native(
10453 &self,
10454 snapshot: Snapshot,
10455 column: Option<u16>,
10456 conditions: &[crate::query::Condition],
10457 agg: NativeAgg,
10458 ) -> Result<Option<NativeAggResult>> {
10459 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
10460 }
10461
10462 pub fn aggregate_native_with_control(
10463 &self,
10464 snapshot: Snapshot,
10465 column: Option<u16>,
10466 conditions: &[crate::query::Condition],
10467 agg: NativeAgg,
10468 control: &crate::ExecutionControl,
10469 ) -> Result<Option<NativeAggResult>> {
10470 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
10471 }
10472
10473 fn aggregate_native_inner(
10474 &self,
10475 snapshot: Snapshot,
10476 column: Option<u16>,
10477 conditions: &[crate::query::Condition],
10478 agg: NativeAgg,
10479 control: Option<&crate::ExecutionControl>,
10480 ) -> Result<Option<NativeAggResult>> {
10481 execution_checkpoint(control, 0)?;
10482 if self.ttl.is_some() {
10483 return Ok(None);
10484 }
10485 if self.run_refs.len() == 1 && conditions.is_empty() {
10487 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
10488 return Ok(Some(res));
10489 }
10490 }
10491 if matches!(agg, NativeAgg::Count) && column.is_none() {
10495 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
10496 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
10497 }
10498 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10499 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
10500 }
10501 let cid = match column {
10504 Some(c) => c,
10505 None => return Ok(None),
10506 };
10507 let ty = self.column_type(cid);
10508 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
10509 execution_checkpoint(control, 0)?;
10510 return match ty {
10511 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10512 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
10513 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10514 }
10515 TypeId::Float64 => {
10516 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
10517 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10518 }
10519 _ => Ok(None),
10520 };
10521 }
10522 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10524 execution_checkpoint(control, 0)?;
10525 match ty {
10526 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10527 let mut count = 0u64;
10528 let mut sum = 0i128;
10529 let mut mn = i64::MAX;
10530 let mut mx = i64::MIN;
10531 for row in &rows {
10532 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
10533 count += 1;
10534 sum += i128::from(*v);
10535 mn = mn.min(*v);
10536 mx = mx.max(*v);
10537 }
10538 }
10539 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10540 }
10541 TypeId::Float64 => {
10542 let mut count = 0u64;
10543 let mut sum = 0.0f64;
10544 let mut mn = f64::INFINITY;
10545 let mut mx = f64::NEG_INFINITY;
10546 for row in &rows {
10547 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
10548 count += 1;
10549 sum += *v;
10550 mn = mn.min(*v);
10551 mx = mx.max(*v);
10552 }
10553 }
10554 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10555 }
10556 _ => Ok(None),
10557 }
10558 }
10559
10560 fn visible_rows_filtered(
10562 &self,
10563 snapshot: Snapshot,
10564 conditions: &[crate::query::Condition],
10565 control: Option<&crate::ExecutionControl>,
10566 ) -> Result<Vec<Row>> {
10567 let rows = if let Some(control) = control {
10568 self.visible_rows_controlled(snapshot, control)?
10569 } else {
10570 self.visible_rows(snapshot)?
10571 };
10572 if conditions.is_empty() {
10573 return Ok(rows);
10574 }
10575 Ok(rows
10576 .into_iter()
10577 .filter(|row| {
10578 conditions
10579 .iter()
10580 .all(|cond| condition_matches_row(cond, row, &self.schema))
10581 })
10582 .collect())
10583 }
10584
10585 fn aggregate_from_stats(
10593 &self,
10594 snapshot: Snapshot,
10595 column: Option<u16>,
10596 agg: NativeAgg,
10597 ) -> Result<Option<NativeAggResult>> {
10598 let cid = match (agg, column) {
10599 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
10600 _ => return Ok(None), };
10602 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
10603 return Ok(None);
10604 };
10605 let Some(cs) = stats.get(&cid) else {
10606 return Ok(None);
10607 };
10608 match agg {
10609 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
10611 self.live_count.saturating_sub(cs.null_count),
10612 ))),
10613 NativeAgg::Min | NativeAgg::Max => {
10614 let bound = if agg == NativeAgg::Min {
10615 &cs.min
10616 } else {
10617 &cs.max
10618 };
10619 match bound {
10620 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
10621 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
10622 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
10627 None => Ok(None),
10628 }
10629 }
10630 _ => Ok(None),
10631 }
10632 }
10633
10634 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
10643 if self.ttl.is_some() {
10644 return Ok(None);
10645 }
10646 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
10647 return Ok(None);
10648 }
10649 self.ensure_indexes_complete()?;
10652 let reader = self.open_reader(self.run_refs[0].run_id)?;
10653 if self.live_count != reader.row_count() as u64 {
10654 return Ok(None);
10655 }
10656 let Some(bm) = self.bitmap.get(&column_id) else {
10657 return Ok(None); };
10659 let mut distinct = bm.value_count() as u64;
10660 if !bm.get(&Value::Null.encode_key()).is_empty() {
10663 distinct = distinct.saturating_sub(1);
10664 }
10665 Ok(Some(distinct))
10666 }
10667
10668 pub fn aggregate_incremental(
10680 &mut self,
10681 cache_key: u64,
10682 conditions: &[crate::query::Condition],
10683 column: Option<u16>,
10684 agg: NativeAgg,
10685 ) -> Result<IncrementalAggResult> {
10686 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
10687 }
10688
10689 pub fn aggregate_incremental_with_control(
10690 &mut self,
10691 cache_key: u64,
10692 conditions: &[crate::query::Condition],
10693 column: Option<u16>,
10694 agg: NativeAgg,
10695 control: &crate::ExecutionControl,
10696 ) -> Result<IncrementalAggResult> {
10697 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
10698 }
10699
10700 fn aggregate_incremental_inner(
10701 &mut self,
10702 cache_key: u64,
10703 conditions: &[crate::query::Condition],
10704 column: Option<u16>,
10705 agg: NativeAgg,
10706 control: Option<&crate::ExecutionControl>,
10707 ) -> Result<IncrementalAggResult> {
10708 execution_checkpoint(control, 0)?;
10709 let snap = self.snapshot();
10710 let cur_wm = self.allocator.current().0;
10711 let cur_epoch = snap.epoch.0;
10712 let incremental_ok = self.ttl.is_none()
10719 && !self.had_deletes
10720 && self.memtable.is_empty()
10721 && self.mutable_run.is_empty();
10722
10723 if incremental_ok {
10726 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10727 if cached.epoch == cur_epoch {
10728 return Ok(IncrementalAggResult {
10729 state: cached.state,
10730 incremental: true,
10731 delta_rows: 0,
10732 });
10733 }
10734 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10735 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10736 let mut delta_rids = Vec::with_capacity(delta_len);
10737 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10738 execution_checkpoint(control, index)?;
10739 delta_rids.push(row_id);
10740 }
10741 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10742 execution_checkpoint(control, 0)?;
10743 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10744 let delta_state = agg_state_from_rows(
10745 &delta_rows,
10746 conditions,
10747 &index_sets,
10748 column,
10749 agg,
10750 &self.schema,
10751 control,
10752 )?;
10753 let merged = cached.state.merge(delta_state);
10754 let delta_n = delta_rids.len() as u64;
10755 Arc::make_mut(&mut self.agg_cache).insert(
10756 cache_key,
10757 CachedAgg {
10758 state: merged.clone(),
10759 watermark: cur_wm,
10760 epoch: cur_epoch,
10761 },
10762 );
10763 return Ok(IncrementalAggResult {
10764 state: merged,
10765 incremental: true,
10766 delta_rows: delta_n,
10767 });
10768 }
10769 }
10770 }
10771
10772 let cursor_ok =
10777 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10778 let state = if cursor_ok && agg != NativeAgg::Avg {
10779 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10780 Some(result) => {
10781 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10782 }
10783 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10784 }
10785 } else {
10786 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10787 };
10788 if incremental_ok {
10790 Arc::make_mut(&mut self.agg_cache).insert(
10791 cache_key,
10792 CachedAgg {
10793 state: state.clone(),
10794 watermark: cur_wm,
10795 epoch: cur_epoch,
10796 },
10797 );
10798 }
10799 Ok(IncrementalAggResult {
10800 state,
10801 incremental: false,
10802 delta_rows: 0,
10803 })
10804 }
10805
10806 fn agg_state_full_scan(
10809 &self,
10810 conditions: &[crate::query::Condition],
10811 column: Option<u16>,
10812 agg: NativeAgg,
10813 snap: Snapshot,
10814 control: Option<&crate::ExecutionControl>,
10815 ) -> Result<AggState> {
10816 execution_checkpoint(control, 0)?;
10817 let rows = self.visible_rows(snap)?;
10818 execution_checkpoint(control, 0)?;
10819 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10820 agg_state_from_rows(
10821 &rows,
10822 conditions,
10823 &index_sets,
10824 column,
10825 agg,
10826 &self.schema,
10827 control,
10828 )
10829 }
10830
10831 fn resolve_index_conditions(
10834 &self,
10835 conditions: &[crate::query::Condition],
10836 snapshot: Snapshot,
10837 ) -> Result<Vec<RowIdSet>> {
10838 use crate::query::Condition;
10839 let mut sets = Vec::new();
10840 for c in conditions {
10841 if matches!(
10842 c,
10843 Condition::Ann { .. }
10844 | Condition::SparseMatch { .. }
10845 | Condition::MinHashSimilar { .. }
10846 ) {
10847 sets.push(self.resolve_condition(c, snapshot)?);
10848 }
10849 }
10850 Ok(sets)
10851 }
10852
10853 fn column_type(&self, cid: u16) -> TypeId {
10854 self.schema
10855 .columns
10856 .iter()
10857 .find(|c| c.id == cid)
10858 .map(|c| c.ty.clone())
10859 .unwrap_or(TypeId::Bytes)
10860 }
10861
10862 pub fn approx_aggregate(
10871 &mut self,
10872 conditions: &[crate::query::Condition],
10873 column: Option<u16>,
10874 agg: ApproxAgg,
10875 z: f64,
10876 ) -> Result<Option<ApproxResult>> {
10877 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10878 }
10879
10880 pub fn approx_aggregate_with_candidate_authorization(
10883 &mut self,
10884 conditions: &[crate::query::Condition],
10885 column: Option<u16>,
10886 agg: ApproxAgg,
10887 z: f64,
10888 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10889 ) -> Result<Option<ApproxResult>> {
10890 use crate::query::Condition;
10891 self.ensure_reservoir_complete()?;
10892 let mut snapshot = self.snapshot();
10897 let n_pop = self.count();
10898 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10899 if sample_rids.is_empty() {
10900 return Ok(None);
10901 }
10902 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10904 if live_sample.is_empty() {
10905 snapshot = Snapshot::unbounded();
10906 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10907 }
10908 let s = live_sample.len();
10909 if s == 0 {
10910 return Ok(None);
10911 }
10912 let authorized = authorization
10913 .map(|authorization| {
10914 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10915 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10916 })
10917 .transpose()?;
10918
10919 let mut index_sets: Vec<RowIdSet> = Vec::new();
10922 for c in conditions {
10923 if matches!(
10924 c,
10925 Condition::Ann { .. }
10926 | Condition::SparseMatch { .. }
10927 | Condition::MinHashSimilar { .. }
10928 ) {
10929 index_sets.push(self.resolve_condition(c, snapshot)?);
10930 }
10931 }
10932
10933 let cid = match (agg, column) {
10935 (ApproxAgg::Count, _) => None,
10936 (_, Some(c)) => Some(c),
10937 _ => return Ok(None),
10938 };
10939 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10940 for r in &live_sample {
10941 if authorized
10942 .as_ref()
10943 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10944 {
10945 continue;
10946 }
10947 if !conditions
10949 .iter()
10950 .all(|c| condition_matches_row(c, r, &self.schema))
10951 {
10952 continue;
10953 }
10954 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10956 continue;
10957 }
10958 if let Some(cid) = cid {
10959 let mut cells = r
10960 .columns
10961 .get(&cid)
10962 .cloned()
10963 .map(|value| vec![(cid, value)])
10964 .unwrap_or_default();
10965 if let Some(authorization) = authorization {
10966 authorization.security.apply_masks_to_cells(
10967 authorization.table,
10968 &mut cells,
10969 authorization.principal,
10970 );
10971 }
10972 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10973 passing_vals.push(v);
10974 } } else {
10976 passing_vals.push(0.0); }
10978 }
10979 let m = passing_vals.len();
10980
10981 let (point, half) = match agg {
10982 ApproxAgg::Count => {
10983 let p = m as f64 / s as f64;
10985 let point = n_pop as f64 * p;
10986 let var = if s > 1 {
10987 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10988 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10989 } else {
10990 0.0
10991 };
10992 (point, z * var.sqrt())
10993 }
10994 ApproxAgg::Sum => {
10995 let y: Vec<f64> = live_sample
10997 .iter()
10998 .map(|r| {
10999 let passes_row = authorized
11000 .as_ref()
11001 .is_none_or(|authorized| authorized.contains(&r.row_id))
11002 && conditions
11003 .iter()
11004 .all(|c| condition_matches_row(c, r, &self.schema))
11005 && index_sets.iter().all(|set| set.contains(r.row_id.0));
11006 if passes_row {
11007 cid.and_then(|cid| {
11008 let mut cells = r
11009 .columns
11010 .get(&cid)
11011 .cloned()
11012 .map(|value| vec![(cid, value)])
11013 .unwrap_or_default();
11014 if let Some(authorization) = authorization {
11015 authorization.security.apply_masks_to_cells(
11016 authorization.table,
11017 &mut cells,
11018 authorization.principal,
11019 );
11020 }
11021 as_f64(cells.first().map(|(_, value)| value))
11022 })
11023 .unwrap_or(0.0)
11024 } else {
11025 0.0
11026 }
11027 })
11028 .collect();
11029 let mean_y = y.iter().sum::<f64>() / s as f64;
11030 let point = n_pop as f64 * mean_y;
11031 let var = if s > 1 {
11032 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
11033 let var_y = ss / (s - 1) as f64;
11034 n_pop as f64 * n_pop as f64 * var_y / s as f64
11035 * (1.0 - s as f64 / n_pop as f64).max(0.0)
11036 } else {
11037 0.0
11038 };
11039 (point, z * var.sqrt())
11040 }
11041 ApproxAgg::Avg => {
11042 if m == 0 {
11043 return Ok(Some(ApproxResult {
11044 point: 0.0,
11045 ci_low: 0.0,
11046 ci_high: 0.0,
11047 n_population: n_pop,
11048 n_sample_live: s,
11049 n_passing: 0,
11050 }));
11051 }
11052 let mean = passing_vals.iter().sum::<f64>() / m as f64;
11053 let half = if m > 1 {
11054 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
11055 let sd = (ss / (m - 1) as f64).sqrt();
11056 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
11057 z * sd / (m as f64).sqrt() * fpc.sqrt()
11058 } else {
11059 0.0
11060 };
11061 (mean, half)
11062 }
11063 };
11064
11065 Ok(Some(ApproxResult {
11066 point,
11067 ci_low: point - half,
11068 ci_high: point + half,
11069 n_population: n_pop,
11070 n_sample_live: s,
11071 n_passing: m,
11072 }))
11073 }
11074
11075 pub fn exact_column_stats(
11083 &self,
11084 _snapshot: Snapshot,
11085 projection: &[u16],
11086 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
11087 if self.ttl.is_some()
11088 || !(self.memtable.is_empty()
11089 && self.mutable_run.is_empty()
11090 && self.run_refs.len() == 1)
11091 {
11092 return Ok(None);
11093 }
11094 let reader = self.open_reader(self.run_refs[0].run_id)?;
11095 if self.live_count != reader.row_count() as u64 {
11096 return Ok(None);
11097 }
11098 let mut out = HashMap::new();
11099 for &cid in projection {
11100 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
11101 Some(c) => c,
11102 None => continue,
11103 };
11104 let Some(stats) = reader.column_page_stats(cid) else {
11106 out.insert(
11107 cid,
11108 ColumnStat {
11109 min: None,
11110 max: None,
11111 null_count: self.live_count,
11112 },
11113 );
11114 continue;
11115 };
11116 let stat = match cdef.ty {
11117 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11118 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
11119 min: mn.map(Value::Int64),
11120 max: mx.map(Value::Int64),
11121 null_count: n,
11122 })
11123 }
11124 TypeId::Float64 => {
11125 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
11126 min: mn.map(Value::Float64),
11127 max: mx.map(Value::Float64),
11128 null_count: n,
11129 })
11130 }
11131 _ => None,
11132 };
11133 if let Some(s) = stat {
11134 out.insert(cid, s);
11135 }
11136 }
11137 Ok(Some(out))
11138 }
11139
11140 pub fn dir(&self) -> &Path {
11141 &self.dir
11142 }
11143
11144 pub fn schema(&self) -> &Schema {
11145 &self.schema
11146 }
11147
11148 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
11149 self.ann.get(&column_id)
11150 }
11151
11152 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
11153 self.ann.get_mut(&column_id)
11154 }
11155
11156 pub(crate) fn set_catalog_name(&mut self, name: String) {
11157 self.name = name;
11158 }
11159
11160 pub(crate) fn prepare_alter_column(
11161 &mut self,
11162 column_name: &str,
11163 change: &AlterColumn,
11164 ) -> Result<(ColumnDef, Option<Schema>)> {
11165 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
11166 return Err(MongrelError::InvalidArgument(
11167 "ALTER COLUMN requires committing staged writes first".into(),
11168 ));
11169 }
11170 let old = self
11171 .schema
11172 .columns
11173 .iter()
11174 .find(|c| c.name == column_name)
11175 .cloned()
11176 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11177 let mut next = old.clone();
11178
11179 if let Some(name) = &change.name {
11180 let trimmed = name.trim();
11181 if trimmed.is_empty() {
11182 return Err(MongrelError::InvalidArgument(
11183 "ALTER COLUMN name must not be empty".into(),
11184 ));
11185 }
11186 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
11187 return Err(MongrelError::Schema(format!(
11188 "column {trimmed} already exists"
11189 )));
11190 }
11191 next.name = trimmed.to_string();
11192 }
11193
11194 if let Some(ty) = &change.ty {
11195 next.ty = ty.clone();
11196 }
11197 if let Some(flags) = change.flags {
11198 validate_alter_column_flags(old.flags, flags)?;
11199 next.flags = flags;
11200 }
11201
11202 if let Some(default_change) = &change.default_value {
11203 next.default_value = default_change.clone();
11204 }
11205 if let Some(source_change) = &change.embedding_source {
11206 next.embedding_source = source_change.clone();
11207 }
11208
11209 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
11210 if old.flags.contains(ColumnFlags::NULLABLE)
11211 && !next.flags.contains(ColumnFlags::NULLABLE)
11212 && self.column_has_nulls(old.id)?
11213 {
11214 return Err(MongrelError::InvalidArgument(format!(
11215 "column '{}' contains NULL values",
11216 old.name
11217 )));
11218 }
11219 if next == old {
11220 return Ok((next, None));
11221 }
11222 let mut schema = self.schema.clone();
11223 let index = schema
11224 .columns
11225 .iter()
11226 .position(|column| column.id == next.id)
11227 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
11228 schema.columns[index] = next.clone();
11229 schema.schema_id = schema
11230 .schema_id
11231 .checked_add(1)
11232 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11233 schema.validate_auto_increment()?;
11234 schema.validate_defaults()?;
11235 Ok((next, Some(schema)))
11236 }
11237
11238 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
11239 self.schema = schema;
11240 self.auto_inc = resolve_auto_inc(&self.schema);
11241 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
11242 self.clear_result_cache();
11243 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
11244 }
11245
11246 pub(crate) fn publish_index_schema_change(
11250 &mut self,
11251 schema: Schema,
11252 artifact: SecondaryIndexArtifact,
11253 ) -> Result<()> {
11254 self.apply_altered_schema_prepared(schema);
11255 self.prune_index_maps_to_schema();
11256 match artifact {
11257 SecondaryIndexArtifact::Bitmap(column_id, index) => {
11258 self.bitmap.insert(column_id, index);
11259 }
11260 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
11261 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
11262 }
11263 SecondaryIndexArtifact::Fm(column_id, index) => {
11264 self.fm.insert(column_id, *index);
11265 }
11266 SecondaryIndexArtifact::Ann(column_id, index) => {
11267 self.ann.insert(column_id, *index);
11268 }
11269 SecondaryIndexArtifact::Sparse(column_id, index) => {
11270 self.sparse.insert(column_id, index);
11271 }
11272 SecondaryIndexArtifact::MinHash(column_id, index) => {
11273 self.minhash.insert(column_id, index);
11274 }
11275 }
11276 self.indexes_complete = true;
11277 self.seal_generations();
11278 let view = Arc::new(self.capture_read_generation());
11279 self.published.store(Arc::clone(&view));
11280 checkpoint_current_schema(self)?;
11281 self.invalidate_index_checkpoint();
11284 Ok(())
11285 }
11286
11287 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
11293 self.apply_altered_schema_prepared(schema);
11294 self.prune_index_maps_to_schema();
11295 self.indexes_complete = true;
11296 self.seal_generations();
11297 let view = Arc::new(self.capture_read_generation());
11298 self.published.store(Arc::clone(&view));
11299 checkpoint_current_schema(self)?;
11300 self.invalidate_index_checkpoint();
11302 Ok(())
11303 }
11304
11305 fn prune_index_maps_to_schema(&mut self) {
11306 let keep_bitmap: std::collections::HashSet<u16> = self
11307 .schema
11308 .indexes
11309 .iter()
11310 .filter(|index| index.kind == IndexKind::Bitmap)
11311 .map(|index| index.column_id)
11312 .collect();
11313 let keep_ann: std::collections::HashSet<u16> = self
11314 .schema
11315 .indexes
11316 .iter()
11317 .filter(|index| index.kind == IndexKind::Ann)
11318 .map(|index| index.column_id)
11319 .collect();
11320 let keep_fm: std::collections::HashSet<u16> = self
11321 .schema
11322 .indexes
11323 .iter()
11324 .filter(|index| index.kind == IndexKind::FmIndex)
11325 .map(|index| index.column_id)
11326 .collect();
11327 let keep_sparse: std::collections::HashSet<u16> = self
11328 .schema
11329 .indexes
11330 .iter()
11331 .filter(|index| index.kind == IndexKind::Sparse)
11332 .map(|index| index.column_id)
11333 .collect();
11334 let keep_minhash: std::collections::HashSet<u16> = self
11335 .schema
11336 .indexes
11337 .iter()
11338 .filter(|index| index.kind == IndexKind::MinHash)
11339 .map(|index| index.column_id)
11340 .collect();
11341 let keep_learned: std::collections::HashSet<u16> = self
11342 .schema
11343 .indexes
11344 .iter()
11345 .filter(|index| index.kind == IndexKind::LearnedRange)
11346 .map(|index| index.column_id)
11347 .collect();
11348 self.bitmap
11349 .retain(|column_id, _| keep_bitmap.contains(column_id));
11350 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
11351 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
11352 self.sparse
11353 .retain(|column_id, _| keep_sparse.contains(column_id));
11354 self.minhash
11355 .retain(|column_id, _| keep_minhash.contains(column_id));
11356 {
11357 let learned = Arc::make_mut(&mut self.learned_range);
11358 learned.retain(|column_id, _| keep_learned.contains(column_id));
11359 }
11360 }
11361
11362 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
11363 checkpoint_current_schema(self)
11364 }
11365
11366 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
11367 self.ensure_writable()?;
11368 let previous_schema = self.schema.clone();
11369 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
11370 if let Some(schema) = schema {
11371 self.apply_altered_schema_prepared(schema);
11372 self.checkpoint_standalone_schema_change(previous_schema)?;
11373 }
11374 Ok(column)
11375 }
11376
11377 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
11378 if self.live_count == 0 {
11379 return Ok(false);
11380 }
11381 let snap = self.snapshot();
11382 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
11383 Ok(columns
11384 .first()
11385 .map(|(_, col)| col.null_count(col.len()) != 0)
11386 .unwrap_or(true))
11387 }
11388
11389 fn has_stored_versions(&self) -> bool {
11390 !self.memtable.is_empty()
11391 || !self.mutable_run.is_empty()
11392 || self.run_refs.iter().any(|r| r.row_count > 0)
11393 || !self.retiring.is_empty()
11394 }
11395
11396 pub fn add_column(
11401 &mut self,
11402 name: &str,
11403 ty: TypeId,
11404 flags: ColumnFlags,
11405 default_value: Option<crate::schema::DefaultExpr>,
11406 ) -> Result<u16> {
11407 self.add_column_with_id(name, ty, flags, default_value, None)
11408 }
11409
11410 pub fn add_column_with_id(
11411 &mut self,
11412 name: &str,
11413 ty: TypeId,
11414 flags: ColumnFlags,
11415 default_value: Option<crate::schema::DefaultExpr>,
11416 requested_id: Option<u16>,
11417 ) -> Result<u16> {
11418 self.ensure_writable()?;
11419 let previous_schema = self.schema.clone();
11420 let (column, schema) =
11421 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
11422 self.apply_altered_schema_prepared(schema);
11423 self.checkpoint_standalone_schema_change(previous_schema)?;
11424 Ok(column.id)
11425 }
11426
11427 pub(crate) fn prepare_add_column(
11428 &mut self,
11429 name: &str,
11430 ty: TypeId,
11431 flags: ColumnFlags,
11432 default_value: Option<crate::schema::DefaultExpr>,
11433 requested_id: Option<u16>,
11434 ) -> Result<(ColumnDef, Schema)> {
11435 self.ensure_writable()?;
11436 if self.schema.columns.iter().any(|c| c.name == name) {
11437 return Err(MongrelError::Schema(format!(
11438 "column {name} already exists"
11439 )));
11440 }
11441 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
11442 if self.schema.columns.iter().any(|c| c.id == id) {
11443 return Err(MongrelError::Schema(format!(
11444 "column id {id} already exists"
11445 )));
11446 }
11447 id
11448 } else {
11449 self.schema
11450 .columns
11451 .iter()
11452 .map(|c| c.id)
11453 .max()
11454 .unwrap_or(0)
11455 .checked_add(1)
11456 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
11457 };
11458 let column = ColumnDef {
11459 id,
11460 name: name.to_string(),
11461 ty,
11462 flags,
11463 default_value,
11464 embedding_source: None,
11465 };
11466 let mut next_schema = self.schema.clone();
11467 next_schema.columns.push(column.clone());
11468 next_schema.schema_id = next_schema
11469 .schema_id
11470 .checked_add(1)
11471 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11472 next_schema.validate_auto_increment()?;
11473 next_schema.validate_defaults()?;
11474 Ok((column, next_schema))
11475 }
11476
11477 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
11486 self.ensure_writable()?;
11487 let cid = self
11488 .schema
11489 .columns
11490 .iter()
11491 .find(|c| c.name == column_name)
11492 .map(|c| c.id)
11493 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11494 let ty = self
11495 .schema
11496 .columns
11497 .iter()
11498 .find(|c| c.id == cid)
11499 .map(|c| c.ty.clone())
11500 .unwrap_or(TypeId::Int64);
11501 if !matches!(
11502 ty,
11503 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
11504 ) {
11505 return Err(MongrelError::Schema(format!(
11506 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
11507 )));
11508 }
11509 if self
11510 .schema
11511 .indexes
11512 .iter()
11513 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
11514 {
11515 return Ok(()); }
11517 let previous_schema = self.schema.clone();
11518 let previous_learned_range = Arc::clone(&self.learned_range);
11519 let mut next_schema = previous_schema.clone();
11520 next_schema.indexes.push(IndexDef {
11521 name: format!("{}_learned_range", column_name),
11522 column_id: cid,
11523 kind: IndexKind::LearnedRange,
11524 predicate: None,
11525 options: Default::default(),
11526 });
11527 next_schema.schema_id = next_schema
11528 .schema_id
11529 .checked_add(1)
11530 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11531 self.apply_altered_schema_prepared(next_schema);
11532 if let Err(error) = self.build_learned_ranges() {
11533 self.apply_altered_schema_prepared(previous_schema);
11534 self.learned_range = previous_learned_range;
11535 return Err(error);
11536 }
11537 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
11538 if !matches!(
11539 &error,
11540 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
11541 ) {
11542 self.learned_range = previous_learned_range;
11543 }
11544 return Err(error);
11545 }
11546 Ok(())
11547 }
11548
11549 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
11550 let mut schema_published = false;
11551 let schema_result = match self._root_guard.as_deref() {
11552 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
11553 schema_published = true;
11554 }),
11555 None => write_schema_with_after(&self.dir, &self.schema, || {
11556 schema_published = true;
11557 }),
11558 };
11559 if schema_result.is_err() && !schema_published {
11560 self.apply_altered_schema_prepared(previous_schema);
11561 return schema_result;
11562 }
11563
11564 let manifest_result = self.persist_manifest(self.current_epoch());
11565 match (schema_result, manifest_result) {
11566 (_, Ok(())) => Ok(()),
11567 (Ok(()), Err(error)) => {
11568 self.poison_after_maintenance_publish_failure();
11569 Err(MongrelError::DurableCommit {
11570 epoch: self.current_epoch().0,
11571 message: format!(
11572 "schema is durable but matching manifest publication failed: {error}"
11573 ),
11574 })
11575 }
11576 (Err(schema_error), Err(manifest_error)) => {
11577 self.poison_after_maintenance_publish_failure();
11578 Err(MongrelError::CommitOutcomeUnknown {
11579 epoch: self.current_epoch().0,
11580 message: format!(
11581 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11582 ),
11583 })
11584 }
11585 }
11586 }
11587
11588 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
11591 self.sync_byte_threshold = threshold;
11592 if let WalSink::Private(w) = &mut self.wal {
11593 w.set_sync_byte_threshold(threshold);
11594 }
11595 }
11596
11597 pub fn page_cache_flush(&self) {
11601 self.page_cache.flush_to_disk();
11602 }
11603
11604 pub fn page_cache_len(&self) -> usize {
11606 self.page_cache.len()
11607 }
11608
11609 pub fn decoded_cache_len(&self) -> usize {
11612 self.decoded_cache.len()
11613 }
11614
11615 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
11618 self.memtable.drain_sorted()
11619 }
11620
11621 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
11622 self.runs_dir().join(format!("r-{run_id}.sr"))
11623 }
11624
11625 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
11626 match self.runs_root.as_deref() {
11627 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
11628 None => Ok(None),
11629 }
11630 }
11631
11632 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
11633 match self.runs_root.as_deref() {
11634 Some(root) => Ok(Some(root.create_regular_new(name)?)),
11635 None => Ok(None),
11636 }
11637 }
11638
11639 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
11640 match self.runs_root.as_deref() {
11641 Some(root) => match root.remove_file(name) {
11642 Ok(()) => Ok(()),
11643 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11644 Err(error) => Err(error.into()),
11645 },
11646 None => match std::fs::remove_file(self.runs_dir().join(name)) {
11647 Ok(()) => Ok(()),
11648 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11649 Err(error) => Err(error.into()),
11650 },
11651 }
11652 }
11653
11654 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
11655 match self.runs_root.as_deref() {
11656 Some(root) => root
11657 .rename_file_new(source, destination)
11658 .map_err(Into::into),
11659 None => crate::durable_file::rename(
11660 &self.runs_dir().join(source),
11661 &self.runs_dir().join(destination),
11662 )
11663 .map_err(Into::into),
11664 }
11665 }
11666
11667 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11668 self.run_refs.iter().map(|run| run.run_id)
11669 }
11670
11671 pub(crate) fn table_dir(&self) -> &Path {
11672 &self.dir
11673 }
11674
11675 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
11676 &self.schema
11677 }
11678
11679 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
11680 let id = self.next_run_id;
11681 self.next_run_id = self
11682 .next_run_id
11683 .checked_add(1)
11684 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
11685 Ok(id)
11686 }
11687
11688 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
11689 self.run_refs.push(run_ref);
11690 }
11691
11692 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
11702 self.retiring.push(crate::manifest::RetiredRun {
11703 run_id,
11704 retire_epoch,
11705 });
11706 }
11707
11708 pub(crate) fn reap_retiring(
11712 &mut self,
11713 min_active: Epoch,
11714 backup_pinned: &std::collections::HashSet<u128>,
11715 ) -> Result<usize> {
11716 if self.retiring.is_empty() {
11717 return Ok(0);
11718 }
11719 let mut reaped = 0;
11720 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11721 for r in std::mem::take(&mut self.retiring) {
11727 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11728 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11729 reaped += 1;
11730 } else {
11731 kept.push(r);
11732 }
11733 }
11734 self.retiring = kept;
11735 if reaped > 0 {
11736 self.persist_manifest(self.current_epoch())?;
11737 }
11738 Ok(reaped)
11739 }
11740
11741 pub(crate) fn has_reapable_retiring(
11742 &self,
11743 min_active: Epoch,
11744 backup_pinned: &std::collections::HashSet<u128>,
11745 ) -> bool {
11746 self.retiring
11747 .iter()
11748 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11749 }
11750
11751 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11752 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11753 return false;
11754 }
11755 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11756 self.run_refs.push(run_ref);
11757 self.indexes_complete = false;
11758 true
11759 }
11760
11761 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11762 self.kek.as_ref()
11763 }
11764
11765 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11766 let mut reader = match self.runs_root.as_deref() {
11767 Some(root) => RunReader::open_file_with_cache(
11768 root.open_regular(format!("r-{run_id}.sr"))?,
11769 self.schema.clone(),
11770 self.kek.clone(),
11771 Some(self.page_cache.clone()),
11772 Some(self.decoded_cache.clone()),
11773 self.table_id,
11774 Some(&self.verified_runs),
11775 None,
11776 )?,
11777 None => RunReader::open_with_cache(
11778 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11779 self.schema.clone(),
11780 self.kek.clone(),
11781 Some(self.page_cache.clone()),
11782 Some(self.decoded_cache.clone()),
11783 self.table_id,
11784 Some(&self.verified_runs),
11785 )?,
11786 };
11787 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11791 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11792 }
11793 Ok(reader)
11794 }
11795
11796 pub(crate) fn run_refs(&self) -> &[RunRef] {
11797 &self.run_refs
11798 }
11799
11800 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11801 self.retiring.iter().map(|run| run.run_id)
11802 }
11803
11804 pub(crate) fn runs_dir(&self) -> PathBuf {
11805 self.runs_root
11806 .as_deref()
11807 .and_then(|root| root.io_path().ok())
11808 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11809 }
11810
11811 pub(crate) fn wal_dir(&self) -> PathBuf {
11812 self.dir.join(WAL_DIR)
11813 }
11814
11815 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11816 self.run_refs = refs;
11817 }
11818
11819 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11820 self.compaction_zstd_level
11821 }
11822
11823 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11824 self.kek.clone()
11825 }
11826
11827 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11831 self.kek.as_ref().map(|k| k.derive_idx_key())
11832 }
11833
11834 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11838 self.kek.as_ref().map(|k| *k.derive_meta_key())
11839 }
11840
11841 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11844 self.column_keys
11845 .iter()
11846 .map(|(&id, &(_, scheme))| (id, scheme))
11847 .collect()
11848 }
11849
11850 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11855 self.tokenize_value_enc(column_id, v)
11856 }
11857
11858 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11859 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11860 let (key, scheme) = self.column_keys.get(&column_id)?;
11861 let token: Vec<u8> = match (*scheme, v) {
11862 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11863 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11864 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11865 _ => hmac_token(key, &v.encode_key()).to_vec(),
11866 };
11867 Some(Value::Bytes(token))
11868 }
11869
11870 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11872 self.index_lookup_key_bytes(column_id, &v.encode_key())
11873 }
11874
11875 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11878 {
11879 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11880 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11881 if *scheme == SCHEME_HMAC_EQ {
11882 return hmac_token(key, encoded).to_vec();
11883 }
11884 }
11885 }
11886 let _ = column_id;
11887 encoded.to_vec()
11888 }
11889}
11890
11891fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11892 let columnar::NativeColumn::Int64 { data, validity } = col else {
11893 return false;
11894 };
11895 if data.len() < n || !columnar::all_non_null(validity, n) {
11896 return false;
11897 }
11898 data.iter()
11899 .take(n)
11900 .zip(data.iter().skip(1))
11901 .all(|(a, b)| a < b)
11902}
11903
11904#[derive(Debug, Clone)]
11908pub struct ColumnStat {
11909 pub min: Option<Value>,
11910 pub max: Option<Value>,
11911 pub null_count: u64,
11912}
11913
11914#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11916pub enum NativeAgg {
11917 Count,
11918 Sum,
11919 Min,
11920 Max,
11921 Avg,
11922}
11923
11924#[derive(Debug, Clone, PartialEq)]
11926pub enum NativeAggResult {
11927 Count(u64),
11928 Int(i64),
11929 Float(f64),
11930 Null,
11932}
11933
11934#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11936pub enum ApproxAgg {
11937 Count,
11938 Sum,
11939 Avg,
11940}
11941
11942#[derive(Debug, Clone)]
11946pub struct ApproxResult {
11947 pub point: f64,
11949 pub ci_low: f64,
11951 pub ci_high: f64,
11953 pub n_population: u64,
11955 pub n_sample_live: usize,
11957 pub n_passing: usize,
11959}
11960
11961#[derive(Debug, Clone, PartialEq)]
11966pub enum AggState {
11967 Count(u64),
11969 SumI {
11971 sum: i128,
11972 count: u64,
11973 },
11974 SumF {
11976 sum: f64,
11977 count: u64,
11978 },
11979 AvgI {
11981 sum: i128,
11982 count: u64,
11983 },
11984 AvgF {
11986 sum: f64,
11987 count: u64,
11988 },
11989 MinI(i64),
11991 MaxI(i64),
11992 MinF(f64),
11994 MaxF(f64),
11995 Empty,
11997}
11998
11999impl AggState {
12000 pub fn merge(self, other: AggState) -> AggState {
12002 use AggState::*;
12003 match (self, other) {
12004 (Empty, x) | (x, Empty) => x,
12005 (Count(a), Count(b)) => Count(a + b),
12006 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
12007 sum: sa + sb,
12008 count: ca + cb,
12009 },
12010 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
12011 sum: sa + sb,
12012 count: ca + cb,
12013 },
12014 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
12015 sum: sa + sb,
12016 count: ca + cb,
12017 },
12018 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
12019 sum: sa + sb,
12020 count: ca + cb,
12021 },
12022 (MinI(a), MinI(b)) => MinI(a.min(b)),
12023 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
12024 (MinF(a), MinF(b)) => MinF(a.min(b)),
12025 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
12026 _ => Empty, }
12028 }
12029
12030 pub fn point(&self) -> Option<f64> {
12032 match self {
12033 AggState::Count(n) => Some(*n as f64),
12034 AggState::SumI { sum, .. } => Some(*sum as f64),
12035 AggState::SumF { sum, .. } => Some(*sum),
12036 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
12037 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
12038 AggState::MinI(n) => Some(*n as f64),
12039 AggState::MaxI(n) => Some(*n as f64),
12040 AggState::MinF(n) => Some(*n),
12041 AggState::MaxF(n) => Some(*n),
12042 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
12043 }
12044 }
12045
12046 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
12050 let is_float = matches!(ty, Some(TypeId::Float64));
12051 match (agg, result) {
12052 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
12053 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
12054 sum: x as i128,
12055 count: 1, },
12057 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
12058 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
12059 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
12060 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
12061 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
12062 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
12063 (NativeAgg::Count, _) => AggState::Empty,
12064 (_, NativeAggResult::Null) => AggState::Empty,
12065 _ => {
12066 let _ = is_float;
12067 AggState::Empty
12068 }
12069 }
12070 }
12071}
12072
12073#[derive(Debug, Clone)]
12076pub struct CachedAgg {
12077 pub state: AggState,
12078 pub watermark: u64,
12079 pub epoch: u64,
12080}
12081
12082#[derive(Debug, Clone)]
12084pub struct IncrementalAggResult {
12085 pub state: AggState,
12087 pub incremental: bool,
12090 pub delta_rows: u64,
12092}
12093
12094fn agg_state_from_rows(
12098 rows: &[Row],
12099 conditions: &[crate::query::Condition],
12100 index_sets: &[RowIdSet],
12101 column: Option<u16>,
12102 agg: NativeAgg,
12103 schema: &Schema,
12104 control: Option<&crate::ExecutionControl>,
12105) -> Result<AggState> {
12106 let mut count: u64 = 0;
12107 let mut sum_i: i128 = 0;
12108 let mut sum_f: f64 = 0.0;
12109 let mut mn_i: i64 = i64::MAX;
12110 let mut mx_i: i64 = i64::MIN;
12111 let mut mn_f: f64 = f64::INFINITY;
12112 let mut mx_f: f64 = f64::NEG_INFINITY;
12113 let mut saw_int = false;
12114 let mut saw_float = false;
12115 for (index, r) in rows.iter().enumerate() {
12116 execution_checkpoint(control, index)?;
12117 if !conditions
12118 .iter()
12119 .all(|c| condition_matches_row(c, r, schema))
12120 {
12121 continue;
12122 }
12123 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
12124 continue;
12125 }
12126 match agg {
12127 NativeAgg::Count => match column {
12128 None => count += 1,
12130 Some(cid) => match r.columns.get(&cid) {
12133 None | Some(Value::Null) => {}
12134 Some(_) => count += 1,
12135 },
12136 },
12137 _ => match column.and_then(|cid| r.columns.get(&cid)) {
12138 Some(Value::Int64(n)) => {
12139 count += 1;
12140 sum_i += *n as i128;
12141 mn_i = mn_i.min(*n);
12142 mx_i = mx_i.max(*n);
12143 saw_int = true;
12144 }
12145 Some(Value::Float64(f)) => {
12146 count += 1;
12147 sum_f += f;
12148 mn_f = mn_f.min(*f);
12149 mx_f = mx_f.max(*f);
12150 saw_float = true;
12151 }
12152 _ => {}
12153 },
12154 }
12155 }
12156 Ok(match agg {
12157 NativeAgg::Count => {
12158 if count == 0 {
12159 AggState::Empty
12160 } else {
12161 AggState::Count(count)
12162 }
12163 }
12164 NativeAgg::Sum => {
12165 if count == 0 {
12166 AggState::Empty
12167 } else if saw_int {
12168 AggState::SumI { sum: sum_i, count }
12169 } else {
12170 AggState::SumF { sum: sum_f, count }
12171 }
12172 }
12173 NativeAgg::Avg => {
12174 if count == 0 {
12175 AggState::Empty
12176 } else if saw_int {
12177 AggState::AvgI { sum: sum_i, count }
12178 } else {
12179 AggState::AvgF { sum: sum_f, count }
12180 }
12181 }
12182 NativeAgg::Min => {
12183 if !saw_int && !saw_float {
12184 AggState::Empty
12185 } else if saw_int {
12186 AggState::MinI(mn_i)
12187 } else {
12188 AggState::MinF(mn_f)
12189 }
12190 }
12191 NativeAgg::Max => {
12192 if !saw_int && !saw_float {
12193 AggState::Empty
12194 } else if saw_int {
12195 AggState::MaxI(mx_i)
12196 } else {
12197 AggState::MaxF(mx_f)
12198 }
12199 }
12200 })
12201}
12202
12203fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
12207 use crate::query::Condition;
12208 match c {
12209 Condition::Pk(key) => match schema.primary_key() {
12210 Some(pk) => row
12211 .columns
12212 .get(&pk.id)
12213 .map(|v| v.encode_key() == *key)
12214 .unwrap_or(false),
12215 None => false,
12216 },
12217 Condition::BitmapEq { column_id, value } => row
12218 .columns
12219 .get(column_id)
12220 .map(|v| v.encode_key() == *value)
12221 .unwrap_or(false),
12222 Condition::BitmapIn { column_id, values } => {
12223 let key = row.columns.get(column_id).map(|v| v.encode_key());
12224 match key {
12225 Some(k) => values.contains(&k),
12226 None => false,
12227 }
12228 }
12229 Condition::BytesPrefix { column_id, prefix } => row
12230 .columns
12231 .get(column_id)
12232 .map(|v| v.encode_key().starts_with(prefix))
12233 .unwrap_or(false),
12234 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
12235 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
12236 _ => false,
12237 },
12238 Condition::RangeF64 {
12239 column_id,
12240 lo,
12241 lo_inclusive,
12242 hi,
12243 hi_inclusive,
12244 } => match row.columns.get(column_id) {
12245 Some(Value::Float64(n)) => {
12246 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
12247 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
12248 lo_ok && hi_ok
12249 }
12250 _ => false,
12251 },
12252 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
12253 Some(Value::Bytes(b)) => {
12254 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
12255 }
12256 _ => false,
12257 },
12258 Condition::FmContainsAll {
12259 column_id,
12260 patterns,
12261 } => match row.columns.get(column_id) {
12262 Some(Value::Bytes(b)) => patterns
12263 .iter()
12264 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
12265 _ => false,
12266 },
12267 Condition::Ann { .. }
12268 | Condition::SparseMatch { .. }
12269 | Condition::MinHashSimilar { .. } => true,
12270 Condition::IsNull { column_id } => {
12271 matches!(row.columns.get(column_id), Some(Value::Null) | None)
12272 }
12273 Condition::IsNotNull { column_id } => {
12274 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
12275 }
12276 }
12277}
12278
12279fn as_f64(v: Option<&Value>) -> Option<f64> {
12281 match v {
12282 Some(Value::Int64(n)) => Some(*n as f64),
12283 Some(Value::Float64(f)) => Some(*f),
12284 _ => None,
12285 }
12286}
12287
12288fn accumulate_int(
12292 cursor: &mut dyn crate::cursor::Cursor,
12293 control: Option<&crate::ExecutionControl>,
12294) -> Result<(u64, i128, i64, i64)> {
12295 let mut count: u64 = 0;
12296 let mut sum: i128 = 0;
12297 let mut mn: i64 = i64::MAX;
12298 let mut mx: i64 = i64::MIN;
12299 while let Some(cols) = cursor.next_batch()? {
12300 execution_checkpoint(control, 0)?;
12301 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
12302 if crate::columnar::all_non_null(validity, data.len()) {
12303 count += data.len() as u64;
12305 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12306 execution_checkpoint(control, chunk_index * 1024)?;
12307 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
12308 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
12309 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
12310 }
12311 } else {
12312 for (i, &v) in data.iter().enumerate() {
12313 execution_checkpoint(control, i)?;
12314 if crate::columnar::validity_bit(validity, i) {
12315 count += 1;
12316 sum += v as i128;
12317 mn = mn.min(v);
12318 mx = mx.max(v);
12319 }
12320 }
12321 }
12322 }
12323 }
12324 Ok((count, sum, mn, mx))
12325}
12326
12327fn accumulate_float(
12329 cursor: &mut dyn crate::cursor::Cursor,
12330 control: Option<&crate::ExecutionControl>,
12331) -> Result<(u64, f64, f64, f64)> {
12332 let mut count: u64 = 0;
12333 let mut sum: f64 = 0.0;
12334 let mut mn: f64 = f64::INFINITY;
12335 let mut mx: f64 = f64::NEG_INFINITY;
12336 while let Some(cols) = cursor.next_batch()? {
12337 execution_checkpoint(control, 0)?;
12338 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
12339 if crate::columnar::all_non_null(validity, data.len()) {
12340 count += data.len() as u64;
12341 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
12342 execution_checkpoint(control, chunk_index * 1024)?;
12343 sum += chunk.iter().sum::<f64>();
12344 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
12345 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
12346 }
12347 } else {
12348 for (i, &v) in data.iter().enumerate() {
12349 execution_checkpoint(control, i)?;
12350 if crate::columnar::validity_bit(validity, i) {
12351 count += 1;
12352 sum += v;
12353 mn = mn.min(v);
12354 mx = mx.max(v);
12355 }
12356 }
12357 }
12358 }
12359 }
12360 Ok((count, sum, mn, mx))
12361}
12362
12363#[inline]
12364fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
12365 if index.is_multiple_of(256) {
12366 control
12367 .map(crate::ExecutionControl::checkpoint)
12368 .transpose()?;
12369 }
12370 Ok(())
12371}
12372
12373fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
12374 if count == 0 && !matches!(agg, NativeAgg::Count) {
12375 return NativeAggResult::Null;
12376 }
12377 match agg {
12378 NativeAgg::Count => NativeAggResult::Count(count),
12379 NativeAgg::Sum => match sum.try_into() {
12382 Ok(v) => NativeAggResult::Int(v),
12383 Err(_) => NativeAggResult::Null,
12384 },
12385 NativeAgg::Min => NativeAggResult::Int(mn),
12386 NativeAgg::Max => NativeAggResult::Int(mx),
12387 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
12388 }
12389}
12390
12391fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
12392 if count == 0 && !matches!(agg, NativeAgg::Count) {
12393 return NativeAggResult::Null;
12394 }
12395 match agg {
12396 NativeAgg::Count => NativeAggResult::Count(count),
12397 NativeAgg::Sum => NativeAggResult::Float(sum),
12398 NativeAgg::Min => NativeAggResult::Float(mn),
12399 NativeAgg::Max => NativeAggResult::Float(mx),
12400 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
12401 }
12402}
12403
12404fn agg_int(
12407 stats: &[crate::page::PageStat],
12408 decode: fn(Option<&[u8]>) -> Option<i64>,
12409) -> Option<(Option<i64>, Option<i64>, u64)> {
12410 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
12411 let mut any = false;
12412 for s in stats {
12413 if let Some(v) = decode(s.min.as_deref()) {
12414 mn = mn.min(v);
12415 any = true;
12416 }
12417 if let Some(v) = decode(s.max.as_deref()) {
12418 mx = mx.max(v);
12419 any = true;
12420 }
12421 nulls += s.null_count;
12422 }
12423 any.then_some((Some(mn), Some(mx), nulls))
12424}
12425
12426fn agg_float(
12428 stats: &[crate::page::PageStat],
12429 decode: fn(Option<&[u8]>) -> Option<f64>,
12430) -> Option<(Option<f64>, Option<f64>, u64)> {
12431 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
12432 let mut any = false;
12433 for s in stats {
12434 if let Some(v) = decode(s.min.as_deref()) {
12435 mn = mn.min(v);
12436 any = true;
12437 }
12438 if let Some(v) = decode(s.max.as_deref()) {
12439 mx = mx.max(v);
12440 any = true;
12441 }
12442 nulls += s.null_count;
12443 }
12444 any.then_some((Some(mn), Some(mx), nulls))
12445}
12446
12447type SecondaryIndexes = (
12449 HashMap<u16, BitmapIndex>,
12450 HashMap<u16, AnnIndex>,
12451 HashMap<u16, FmIndex>,
12452 HashMap<u16, SparseIndex>,
12453 HashMap<u16, MinHashIndex>,
12454);
12455
12456fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
12457 let mut bitmap = HashMap::new();
12458 let mut ann = HashMap::new();
12459 let mut fm = HashMap::new();
12460 let mut sparse = HashMap::new();
12461 let mut minhash = HashMap::new();
12462 for idef in &schema.indexes {
12463 match idef.kind {
12464 IndexKind::Bitmap => {
12465 bitmap.insert(idef.column_id, BitmapIndex::new());
12466 }
12467 IndexKind::Ann => {
12468 let dim = schema
12469 .columns
12470 .iter()
12471 .find(|c| c.id == idef.column_id)
12472 .and_then(|c| match c.ty {
12473 TypeId::Embedding { dim } => Some(dim as usize),
12474 _ => None,
12475 })
12476 .unwrap_or(0);
12477 let options = idef.options.ann.clone().unwrap_or_default();
12478 ann.insert(
12479 idef.column_id,
12480 AnnIndex::with_full_options(
12481 dim,
12482 options.m,
12483 options.ef_construction,
12484 options.ef_search,
12485 &options,
12486 ),
12487 );
12488 }
12489 IndexKind::FmIndex => {
12490 fm.insert(idef.column_id, FmIndex::new());
12491 }
12492 IndexKind::Sparse => {
12493 sparse.insert(idef.column_id, SparseIndex::new());
12494 }
12495 IndexKind::MinHash => {
12496 let options = idef.options.minhash.clone().unwrap_or_default();
12497 minhash.insert(
12498 idef.column_id,
12499 MinHashIndex::with_options(options.permutations, options.bands),
12500 );
12501 }
12502 _ => {}
12503 }
12504 }
12505 (bitmap, ann, fm, sparse, minhash)
12506}
12507
12508const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
12509 | ColumnFlags::AUTO_INCREMENT
12510 | ColumnFlags::ENCRYPTED
12511 | ColumnFlags::ENCRYPTED_INDEXABLE
12512 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
12513
12514fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
12515 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
12516 return Err(MongrelError::Schema(
12517 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
12518 ));
12519 }
12520 Ok(())
12521}
12522
12523fn validate_alter_column_type(
12524 schema: &Schema,
12525 old: &ColumnDef,
12526 next: &ColumnDef,
12527 has_stored_versions: bool,
12528) -> Result<()> {
12529 if old.ty == next.ty {
12530 return Ok(());
12531 }
12532 if schema.indexes.iter().any(|i| i.column_id == old.id) {
12533 return Err(MongrelError::Schema(format!(
12534 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
12535 old.name
12536 )));
12537 }
12538 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
12539 return Ok(());
12540 }
12541 Err(MongrelError::Schema(format!(
12542 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
12543 old.ty, next.ty
12544 )))
12545}
12546
12547fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
12548 matches!(
12549 (old, new),
12550 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
12551 )
12552}
12553
12554fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
12560 let mut prev: Option<i64> = None;
12561 for r in rows {
12562 match r.columns.get(&pk_id) {
12563 Some(Value::Int64(v)) => {
12564 if prev.is_some_and(|p| p >= *v) {
12565 return false;
12566 }
12567 prev = Some(*v);
12568 }
12569 _ => return false,
12570 }
12571 }
12572 true
12573}
12574
12575#[allow(clippy::too_many_arguments)]
12576fn index_into(
12577 schema: &Schema,
12578 row: &Row,
12579 hot: &mut HotIndex,
12580 bitmap: &mut HashMap<u16, BitmapIndex>,
12581 ann: &mut HashMap<u16, AnnIndex>,
12582 fm: &mut HashMap<u16, FmIndex>,
12583 sparse: &mut HashMap<u16, SparseIndex>,
12584 minhash: &mut HashMap<u16, MinHashIndex>,
12585) {
12586 for idef in &schema.indexes {
12587 let Some(val) = row.columns.get(&idef.column_id) else {
12588 continue;
12589 };
12590 match idef.kind {
12591 IndexKind::Bitmap => {
12592 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12593 b.insert(val.encode_key(), row.row_id);
12594 }
12595 }
12596 IndexKind::Ann => {
12597 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12598 if let Some(meta) = val.generated_embedding_metadata() {
12599 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12601 continue;
12602 }
12603 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12604 .is_err()
12605 {
12606 continue;
12607 }
12608 }
12609 a.insert_validated(v, row.row_id);
12610 }
12611 }
12612 IndexKind::FmIndex => {
12613 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12614 f.insert(b.clone(), row.row_id);
12615 }
12616 }
12617 IndexKind::Sparse => {
12618 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12619 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12622 s.insert(&terms, row.row_id);
12623 }
12624 }
12625 }
12626 IndexKind::MinHash => {
12627 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12628 let tokens = crate::index::token_hashes_from_bytes(b);
12631 mh.insert(&tokens, row.row_id);
12632 }
12633 }
12634 _ => {}
12635 }
12636 }
12637 if let Some(pk_col) = schema.primary_key() {
12638 if let Some(pk_val) = row.columns.get(&pk_col.id) {
12639 hot.insert(pk_val.encode_key(), row.row_id);
12640 }
12641 }
12642}
12643
12644#[allow(clippy::too_many_arguments)]
12647fn index_into_single(
12648 idef: &IndexDef,
12649 _schema: &Schema,
12650 row: &Row,
12651 _hot: &mut HotIndex,
12652 bitmap: &mut HashMap<u16, BitmapIndex>,
12653 ann: &mut HashMap<u16, AnnIndex>,
12654 fm: &mut HashMap<u16, FmIndex>,
12655 sparse: &mut HashMap<u16, SparseIndex>,
12656 minhash: &mut HashMap<u16, MinHashIndex>,
12657) {
12658 let Some(val) = row.columns.get(&idef.column_id) else {
12659 return;
12660 };
12661 match idef.kind {
12662 IndexKind::Bitmap => {
12663 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12664 b.insert(val.encode_key(), row.row_id);
12665 }
12666 }
12667 IndexKind::Ann => {
12668 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12669 if let Some(meta) = val.generated_embedding_metadata() {
12670 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12672 return;
12673 }
12674 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12675 .is_err()
12676 {
12677 return;
12678 }
12679 }
12680 a.insert_validated(v, row.row_id);
12681 }
12682 }
12683 IndexKind::FmIndex => {
12684 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12685 f.insert(b.clone(), row.row_id);
12686 }
12687 }
12688 IndexKind::Sparse => {
12689 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12690 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12691 s.insert(&terms, row.row_id);
12692 }
12693 }
12694 }
12695 IndexKind::MinHash => {
12696 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12697 let tokens = crate::index::token_hashes_from_bytes(b);
12698 mh.insert(&tokens, row.row_id);
12699 }
12700 }
12701 _ => {}
12702 }
12703}
12704
12705fn eval_partial_predicate(
12711 pred: &str,
12712 columns_map: &HashMap<u16, &Value>,
12713 name_to_id: &HashMap<&str, u16>,
12714) -> bool {
12715 let lower = pred.trim().to_ascii_lowercase();
12716 if let Some(rest) = lower.strip_suffix(" is not null") {
12718 let col_name = rest.trim();
12719 if let Some(col_id) = name_to_id.get(col_name) {
12720 return columns_map
12721 .get(col_id)
12722 .is_some_and(|v| !matches!(v, Value::Null));
12723 }
12724 }
12725 if let Some(rest) = lower.strip_suffix(" is null") {
12727 let col_name = rest.trim();
12728 if let Some(col_id) = name_to_id.get(col_name) {
12729 return columns_map
12730 .get(col_id)
12731 .is_none_or(|v| matches!(v, Value::Null));
12732 }
12733 }
12734 true
12737}
12738
12739#[allow(dead_code)]
12745fn bulk_index_key(
12746 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12747 column_id: u16,
12748 ty: TypeId,
12749 col: &columnar::NativeColumn,
12750 i: usize,
12751) -> Option<Vec<u8>> {
12752 let encoded = columnar::encode_key_native(ty, col, i)?;
12753 {
12754 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12755 if let Some((key, scheme)) = column_keys.get(&column_id) {
12756 return Some(match (*scheme, col) {
12757 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12758 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12759 ope_token_i64(key, data[i]).to_vec()
12760 }
12761 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12762 ope_token_f64(key, data[i]).to_vec()
12763 }
12764 _ => hmac_token(key, &encoded).to_vec(),
12765 });
12766 }
12767 }
12768 Some(encoded)
12769}
12770
12771pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12772 write_schema_with_after(dir, schema, || {})
12773}
12774
12775pub(crate) fn write_schema_durable(
12776 root: &crate::durable_file::DurableRoot,
12777 schema: &Schema,
12778) -> Result<()> {
12779 write_schema_durable_with_after(root, schema, || {})
12780}
12781
12782fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12783where
12784 F: FnOnce(),
12785{
12786 let json = serde_json::to_string_pretty(schema)
12787 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12788 crate::durable_file::write_atomic_with_after(
12789 &dir.join(SCHEMA_FILENAME),
12790 json.as_bytes(),
12791 after_publish,
12792 )?;
12793 Ok(())
12794}
12795
12796fn write_schema_durable_with_after<F>(
12797 root: &crate::durable_file::DurableRoot,
12798 schema: &Schema,
12799 after_publish: F,
12800) -> Result<()>
12801where
12802 F: FnOnce(),
12803{
12804 let json = serde_json::to_string_pretty(schema)
12805 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12806 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12807 Ok(())
12808}
12809
12810fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12811 let mut schema_published = false;
12812 let schema_result = match table._root_guard.as_deref() {
12813 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12814 schema_published = true;
12815 }),
12816 None => write_schema_with_after(&table.dir, &table.schema, || {
12817 schema_published = true;
12818 }),
12819 };
12820 if schema_result.is_err() && !schema_published {
12821 return schema_result;
12822 }
12823 match table.persist_manifest(table.current_epoch()) {
12824 Ok(()) => Ok(()),
12825 Err(manifest_error) => Err(match schema_result {
12826 Ok(()) => manifest_error,
12827 Err(schema_error) => MongrelError::Other(format!(
12828 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12829 )),
12830 }),
12831 }
12832}
12833
12834fn read_schema(dir: &Path) -> Result<Schema> {
12835 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12836 read_schema_file(file)
12837}
12838
12839fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12840 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12841 use std::io::Read;
12842
12843 let length = file.metadata()?.len();
12844 if length > MAX_SCHEMA_BYTES {
12845 return Err(MongrelError::ResourceLimitExceeded {
12846 resource: "schema bytes",
12847 requested: usize::try_from(length).unwrap_or(usize::MAX),
12848 limit: MAX_SCHEMA_BYTES as usize,
12849 });
12850 }
12851 let mut bytes = Vec::with_capacity(length as usize);
12852 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12853 if bytes.len() as u64 != length {
12854 return Err(MongrelError::Schema(
12855 "schema length changed while reading".into(),
12856 ));
12857 }
12858 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12859}
12860
12861fn preflight_standalone_open(
12862 dir: &Path,
12863 runs_root: Option<&crate::durable_file::DurableRoot>,
12864 idx_root: Option<&crate::durable_file::DurableRoot>,
12865 manifest: &Manifest,
12866 schema: &Schema,
12867 records: &[crate::wal::Record],
12868 kek: Option<Arc<Kek>>,
12869) -> Result<()> {
12870 crate::wal::validate_shared_transaction_framing(records)?;
12871 if manifest.schema_id > schema.schema_id
12872 || manifest.flushed_epoch > manifest.current_epoch
12873 || manifest.global_idx_epoch > manifest.current_epoch
12874 || manifest.next_row_id == u64::MAX
12875 || manifest.auto_inc_next < 0
12876 || manifest.auto_inc_next == i64::MAX
12877 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12878 {
12879 return Err(MongrelError::InvalidArgument(
12880 "manifest counters or schema identity are invalid".into(),
12881 ));
12882 }
12883 let mut run_ids = HashSet::new();
12884 let mut maximum_row_id = None::<u64>;
12885 for run in &manifest.runs {
12886 if run.run_id >= u64::MAX as u128
12887 || !run_ids.insert(run.run_id)
12888 || run.epoch_created > manifest.current_epoch
12889 {
12890 return Err(MongrelError::InvalidArgument(
12891 "manifest contains an invalid or duplicate active run".into(),
12892 ));
12893 }
12894 let mut reader = match runs_root {
12895 Some(root) => RunReader::open_file(
12896 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12897 schema.clone(),
12898 kek.clone(),
12899 )?,
12900 None => RunReader::open(
12901 dir.join(RUNS_DIR)
12902 .join(format!("r-{}.sr", run.run_id as u64)),
12903 schema.clone(),
12904 kek.clone(),
12905 )?,
12906 };
12907 let header = reader.header();
12908 if header.run_id != run.run_id
12909 || header.level != run.level
12910 || header.row_count != run.row_count
12911 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12912 || header.is_uniform_epoch() && header.epoch_created != 0
12913 || header.schema_id > schema.schema_id
12914 {
12915 return Err(MongrelError::InvalidArgument(format!(
12916 "run {} differs from its manifest",
12917 run.run_id
12918 )));
12919 }
12920 if header.row_count != 0 {
12921 maximum_row_id = Some(
12922 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12923 );
12924 }
12925 reader.validate_all_pages()?;
12926 }
12927 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12928 return Err(MongrelError::InvalidArgument(
12929 "manifest next_row_id does not advance beyond persisted rows".into(),
12930 ));
12931 }
12932 for run in &manifest.retiring {
12933 if run.run_id >= u64::MAX as u128
12934 || run.retire_epoch > manifest.current_epoch
12935 || !run_ids.insert(run.run_id)
12936 {
12937 return Err(MongrelError::InvalidArgument(
12938 "manifest contains an invalid or duplicate retired run".into(),
12939 ));
12940 }
12941 }
12942 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12943 match idx_root {
12944 Some(root) => {
12945 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12946 }
12947 None => {
12948 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12949 }
12950 }
12951
12952 let committed = records
12953 .iter()
12954 .filter_map(|record| match record.op {
12955 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12956 _ => None,
12957 })
12958 .collect::<HashMap<_, _>>();
12959 for record in records {
12960 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12961 continue;
12962 };
12963 match &record.op {
12964 Op::Put { table_id, rows } => {
12965 if *table_id != manifest.table_id {
12966 return Err(MongrelError::CorruptWal {
12967 offset: record.seq.0,
12968 reason: format!(
12969 "private WAL record references table {table_id}, expected {}",
12970 manifest.table_id
12971 ),
12972 });
12973 }
12974 let rows: Vec<Row> =
12975 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12976 offset: record.seq.0,
12977 reason: format!("committed Put payload could not be decoded: {error}"),
12978 })?;
12979 for row in rows {
12980 if row.deleted || row.row_id.0 == u64::MAX {
12981 return Err(MongrelError::CorruptWal {
12982 offset: record.seq.0,
12983 reason: "committed Put contains an invalid row identity".into(),
12984 });
12985 }
12986 let cells = row.columns.into_iter().collect::<Vec<_>>();
12987 schema
12988 .validate_values(&cells)
12989 .map_err(|error| MongrelError::CorruptWal {
12990 offset: record.seq.0,
12991 reason: format!("committed Put violates table schema: {error}"),
12992 })?;
12993 if schema.auto_increment_column().is_some_and(|column| {
12994 matches!(
12995 cells.iter().find(|(id, _)| *id == column.id),
12996 Some((_, Value::Int64(value))) if *value == i64::MAX
12997 )
12998 }) {
12999 return Err(MongrelError::CorruptWal {
13000 offset: record.seq.0,
13001 reason: "committed Put exhausts AUTO_INCREMENT".into(),
13002 });
13003 }
13004 }
13005 }
13006 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
13007 if *table_id != manifest.table_id =>
13008 {
13009 return Err(MongrelError::CorruptWal {
13010 offset: record.seq.0,
13011 reason: format!(
13012 "private WAL record references table {table_id}, expected {}",
13013 manifest.table_id
13014 ),
13015 });
13016 }
13017 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
13018 return Err(MongrelError::CorruptWal {
13019 offset: record.seq.0,
13020 reason: "private WAL contains shared spilled-run metadata".into(),
13021 });
13022 }
13023 _ => {}
13024 }
13025 }
13026 Ok(())
13027}
13028
13029fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
13030 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
13031}
13032
13033fn wal_segment_number(path: &Path) -> Option<u64> {
13034 path.file_stem()
13035 .and_then(|stem| stem.to_str())
13036 .and_then(|stem| stem.strip_prefix("seg-"))
13037 .and_then(|number| number.parse().ok())
13038}
13039
13040fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
13041 let n = list_wal_numbers(wal_dir)?;
13042 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
13043}
13044
13045fn next_wal_number(wal_dir: &Path) -> Result<u32> {
13046 list_wal_numbers(wal_dir)?
13047 .map(|maximum| {
13048 maximum
13049 .checked_add(1)
13050 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
13051 })
13052 .unwrap_or(Ok(0))
13053}
13054
13055fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
13056 let mut max_n = None;
13057 let entries = match std::fs::read_dir(wal_dir) {
13058 Ok(entries) => entries,
13059 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
13060 Err(error) => return Err(error.into()),
13061 };
13062 for entry in entries {
13063 let entry = entry?;
13064 let fname = entry.file_name();
13065 let Some(s) = fname.to_str() else {
13066 continue;
13067 };
13068 let Some(stripped) = s.strip_prefix("seg-") else {
13069 continue;
13070 };
13071 let Some(number) = stripped.strip_suffix(".wal") else {
13072 return Err(MongrelError::CorruptWal {
13073 offset: 0,
13074 reason: format!("malformed WAL segment name {s:?}"),
13075 });
13076 };
13077 let n = number
13078 .parse::<u32>()
13079 .map_err(|_| MongrelError::CorruptWal {
13080 offset: 0,
13081 reason: format!("malformed WAL segment name {s:?}"),
13082 })?;
13083 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
13084 return Err(MongrelError::CorruptWal {
13085 offset: n as u64,
13086 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
13087 });
13088 }
13089 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
13090 }
13091 Ok(max_n)
13092}