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, 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, 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: std::collections::VecDeque<u64>,
948 bytes: u64,
949 max_bytes: u64,
950 dir: Option<std::path::PathBuf>,
951 #[allow(dead_code)]
952 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
953}
954
955#[derive(serde::Serialize, serde::Deserialize)]
957struct SerializedEntry {
958 condition_cols: Vec<u16>,
959 footprint_bits: Vec<u32>,
960 data: SerializedData,
961}
962
963#[derive(serde::Serialize, serde::Deserialize)]
964enum SerializedData {
965 Rows(Vec<Row>),
966 Columns(Vec<(u16, columnar::NativeColumn)>),
967}
968
969impl SerializedEntry {
970 fn from_entry(entry: &CachedEntry) -> Self {
971 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
972 let data = match &entry.data {
973 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
974 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
975 };
976 Self {
977 condition_cols: entry.condition_cols.clone(),
978 footprint_bits,
979 data,
980 }
981 }
982
983 fn into_entry(self) -> Option<CachedEntry> {
984 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
985 let data = match self.data {
986 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
987 SerializedData::Columns(c) => {
988 if !c.iter().all(|(_, col)| col.validate()) {
991 return None;
992 }
993 CachedData::Columns(Arc::new(c))
994 }
995 };
996 Some(CachedEntry {
997 data,
998 footprint,
999 condition_cols: self.condition_cols,
1000 })
1001 }
1002}
1003
1004impl ResultCache {
1005 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1006
1007 fn new() -> Self {
1008 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1009 }
1010
1011 fn with_max_bytes(max_bytes: u64) -> Self {
1012 Self {
1013 entries: std::collections::HashMap::new(),
1014 order: std::collections::VecDeque::new(),
1015 bytes: 0,
1016 max_bytes,
1017 dir: None,
1018 cache_dek: None,
1019 }
1020 }
1021
1022 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1023 let _ = std::fs::create_dir_all(&dir);
1024 self.dir = Some(dir);
1025 self
1026 }
1027
1028 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1029 self.cache_dek = dek;
1030 self
1031 }
1032
1033 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1034 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1035 }
1036
1037 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1041 let Some(path) = self.disk_path(key) else {
1042 return;
1043 };
1044 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
1045 Ok(s) => s,
1046 Err(_) => return,
1047 };
1048 let on_disk = if let Some(dek) = &self.cache_dek {
1050 match self.encrypt_cache(&serialized, dek) {
1051 Some(b) => b,
1052 None => return,
1053 }
1054 } else {
1055 serialized
1056 };
1057 let tmp = path.with_extension("tmp");
1058 use std::io::Write;
1059 let write = || -> std::io::Result<()> {
1060 let mut f = std::fs::File::create(&tmp)?;
1061 f.write_all(&on_disk)?;
1062 f.flush()?;
1063 Ok(())
1064 };
1065 if write().is_err() {
1066 let _ = std::fs::remove_file(&tmp);
1067 return;
1068 }
1069 let _ = std::fs::rename(&tmp, &path);
1070 }
1071
1072 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1074 let path = self.disk_path(key)?;
1075 let bytes = std::fs::read(&path).ok()?;
1076 let plaintext = if let Some(dek) = &self.cache_dek {
1077 self.decrypt_cache(&bytes, dek)?
1078 } else {
1079 bytes
1080 };
1081 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1082 serialized.into_entry()
1083 }
1084
1085 fn remove_from_disk(&self, key: u64) {
1087 if let Some(path) = self.disk_path(key) {
1088 let _ = std::fs::remove_file(&path);
1089 }
1090 }
1091
1092 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1094 use crate::encryption::Cipher;
1095 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1096 let mut nonce = [0u8; 12];
1097 crate::encryption::fill_random(&mut nonce).ok()?;
1098 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1099 let mut out = Vec::with_capacity(12 + ct.len());
1100 out.extend_from_slice(&nonce);
1101 out.extend_from_slice(&ct);
1102 Some(out)
1103 }
1104
1105 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1107 use crate::encryption::Cipher;
1108 if bytes.len() < 28 {
1109 return None;
1110 }
1111 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1112 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1113 let ct = &bytes[12..];
1114 cipher.decrypt_page(&nonce, ct).ok()
1115 }
1116
1117 fn load_persistent(&mut self) {
1120 let Some(dir) = self.dir.as_ref().cloned() else {
1121 return;
1122 };
1123 let entries = match std::fs::read_dir(&dir) {
1124 Ok(e) => e,
1125 Err(_) => return,
1126 };
1127 for entry in entries.flatten() {
1128 let path = entry.path();
1129 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1131 let _ = std::fs::remove_file(&path);
1132 continue;
1133 }
1134 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1135 continue;
1136 }
1137 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1138 Some(s) => s,
1139 None => continue,
1140 };
1141 let key = match u64::from_str_radix(stem, 16) {
1142 Ok(k) => k,
1143 Err(_) => continue,
1144 };
1145 let bytes = match std::fs::read(&path) {
1146 Ok(b) => b,
1147 Err(_) => continue,
1148 };
1149 let plaintext = if let Some(dek) = &self.cache_dek {
1151 match self.decrypt_cache(&bytes, dek) {
1152 Some(p) => p,
1153 None => {
1154 let _ = std::fs::remove_file(&path);
1155 continue;
1156 }
1157 }
1158 } else {
1159 bytes
1160 };
1161 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1162 Ok(serialized) => {
1163 if let Some(entry) = serialized.into_entry() {
1164 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1165 self.entries.insert(key, entry);
1166 self.order.push_back(key);
1167 } else {
1168 let _ = std::fs::remove_file(&path);
1169 }
1170 }
1171 Err(_) => {
1172 let _ = std::fs::remove_file(&path);
1173 }
1174 }
1175 }
1176 self.evict();
1177 }
1178
1179 fn set_max_bytes(&mut self, max_bytes: u64) {
1180 self.max_bytes = max_bytes;
1181 self.evict();
1182 }
1183
1184 fn touch(&mut self, key: u64) {
1186 self.order.retain(|k| *k != key);
1187 self.order.push_back(key);
1188 }
1189
1190 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1191 let res = self.entries.get(&key).and_then(|e| match &e.data {
1192 CachedData::Rows(r) => Some(r.clone()),
1193 CachedData::Columns(_) => None,
1194 });
1195 if res.is_some() {
1196 self.touch(key);
1197 return res;
1198 }
1199 if let Some(entry) = self.load_from_disk(key) {
1201 let res = match &entry.data {
1202 CachedData::Rows(r) => Some(r.clone()),
1203 CachedData::Columns(_) => None,
1204 };
1205 if res.is_some() {
1206 let approx = entry.data.approx_bytes();
1207 self.bytes = self.bytes.saturating_add(approx);
1208 self.entries.insert(key, entry);
1209 self.order.push_back(key);
1210 self.evict();
1211 return res;
1212 }
1213 }
1214 None
1215 }
1216
1217 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1218 let res = self.entries.get(&key).and_then(|e| match &e.data {
1219 CachedData::Columns(c) => Some(c.clone()),
1220 CachedData::Rows(_) => None,
1221 });
1222 if res.is_some() {
1223 self.touch(key);
1224 return res;
1225 }
1226 if let Some(entry) = self.load_from_disk(key) {
1228 let res = match &entry.data {
1229 CachedData::Columns(c) => Some(c.clone()),
1230 CachedData::Rows(_) => None,
1231 };
1232 if res.is_some() {
1233 let approx = entry.data.approx_bytes();
1234 self.bytes = self.bytes.saturating_add(approx);
1235 self.entries.insert(key, entry);
1236 self.order.push_back(key);
1237 self.evict();
1238 return res;
1239 }
1240 }
1241 None
1242 }
1243
1244 fn insert(&mut self, key: u64, entry: CachedEntry) {
1245 let approx = entry.data.approx_bytes();
1246 if self.entries.remove(&key).is_some() {
1247 self.order.retain(|k| *k != key);
1248 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1249 }
1250 self.store_to_disk(key, &entry);
1252 self.bytes = self.bytes.saturating_add(approx);
1253 self.entries.insert(key, entry);
1254 self.order.push_back(key);
1255 self.evict();
1256 }
1257
1258 fn invalidate(
1267 &mut self,
1268 delete_rids: &roaring::RoaringBitmap,
1269 put_cols: &std::collections::HashSet<u16>,
1270 ) {
1271 if self.entries.is_empty() {
1272 return;
1273 }
1274 let has_deletes = !delete_rids.is_empty();
1275 let to_remove: std::collections::HashSet<u64> = self
1276 .entries
1277 .iter()
1278 .filter(|(_, e)| {
1279 let delete_hit = if e.footprint.is_empty() {
1280 has_deletes
1281 } else {
1282 e.footprint.intersection_len(delete_rids) > 0
1283 };
1284 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1285 delete_hit || col_hit
1286 })
1287 .map(|(&k, _)| k)
1288 .collect();
1289 for key in &to_remove {
1290 if let Some(e) = self.entries.remove(key) {
1291 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1292 }
1293 self.remove_from_disk(*key);
1294 }
1295 if !to_remove.is_empty() {
1296 self.order.retain(|k| !to_remove.contains(k));
1297 }
1298 }
1299
1300 fn clear(&mut self) {
1301 if let Some(dir) = &self.dir {
1303 if let Ok(entries) = std::fs::read_dir(dir) {
1304 for entry in entries.flatten() {
1305 let path = entry.path();
1306 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1307 let _ = std::fs::remove_file(&path);
1308 }
1309 }
1310 }
1311 }
1312 self.entries.clear();
1313 self.order.clear();
1314 self.bytes = 0;
1315 }
1316
1317 fn evict(&mut self) {
1318 while self.bytes > self.max_bytes {
1319 let Some(k) = self.order.pop_front() else {
1320 break;
1321 };
1322 if let Some(e) = self.entries.remove(&k) {
1323 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1324 self.remove_from_disk(k);
1328 }
1329 }
1330 }
1331}
1332
1333type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1340
1341fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1342 let _ = kek;
1343 {
1344 if let Some(k) = kek {
1345 return (
1346 Some(k.derive_table_wal_key(_table_id)),
1347 Some(k.derive_cache_key()),
1348 );
1349 }
1350 }
1351 (None, None)
1352}
1353
1354fn read_table_encryption_salt_root(
1355 root: &crate::durable_file::DurableRoot,
1356) -> Result<[u8; crate::encryption::SALT_LEN]> {
1357 use std::io::Read;
1358
1359 let mut file = root
1360 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1361 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1362 let length = file.metadata()?.len();
1363 if length != crate::encryption::SALT_LEN as u64 {
1364 return Err(MongrelError::InvalidArgument(format!(
1365 "salt file is {length} bytes, expected {}",
1366 crate::encryption::SALT_LEN
1367 )));
1368 }
1369 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1370 file.read_exact(&mut salt)?;
1371 Ok(salt)
1372}
1373
1374fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1376 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1377}
1378
1379fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1380 let Some(kek) = kek else {
1381 return HashMap::new();
1382 };
1383 {
1384 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1385 schema
1386 .columns
1387 .iter()
1388 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1389 .map(|c| {
1390 let scheme = if schema
1391 .indexes
1392 .iter()
1393 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1394 {
1395 SCHEME_OPE_RANGE
1396 } else {
1397 SCHEME_HMAC_EQ
1398 };
1399 let key: [u8; 32] = *kek.derive_column_key(c.id);
1400 (c.id, (key, scheme))
1401 })
1402 .collect()
1403 }
1404}
1405
1406pub(crate) struct SharedCtx {
1411 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1412 pub epoch: Arc<EpochAuthority>,
1413 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1414 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1415 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1416 pub kek: Option<Arc<Kek>>,
1417 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1423 pub shared: Option<SharedWalCtx>,
1427 pub table_name: Option<String>,
1430 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1433 pub read_only: bool,
1435}
1436
1437#[derive(Clone)]
1443pub(crate) struct SharedWalCtx {
1444 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1445 pub group: Arc<GroupCommit>,
1446 pub poisoned: Arc<AtomicBool>,
1447 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1448 pub change_wake: tokio::sync::broadcast::Sender<()>,
1449 pub lifecycle: Arc<crate::core::LifecycleController>,
1452}
1453
1454enum WalSink {
1457 Private(Wal),
1458 Shared(SharedWalCtx),
1459 ReadOnly,
1460}
1461
1462impl Clone for WalSink {
1463 fn clone(&self) -> Self {
1464 match self {
1465 Self::Shared(shared) => Self::Shared(shared.clone()),
1466 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1467 }
1468 }
1469}
1470
1471impl SharedCtx {
1472 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1476 let n_shards = if cache_dir.is_some() {
1480 1
1481 } else {
1482 crate::cache::CACHE_SHARDS
1483 };
1484 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1485 let page_cache = if let Some(d) = cache_dir {
1486 Arc::new(crate::cache::Sharded::new(1, || {
1487 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1488 }))
1489 } else {
1490 Arc::new(crate::cache::Sharded::new(n_shards, || {
1491 crate::cache::PageCache::new(per_shard)
1492 }))
1493 };
1494 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1495 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1496 crate::cache::CACHE_SHARDS,
1497 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1498 ));
1499 Self {
1500 root_guard: None,
1501 epoch: Arc::new(EpochAuthority::new(0)),
1502 page_cache,
1503 decoded_cache,
1504 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1505 kek,
1506 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1507 shared: None,
1508 table_name: None,
1509 auth: None,
1510 read_only: false,
1511 }
1512 }
1513}
1514
1515fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1519 use crate::query::Condition;
1520 match c {
1521 Condition::Pk(_)
1523 | Condition::BitmapEq { .. }
1524 | Condition::BitmapIn { .. }
1525 | Condition::BytesPrefix { .. }
1526 | Condition::IsNull { .. }
1527 | Condition::IsNotNull { .. } => 0,
1528 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1530 1
1531 }
1532 Condition::FmContains { .. }
1534 | Condition::FmContainsAll { .. }
1535 | Condition::Ann { .. }
1536 | Condition::SparseMatch { .. } => 2,
1537 }
1538}
1539
1540impl Table {
1541 pub(crate) fn build_secondary_index_artifact<F>(
1547 &self,
1548 definition: &IndexDef,
1549 rows: &[Row],
1550 mut checkpoint: F,
1551 ) -> Result<SecondaryIndexArtifact>
1552 where
1553 F: FnMut(usize, usize) -> Result<()>,
1554 {
1555 let mut build_schema = self.schema.clone();
1556 build_schema.indexes = vec![definition.clone()];
1557 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
1558 let name_to_id: HashMap<&str, u16> = self
1559 .schema
1560 .columns
1561 .iter()
1562 .map(|column| (column.name.as_str(), column.id))
1563 .collect();
1564 let mut hot = HotIndex::new();
1565 let mut accepted = Vec::with_capacity(rows.len());
1566
1567 for (offset, row) in rows.iter().enumerate() {
1568 checkpoint(offset, rows.len())?;
1569 if row.deleted {
1570 continue;
1571 }
1572 if let Some(predicate) = &definition.predicate {
1573 let columns: HashMap<u16, &Value> =
1574 row.columns.iter().map(|(id, value)| (*id, value)).collect();
1575 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
1576 continue;
1577 }
1578 }
1579 accepted.push(row);
1580 if definition.kind == IndexKind::LearnedRange {
1581 continue;
1582 }
1583 let effective = if self.column_keys.is_empty() {
1584 row.clone()
1585 } else {
1586 self.tokenized_for_indexes(row)
1587 };
1588 if definition.kind == IndexKind::Ann {
1589 if let (Some(index), Some(vector)) = (
1590 ann.get_mut(&definition.column_id),
1591 effective
1592 .columns
1593 .get(&definition.column_id)
1594 .and_then(Value::as_embedding),
1595 ) {
1596 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
1597 checkpoint(offset, rows.len())
1598 })?;
1599 }
1600 } else {
1601 index_into_single(
1602 definition,
1603 &build_schema,
1604 &effective,
1605 &mut hot,
1606 &mut bitmap,
1607 &mut ann,
1608 &mut fm,
1609 &mut sparse,
1610 &mut minhash,
1611 );
1612 }
1613 }
1614 checkpoint(rows.len(), rows.len())?;
1615
1616 if let Some(index) = ann.get_mut(&definition.column_id) {
1617 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
1618 }
1619
1620 let column_id = definition.column_id;
1621 match definition.kind {
1622 IndexKind::Bitmap => bitmap
1623 .remove(&column_id)
1624 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
1625 IndexKind::Ann => ann
1626 .remove(&column_id)
1627 .map(|index| SecondaryIndexArtifact::Ann(column_id, index)),
1628 IndexKind::FmIndex => fm
1629 .remove(&column_id)
1630 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
1631 IndexKind::Sparse => sparse
1632 .remove(&column_id)
1633 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
1634 IndexKind::MinHash => minhash
1635 .remove(&column_id)
1636 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
1637 IndexKind::LearnedRange => {
1638 let epsilon = definition
1639 .options
1640 .learned_range
1641 .as_ref()
1642 .map(|options| options.epsilon)
1643 .unwrap_or(16);
1644 let column = self
1645 .schema
1646 .columns
1647 .iter()
1648 .find(|column| column.id == column_id)
1649 .ok_or_else(|| {
1650 MongrelError::Schema(format!(
1651 "index {} references unknown column {column_id}",
1652 definition.name
1653 ))
1654 })?;
1655 let index = match column.ty {
1656 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1657 let pairs: Vec<_> = accepted
1658 .iter()
1659 .filter_map(|row| match row.columns.get(&column_id) {
1660 Some(Value::Int64(value)) if !row.deleted => {
1661 Some((*value, row.row_id.0))
1662 }
1663 _ => None,
1664 })
1665 .collect();
1666 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
1667 }
1668 TypeId::Float32 | TypeId::Float64 => {
1669 let pairs: Vec<_> = accepted
1670 .iter()
1671 .filter_map(|row| match row.columns.get(&column_id) {
1672 Some(Value::Float64(value)) if !row.deleted => {
1673 Some((*value, row.row_id.0))
1674 }
1675 _ => None,
1676 })
1677 .collect();
1678 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
1679 }
1680 ref ty => {
1681 return Err(MongrelError::Schema(format!(
1682 "LearnedRange index {} does not support {ty:?}",
1683 definition.name
1684 )));
1685 }
1686 };
1687 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
1688 }
1689 }
1690 .ok_or_else(|| {
1691 MongrelError::Other(format!(
1692 "failed to construct hidden index {}",
1693 definition.name
1694 ))
1695 })
1696 }
1697
1698 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1699 let dir = dir.as_ref().to_path_buf();
1700 crate::durable_file::create_directory_all(&dir)?;
1701 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1702 let pinned = root.io_path()?;
1703 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1704 ctx.root_guard = Some(root);
1705 Self::create_in(&pinned, schema, table_id, ctx)
1706 }
1707
1708 pub fn create_encrypted(
1719 dir: impl AsRef<Path>,
1720 schema: Schema,
1721 table_id: u64,
1722 passphrase: &str,
1723 ) -> Result<Self> {
1724 let dir = dir.as_ref().to_path_buf();
1725 crate::durable_file::create_directory_all(&dir)?;
1726 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1727 root.create_directory_all(META_DIR)?;
1728 let salt = crate::encryption::random_salt()?;
1729 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1730 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1731 let pinned = root.io_path()?;
1732 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1733 ctx.root_guard = Some(root);
1734 Self::create_in(&pinned, schema, table_id, ctx)
1735 }
1736
1737 pub fn create_with_key(
1742 dir: impl AsRef<Path>,
1743 schema: Schema,
1744 table_id: u64,
1745 key: &[u8],
1746 ) -> Result<Self> {
1747 let dir = dir.as_ref().to_path_buf();
1748 crate::durable_file::create_directory_all(&dir)?;
1749 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1750 root.create_directory_all(META_DIR)?;
1751 let salt = crate::encryption::random_salt()?;
1752 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1753 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1754 let pinned = root.io_path()?;
1755 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1756 ctx.root_guard = Some(root);
1757 Self::create_in(&pinned, schema, table_id, ctx)
1758 }
1759
1760 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1762 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1763 let salt = read_table_encryption_salt_root(&root)?;
1764 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1765 let pinned = root.io_path()?;
1766 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1767 ctx.root_guard = Some(root);
1768 Self::open_in(&pinned, ctx)
1769 }
1770
1771 pub(crate) fn create_in(
1772 dir: impl AsRef<Path>,
1773 schema: Schema,
1774 table_id: u64,
1775 ctx: SharedCtx,
1776 ) -> Result<Self> {
1777 schema.validate_auto_increment()?;
1778 schema.validate_defaults()?;
1779 schema.validate_ai()?;
1780 for index in &schema.indexes {
1781 index.validate_options()?;
1782 }
1783 let dir = dir.as_ref().to_path_buf();
1784 let runs_root = match ctx.root_guard.as_ref() {
1785 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1786 None => {
1787 crate::durable_file::create_directory_all(&dir)?;
1788 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1789 None
1790 }
1791 };
1792 match ctx.root_guard.as_deref() {
1793 Some(root) => write_schema_durable(root, &schema)?,
1794 None => write_schema(&dir, &schema)?,
1795 }
1796 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1797 let (wal, current_txn_id) = match ctx.shared.clone() {
1800 Some(s) => (WalSink::Shared(s), 0),
1801 None => {
1802 let pinned_wal_root = match ctx.root_guard.as_deref() {
1803 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1804 None => None,
1805 };
1806 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1807 root.io_path()?
1808 } else {
1809 let wal_dir = dir.join(WAL_DIR);
1810 crate::durable_file::create_directory_all(&wal_dir)?;
1811 wal_dir
1812 };
1813 let mut w = if let Some(ref dk) = wal_dek {
1814 Wal::create_with_cipher(
1815 wal_dir.join("seg-000000.wal"),
1816 Epoch(0),
1817 Some(make_cipher(dk)),
1818 0,
1819 )?
1820 } else {
1821 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1822 };
1823 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1824 (WalSink::Private(w), 1)
1825 }
1826 };
1827 let mut manifest = Manifest::new(table_id, schema.schema_id);
1828 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1833 match ctx.root_guard.as_deref() {
1834 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1835 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1836 }
1837 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1838 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1839 let auto_inc = resolve_auto_inc(&schema);
1840 let rcache_dir = dir.join(RCACHE_DIR);
1841 let initial_view = ReadGeneration::empty(&schema);
1842 Ok(Self {
1843 dir,
1844 _root_guard: ctx.root_guard,
1845 runs_root,
1846 idx_root: None,
1847 table_id,
1848 name: ctx.table_name.unwrap_or_default(),
1849 auth: ctx.auth,
1850 read_only: ctx.read_only,
1851 durable_commit_failed: false,
1852 wal,
1853 memtable: Memtable::new(),
1854 mutable_run: MutableRun::new(),
1855 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1856 compaction_zstd_level: 3,
1857 allocator: RowIdAllocator::new(0),
1858 epoch: ctx.epoch,
1859 data_generation: 0,
1860 schema,
1861 hot: HotIndex::new(),
1862 kek: ctx.kek,
1863 column_keys,
1864 run_refs: Vec::new(),
1865 retiring: Vec::new(),
1866 next_run_id: 1,
1867 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1868 current_txn_id,
1869 pending_private_mutations: false,
1870 bitmap,
1871 ann,
1872 fm,
1873 sparse,
1874 minhash,
1875 learned_range: Arc::new(HashMap::new()),
1876 pk_by_row: ReversePkMap::new(),
1877 pinned: BTreeMap::new(),
1878 live_count: 0,
1879 reservoir: crate::reservoir::Reservoir::default(),
1880 reservoir_complete: true,
1881 had_deletes: false,
1882 agg_cache: Arc::new(HashMap::new()),
1883 global_idx_epoch: 0,
1884 indexes_complete: true,
1885 index_build_policy: IndexBuildPolicy::default(),
1886 pk_by_row_complete: false,
1887 flushed_epoch: 0,
1888 page_cache: ctx.page_cache,
1889 decoded_cache: ctx.decoded_cache,
1890 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1891 snapshots: ctx.snapshots,
1892 commit_lock: ctx.commit_lock,
1893 result_cache: Arc::new(parking_lot::Mutex::new(
1894 ResultCache::new()
1895 .with_dir(rcache_dir)
1896 .with_cache_dek(cache_dek.clone()),
1897 )),
1898 pending_delete_rids: roaring::RoaringBitmap::new(),
1899 pending_put_cols: std::collections::HashSet::new(),
1900 pending_rows: Vec::new(),
1901 pending_rows_auto_inc: Vec::new(),
1902 pending_dels: Vec::new(),
1903 pending_truncate: None,
1904 wal_dek,
1905 auto_inc,
1906 ttl: None,
1907 pins: Arc::new(crate::retention::PinRegistry::new()),
1908 published: Arc::new(ArcSwap::from_pointee(initial_view)),
1909 read_generation_pin: None,
1910 })
1911 }
1912
1913 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1917 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1918 let pinned = root.io_path()?;
1919 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1920 ctx.root_guard = Some(root);
1921 Self::open_in(&pinned, ctx)
1922 }
1923
1924 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1927 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1928 let salt = read_table_encryption_salt_root(&root)?;
1929 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1930 let pinned = root.io_path()?;
1931 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1932 ctx.root_guard = Some(root);
1933 let t = Self::open_in(&pinned, ctx)?;
1934 Ok(t)
1935 }
1936
1937 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1938 let dir = dir.as_ref().to_path_buf();
1939 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1940 let mut manifest = match ctx.root_guard.as_ref() {
1941 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1942 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1943 };
1944 let schema: Schema = match ctx.root_guard.as_ref() {
1945 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1946 None => read_schema(&dir)?,
1947 };
1948 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1954 let runs_root = match ctx.root_guard.as_ref() {
1955 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1956 None => None,
1957 };
1958 let idx_root = match ctx.root_guard.as_ref() {
1959 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1960 Ok(root) => Some(Arc::new(root)),
1961 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1962 Err(error) => return Err(error.into()),
1963 },
1964 None => None,
1965 };
1966 schema.validate_auto_increment()?;
1967 schema.validate_defaults()?;
1968 schema.validate_ai()?;
1969 for index in &schema.indexes {
1970 index.validate_options()?;
1971 }
1972 let replay_epoch = Epoch(manifest.current_epoch);
1973 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1974 let private_replayed = if ctx.shared.is_none() {
1975 match latest_wal_segment(&dir.join(WAL_DIR))? {
1976 Some(path) => {
1977 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1978 crate::wal::replay_with_cipher(path, cipher)?
1979 }
1980 None => Vec::new(),
1981 }
1982 } else {
1983 Vec::new()
1984 };
1985 if ctx.shared.is_none() {
1986 preflight_standalone_open(
1987 &dir,
1988 runs_root.as_deref(),
1989 idx_root.as_deref(),
1990 &manifest,
1991 &schema,
1992 &private_replayed,
1993 ctx.kek.clone(),
1994 )?;
1995 }
1996 let next_run_id = derive_next_run_id(
1997 &dir,
1998 runs_root.as_deref(),
1999 &manifest.runs,
2000 &manifest.retiring,
2001 )?;
2002 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2006 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2007 None => {
2008 let replayed = private_replayed;
2009 let wal_dir = dir.join(WAL_DIR);
2015 crate::durable_file::create_directory_all(&wal_dir)?;
2016 let segment = next_wal_segment(&wal_dir)?;
2017 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2018 let temporary = wal_dir.join(format!(
2019 ".recovery-{}-{}-{segment_no:06}.tmp",
2020 std::process::id(),
2021 std::time::SystemTime::now()
2022 .duration_since(std::time::UNIX_EPOCH)
2023 .unwrap_or_default()
2024 .as_nanos()
2025 ));
2026 let mut w = Wal::create_with_cipher(
2027 &temporary,
2028 replay_epoch,
2029 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2030 segment_no,
2031 )?;
2032 for record in &replayed {
2033 w.append_txn(record.txn_id, record.op.clone())?;
2034 }
2035 let mut w = w.publish_as(segment)?;
2036 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2037 let next_txn_id = replayed
2038 .iter()
2039 .map(|record| record.txn_id)
2040 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2041 .max()
2042 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2043 .unwrap_or(1);
2044 (WalSink::Private(w), replayed, next_txn_id)
2045 }
2046 };
2047
2048 let mut memtable = Memtable::new();
2049 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2050 let persisted_epoch = manifest.current_epoch;
2051 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2058 s.next = manifest.auto_inc_next;
2059 s.seeded = manifest.auto_inc_next > 0;
2060 s
2061 });
2062
2063 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2070 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2071 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2072 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2073 std::collections::BTreeMap::new();
2074 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2075 let mut recovered_epoch = manifest.current_epoch;
2076 let mut recovered_manifest_dirty = schema_manifest_repair;
2077 let mut saw_delete = false;
2078 for record in replayed {
2079 let txn_id = record.txn_id;
2080 match record.op {
2081 Op::Put { rows, .. } => {
2082 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2083 for row in &rows {
2084 allocator.advance_to(row.row_id)?;
2085 if let Some(ai) = auto_inc.as_mut() {
2086 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2087 let next = n.checked_add(1).ok_or_else(|| {
2088 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2089 })?;
2090 if next > ai.next {
2091 ai.next = next;
2092 }
2093 }
2094 }
2095 }
2096 staged_puts.entry(txn_id).or_default().extend(rows);
2097 }
2098 Op::Delete { row_ids, .. } => {
2099 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2100 }
2101 Op::TxnCommit { epoch, .. } => {
2102 let commit_epoch = Epoch(epoch);
2103 recovered_epoch = recovered_epoch.max(epoch);
2104 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2105 memtable = Memtable::new();
2106 replayed_puts.clear();
2107 replayed_deletes.clear();
2108 manifest.runs.clear();
2109 manifest.retiring.clear();
2110 manifest.live_count = 0;
2111 manifest.global_idx_epoch = 0;
2112 manifest.current_epoch = manifest.current_epoch.max(epoch);
2113 recovered_manifest_dirty = true;
2114 saw_delete = true;
2115 }
2116 if let Some(puts) = staged_puts.remove(&txn_id) {
2117 if commit_epoch.0 > manifest.flushed_epoch {
2118 for row in &puts {
2119 memtable.upsert(row.clone());
2120 }
2121 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2122 }
2123 }
2124 if let Some(dels) = staged_deletes.remove(&txn_id) {
2125 saw_delete = true;
2126 if commit_epoch.0 > manifest.flushed_epoch {
2127 for rid in dels {
2128 memtable.tombstone(rid, commit_epoch);
2129 replayed_deletes.push((rid, commit_epoch));
2130 }
2131 }
2132 }
2133 }
2134 Op::TxnAbort => {
2135 staged_puts.remove(&txn_id);
2136 staged_deletes.remove(&txn_id);
2137 staged_truncates.remove(&txn_id);
2138 }
2139 Op::TruncateTable { .. } => {
2140 staged_puts.remove(&txn_id);
2141 staged_deletes.remove(&txn_id);
2142 staged_truncates.insert(txn_id);
2143 }
2144 Op::ExternalTableState { .. }
2145 | Op::Flush { .. }
2146 | Op::Ddl(_)
2147 | Op::BeforeImage { .. }
2148 | Op::CommitTimestamp { .. }
2149 | Op::SpilledRows { .. } => {}
2150 }
2151 }
2152
2153 let rcache_dir = dir.join(RCACHE_DIR);
2154 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2155 let initial_view = ReadGeneration::empty(&schema);
2156 let mut db = Self {
2157 dir,
2158 _root_guard: ctx.root_guard,
2159 runs_root,
2160 idx_root,
2161 table_id: manifest.table_id,
2162 name: ctx.table_name.unwrap_or_default(),
2163 auth: ctx.auth,
2164 read_only: ctx.read_only,
2165 durable_commit_failed: false,
2166 wal,
2167 memtable,
2168 mutable_run: MutableRun::new(),
2169 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2170 compaction_zstd_level: 3,
2171 allocator,
2172 epoch: ctx.epoch,
2173 data_generation: persisted_epoch,
2174 schema,
2175 hot: HotIndex::new(),
2176 kek: ctx.kek,
2177 column_keys,
2178 run_refs: manifest.runs.clone(),
2179 retiring: manifest.retiring.clone(),
2180 next_run_id,
2181 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2182 current_txn_id,
2183 pending_private_mutations: false,
2184 bitmap: HashMap::new(),
2185 ann: HashMap::new(),
2186 fm: HashMap::new(),
2187 sparse: HashMap::new(),
2188 minhash: HashMap::new(),
2189 learned_range: Arc::new(HashMap::new()),
2190 pk_by_row: ReversePkMap::new(),
2191 pinned: BTreeMap::new(),
2192 live_count: manifest.live_count,
2193 reservoir: crate::reservoir::Reservoir::default(),
2194 reservoir_complete: false,
2195 had_deletes: saw_delete
2196 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2197 != manifest.live_count,
2198 agg_cache: Arc::new(HashMap::new()),
2199 global_idx_epoch: manifest.global_idx_epoch,
2200 indexes_complete: true,
2201 index_build_policy: IndexBuildPolicy::default(),
2202 pk_by_row_complete: false,
2203 flushed_epoch: manifest.flushed_epoch,
2204 page_cache: ctx.page_cache,
2205 decoded_cache: ctx.decoded_cache,
2206 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2207 snapshots: ctx.snapshots,
2208 commit_lock: ctx.commit_lock,
2209 result_cache: Arc::new(parking_lot::Mutex::new(
2210 ResultCache::new()
2211 .with_dir(rcache_dir)
2212 .with_cache_dek(cache_dek.clone()),
2213 )),
2214 pending_delete_rids: roaring::RoaringBitmap::new(),
2215 pending_put_cols: std::collections::HashSet::new(),
2216 pending_rows: Vec::new(),
2217 pending_rows_auto_inc: Vec::new(),
2218 pending_dels: Vec::new(),
2219 pending_truncate: None,
2220 wal_dek,
2221 auto_inc,
2222 ttl: manifest.ttl,
2223 pins: Arc::new(crate::retention::PinRegistry::new()),
2224 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2225 read_generation_pin: None,
2226 };
2227
2228 db.epoch.advance_recovered(Epoch(recovered_epoch));
2231
2232 let checkpoint = match db.idx_root.as_deref() {
2237 Some(root) => {
2238 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2239 }
2240 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2241 };
2242 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2243 c.epoch_built == manifest.global_idx_epoch
2244 && manifest.global_idx_epoch > 0
2245 && manifest
2246 .runs
2247 .iter()
2248 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2249 });
2250 if let Some(loaded) = checkpoint {
2251 if checkpoint_valid {
2252 db.hot = loaded.hot;
2253 db.bitmap = loaded.bitmap;
2254 db.ann = loaded.ann;
2255 db.fm = loaded.fm;
2256 db.sparse = loaded.sparse;
2257 db.minhash = loaded.minhash;
2258 db.learned_range = Arc::new(loaded.learned_range);
2259 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2264 for (cid, idx) in bitmap0 {
2265 db.bitmap.entry(cid).or_insert(idx);
2266 }
2267 for (cid, idx) in ann0 {
2268 db.ann.entry(cid).or_insert(idx);
2269 }
2270 for (cid, idx) in fm0 {
2271 db.fm.entry(cid).or_insert(idx);
2272 }
2273 for (cid, idx) in sparse0 {
2274 db.sparse.entry(cid).or_insert(idx);
2275 }
2276 for (cid, idx) in minhash0 {
2277 db.minhash.entry(cid).or_insert(idx);
2278 }
2279 }
2282 }
2283 if !checkpoint_valid {
2284 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2285 db.bitmap = bitmap;
2286 db.ann = ann;
2287 db.fm = fm;
2288 db.sparse = sparse;
2289 db.minhash = minhash;
2290 db.rebuild_indexes_from_runs()?;
2291 db.build_learned_ranges()?;
2292 }
2293
2294 for (epoch, group) in replayed_puts {
2299 let (losers, winner_pks) = db.partition_pk_winners(&group);
2300 for (key, &row_id) in &winner_pks {
2301 if let Some(old_rid) = db.hot.get(key) {
2302 if old_rid != row_id {
2303 db.tombstone_row(old_rid, epoch, false);
2304 }
2305 }
2306 }
2307 for &loser_rid in &losers {
2308 db.tombstone_row(loser_rid, epoch, false);
2309 }
2310 for (key, row_id) in winner_pks {
2311 db.insert_hot_pk(key, row_id);
2312 }
2313 if db.schema.primary_key().is_none() {
2314 for r in &group {
2315 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2316 }
2317 }
2318 for r in &group {
2319 if !losers.contains(&r.row_id) {
2320 db.index_row(r);
2321 }
2322 }
2323 }
2324 for (rid, epoch) in &replayed_deletes {
2328 db.remove_hot_for_row(*rid, *epoch);
2329 }
2330
2331 if recovered_manifest_dirty {
2332 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2333 db.live_count = rows.len() as u64;
2334 db.persist_manifest(Epoch(recovered_epoch))?;
2335 }
2336
2337 db.result_cache.lock().load_persistent();
2344 Ok(db)
2345 }
2346
2347 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2353 if self.reservoir_complete {
2354 return Ok(());
2355 }
2356 self.rebuild_reservoir()?;
2357 self.reservoir_complete = true;
2358 Ok(())
2359 }
2360
2361 fn rebuild_reservoir(&mut self) -> Result<()> {
2364 let snap = self.snapshot();
2365 let rows = self.visible_rows(snap)?;
2366 self.reservoir.reset();
2367 for r in rows {
2368 self.reservoir.offer(r.row_id.0);
2369 }
2370 Ok(())
2371 }
2372
2373 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2374 self.rebuild_indexes_from_runs_inner(None)
2375 }
2376
2377 fn rebuild_indexes_from_runs_inner(
2378 &mut self,
2379 control: Option<&crate::ExecutionControl>,
2380 ) -> Result<()> {
2381 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2384 crate::retention::PinSource::OnlineIndexBuild,
2385 self.current_epoch(),
2386 );
2387 self.hot = HotIndex::new();
2388 self.pk_by_row.clear();
2389 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2390 self.bitmap = bitmap;
2391 self.ann = ann;
2392 self.fm = fm;
2393 self.sparse = sparse;
2394 self.minhash = minhash;
2395 let snapshot = Epoch(u64::MAX);
2396 let ttl_now = unix_nanos_now();
2397 let mut scanned = 0_usize;
2398 for rr in self.run_refs.clone() {
2399 if let Some(control) = control {
2400 control.checkpoint()?;
2401 }
2402 let mut reader = self.open_reader(rr.run_id)?;
2403 for row in reader.visible_rows(snapshot)? {
2404 if scanned.is_multiple_of(256) {
2405 if let Some(control) = control {
2406 control.checkpoint()?;
2407 }
2408 }
2409 scanned += 1;
2410 if self.row_expired_at(&row, ttl_now) {
2411 continue;
2412 }
2413 let tok_row = self.tokenized_for_indexes(&row);
2414 index_into(
2415 &self.schema,
2416 &tok_row,
2417 &mut self.hot,
2418 &mut self.bitmap,
2419 &mut self.ann,
2420 &mut self.fm,
2421 &mut self.sparse,
2422 &mut self.minhash,
2423 );
2424 }
2425 }
2426 for row in self.mutable_run.visible_versions(snapshot) {
2427 if scanned.is_multiple_of(256) {
2428 if let Some(control) = control {
2429 control.checkpoint()?;
2430 }
2431 }
2432 scanned += 1;
2433 if row.deleted {
2434 self.remove_hot_for_row(row.row_id, snapshot);
2435 } else if !self.row_expired_at(&row, ttl_now) {
2436 self.index_row(&row);
2437 }
2438 }
2439 for row in self.memtable.visible_versions(snapshot) {
2440 if scanned.is_multiple_of(256) {
2441 if let Some(control) = control {
2442 control.checkpoint()?;
2443 }
2444 }
2445 scanned += 1;
2446 if row.deleted {
2447 self.remove_hot_for_row(row.row_id, snapshot);
2448 } else if !self.row_expired_at(&row, ttl_now) {
2449 self.index_row(&row);
2450 }
2451 }
2452 self.refresh_pk_by_row_from_hot();
2453 Ok(())
2454 }
2455
2456 fn refresh_pk_by_row_from_hot(&mut self) {
2457 self.pk_by_row_complete = true;
2458 if self.schema.primary_key().is_none() {
2459 self.pk_by_row.clear();
2460 return;
2461 }
2462 self.pk_by_row = ReversePkMap::from_entries(
2468 self.hot
2469 .entries()
2470 .into_iter()
2471 .map(|(key, row_id)| (row_id, key)),
2472 );
2473 }
2474
2475 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2476 if self.schema.primary_key().is_some() {
2477 self.pk_by_row.insert(row_id, key.clone());
2478 }
2479 self.hot.insert(key, row_id);
2480 }
2481
2482 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2486 self.build_learned_ranges_inner(None)
2487 }
2488
2489 fn build_learned_ranges_inner(
2490 &mut self,
2491 control: Option<&crate::ExecutionControl>,
2492 ) -> Result<()> {
2493 self.learned_range = Arc::new(HashMap::new());
2494 if self.run_refs.len() != 1 {
2495 return Ok(());
2496 }
2497 let cols: Vec<(u16, usize)> = self
2498 .schema
2499 .indexes
2500 .iter()
2501 .filter(|i| i.kind == IndexKind::LearnedRange)
2502 .map(|i| {
2503 (
2504 i.column_id,
2505 i.options
2506 .learned_range
2507 .as_ref()
2508 .map(|options| options.epsilon)
2509 .unwrap_or(16),
2510 )
2511 })
2512 .collect();
2513 if cols.is_empty() {
2514 return Ok(());
2515 }
2516 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2517 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2518 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2519 _ => return Ok(()),
2520 };
2521 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2522 if column_index % 256 == 0 {
2523 if let Some(control) = control {
2524 control.checkpoint()?;
2525 }
2526 }
2527 let ty = self
2528 .schema
2529 .columns
2530 .iter()
2531 .find(|c| c.id == cid)
2532 .map(|c| c.ty.clone())
2533 .unwrap_or(TypeId::Int64);
2534 match ty {
2535 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2536 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2537 let pairs: Vec<(i64, u64)> = data
2538 .iter()
2539 .zip(row_ids.iter())
2540 .map(|(v, r)| (*v, *r))
2541 .collect();
2542 Arc::make_mut(&mut self.learned_range).insert(
2543 cid,
2544 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2545 );
2546 }
2547 }
2548 TypeId::Float64 => {
2549 if let columnar::NativeColumn::Float64 { data, .. } =
2550 reader.column_native(cid)?
2551 {
2552 let pairs: Vec<(f64, u64)> = data
2553 .iter()
2554 .zip(row_ids.iter())
2555 .map(|(v, r)| (*v, *r))
2556 .collect();
2557 Arc::make_mut(&mut self.learned_range).insert(
2558 cid,
2559 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2560 );
2561 }
2562 }
2563 _ => {}
2564 }
2565 }
2566 Ok(())
2567 }
2568
2569 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2576 if self.indexes_complete {
2577 crate::trace::QueryTrace::record(|t| {
2578 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2579 });
2580 return Ok(());
2581 }
2582 crate::trace::QueryTrace::record(|t| {
2583 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2584 });
2585 self.rebuild_indexes_from_runs()?;
2586 self.build_learned_ranges()?;
2587 self.indexes_complete = true;
2588 let epoch = self.current_epoch();
2589 self.checkpoint_indexes(epoch);
2590 Ok(())
2591 }
2592
2593 #[doc(hidden)]
2596 pub fn ensure_indexes_complete_controlled<F>(
2597 &mut self,
2598 control: &crate::ExecutionControl,
2599 before_publish: F,
2600 ) -> Result<bool>
2601 where
2602 F: FnOnce() -> bool,
2603 {
2604 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2605 .map(|(changed, _)| changed)
2606 }
2607
2608 #[doc(hidden)]
2611 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2612 &mut self,
2613 control: &crate::ExecutionControl,
2614 before_publish: F,
2615 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2616 where
2617 F: FnOnce() -> bool,
2618 {
2619 if self.indexes_complete {
2620 crate::trace::QueryTrace::record(|trace| {
2621 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2622 });
2623 return Ok((false, None));
2624 }
2625 crate::trace::QueryTrace::record(|trace| {
2626 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2627 });
2628 control.checkpoint()?;
2629 let maintenance_epoch = self.current_epoch();
2630 self.rebuild_indexes_from_runs_inner(Some(control))?;
2631 self.build_learned_ranges_inner(Some(control))?;
2632 control.checkpoint()?;
2633 if !before_publish() {
2634 return Err(MongrelError::Cancelled);
2635 }
2636 self.indexes_complete = true;
2637 self.checkpoint_indexes(maintenance_epoch);
2638 Ok((
2639 true,
2640 Some(MaintenanceReceipt {
2641 epoch: maintenance_epoch,
2642 }),
2643 ))
2644 }
2645
2646 fn pending_epoch(&self) -> Epoch {
2647 Epoch(self.epoch.visible().0 + 1)
2648 }
2649
2650 fn is_shared(&self) -> bool {
2653 matches!(self.wal, WalSink::Shared(_))
2654 }
2655
2656 fn ensure_txn_id(&mut self) -> Result<u64> {
2660 if self.current_txn_id == 0 {
2661 let id = match &self.wal {
2662 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2663 WalSink::Private(_) => {
2664 return Err(MongrelError::Full(
2665 "standalone transaction id namespace exhausted".into(),
2666 ))
2667 }
2668 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2669 };
2670 self.current_txn_id = id;
2671 }
2672 Ok(self.current_txn_id)
2673 }
2674
2675 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2678 self.ensure_writable()?;
2679 let txn_id = self.ensure_txn_id()?;
2680 let table_id = self.table_id;
2681 match &mut self.wal {
2682 WalSink::Private(w) => {
2683 w.append_txn(txn_id, op)?;
2684 self.pending_private_mutations = true;
2685 }
2686 WalSink::Shared(s) => {
2687 s.wal.lock().append(txn_id, table_id, op)?;
2688 }
2689 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2690 }
2691 Ok(())
2692 }
2693
2694 fn ensure_writable(&self) -> Result<()> {
2695 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2696 return Err(MongrelError::ReadOnlyReplica);
2697 }
2698 if self.durable_commit_failed {
2699 return Err(MongrelError::Other(
2700 "table poisoned by post-commit failure; reopen required".into(),
2701 ));
2702 }
2703 Ok(())
2704 }
2705
2706 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2717 match &self.auth {
2718 Some(checker) => checker.check(&self.name, perm),
2719 None => Ok(()),
2720 }
2721 }
2722 pub fn require_select(&self) -> Result<()> {
2727 self.require(crate::auth_state::RequiredPermission::Select)
2728 }
2729 fn require_insert(&self) -> Result<()> {
2730 self.require(crate::auth_state::RequiredPermission::Insert)
2731 }
2732 #[allow(dead_code)]
2736 fn require_update(&self) -> Result<()> {
2737 self.require(crate::auth_state::RequiredPermission::Update)
2738 }
2739 fn require_delete(&self) -> Result<()> {
2740 self.require(crate::auth_state::RequiredPermission::Delete)
2741 }
2742
2743 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2746 self.require_insert()?;
2747 Ok(self.put_returning(columns)?.0)
2748 }
2749
2750 pub fn put_returning(
2755 &mut self,
2756 mut columns: Vec<(u16, Value)>,
2757 ) -> Result<(RowId, Option<i64>)> {
2758 self.require_insert()?;
2759 let assigned = self.fill_auto_inc(&mut columns)?;
2760 self.apply_defaults(&mut columns)?;
2761 self.schema.validate_values(&columns)?;
2762 let row_id = if self.schema.clustered {
2767 self.derive_clustered_row_id(&columns)?
2768 } else {
2769 self.allocator.alloc()?
2770 };
2771 let epoch = self.pending_epoch();
2772 let mut row = Row::new(row_id, epoch);
2773 for (col_id, val) in columns {
2774 row.columns.insert(col_id, val);
2775 }
2776 self.commit_rows(vec![row], assigned.is_some())?;
2777 Ok((row_id, assigned))
2778 }
2779
2780 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2783 self.require_insert()?;
2784 Ok(self
2785 .put_batch_returning(batch)?
2786 .into_iter()
2787 .map(|(r, _)| r)
2788 .collect())
2789 }
2790
2791 pub fn put_batch_returning(
2794 &mut self,
2795 batch: Vec<Vec<(u16, Value)>>,
2796 ) -> Result<Vec<(RowId, Option<i64>)>> {
2797 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2798 for mut cols in batch {
2799 let assigned = self.fill_auto_inc(&mut cols)?;
2800 self.apply_defaults(&mut cols)?;
2801 filled.push((cols, assigned));
2802 }
2803 for (cols, _) in &filled {
2804 self.schema.validate_values(cols)?;
2805 }
2806 let epoch = self.pending_epoch();
2807 let mut rows = Vec::with_capacity(filled.len());
2808 let mut ids = Vec::with_capacity(filled.len());
2809 let first_row_id = if self.schema.clustered {
2810 None
2811 } else {
2812 let count = u64::try_from(filled.len())
2813 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2814 Some(self.allocator.alloc_range(count)?.0)
2815 };
2816 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2817 let row_id = match first_row_id {
2818 Some(first) => RowId(first + row_index as u64),
2819 None => self.derive_clustered_row_id(&cols)?,
2820 };
2821 let mut row = Row::new(row_id, epoch);
2822 for (c, v) in cols {
2823 row.columns.insert(c, v);
2824 }
2825 ids.push((row_id, assigned));
2826 rows.push(row);
2827 }
2828 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2829 self.commit_rows(rows, all_auto_generated)?;
2830 Ok(ids)
2831 }
2832
2833 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2839 self.ensure_writable()?;
2840 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2841 return Ok(None);
2842 };
2843 let pos = columns.iter().position(|(c, _)| *c == cid);
2844 let assigned = match pos {
2845 Some(i) => match &columns[i].1 {
2846 Value::Null => {
2847 let next = self.alloc_auto_inc_value()?;
2848 columns[i].1 = Value::Int64(next);
2849 Some(next)
2850 }
2851 Value::Int64(n) => {
2852 self.advance_auto_inc_past(*n)?;
2853 None
2854 }
2855 other => {
2856 return Err(MongrelError::InvalidArgument(format!(
2857 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2858 other
2859 )))
2860 }
2861 },
2862 None => {
2863 let next = self.alloc_auto_inc_value()?;
2864 columns.push((cid, Value::Int64(next)));
2865 Some(next)
2866 }
2867 };
2868 Ok(assigned)
2869 }
2870
2871 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2877 for col in &self.schema.columns {
2878 let Some(expr) = &col.default_value else {
2879 continue;
2880 };
2881 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2883 continue;
2884 }
2885 let pos = columns.iter().position(|(c, _)| *c == col.id);
2886 let needs_default = match pos {
2887 None => true,
2888 Some(i) => matches!(columns[i].1, Value::Null),
2889 };
2890 if !needs_default {
2891 continue;
2892 }
2893 let v = match expr {
2894 crate::schema::DefaultExpr::Static(v) => v.clone(),
2895 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2896 crate::schema::DefaultExpr::Uuid => {
2897 let mut buf = [0u8; 16];
2898 getrandom::getrandom(&mut buf)
2899 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2900 Value::Uuid(buf)
2901 }
2902 };
2903 match pos {
2904 None => columns.push((col.id, v)),
2905 Some(i) => columns[i].1 = v,
2906 }
2907 }
2908 Ok(())
2909 }
2910
2911 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2913 self.ensure_auto_inc_seeded()?;
2914 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2916 let v = ai.next;
2917 ai.next = ai
2918 .next
2919 .checked_add(1)
2920 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2921 Ok(v)
2922 }
2923
2924 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2927 self.ensure_auto_inc_seeded()?;
2928 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2929 let floor = used
2930 .checked_add(1)
2931 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2932 .max(1);
2933 if ai.next < floor {
2934 ai.next = floor;
2935 }
2936 Ok(())
2937 }
2938
2939 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2944 let needs_seed = match self.auto_inc {
2945 Some(ai) => !ai.seeded,
2946 None => return Ok(()),
2947 };
2948 if !needs_seed {
2949 return Ok(());
2950 }
2951 if self.seed_empty_auto_inc() {
2952 return Ok(());
2953 }
2954 let cid = self
2955 .auto_inc
2956 .as_ref()
2957 .expect("auto-inc column present")
2958 .column_id;
2959 let max = self.scan_max_int64(cid)?;
2960 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2961 let floor = max
2962 .checked_add(1)
2963 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2964 .max(1);
2965 if ai.next < floor {
2966 ai.next = floor;
2967 }
2968 ai.seeded = true;
2969 Ok(())
2970 }
2971
2972 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2973 if n == 0 || self.auto_inc.is_none() {
2974 return Ok(None);
2975 }
2976 self.ensure_auto_inc_seeded()?;
2977 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2978 let start = ai.next;
2979 let count = i64::try_from(n)
2980 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2981 ai.next = ai
2982 .next
2983 .checked_add(count)
2984 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2985 Ok(Some(start))
2986 }
2987
2988 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2992 let mut max: i64 = 0;
2993 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2994 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2995 if *n > max {
2996 max = *n;
2997 }
2998 }
2999 }
3000 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3001 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3002 if *n > max {
3003 max = *n;
3004 }
3005 }
3006 }
3007 for rr in self.run_refs.clone() {
3008 let reader = self.open_reader(rr.run_id)?;
3009 if let Some(stats) = reader.column_page_stats(column_id) {
3010 for s in stats {
3011 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3012 if n > max {
3013 max = n;
3014 }
3015 }
3016 }
3017 } else if reader.has_column(column_id) {
3018 if let columnar::NativeColumn::Int64 { data, validity } =
3019 reader.column_native_shared(column_id)?
3020 {
3021 for (i, n) in data.iter().enumerate() {
3022 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3023 {
3024 max = *n;
3025 }
3026 }
3027 }
3028 }
3029 }
3030 Ok(max)
3031 }
3032
3033 fn seed_empty_auto_inc(&mut self) -> bool {
3034 let Some(ai) = self.auto_inc.as_mut() else {
3035 return false;
3036 };
3037 if ai.seeded || self.live_count != 0 {
3038 return false;
3039 }
3040 if ai.next < 1 {
3041 ai.next = 1;
3042 }
3043 ai.seeded = true;
3044 true
3045 }
3046
3047 fn advance_auto_inc_from_native_columns(
3048 &mut self,
3049 columns: &[(u16, columnar::NativeColumn)],
3050 n: usize,
3051 live_before: u64,
3052 ) -> Result<()> {
3053 let Some(ai) = self.auto_inc.as_mut() else {
3054 return Ok(());
3055 };
3056 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3057 return Ok(());
3058 };
3059 let columnar::NativeColumn::Int64 { data, validity } = col else {
3060 return Err(MongrelError::InvalidArgument(format!(
3061 "AUTO_INCREMENT column {} must be Int64",
3062 ai.column_id
3063 )));
3064 };
3065 let max = if native_int64_strictly_increasing(col, n) {
3066 data.get(n.saturating_sub(1)).copied()
3067 } else {
3068 data.iter()
3069 .take(n)
3070 .enumerate()
3071 .filter_map(|(i, v)| {
3072 if validity.is_empty() || columnar::validity_bit(validity, i) {
3073 Some(*v)
3074 } else {
3075 None
3076 }
3077 })
3078 .max()
3079 };
3080 if let Some(max) = max {
3081 let floor = max
3082 .checked_add(1)
3083 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3084 .max(1);
3085 if ai.next < floor {
3086 ai.next = floor;
3087 }
3088 if ai.seeded || live_before == 0 {
3089 ai.seeded = true;
3090 }
3091 }
3092 Ok(())
3093 }
3094
3095 fn fill_auto_inc_native_columns(
3096 &mut self,
3097 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3098 n: usize,
3099 ) -> Result<()> {
3100 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3101 return Ok(());
3102 };
3103 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3104 if let Some(start) = self.alloc_auto_inc_range(n)? {
3105 columns.push((
3106 cid,
3107 columnar::NativeColumn::Int64 {
3108 data: (start..start.saturating_add(n as i64)).collect(),
3109 validity: vec![0xFF; n.div_ceil(8)],
3110 },
3111 ));
3112 }
3113 return Ok(());
3114 };
3115
3116 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3117 return Err(MongrelError::InvalidArgument(format!(
3118 "AUTO_INCREMENT column {cid} must be Int64"
3119 )));
3120 };
3121 if data.len() < n {
3122 return Err(MongrelError::InvalidArgument(format!(
3123 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3124 data.len()
3125 )));
3126 }
3127 if columnar::all_non_null(validity, n) {
3128 return Ok(());
3129 }
3130 if validity.iter().all(|b| *b == 0) {
3131 if let Some(start) = self.alloc_auto_inc_range(n)? {
3132 for (i, slot) in data.iter_mut().take(n).enumerate() {
3133 *slot = start.saturating_add(i as i64);
3134 }
3135 *validity = vec![0xFF; n.div_ceil(8)];
3136 }
3137 return Ok(());
3138 }
3139
3140 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3141 for (i, slot) in data.iter_mut().enumerate().take(n) {
3142 if columnar::validity_bit(validity, i) {
3143 self.advance_auto_inc_past(*slot)?;
3144 } else {
3145 *slot = self.alloc_auto_inc_value()?;
3146 }
3147 }
3148 *validity = new_validity;
3149 Ok(())
3150 }
3151
3152 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3166 self.ensure_writable()?;
3167 if self.auto_inc.is_none() {
3168 return Ok(None);
3169 }
3170 Ok(Some(self.alloc_auto_inc_value()?))
3171 }
3172
3173 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3179 let payload = bincode::serialize(&rows)?;
3180 self.wal_append_data(Op::Put {
3181 table_id: self.table_id,
3182 rows: payload,
3183 })?;
3184 if self.is_shared() {
3185 self.pending_rows_auto_inc
3186 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3187 self.pending_rows.extend(rows);
3188 } else {
3189 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3190 }
3191 Ok(())
3192 }
3193
3194 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3197 self.ensure_indexes_complete()
3198 }
3199
3200 pub(crate) fn prepare_durable_publish_controlled(
3201 &mut self,
3202 control: &crate::ExecutionControl,
3203 ) -> Result<()> {
3204 self.ensure_indexes_complete_controlled(control, || true)?;
3205 Ok(())
3206 }
3207
3208 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3209 self.apply_put_rows_inner_prepared(rows, true);
3210 }
3211
3212 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3213 if check_existing_pk {
3214 self.ensure_indexes_complete()?;
3215 }
3216 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3217 Ok(())
3218 }
3219
3220 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3224 if rows.len() == 1 {
3228 let row = rows.into_iter().next().expect("len checked");
3229 self.apply_put_row_single(row, check_existing_pk);
3230 return;
3231 }
3232 let pk_id = self.schema.primary_key().map(|c| c.id);
3249 let probe = match pk_id {
3250 Some(pid) => {
3251 check_existing_pk
3252 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3253 }
3254 None => false,
3255 };
3256 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3259 for r in rows {
3260 for &cid in r.columns.keys() {
3261 self.pending_put_cols.insert(cid);
3262 }
3263 match pk_id {
3264 Some(pid) if probe || maintain_pk_by_row => {
3265 if let Some(pk_val) = r.columns.get(&pid) {
3266 let key = self.index_lookup_key(pid, pk_val);
3267 if probe {
3268 if let Some(old_rid) = self.hot.get(&key) {
3269 if old_rid != r.row_id {
3270 self.tombstone_row(old_rid, r.committed_epoch, true);
3271 }
3272 }
3273 }
3274 if maintain_pk_by_row {
3275 self.pk_by_row.insert(r.row_id, key);
3276 }
3277 }
3278 }
3279 Some(_) => {}
3280 None => {
3281 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3282 }
3283 }
3284 self.index_row(&r);
3285 self.reservoir.offer(r.row_id.0);
3286 self.memtable.upsert(r);
3287 self.live_count = self.live_count.saturating_add(1);
3290 }
3291 self.data_generation = self.data_generation.wrapping_add(1);
3292 }
3293
3294 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3298 for &cid in row.columns.keys() {
3299 self.pending_put_cols.insert(cid);
3300 }
3301 let epoch = row.committed_epoch;
3302 if let Some(pk_col) = self.schema.primary_key() {
3303 let pk_id = pk_col.id;
3304 if let Some(pk_val) = row.columns.get(&pk_id) {
3305 let maintain_pk_by_row = self.pk_by_row_complete;
3309 if check_existing_pk || maintain_pk_by_row {
3310 let key = self.index_lookup_key(pk_id, pk_val);
3311 if check_existing_pk {
3312 if let Some(old_rid) = self.hot.get(&key) {
3313 if old_rid != row.row_id {
3314 self.tombstone_row(old_rid, epoch, true);
3315 }
3316 }
3317 }
3318 if maintain_pk_by_row {
3319 self.pk_by_row.insert(row.row_id, key);
3320 }
3321 }
3322 }
3323 } else {
3324 self.hot
3325 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3326 }
3327 self.index_row(&row);
3328 self.reservoir.offer(row.row_id.0);
3329 self.memtable.upsert(row);
3330 self.live_count = self.live_count.saturating_add(1);
3331 self.data_generation = self.data_generation.wrapping_add(1);
3332 }
3333
3334 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3337 self.allocator.alloc()
3338 }
3339
3340 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3346 let pk = self.schema.primary_key().ok_or_else(|| {
3347 MongrelError::Schema("clustered table requires a single-column primary key".into())
3348 })?;
3349 let pk_val = columns
3350 .iter()
3351 .find(|(id, _)| *id == pk.id)
3352 .map(|(_, v)| v)
3353 .ok_or_else(|| {
3354 MongrelError::Schema(format!(
3355 "clustered table missing primary key column {} ({})",
3356 pk.id, pk.name
3357 ))
3358 })?;
3359 Ok(clustered_row_id(pk_val))
3360 }
3361
3362 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3370 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3371 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3372 }
3373 let n = rows.len();
3374 for r in rows {
3375 for &cid in r.columns.keys() {
3376 self.pending_put_cols.insert(cid);
3377 }
3378 }
3379 let (losers, winner_pks) = self.partition_pk_winners(rows);
3380 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3381 for (key, &row_id) in &winner_pks {
3383 if let Some(old_rid) = self.hot.get(key) {
3384 if old_rid != row_id {
3385 self.tombstone_row(old_rid, epoch, true);
3386 }
3387 }
3388 }
3389 for &loser_rid in &losers {
3392 self.tombstone_row(loser_rid, epoch, false);
3393 }
3394 for (key, row_id) in winner_pks {
3396 self.insert_hot_pk(key, row_id);
3397 }
3398 if self.schema.primary_key().is_none() {
3399 for r in rows {
3400 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3401 }
3402 }
3403 for r in rows {
3404 self.allocator.advance_to(r.row_id)?;
3405 if !losers.contains(&r.row_id) {
3406 self.index_row(r);
3407 }
3408 }
3409 for r in rows {
3410 if !losers.contains(&r.row_id) {
3411 self.reservoir.offer(r.row_id.0);
3412 }
3413 }
3414 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3415 self.data_generation = self.data_generation.wrapping_add(1);
3416 Ok(())
3417 }
3418
3419 pub(crate) fn recover_apply(
3424 &mut self,
3425 rows: Vec<Row>,
3426 deletes: Vec<(RowId, Epoch)>,
3427 ) -> Result<()> {
3428 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3432 std::collections::BTreeMap::new();
3433 for row in rows {
3434 if row.row_id.0 >= u64::MAX - 1 {
3435 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3436 }
3437 self.allocator.advance_to(row.row_id)?;
3438 if let Some(ai) = self.auto_inc.as_mut() {
3443 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3444 let next = n.checked_add(1).ok_or_else(|| {
3445 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3446 })?;
3447 if next > ai.next {
3448 ai.next = next;
3449 }
3450 }
3451 }
3452 by_epoch.entry(row.committed_epoch).or_default().push(row);
3453 }
3454 for (epoch, group) in by_epoch {
3455 let (losers, winner_pks) = self.partition_pk_winners(&group);
3456 for (key, &row_id) in &winner_pks {
3458 if let Some(old_rid) = self.hot.get(key) {
3459 if old_rid != row_id {
3460 self.tombstone_row(old_rid, epoch, false);
3461 }
3462 }
3463 }
3464 for (key, row_id) in winner_pks {
3465 self.insert_hot_pk(key, row_id);
3466 }
3467 if self.schema.primary_key().is_none() {
3468 for r in &group {
3469 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3470 }
3471 }
3472 for r in &group {
3473 if !losers.contains(&r.row_id) {
3474 self.memtable.upsert(r.clone());
3475 self.index_row(r);
3476 }
3477 }
3478 }
3479 for (rid, epoch) in deletes {
3480 self.memtable.tombstone(rid, epoch);
3481 self.remove_hot_for_row(rid, epoch);
3482 }
3483 self.reservoir_complete = false;
3486 Ok(())
3487 }
3488
3489 pub(crate) fn flushed_epoch(&self) -> u64 {
3491 self.flushed_epoch
3492 }
3493
3494 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3495 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3496 }
3497
3498 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3500 self.schema.validate_values(cells)
3501 }
3502
3503 fn validate_columns_not_null(
3507 &self,
3508 columns: &[(u16, columnar::NativeColumn)],
3509 n: usize,
3510 ) -> Result<()> {
3511 let by_id: HashMap<u16, &columnar::NativeColumn> =
3512 columns.iter().map(|(id, c)| (*id, c)).collect();
3513 for col in &self.schema.columns {
3514 if !col.flags.contains(ColumnFlags::NULLABLE) {
3515 match by_id.get(&col.id) {
3516 None => {
3517 return Err(MongrelError::InvalidArgument(format!(
3518 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3519 col.name, col.id
3520 )));
3521 }
3522 Some(c) => {
3523 if c.null_count(n) != 0 {
3524 return Err(MongrelError::InvalidArgument(format!(
3525 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3526 col.name, col.id
3527 )));
3528 }
3529 }
3530 }
3531 }
3532 if let TypeId::Enum { variants } = &col.ty {
3533 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3534 if by_id.contains_key(&col.id) {
3535 return Err(MongrelError::InvalidArgument(format!(
3536 "column '{}' ({}) enum requires a bytes column",
3537 col.name, col.id
3538 )));
3539 }
3540 continue;
3541 };
3542 for index in 0..n {
3543 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3544 continue;
3545 };
3546 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3547 return Err(MongrelError::InvalidArgument(format!(
3548 "column '{}' ({}) enum value {:?} is not one of {:?}",
3549 col.name,
3550 col.id,
3551 String::from_utf8_lossy(value),
3552 variants
3553 )));
3554 }
3555 }
3556 }
3557 }
3558 Ok(())
3559 }
3560
3561 fn bulk_pk_winner_indices(
3566 &self,
3567 columns: &[(u16, columnar::NativeColumn)],
3568 n: usize,
3569 ) -> Option<Vec<usize>> {
3570 let pk_col = self.schema.primary_key()?;
3571 let pk_id = pk_col.id;
3572 let pk_ty = pk_col.ty.clone();
3573 let by_id: HashMap<u16, &columnar::NativeColumn> =
3574 columns.iter().map(|(id, c)| (*id, c)).collect();
3575 let pk_native = by_id.get(&pk_id)?;
3576 if native_int64_strictly_increasing(pk_native, n) {
3577 return None;
3578 }
3579 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3581 let mut null_pk_rows: Vec<usize> = Vec::new();
3582 for i in 0..n {
3583 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3584 Some(key) => {
3585 last.insert(key, i);
3586 }
3587 None => null_pk_rows.push(i),
3588 }
3589 }
3590 let mut winners: HashSet<usize> = last.values().copied().collect();
3591 for i in null_pk_rows {
3592 winners.insert(i);
3593 }
3594 Some((0..n).filter(|i| winners.contains(i)).collect())
3595 }
3596
3597 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3599 self.require_delete()?;
3600 let epoch = self.pending_epoch();
3601 self.wal_append_data(Op::Delete {
3602 table_id: self.table_id,
3603 row_ids: vec![row_id],
3604 })?;
3605 if self.is_shared() {
3606 self.pending_dels.push(row_id);
3607 } else {
3608 self.apply_delete(row_id, epoch);
3609 }
3610 Ok(())
3611 }
3612
3613 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3614 let pre = self.get(row_id, self.snapshot());
3615 self.delete(row_id)?;
3616 Ok(pre.map(|row| {
3617 let mut columns: Vec<_> = row.columns.into_iter().collect();
3618 columns.sort_by_key(|(id, _)| *id);
3619 OwnedRow { columns }
3620 }))
3621 }
3622
3623 pub fn truncate(&mut self) -> Result<()> {
3625 self.require_delete()?;
3626 let epoch = self.pending_epoch();
3627 self.wal_append_data(Op::TruncateTable {
3628 table_id: self.table_id,
3629 })?;
3630 self.pending_rows.clear();
3631 self.pending_rows_auto_inc.clear();
3632 self.pending_dels.clear();
3633 self.pending_truncate = Some(epoch);
3634 Ok(())
3635 }
3636
3637 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3639 self.run_refs.clear();
3644 self.retiring.clear();
3645 self.memtable = Memtable::new();
3646 self.mutable_run = MutableRun::new();
3647 self.hot = HotIndex::new();
3648 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3649 self.bitmap = bitmap;
3650 self.ann = ann;
3651 self.fm = fm;
3652 self.sparse = sparse;
3653 self.minhash = minhash;
3654 self.learned_range = Arc::new(HashMap::new());
3655 self.pk_by_row.clear();
3656 self.pk_by_row_complete = false;
3657 self.live_count = 0;
3658 self.reservoir = crate::reservoir::Reservoir::default();
3659 self.reservoir_complete = true;
3660 self.had_deletes = true;
3661 self.agg_cache = Arc::new(HashMap::new());
3662 self.global_idx_epoch = 0;
3663 self.indexes_complete = true;
3664 self.pending_delete_rids.clear();
3665 self.pending_put_cols.clear();
3666 self.pending_rows.clear();
3667 self.pending_rows_auto_inc.clear();
3668 self.pending_dels.clear();
3669 self.clear_result_cache();
3670 self.invalidate_index_checkpoint();
3671 self.data_generation = self.data_generation.wrapping_add(1);
3672 }
3673
3674 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3677 self.remove_hot_for_row(row_id, epoch);
3678 self.tombstone_row(row_id, epoch, true);
3679 self.data_generation = self.data_generation.wrapping_add(1);
3680 }
3681
3682 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3686 let tombstone = Row {
3687 row_id,
3688 committed_epoch: epoch,
3689 columns: std::collections::HashMap::new(),
3690 deleted: true,
3691 };
3692 self.memtable.upsert(tombstone);
3693 self.pk_by_row.remove(&row_id);
3694 if adjust_live_count {
3695 self.live_count = self.live_count.saturating_sub(1);
3696 }
3697 self.pending_delete_rids.insert(row_id.0 as u32);
3699 self.had_deletes = true;
3702 self.agg_cache = Arc::new(HashMap::new());
3703 }
3704
3705 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3709 let Some(pk_col) = self.schema.primary_key() else {
3710 return;
3711 };
3712 if self.pk_by_row_complete {
3715 if let Some(key) = self.pk_by_row.remove(&row_id) {
3716 if self.hot.get(&key) == Some(row_id) {
3717 self.hot.remove(&key);
3718 }
3719 }
3720 return;
3721 }
3722 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3741 if self.indexes_complete {
3742 let pk_val = self
3743 .memtable
3744 .get_version(row_id, lookup_epoch)
3745 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3746 .or_else(|| {
3747 self.mutable_run
3748 .get_version(row_id, lookup_epoch)
3749 .filter(|(_, r)| !r.deleted)
3750 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3751 })
3752 .or_else(|| {
3753 self.run_refs.iter().find_map(|rr| {
3754 let mut reader = self.open_reader(rr.run_id).ok()?;
3755 let (_, deleted, val) = reader
3756 .get_version_column(row_id, lookup_epoch, pk_col.id)
3757 .ok()??;
3758 if deleted {
3759 return None;
3760 }
3761 val
3762 })
3763 });
3764 if let Some(pk_val) = pk_val {
3765 let key = self.index_lookup_key(pk_col.id, &pk_val);
3766 if self.hot.get(&key) == Some(row_id) {
3767 self.hot.remove(&key);
3768 }
3769 return;
3770 }
3771 }
3772 self.refresh_pk_by_row_from_hot();
3777 if let Some(key) = self.pk_by_row.remove(&row_id) {
3778 if self.hot.get(&key) == Some(row_id) {
3779 self.hot.remove(&key);
3780 }
3781 }
3782 }
3783
3784 fn partition_pk_winners(
3789 &self,
3790 rows: &[Row],
3791 ) -> (
3792 std::collections::HashSet<RowId>,
3793 std::collections::HashMap<Vec<u8>, RowId>,
3794 ) {
3795 let mut losers = std::collections::HashSet::new();
3796 let Some(pk_col) = self.schema.primary_key() else {
3797 return (losers, std::collections::HashMap::new());
3798 };
3799 let pk_id = pk_col.id;
3800 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3801 std::collections::HashMap::new();
3802 for r in rows {
3803 let Some(pk_val) = r.columns.get(&pk_id) else {
3804 continue;
3805 };
3806 let key = self.index_lookup_key(pk_id, pk_val);
3807 if let Some(&old_rid) = winners.get(&key) {
3808 losers.insert(old_rid);
3809 }
3810 winners.insert(key, r.row_id);
3811 }
3812 (losers, winners)
3813 }
3814
3815 fn index_row(&mut self, row: &Row) {
3816 if row.deleted {
3817 return;
3818 }
3819 let any_predicate = self
3827 .schema
3828 .indexes
3829 .iter()
3830 .any(|idx| idx.predicate.is_some());
3831 if any_predicate {
3832 let columns_map: HashMap<u16, &Value> =
3833 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3834 let name_to_id: HashMap<&str, u16> = self
3835 .schema
3836 .columns
3837 .iter()
3838 .map(|c| (c.name.as_str(), c.id))
3839 .collect();
3840 for idx in &self.schema.indexes {
3841 if let Some(pred) = &idx.predicate {
3842 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3843 continue; }
3845 }
3846 index_into_single(
3848 idx,
3849 &self.schema,
3850 row,
3851 &mut self.hot,
3852 &mut self.bitmap,
3853 &mut self.ann,
3854 &mut self.fm,
3855 &mut self.sparse,
3856 &mut self.minhash,
3857 );
3858 }
3859 return;
3860 }
3861 if self.column_keys.is_empty() {
3865 index_into(
3866 &self.schema,
3867 row,
3868 &mut self.hot,
3869 &mut self.bitmap,
3870 &mut self.ann,
3871 &mut self.fm,
3872 &mut self.sparse,
3873 &mut self.minhash,
3874 );
3875 return;
3876 }
3877 let effective_row = self.tokenized_for_indexes(row);
3878 index_into(
3879 &self.schema,
3880 &effective_row,
3881 &mut self.hot,
3882 &mut self.bitmap,
3883 &mut self.ann,
3884 &mut self.fm,
3885 &mut self.sparse,
3886 &mut self.minhash,
3887 );
3888 }
3889
3890 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3896 if self.column_keys.is_empty() {
3897 return row.clone();
3898 }
3899 {
3900 use crate::encryption::SCHEME_HMAC_EQ;
3901 let mut tok = row.clone();
3902 for (&cid, &(_, scheme)) in &self.column_keys {
3903 if scheme != SCHEME_HMAC_EQ {
3904 continue;
3905 }
3906 if let Some(v) = tok.columns.get(&cid).cloned() {
3907 if let Some(t) = self.tokenize_value(cid, &v) {
3908 tok.columns.insert(cid, t);
3909 }
3910 }
3911 }
3912 tok
3913 }
3914 }
3915
3916 pub fn commit(&mut self) -> Result<Epoch> {
3921 self.commit_inner(None)
3922 }
3923
3924 #[doc(hidden)]
3927 pub fn commit_controlled<F>(
3928 &mut self,
3929 control: &crate::ExecutionControl,
3930 mut before_commit: F,
3931 ) -> Result<Epoch>
3932 where
3933 F: FnMut() -> Result<()>,
3934 {
3935 self.commit_inner(Some((control, &mut before_commit)))
3936 }
3937
3938 fn commit_inner(
3939 &mut self,
3940 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3941 ) -> Result<Epoch> {
3942 self.ensure_writable()?;
3943 if !self.has_pending_mutations() {
3944 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3945 return Err(MongrelError::Full(
3946 "standalone transaction id namespace exhausted".into(),
3947 ));
3948 }
3949 return Ok(self.epoch.visible());
3950 }
3951 self.commit_new_epoch_inner(controlled)
3952 }
3953
3954 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3958 self.commit_new_epoch_inner(None)
3959 }
3960
3961 fn commit_new_epoch_inner(
3962 &mut self,
3963 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3964 ) -> Result<Epoch> {
3965 self.ensure_writable()?;
3966 if self.is_shared() {
3967 self.commit_shared(controlled)
3968 } else {
3969 self.commit_private(controlled)
3970 }
3971 }
3972
3973 fn commit_private(
3975 &mut self,
3976 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3977 ) -> Result<Epoch> {
3978 let commit_lock = Arc::clone(&self.commit_lock);
3982 let _g = commit_lock.lock();
3983 let txn_id = self.ensure_txn_id()?;
3986 if let Some((control, before_commit)) = controlled {
3987 control.checkpoint()?;
3988 before_commit()?;
3989 }
3990 let new_epoch = self.epoch.bump_assigned();
3991 let epoch_authority = Arc::clone(&self.epoch);
3992 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3993 let wal_result = match &mut self.wal {
3997 WalSink::Private(w) => w
3998 .append_txn(
3999 txn_id,
4000 Op::TxnCommit {
4001 epoch: new_epoch.0,
4002 added_runs: Vec::new(),
4003 },
4004 )
4005 .and_then(|_| w.sync()),
4006 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4007 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4008 };
4009 if let Err(error) = wal_result {
4010 self.durable_commit_failed = true;
4011 return Err(MongrelError::CommitOutcomeUnknown {
4012 epoch: new_epoch.0,
4013 message: error.to_string(),
4014 });
4015 }
4016 if let Some(epoch) = self.pending_truncate.take() {
4019 self.apply_truncate(epoch);
4020 }
4021 self.invalidate_pending_cache();
4022 let publish_result = self.persist_manifest(new_epoch);
4023 self.epoch.publish_in_order(new_epoch);
4027 epoch_guard.disarm();
4028 if let Err(error) = publish_result {
4029 self.durable_commit_failed = true;
4030 return Err(MongrelError::DurableCommit {
4031 epoch: new_epoch.0,
4032 message: error.to_string(),
4033 });
4034 }
4035 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4036 self.pending_private_mutations = false;
4037 self.data_generation = self.data_generation.wrapping_add(1);
4038 Ok(new_epoch)
4039 }
4040
4041 fn commit_shared(
4047 &mut self,
4048 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4049 ) -> Result<Epoch> {
4050 use std::sync::atomic::Ordering;
4051 let s = match &self.wal {
4052 WalSink::Shared(s) => s.clone(),
4053 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4054 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4055 };
4056 if s.poisoned.load(Ordering::Relaxed) {
4057 return Err(MongrelError::Other(
4058 "database poisoned by fsync error".into(),
4059 ));
4060 }
4061 let commit_lock = Arc::clone(&self.commit_lock);
4068 let _g = commit_lock.lock();
4069 if !self.pending_rows.is_empty() {
4070 match controlled.as_ref() {
4071 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4072 None => self.prepare_durable_publish()?,
4073 }
4074 }
4075 let txn_id = self.ensure_txn_id()?;
4078 let mut wal = s.wal.lock();
4079 if let Some((control, before_commit)) = controlled {
4080 control.checkpoint()?;
4081 before_commit()?;
4082 }
4083 let new_epoch = self.epoch.bump_assigned();
4084 let epoch_authority = Arc::clone(&self.epoch);
4085 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4086 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
4087 Ok(commit_seq) => commit_seq,
4088 Err(error) => {
4089 s.poisoned.store(true, Ordering::Relaxed);
4090 s.lifecycle.poison();
4091 return Err(MongrelError::CommitOutcomeUnknown {
4092 epoch: new_epoch.0,
4093 message: error.to_string(),
4094 });
4095 }
4096 };
4097 drop(wal);
4098 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4099 s.poisoned.store(true, Ordering::Relaxed);
4100 s.lifecycle.poison();
4101 return Err(MongrelError::CommitOutcomeUnknown {
4102 epoch: new_epoch.0,
4103 message: error.to_string(),
4104 });
4105 }
4106
4107 if self.pending_truncate.take().is_some() {
4110 self.apply_truncate(new_epoch);
4111 }
4112 let mut rows = std::mem::take(&mut self.pending_rows);
4113 if !rows.is_empty() {
4114 for r in &mut rows {
4115 r.committed_epoch = new_epoch;
4116 }
4117 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4118 let all_auto_generated =
4119 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4120 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4121 } else {
4122 self.pending_rows_auto_inc.clear();
4123 }
4124 let dels = std::mem::take(&mut self.pending_dels);
4125 for rid in dels {
4126 self.apply_delete(rid, new_epoch);
4127 }
4128
4129 self.invalidate_pending_cache();
4130 let publish_result = self.persist_manifest(new_epoch);
4131 self.epoch.publish_in_order(new_epoch);
4132 epoch_guard.disarm();
4133 let _ = s.change_wake.send(());
4134 if let Err(error) = publish_result {
4135 self.durable_commit_failed = true;
4136 s.poisoned.store(true, Ordering::Relaxed);
4137 s.lifecycle.poison();
4138 return Err(MongrelError::DurableCommit {
4139 epoch: new_epoch.0,
4140 message: error.to_string(),
4141 });
4142 }
4143 self.current_txn_id = 0;
4145 self.data_generation = self.data_generation.wrapping_add(1);
4146 Ok(new_epoch)
4147 }
4148
4149 pub fn flush(&mut self) -> Result<Epoch> {
4157 self.flush_with_outcome().map(|(epoch, _)| epoch)
4158 }
4159
4160 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4162 self.flush_with_outcome_inner(None)
4163 }
4164
4165 #[doc(hidden)]
4168 pub fn flush_with_outcome_controlled<F>(
4169 &mut self,
4170 control: &crate::ExecutionControl,
4171 mut before_commit: F,
4172 ) -> Result<(Epoch, bool)>
4173 where
4174 F: FnMut() -> Result<()>,
4175 {
4176 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4177 }
4178
4179 fn flush_with_outcome_inner(
4180 &mut self,
4181 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4182 ) -> Result<(Epoch, bool)> {
4183 match controlled.as_ref() {
4184 Some((control, _)) => {
4185 self.ensure_indexes_complete_controlled(control, || true)?;
4186 }
4187 None => self.ensure_indexes_complete()?,
4188 }
4189 let committed = self.has_pending_mutations();
4190 let epoch = self.commit_inner(controlled)?;
4191 let finish: Result<(Epoch, bool)> = (|| {
4192 let rows = self.memtable.drain_sorted();
4193 if !rows.is_empty() {
4194 self.mutable_run.insert_many(rows);
4195 }
4196 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4197 self.spill_mutable_run(epoch)?;
4198 self.mark_flushed(epoch)?;
4202 self.persist_manifest(epoch)?;
4203 self.build_learned_ranges()?;
4204 self.checkpoint_indexes(epoch);
4207 }
4208 Ok((epoch, committed))
4211 })();
4212 let outcome = match finish {
4213 Err(error) if committed => Err(MongrelError::DurableCommit {
4214 epoch: epoch.0,
4215 message: error.to_string(),
4216 }),
4217 result => result,
4218 };
4219 if outcome.is_ok() {
4220 let _ = self.publish_read_generation();
4226 }
4227 outcome
4228 }
4229
4230 fn has_pending_mutations(&self) -> bool {
4231 self.pending_private_mutations
4232 || !self.pending_rows.is_empty()
4233 || !self.pending_dels.is_empty()
4234 || self.pending_truncate.is_some()
4235 }
4236
4237 pub fn has_pending_writes(&self) -> bool {
4238 self.has_pending_mutations()
4239 }
4240
4241 pub fn force_flush(&mut self) -> Result<Epoch> {
4250 let saved = self.mutable_run_spill_bytes;
4251 self.mutable_run_spill_bytes = 1;
4252 let result = self.flush();
4253 self.mutable_run_spill_bytes = saved;
4254 result
4255 }
4256
4257 pub fn close(&mut self) -> Result<()> {
4264 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4265 self.force_flush()?;
4266 }
4267 Ok(())
4268 }
4269
4270 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4277 let op = Op::Flush {
4278 table_id: self.table_id,
4279 flushed_epoch: epoch.0,
4280 };
4281 match &mut self.wal {
4282 WalSink::Private(w) => {
4283 w.append_system(op)?;
4284 w.sync()?;
4285 }
4286 WalSink::Shared(s) => {
4287 s.wal.lock().append_system(op)?;
4292 }
4293 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4294 }
4295 self.flushed_epoch = epoch.0;
4296 if matches!(self.wal, WalSink::Private(_)) {
4297 self.rotate_wal(epoch)?;
4298 }
4299 Ok(())
4300 }
4301
4302 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4306 if self.mutable_run.is_empty() {
4307 return Ok(());
4308 }
4309 let run_id = self.alloc_run_id()?;
4310 let rows = self.mutable_run.drain_sorted();
4311 if rows.is_empty() {
4312 return Ok(());
4313 }
4314 let path = self.run_path(run_id);
4315 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4316 if let Some(kek) = &self.kek {
4317 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4318 }
4319 let header = match self.create_run_file(run_id)? {
4320 Some(file) => writer.write_file(file, &rows)?,
4321 None => writer.write(&path, &rows)?,
4322 };
4323 self.run_refs.push(RunRef {
4324 run_id: run_id as u128,
4325 level: 0,
4326 epoch_created: epoch.0,
4327 row_count: header.row_count,
4328 });
4329 Ok(())
4330 }
4331
4332 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4336 self.mutable_run_spill_bytes = bytes.max(1);
4337 }
4338
4339 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4343 self.compaction_zstd_level = level;
4344 }
4345
4346 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4350 self.result_cache.lock().set_max_bytes(max_bytes);
4351 }
4352
4353 pub(crate) fn clear_result_cache(&mut self) {
4357 self.result_cache.lock().clear();
4358 }
4359
4360 pub fn mutable_run_len(&self) -> usize {
4362 self.mutable_run.len()
4363 }
4364
4365 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4368 self.mutable_run.drain_sorted()
4369 }
4370
4371 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4373 let mut snapshot = self.mutable_run.clone();
4374 snapshot.drain_sorted()
4375 }
4376
4377 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4382 self.ensure_writable()?;
4383 let n = batch.len();
4384 if n == 0 {
4385 return Ok(self.current_epoch());
4386 }
4387 for row in &batch {
4388 self.schema.validate_values(row)?;
4389 }
4390 let epoch = self.commit_new_epoch()?;
4391 let live_before = self.live_count;
4392 self.spill_mutable_run(epoch)?;
4396 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4397 && self.indexes_complete
4398 && self.run_refs.is_empty()
4399 && self.memtable.is_empty()
4400 && self.mutable_run.is_empty();
4401 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4407 use rayon::prelude::*;
4408 self.schema
4409 .columns
4410 .par_iter()
4411 .map(|cdef| {
4412 (
4413 cdef.id,
4414 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4415 )
4416 })
4417 .collect::<Vec<_>>()
4418 };
4419 drop(batch);
4420 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4425 self.validate_columns_not_null(&user_columns, n)?;
4426 let winner_idx = self
4427 .bulk_pk_winner_indices(&user_columns, n)
4428 .filter(|idx| idx.len() != n);
4429 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4430 match winner_idx.as_deref() {
4431 Some(idx) => {
4432 let compacted = user_columns
4433 .iter()
4434 .map(|(id, c)| (*id, c.gather(idx)))
4435 .collect();
4436 (compacted, idx.len())
4437 }
4438 None => (std::mem::take(&mut user_columns), n),
4439 };
4440 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4441 let first = self.allocator.alloc_range(write_n as u64)?.0;
4442 for rid in first..first + write_n as u64 {
4443 self.reservoir.offer(rid);
4444 }
4445 let run_id = self.alloc_run_id()?;
4446 let path = self.run_path(run_id);
4447 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4448 .clean(true)
4449 .with_lz4()
4450 .with_native_endian();
4451 if let Some(kek) = &self.kek {
4452 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4453 }
4454 let header = match self.create_run_file(run_id)? {
4455 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4456 None => writer.write_native(&path, &write_columns, write_n, first)?,
4457 };
4458 self.run_refs.push(RunRef {
4459 run_id: run_id as u128,
4460 level: 0,
4461 epoch_created: epoch.0,
4462 row_count: header.row_count,
4463 });
4464 self.live_count = self.live_count.saturating_add(write_n as u64);
4465 if eager_index_build {
4466 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4467 self.index_columns_bulk(&write_columns, &row_ids);
4468 self.indexes_complete = true;
4469 self.build_learned_ranges()?;
4470 } else {
4471 self.indexes_complete = false;
4472 }
4473 self.mark_flushed(epoch)?;
4474 self.persist_manifest(epoch)?;
4475 if eager_index_build {
4476 self.checkpoint_indexes(epoch);
4477 }
4478 self.clear_result_cache();
4479 Ok(epoch)
4480 }
4481
4482 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4485 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4486 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4487 let segment_no = segment
4490 .file_stem()
4491 .and_then(|s| s.to_str())
4492 .and_then(|s| s.strip_prefix("seg-"))
4493 .and_then(|s| s.parse::<u64>().ok())
4494 .unwrap_or(0);
4495 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4496 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4497 wal.sync()?;
4498 self.wal = WalSink::Private(wal);
4499 Ok(())
4500 }
4501
4502 pub(crate) fn invalidate_pending_cache(&mut self) {
4507 self.result_cache
4508 .lock()
4509 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4510 self.pending_delete_rids.clear();
4511 self.pending_put_cols.clear();
4512 }
4513
4514 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4515 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4516 m.current_epoch = epoch.0;
4517 m.next_row_id = self.allocator.current().0;
4518 m.runs = self.run_refs.clone();
4519 m.live_count = self.live_count;
4520 m.global_idx_epoch = self.global_idx_epoch;
4521 m.flushed_epoch = self.flushed_epoch;
4522 m.retiring = self.retiring.clone();
4523 m.auto_inc_next = match self.auto_inc {
4527 Some(ai) if ai.seeded => ai.next,
4528 _ => 0,
4529 };
4530 m.ttl = self.ttl;
4531 let meta_dek = self.manifest_meta_dek();
4532 match self._root_guard.as_deref() {
4533 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4534 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4535 }
4536 Ok(())
4537 }
4538
4539 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4540 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4544 let live_count = u64::try_from(rows.len())
4545 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4546 let auto_inc = match self.auto_inc {
4547 Some(mut state) => {
4548 let maximum = self.scan_max_int64(state.column_id)?;
4549 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4550 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4551 })?;
4552 state.next = state.next.max(after_maximum).max(1);
4553 state.seeded = true;
4554 Some(state)
4555 }
4556 None => None,
4557 };
4558 Ok(RecoveryMetadataPlan {
4559 live_count,
4560 auto_inc,
4561 changed: live_count != self.live_count
4562 || auto_inc.is_some_and(|planned| {
4563 self.auto_inc.is_none_or(|current| {
4564 current.next != planned.next || current.seeded != planned.seeded
4565 })
4566 }),
4567 })
4568 }
4569
4570 pub(crate) fn apply_recovered_metadata(
4571 &mut self,
4572 plan: RecoveryMetadataPlan,
4573 epoch: Epoch,
4574 ) -> Result<()> {
4575 if !plan.changed {
4576 return Ok(());
4577 }
4578 self.live_count = plan.live_count;
4579 self.auto_inc = plan.auto_inc;
4580 self.persist_manifest(epoch)
4581 }
4582
4583 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4589 if !self.indexes_complete {
4592 return;
4593 }
4594 if crate::catalog::inject_hook("index.publish.before").is_err() {
4597 return;
4598 }
4599 if self.idx_root.is_none() {
4600 if let Some(root) = self._root_guard.as_ref() {
4601 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4602 return;
4603 };
4604 self.idx_root = Some(Arc::new(idx_root));
4605 }
4606 }
4607 let snap = global_idx::IndexSnapshot {
4608 hot: &self.hot,
4609 bitmap: &self.bitmap,
4610 ann: &self.ann,
4611 fm: &self.fm,
4612 sparse: &self.sparse,
4613 minhash: &self.minhash,
4614 learned_range: &self.learned_range,
4615 };
4616 let idx_dek = self.idx_dek();
4618 let written = match self.idx_root.as_deref() {
4619 Some(root) => global_idx::write_atomic_root(
4620 root,
4621 self.table_id,
4622 epoch.0,
4623 snap,
4624 idx_dek.as_deref(),
4625 ),
4626 None => global_idx::write_atomic(
4627 &self.dir,
4628 self.table_id,
4629 epoch.0,
4630 snap,
4631 idx_dek.as_deref(),
4632 ),
4633 };
4634 if written.is_ok() {
4635 self.global_idx_epoch = epoch.0;
4636 let _ = self.persist_manifest(epoch);
4637 let _ = crate::catalog::inject_hook("index.publish.after");
4639 }
4640 }
4641
4642 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4645 self.global_idx_epoch = 0;
4646 if let Some(root) = self.idx_root.as_deref() {
4647 let _ = root.remove_file(global_idx::IDX_FILENAME);
4648 } else {
4649 global_idx::remove(&self.dir);
4650 }
4651 let _ = self.persist_manifest(self.epoch.visible());
4652 }
4653
4654 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4659 self.indexes_complete = false;
4660 self.global_idx_epoch = 0;
4661 if let Some(root) = self.idx_root.as_deref() {
4662 let _ = root.remove_file(global_idx::IDX_FILENAME);
4663 } else {
4664 global_idx::remove(&self.dir);
4665 }
4666 }
4667
4668 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4669 self.indexes_complete = true;
4670 }
4671
4672 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4678 self.durable_commit_failed = true;
4679 if let WalSink::Shared(shared) = &self.wal {
4680 shared
4681 .poisoned
4682 .store(true, std::sync::atomic::Ordering::Relaxed);
4683 }
4684 }
4685
4686 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4690 self.durable_commit_failed = true;
4691 }
4692
4693 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4696 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4697 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4698 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4699 best = Some((epoch, row));
4700 }
4701 }
4702 for rr in &self.run_refs {
4703 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4704 continue;
4705 };
4706 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4707 continue;
4708 };
4709 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4710 best = Some((epoch, row));
4711 }
4712 }
4713 let now_nanos = unix_nanos_now();
4714 match best {
4715 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4716 Some((_, r)) => Some(r),
4717 None => None,
4718 }
4719 }
4720
4721 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4725 self.visible_rows_at_time(snapshot, unix_nanos_now())
4726 }
4727
4728 #[doc(hidden)]
4731 pub fn visible_rows_controlled(
4732 &self,
4733 snapshot: Snapshot,
4734 control: &crate::ExecutionControl,
4735 ) -> Result<Vec<Row>> {
4736 let mut out = Vec::new();
4737 self.for_each_visible_row_controlled(snapshot, control, |row| {
4738 out.push(row);
4739 Ok(())
4740 })?;
4741 Ok(out)
4742 }
4743
4744 #[doc(hidden)]
4747 pub fn for_each_visible_row_controlled<F>(
4748 &self,
4749 snapshot: Snapshot,
4750 control: &crate::ExecutionControl,
4751 visit: F,
4752 ) -> Result<()>
4753 where
4754 F: FnMut(Row) -> Result<()>,
4755 {
4756 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4757 control.checkpoint()?;
4758 let memtable = self.memtable.visible_versions(snapshot.epoch);
4759 if !memtable.is_empty() {
4760 sources.push(ControlledVisibleSource::memory(memtable));
4761 }
4762 control.checkpoint()?;
4763 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4764 if !mutable.is_empty() {
4765 sources.push(ControlledVisibleSource::memory(mutable));
4766 }
4767 for run in &self.run_refs {
4768 control.checkpoint()?;
4769 let reader = self.open_reader(run.run_id)?;
4770 sources.push(ControlledVisibleSource::run(
4771 reader.into_visible_version_cursor(snapshot.epoch)?,
4772 ));
4773 }
4774 let now_nanos = unix_nanos_now();
4775 merge_controlled_visible_sources(
4776 &mut sources,
4777 control,
4778 |row| self.row_expired_at(row, now_nanos),
4779 visit,
4780 )
4781 }
4782
4783 #[doc(hidden)]
4784 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4785 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4786 let mut fold = |row: Row| {
4787 best.entry(row.row_id.0)
4788 .and_modify(|e| {
4789 if row.committed_epoch > e.0 {
4790 *e = (row.committed_epoch, row.clone());
4791 }
4792 })
4793 .or_insert_with(|| (row.committed_epoch, row));
4794 };
4795 for row in self.memtable.visible_versions(snapshot.epoch) {
4796 fold(row);
4797 }
4798 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4799 fold(row);
4800 }
4801 for rr in &self.run_refs {
4802 let mut reader = self.open_reader(rr.run_id)?;
4803 for row in reader.visible_versions(snapshot.epoch)? {
4804 fold(row);
4805 }
4806 }
4807 let mut out: Vec<Row> = best
4808 .into_values()
4809 .filter_map(|(_, r)| {
4810 if r.deleted || self.row_expired_at(&r, now_nanos) {
4811 None
4812 } else {
4813 Some(r)
4814 }
4815 })
4816 .collect();
4817 out.sort_by_key(|r| r.row_id);
4818 Ok(out)
4819 }
4820
4821 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4828 if self.ttl.is_none()
4829 && self.memtable.is_empty()
4830 && self.mutable_run.is_empty()
4831 && self.run_refs.len() == 1
4832 {
4833 let rr = self.run_refs[0].clone();
4834 let mut reader = self.open_reader(rr.run_id)?;
4835 let idxs = reader.visible_indices(snapshot.epoch)?;
4836 let mut cols = Vec::with_capacity(self.schema.columns.len());
4837 for cdef in &self.schema.columns {
4838 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4839 }
4840 return Ok(cols);
4841 }
4842 let rows = self.visible_rows(snapshot)?;
4844 let mut cols: Vec<(u16, Vec<Value>)> = self
4845 .schema
4846 .columns
4847 .iter()
4848 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4849 .collect();
4850 for r in &rows {
4851 for (cid, vec) in cols.iter_mut() {
4852 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4853 }
4854 }
4855 Ok(cols)
4856 }
4857
4858 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4860 let row_id = self.hot.get(key)?;
4861 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4862 Some(row_id)
4863 } else {
4864 None
4865 }
4866 }
4867
4868 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4873 self.query_at_with_allowed(q, self.snapshot(), None)
4874 }
4875
4876 pub fn query_controlled(
4879 &mut self,
4880 q: &crate::query::Query,
4881 control: &crate::ExecutionControl,
4882 ) -> Result<Vec<Row>> {
4883 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4884 }
4885
4886 pub fn query_at_with_allowed(
4889 &mut self,
4890 q: &crate::query::Query,
4891 snapshot: Snapshot,
4892 allowed: Option<&std::collections::HashSet<RowId>>,
4893 ) -> Result<Vec<Row>> {
4894 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4895 }
4896
4897 #[doc(hidden)]
4898 pub fn query_at_with_allowed_controlled(
4899 &mut self,
4900 q: &crate::query::Query,
4901 snapshot: Snapshot,
4902 allowed: Option<&std::collections::HashSet<RowId>>,
4903 control: &crate::ExecutionControl,
4904 ) -> Result<Vec<Row>> {
4905 self.require_select()?;
4906 self.ensure_indexes_complete_controlled(control, || true)?;
4907 self.validate_native_query(q)?;
4908 self.query_conditions_at(
4909 &q.conditions,
4910 snapshot,
4911 allowed,
4912 q.limit,
4913 q.offset,
4914 None,
4915 unix_nanos_now(),
4916 Some(control),
4917 )
4918 }
4919
4920 #[doc(hidden)]
4921 pub fn query_at_with_allowed_after(
4922 &mut self,
4923 q: &crate::query::Query,
4924 snapshot: Snapshot,
4925 allowed: Option<&std::collections::HashSet<RowId>>,
4926 after_row_id: Option<RowId>,
4927 ) -> Result<Vec<Row>> {
4928 self.query_at_with_allowed_after_at_time(
4929 q,
4930 snapshot,
4931 allowed,
4932 after_row_id,
4933 unix_nanos_now(),
4934 )
4935 }
4936
4937 #[doc(hidden)]
4938 pub fn query_at_with_allowed_after_at_time(
4939 &mut self,
4940 q: &crate::query::Query,
4941 snapshot: Snapshot,
4942 allowed: Option<&std::collections::HashSet<RowId>>,
4943 after_row_id: Option<RowId>,
4944 query_time_nanos: i64,
4945 ) -> Result<Vec<Row>> {
4946 self.require_select()?;
4947 self.ensure_indexes_complete()?;
4948 self.validate_native_query(q)?;
4949 self.query_conditions_at(
4950 &q.conditions,
4951 snapshot,
4952 allowed,
4953 q.limit,
4954 q.offset,
4955 after_row_id,
4956 query_time_nanos,
4957 None,
4958 )
4959 }
4960
4961 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4962 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4963 return Err(MongrelError::InvalidArgument(format!(
4964 "query exceeds {} conditions",
4965 crate::query::MAX_HARD_CONDITIONS
4966 )));
4967 }
4968 if let Some(limit) = q.limit {
4969 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4970 return Err(MongrelError::InvalidArgument(format!(
4971 "query limit must be between 1 and {}",
4972 crate::query::MAX_FINAL_LIMIT
4973 )));
4974 }
4975 }
4976 if q.offset > crate::query::MAX_QUERY_OFFSET {
4977 return Err(MongrelError::InvalidArgument(format!(
4978 "query offset exceeds {}",
4979 crate::query::MAX_QUERY_OFFSET
4980 )));
4981 }
4982 Ok(())
4983 }
4984
4985 #[doc(hidden)]
4988 pub fn query_all_at(
4989 &mut self,
4990 conditions: &[crate::query::Condition],
4991 snapshot: Snapshot,
4992 ) -> Result<Vec<Row>> {
4993 self.require_select()?;
4994 self.ensure_indexes_complete()?;
4995 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4996 return Err(MongrelError::InvalidArgument(format!(
4997 "query exceeds {} conditions",
4998 crate::query::MAX_HARD_CONDITIONS
4999 )));
5000 }
5001 self.query_conditions_at(
5002 conditions,
5003 snapshot,
5004 None,
5005 None,
5006 0,
5007 None,
5008 unix_nanos_now(),
5009 None,
5010 )
5011 }
5012
5013 #[allow(clippy::too_many_arguments)]
5014 fn query_conditions_at(
5015 &self,
5016 conditions: &[crate::query::Condition],
5017 snapshot: Snapshot,
5018 allowed: Option<&std::collections::HashSet<RowId>>,
5019 limit: Option<usize>,
5020 offset: usize,
5021 after_row_id: Option<RowId>,
5022 query_time_nanos: i64,
5023 control: Option<&crate::ExecutionControl>,
5024 ) -> Result<Vec<Row>> {
5025 control
5026 .map(crate::ExecutionControl::checkpoint)
5027 .transpose()?;
5028 crate::trace::QueryTrace::record(|t| {
5029 t.run_count = self.run_refs.len();
5030 t.memtable_rows = self.memtable.len();
5031 t.mutable_run_rows = self.mutable_run.len();
5032 });
5033 if conditions.is_empty() {
5037 crate::trace::QueryTrace::record(|t| {
5038 t.scan_mode = crate::trace::ScanMode::Materialized;
5039 t.row_materialized = true;
5040 });
5041 let mut rows = match control {
5042 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5043 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5044 };
5045 if let Some(allowed) = allowed {
5046 let mut filtered = Vec::with_capacity(rows.len());
5047 for (index, row) in rows.into_iter().enumerate() {
5048 if index & 255 == 0 {
5049 control
5050 .map(crate::ExecutionControl::checkpoint)
5051 .transpose()?;
5052 }
5053 if allowed.contains(&row.row_id) {
5054 filtered.push(row);
5055 }
5056 }
5057 rows = filtered;
5058 }
5059 if let Some(after_row_id) = after_row_id {
5060 rows.retain(|row| row.row_id > after_row_id);
5061 }
5062 rows.drain(..offset.min(rows.len()));
5063 if let Some(limit) = limit {
5064 rows.truncate(limit);
5065 }
5066 return Ok(rows);
5067 }
5068 crate::trace::QueryTrace::record(|t| {
5069 t.conditions_pushed = conditions.len();
5070 t.scan_mode = crate::trace::ScanMode::Materialized;
5071 t.row_materialized = true;
5072 });
5073 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5080 ordered.sort_by_key(|c| condition_cost_rank(c));
5081 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5082 for c in &ordered {
5083 control
5084 .map(crate::ExecutionControl::checkpoint)
5085 .transpose()?;
5086 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5087 let empty = s.is_empty();
5088 sets.push(s);
5089 if empty {
5090 break;
5091 }
5092 }
5093 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5094 if let Some(allowed) = allowed {
5095 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5096 }
5097 if let Some(after_row_id) = after_row_id {
5098 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5099 rids.drain(..first);
5100 }
5101 rids.drain(..offset.min(rids.len()));
5102 if let Some(limit) = limit {
5103 rids.truncate(limit);
5104 }
5105 control
5106 .map(crate::ExecutionControl::checkpoint)
5107 .transpose()?;
5108 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5109 }
5110
5111 pub fn retrieve(
5113 &mut self,
5114 retriever: &crate::query::Retriever,
5115 ) -> Result<Vec<crate::query::RetrieverHit>> {
5116 self.retrieve_with_allowed(retriever, None)
5117 }
5118
5119 pub fn retrieve_at(
5120 &mut self,
5121 retriever: &crate::query::Retriever,
5122 snapshot: Snapshot,
5123 allowed: Option<&std::collections::HashSet<RowId>>,
5124 ) -> Result<Vec<crate::query::RetrieverHit>> {
5125 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5126 }
5127
5128 pub fn retrieve_with_allowed(
5131 &mut self,
5132 retriever: &crate::query::Retriever,
5133 allowed: Option<&std::collections::HashSet<RowId>>,
5134 ) -> Result<Vec<crate::query::RetrieverHit>> {
5135 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5136 }
5137
5138 pub fn retrieve_at_with_allowed(
5139 &mut self,
5140 retriever: &crate::query::Retriever,
5141 snapshot: Snapshot,
5142 allowed: Option<&std::collections::HashSet<RowId>>,
5143 ) -> Result<Vec<crate::query::RetrieverHit>> {
5144 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5145 }
5146
5147 pub fn retrieve_at_with_allowed_and_context(
5148 &mut self,
5149 retriever: &crate::query::Retriever,
5150 snapshot: Snapshot,
5151 allowed: Option<&std::collections::HashSet<RowId>>,
5152 context: Option<&crate::query::AiExecutionContext>,
5153 ) -> Result<Vec<crate::query::RetrieverHit>> {
5154 self.require_select()?;
5155 self.ensure_indexes_complete()?;
5156 self.validate_retriever(retriever)?;
5157 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5158 }
5159
5160 pub fn retrieve_at_with_candidate_authorization_and_context(
5161 &mut self,
5162 retriever: &crate::query::Retriever,
5163 snapshot: Snapshot,
5164 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5165 context: Option<&crate::query::AiExecutionContext>,
5166 ) -> Result<Vec<crate::query::RetrieverHit>> {
5167 self.require_select()?;
5168 self.ensure_indexes_complete()?;
5169 self.retrieve_at_with_candidate_authorization_on_generation(
5170 retriever,
5171 snapshot,
5172 authorization,
5173 context,
5174 )
5175 }
5176
5177 #[doc(hidden)]
5178 pub fn retrieve_at_with_candidate_authorization_on_generation(
5179 &self,
5180 retriever: &crate::query::Retriever,
5181 snapshot: Snapshot,
5182 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5183 context: Option<&crate::query::AiExecutionContext>,
5184 ) -> Result<Vec<crate::query::RetrieverHit>> {
5185 self.require_select()?;
5186 self.validate_retriever(retriever)?;
5187 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5188 }
5189
5190 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5191 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5192 let (column_id, k) = match retriever {
5193 Retriever::Ann {
5194 column_id,
5195 query,
5196 k,
5197 } => {
5198 let index = self.ann.get(column_id).ok_or_else(|| {
5199 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5200 })?;
5201 if query.len() != index.dim() {
5202 return Err(MongrelError::InvalidArgument(format!(
5203 "ANN query dimension must be {}, got {}",
5204 index.dim(),
5205 query.len()
5206 )));
5207 }
5208 if query.iter().any(|value| !value.is_finite()) {
5209 return Err(MongrelError::InvalidArgument(
5210 "ANN query values must be finite".into(),
5211 ));
5212 }
5213 (*column_id, *k)
5214 }
5215 Retriever::Sparse {
5216 column_id,
5217 query,
5218 k,
5219 } => {
5220 if !self.sparse.contains_key(column_id) {
5221 return Err(MongrelError::InvalidArgument(format!(
5222 "column {column_id} has no Sparse index"
5223 )));
5224 }
5225 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5226 return Err(MongrelError::InvalidArgument(
5227 "Sparse query must be non-empty with finite weights".into(),
5228 ));
5229 }
5230 if query.len() > MAX_SPARSE_TERMS {
5231 return Err(MongrelError::InvalidArgument(format!(
5232 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5233 )));
5234 }
5235 (*column_id, *k)
5236 }
5237 Retriever::MinHash {
5238 column_id,
5239 members,
5240 k,
5241 } => {
5242 if !self.minhash.contains_key(column_id) {
5243 return Err(MongrelError::InvalidArgument(format!(
5244 "column {column_id} has no MinHash index"
5245 )));
5246 }
5247 if members.is_empty() {
5248 return Err(MongrelError::InvalidArgument(
5249 "MinHash members must not be empty".into(),
5250 ));
5251 }
5252 if members.len() > MAX_SET_MEMBERS {
5253 return Err(MongrelError::InvalidArgument(format!(
5254 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5255 )));
5256 }
5257 let mut total_bytes = 0usize;
5258 for member in members {
5259 let bytes = member.encoded_len();
5260 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5261 return Err(MongrelError::InvalidArgument(format!(
5262 "MinHash member exceeds {} bytes",
5263 crate::query::MAX_SET_MEMBER_BYTES
5264 )));
5265 }
5266 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5267 MongrelError::InvalidArgument("MinHash input size overflow".into())
5268 })?;
5269 }
5270 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5271 return Err(MongrelError::InvalidArgument(format!(
5272 "MinHash input exceeds {} bytes",
5273 crate::query::MAX_SET_INPUT_BYTES
5274 )));
5275 }
5276 (*column_id, *k)
5277 }
5278 };
5279 if k == 0 {
5280 return Err(MongrelError::InvalidArgument(
5281 "retriever k must be > 0".into(),
5282 ));
5283 }
5284 if k > MAX_RETRIEVER_K {
5285 return Err(MongrelError::InvalidArgument(format!(
5286 "retriever k exceeds {MAX_RETRIEVER_K}"
5287 )));
5288 }
5289 debug_assert!(self
5290 .schema
5291 .columns
5292 .iter()
5293 .any(|column| column.id == column_id));
5294 Ok(())
5295 }
5296
5297 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5298 use crate::query::Condition;
5299 match condition {
5300 Condition::Ann {
5301 column_id,
5302 query,
5303 k,
5304 } => self.validate_retriever(&crate::query::Retriever::Ann {
5305 column_id: *column_id,
5306 query: query.clone(),
5307 k: *k,
5308 }),
5309 Condition::SparseMatch {
5310 column_id,
5311 query,
5312 k,
5313 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5314 column_id: *column_id,
5315 query: query.clone(),
5316 k: *k,
5317 }),
5318 Condition::MinHashSimilar {
5319 column_id,
5320 query,
5321 k,
5322 } => {
5323 if !self.minhash.contains_key(column_id) {
5324 return Err(MongrelError::InvalidArgument(format!(
5325 "column {column_id} has no MinHash index"
5326 )));
5327 }
5328 if query.is_empty() || *k == 0 {
5329 return Err(MongrelError::InvalidArgument(
5330 "MinHash query must be non-empty and k must be > 0".into(),
5331 ));
5332 }
5333 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5334 {
5335 return Err(MongrelError::InvalidArgument(format!(
5336 "MinHash query must have <= {} members and k <= {}",
5337 crate::query::MAX_SET_MEMBERS,
5338 crate::query::MAX_RETRIEVER_K
5339 )));
5340 }
5341 Ok(())
5342 }
5343 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5344 Err(MongrelError::InvalidArgument(format!(
5345 "bitmap IN exceeds {} values",
5346 crate::query::MAX_SET_MEMBERS
5347 )))
5348 }
5349 Condition::FmContainsAll { patterns, .. }
5350 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5351 {
5352 Err(MongrelError::InvalidArgument(format!(
5353 "FM query exceeds {} patterns",
5354 crate::query::MAX_HARD_CONDITIONS
5355 )))
5356 }
5357 _ => Ok(()),
5358 }
5359 }
5360
5361 fn retrieve_filtered(
5362 &self,
5363 retriever: &crate::query::Retriever,
5364 snapshot: Snapshot,
5365 hard_filter: Option<&RowIdSet>,
5366 allowed: Option<&std::collections::HashSet<RowId>>,
5367 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5368 context: Option<&crate::query::AiExecutionContext>,
5369 ) -> Result<Vec<crate::query::RetrieverHit>> {
5370 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5371 let started = std::time::Instant::now();
5372 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5373 Retriever::Ann {
5374 column_id,
5375 query,
5376 k,
5377 } => {
5378 let Some(index) = self.ann.get(column_id) else {
5379 return Ok(Vec::new());
5380 };
5381 let cap = ann_candidate_cap(index.len(), context);
5382 if cap == 0 {
5383 return Ok(Vec::new());
5384 }
5385 let mut breadth = (*k).max(1).min(cap);
5386 let mut eligibility = std::collections::HashMap::new();
5387 let mut filtered = loop {
5388 let mut seen = std::collections::HashSet::new();
5389 if let Some(context) = context {
5390 context.checkpoint()?;
5391 }
5392 let raw = index.search_with_context(query, breadth, context)?;
5393 let unchecked: Vec<_> = raw
5394 .iter()
5395 .map(|(row_id, _)| *row_id)
5396 .filter(|row_id| !eligibility.contains_key(row_id))
5397 .filter(|row_id| {
5398 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5399 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5400 })
5401 .collect();
5402 let eligible = self.eligible_and_authorized_candidate_ids(
5403 &unchecked,
5404 *column_id,
5405 snapshot,
5406 candidate_authorization,
5407 context,
5408 )?;
5409 for row_id in unchecked {
5410 eligibility.insert(row_id, eligible.contains(&row_id));
5411 }
5412 let filtered: Vec<_> = raw
5413 .into_iter()
5414 .filter(|(row_id, _)| {
5415 seen.insert(*row_id)
5416 && eligibility.get(row_id).copied().unwrap_or(false)
5417 })
5418 .map(|(row_id, score)| {
5419 let score = match score {
5420 crate::index::AnnDistance::Hamming(d) => {
5421 RetrieverScore::AnnHammingDistance(d)
5422 }
5423 crate::index::AnnDistance::Cosine(d) => {
5424 RetrieverScore::AnnCosineDistance(d)
5425 }
5426 };
5427 (row_id, score)
5428 })
5429 .collect();
5430 if filtered.len() >= *k || breadth >= cap {
5431 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5432 crate::trace::QueryTrace::record(|trace| {
5433 trace.ann_candidate_cap_hit = true;
5434 });
5435 }
5436 break filtered;
5437 }
5438 breadth = breadth.saturating_mul(2).min(cap);
5439 };
5440 filtered.truncate(*k);
5441 filtered
5442 }
5443 Retriever::Sparse {
5444 column_id,
5445 query,
5446 k,
5447 } => self
5448 .sparse
5449 .get(column_id)
5450 .map(|index| -> Result<Vec<_>> {
5451 let mut breadth = (*k).max(1);
5452 let mut eligibility = std::collections::HashMap::new();
5453 loop {
5454 if let Some(context) = context {
5455 context.checkpoint()?;
5456 }
5457 let raw = index.search_with_context(query, breadth, context)?;
5458 let unchecked: Vec<_> = raw
5459 .iter()
5460 .map(|(row_id, _)| *row_id)
5461 .filter(|row_id| !eligibility.contains_key(row_id))
5462 .filter(|row_id| {
5463 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5464 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5465 })
5466 .collect();
5467 let eligible = self.eligible_and_authorized_candidate_ids(
5468 &unchecked,
5469 *column_id,
5470 snapshot,
5471 candidate_authorization,
5472 context,
5473 )?;
5474 for row_id in unchecked {
5475 eligibility.insert(row_id, eligible.contains(&row_id));
5476 }
5477 let filtered: Vec<_> = raw
5478 .iter()
5479 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5480 .take(*k)
5481 .map(|(row_id, score)| {
5482 (*row_id, RetrieverScore::SparseDotProduct(*score))
5483 })
5484 .collect();
5485 if filtered.len() >= *k || raw.len() < breadth {
5486 break Ok(filtered);
5487 }
5488 let next = breadth.saturating_mul(2);
5489 if next == breadth {
5490 break Ok(filtered);
5491 }
5492 breadth = next;
5493 }
5494 })
5495 .transpose()?
5496 .unwrap_or_default(),
5497 Retriever::MinHash {
5498 column_id,
5499 members,
5500 k,
5501 } => self
5502 .minhash
5503 .get(column_id)
5504 .map(|index| -> Result<Vec<_>> {
5505 let mut hashes = Vec::with_capacity(members.len());
5506 for member in members {
5507 if let Some(context) = context {
5508 context.consume(crate::query::work_units(
5509 member.encoded_len(),
5510 crate::query::PARSE_WORK_QUANTUM,
5511 ))?;
5512 }
5513 hashes.push(member.hash_v1());
5514 }
5515 let mut breadth = (*k).max(1);
5516 let mut eligibility = std::collections::HashMap::new();
5517 loop {
5518 if let Some(context) = context {
5519 context.checkpoint()?;
5520 }
5521 let raw = index.search_with_context(&hashes, breadth, context)?;
5522 let unchecked: Vec<_> = raw
5523 .iter()
5524 .map(|(row_id, _)| *row_id)
5525 .filter(|row_id| !eligibility.contains_key(row_id))
5526 .filter(|row_id| {
5527 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5528 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5529 })
5530 .collect();
5531 let eligible = self.eligible_and_authorized_candidate_ids(
5532 &unchecked,
5533 *column_id,
5534 snapshot,
5535 candidate_authorization,
5536 context,
5537 )?;
5538 for row_id in unchecked {
5539 eligibility.insert(row_id, eligible.contains(&row_id));
5540 }
5541 let filtered: Vec<_> = raw
5542 .iter()
5543 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5544 .take(*k)
5545 .map(|(row_id, score)| {
5546 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5547 })
5548 .collect();
5549 if filtered.len() >= *k || raw.len() < breadth {
5550 break Ok(filtered);
5551 }
5552 let next = breadth.saturating_mul(2);
5553 if next == breadth {
5554 break Ok(filtered);
5555 }
5556 breadth = next;
5557 }
5558 })
5559 .transpose()?
5560 .unwrap_or_default(),
5561 };
5562 let elapsed = started.elapsed().as_nanos() as u64;
5563 crate::trace::QueryTrace::record(|trace| {
5564 match retriever {
5565 Retriever::Ann { .. } => {
5566 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
5567 if let Retriever::Ann { column_id, .. } = retriever {
5568 if let Some(index) = self.ann.get(column_id) {
5569 trace.ann_algorithm = Some(index.algorithm());
5570 trace.ann_quantization = Some(index.quantization());
5571 trace.ann_backend = Some(index.backend_name());
5572 }
5573 }
5574 }
5575 Retriever::Sparse { .. } => {
5576 trace.sparse_candidate_nanos =
5577 trace.sparse_candidate_nanos.saturating_add(elapsed)
5578 }
5579 Retriever::MinHash { .. } => {
5580 trace.minhash_candidate_nanos =
5581 trace.minhash_candidate_nanos.saturating_add(elapsed)
5582 }
5583 }
5584 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5585 });
5586 Ok(scored
5587 .into_iter()
5588 .enumerate()
5589 .map(|(rank, (row_id, score))| RetrieverHit {
5590 row_id,
5591 rank: rank + 1,
5592 score,
5593 })
5594 .collect())
5595 }
5596
5597 fn eligible_candidate_ids(
5598 &self,
5599 candidates: &[RowId],
5600 _column_id: u16,
5601 snapshot: Snapshot,
5602 context: Option<&crate::query::AiExecutionContext>,
5603 ) -> Result<std::collections::HashSet<RowId>> {
5604 if !self.had_deletes
5605 && self.ttl.is_none()
5606 && self.pending_put_cols.is_empty()
5607 && snapshot.epoch == self.snapshot().epoch
5608 {
5609 return Ok(candidates.iter().copied().collect());
5610 }
5611 let mut readers: Vec<_> = self
5612 .run_refs
5613 .iter()
5614 .map(|run| self.open_reader(run.run_id))
5615 .collect::<Result<_>>()?;
5616 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5617 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5618 for &row_id in candidates {
5619 if let Some(context) = context {
5620 context.consume(1)?;
5621 }
5622 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5623 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5624 let overlay = match (mem, mutable) {
5625 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5626 (Some(value), None) | (None, Some(value)) => Some(value),
5627 (None, None) => None,
5628 };
5629 if let Some((_, row)) = overlay {
5630 if !row.deleted && !self.row_expired_at(&row, now) {
5631 eligible.insert(row_id);
5632 }
5633 continue;
5634 }
5635 let mut best: Option<(Epoch, bool, usize)> = None;
5636 for (index, reader) in readers.iter_mut().enumerate() {
5637 if let Some((epoch, deleted)) =
5638 reader.get_version_visibility(row_id, snapshot.epoch)?
5639 {
5640 if best
5641 .as_ref()
5642 .map(|(best_epoch, ..)| epoch > *best_epoch)
5643 .unwrap_or(true)
5644 {
5645 best = Some((epoch, deleted, index));
5646 }
5647 }
5648 }
5649 let Some((_, false, reader_index)) = best else {
5650 continue;
5651 };
5652 if let Some(ttl) = self.ttl {
5653 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5654 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5655 {
5656 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5657 continue;
5658 }
5659 }
5660 }
5661 eligible.insert(row_id);
5662 }
5663 Ok(eligible)
5664 }
5665
5666 fn eligible_and_authorized_candidate_ids(
5667 &self,
5668 candidates: &[RowId],
5669 column_id: u16,
5670 snapshot: Snapshot,
5671 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5672 context: Option<&crate::query::AiExecutionContext>,
5673 ) -> Result<std::collections::HashSet<RowId>> {
5674 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5675 let Some(authorization) = authorization else {
5676 return Ok(eligible);
5677 };
5678 let candidates: Vec<_> = eligible.into_iter().collect();
5679 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5680 }
5681
5682 fn policy_allowed_candidate_ids(
5683 &self,
5684 candidates: &[RowId],
5685 snapshot: Snapshot,
5686 authorization: &crate::security::CandidateAuthorization<'_>,
5687 context: Option<&crate::query::AiExecutionContext>,
5688 ) -> Result<std::collections::HashSet<RowId>> {
5689 let started = std::time::Instant::now();
5690 if candidates.is_empty()
5691 || authorization.principal.is_admin
5692 || !authorization.security.rls_enabled(authorization.table)
5693 {
5694 return Ok(candidates.iter().copied().collect());
5695 }
5696 if let Some(context) = context {
5697 context.checkpoint()?;
5698 }
5699 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5700 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5701 .iter()
5702 .map(|row_id| {
5703 (
5704 *row_id,
5705 Row {
5706 row_id: *row_id,
5707 committed_epoch: snapshot.epoch,
5708 columns: std::collections::HashMap::new(),
5709 deleted: false,
5710 },
5711 )
5712 })
5713 .collect();
5714 let columns = authorization
5715 .security
5716 .select_policy_columns(authorization.table, authorization.principal);
5717 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5718 let mut decoded = 0usize;
5719 for column_id in &columns {
5720 if let Some(context) = context {
5721 context.checkpoint()?;
5722 }
5723 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5724 &row_ids, *column_id, snapshot, query_now, context,
5725 )? {
5726 if let Some(row) = rows.get_mut(&row_id) {
5727 row.columns.insert(*column_id, value);
5728 decoded = decoded.saturating_add(1);
5729 }
5730 }
5731 }
5732 if let Some(context) = context {
5733 context.consume(candidates.len().saturating_add(decoded))?;
5734 }
5735 let allowed = rows
5736 .into_values()
5737 .filter_map(|row| {
5738 authorization
5739 .security
5740 .row_allowed(
5741 authorization.table,
5742 crate::security::PolicyCommand::Select,
5743 &row,
5744 authorization.principal,
5745 false,
5746 )
5747 .then_some(row.row_id)
5748 })
5749 .collect();
5750 crate::trace::QueryTrace::record(|trace| {
5751 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5752 trace.rls_policy_columns_decoded =
5753 trace.rls_policy_columns_decoded.saturating_add(decoded);
5754 trace.authorization_nanos = trace
5755 .authorization_nanos
5756 .saturating_add(started.elapsed().as_nanos() as u64);
5757 });
5758 Ok(allowed)
5759 }
5760
5761 pub fn search(
5763 &mut self,
5764 request: &crate::query::SearchRequest,
5765 ) -> Result<Vec<crate::query::SearchHit>> {
5766 self.search_with_allowed(request, None)
5767 }
5768
5769 pub fn search_at(
5770 &mut self,
5771 request: &crate::query::SearchRequest,
5772 snapshot: Snapshot,
5773 authorized: Option<&std::collections::HashSet<RowId>>,
5774 ) -> Result<Vec<crate::query::SearchHit>> {
5775 self.search_at_with_allowed(request, snapshot, authorized)
5776 }
5777
5778 pub fn search_with_allowed(
5779 &mut self,
5780 request: &crate::query::SearchRequest,
5781 authorized: Option<&std::collections::HashSet<RowId>>,
5782 ) -> Result<Vec<crate::query::SearchHit>> {
5783 self.search_at_with_allowed(request, self.snapshot(), authorized)
5784 }
5785
5786 pub fn search_at_with_allowed(
5787 &mut self,
5788 request: &crate::query::SearchRequest,
5789 snapshot: Snapshot,
5790 authorized: Option<&std::collections::HashSet<RowId>>,
5791 ) -> Result<Vec<crate::query::SearchHit>> {
5792 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5793 }
5794
5795 pub fn search_at_with_allowed_and_context(
5796 &mut self,
5797 request: &crate::query::SearchRequest,
5798 snapshot: Snapshot,
5799 authorized: Option<&std::collections::HashSet<RowId>>,
5800 context: Option<&crate::query::AiExecutionContext>,
5801 ) -> Result<Vec<crate::query::SearchHit>> {
5802 self.ensure_indexes_complete()?;
5803 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5804 }
5805
5806 pub fn search_at_with_candidate_authorization_and_context(
5807 &mut self,
5808 request: &crate::query::SearchRequest,
5809 snapshot: Snapshot,
5810 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5811 context: Option<&crate::query::AiExecutionContext>,
5812 ) -> Result<Vec<crate::query::SearchHit>> {
5813 self.ensure_indexes_complete()?;
5814 self.search_at_with_filters_and_context(
5815 request,
5816 snapshot,
5817 None,
5818 authorization,
5819 context,
5820 None,
5821 )
5822 }
5823
5824 #[doc(hidden)]
5825 pub fn search_at_with_candidate_authorization_on_generation(
5826 &self,
5827 request: &crate::query::SearchRequest,
5828 snapshot: Snapshot,
5829 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5830 context: Option<&crate::query::AiExecutionContext>,
5831 ) -> Result<Vec<crate::query::SearchHit>> {
5832 self.search_at_with_filters_and_context(
5833 request,
5834 snapshot,
5835 None,
5836 authorization,
5837 context,
5838 None,
5839 )
5840 }
5841
5842 #[doc(hidden)]
5843 pub fn search_at_with_candidate_authorization_on_generation_after(
5844 &self,
5845 request: &crate::query::SearchRequest,
5846 snapshot: Snapshot,
5847 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5848 context: Option<&crate::query::AiExecutionContext>,
5849 after: Option<crate::query::SearchAfter>,
5850 ) -> Result<Vec<crate::query::SearchHit>> {
5851 self.search_at_with_filters_and_context(
5852 request,
5853 snapshot,
5854 None,
5855 authorization,
5856 context,
5857 after,
5858 )
5859 }
5860
5861 fn search_at_with_filters_and_context(
5862 &self,
5863 request: &crate::query::SearchRequest,
5864 snapshot: Snapshot,
5865 authorized: Option<&std::collections::HashSet<RowId>>,
5866 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5867 context: Option<&crate::query::AiExecutionContext>,
5868 after: Option<crate::query::SearchAfter>,
5869 ) -> Result<Vec<crate::query::SearchHit>> {
5870 use crate::query::{
5871 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5872 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5873 };
5874 let total_started = std::time::Instant::now();
5875 let rank_offset = after.map_or(0, |after| after.returned_count);
5876 self.require_select()?;
5877 if request.limit == 0 {
5878 return Err(MongrelError::InvalidArgument(
5879 "search limit must be > 0".into(),
5880 ));
5881 }
5882 if request.limit > MAX_FINAL_LIMIT {
5883 return Err(MongrelError::InvalidArgument(format!(
5884 "search limit exceeds {MAX_FINAL_LIMIT}"
5885 )));
5886 }
5887 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5888 return Err(MongrelError::InvalidArgument(
5889 "search-after score must be finite".into(),
5890 ));
5891 }
5892 if request.retrievers.is_empty() {
5893 return Err(MongrelError::InvalidArgument(
5894 "search requires at least one retriever".into(),
5895 ));
5896 }
5897 if request.retrievers.len() > MAX_RETRIEVERS {
5898 return Err(MongrelError::InvalidArgument(format!(
5899 "search exceeds {MAX_RETRIEVERS} retrievers"
5900 )));
5901 }
5902 if request.must.len() > MAX_HARD_CONDITIONS {
5903 return Err(MongrelError::InvalidArgument(format!(
5904 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5905 )));
5906 }
5907 for condition in &request.must {
5908 self.validate_condition(condition)?;
5909 }
5910 if request.must.iter().any(|condition| {
5911 matches!(
5912 condition,
5913 Condition::Ann { .. }
5914 | Condition::SparseMatch { .. }
5915 | Condition::MinHashSimilar { .. }
5916 )
5917 }) {
5918 return Err(MongrelError::InvalidArgument(
5919 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5920 .into(),
5921 ));
5922 }
5923 let mut names = std::collections::HashSet::new();
5924 for named in &request.retrievers {
5925 if named.name.is_empty()
5926 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5927 || !names.insert(named.name.as_str())
5928 {
5929 return Err(MongrelError::InvalidArgument(format!(
5930 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5931 crate::query::MAX_RETRIEVER_NAME_BYTES
5932 )));
5933 }
5934 if !named.weight.is_finite()
5935 || named.weight < 0.0
5936 || named.weight > MAX_RETRIEVER_WEIGHT
5937 {
5938 return Err(MongrelError::InvalidArgument(format!(
5939 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5940 )));
5941 }
5942 self.validate_retriever(&named.retriever)?;
5943 }
5944 let projection = request
5945 .projection
5946 .clone()
5947 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5948 if projection.len() > MAX_PROJECTION_COLUMNS {
5949 return Err(MongrelError::InvalidArgument(format!(
5950 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5951 )));
5952 }
5953 for column_id in &projection {
5954 if !self
5955 .schema
5956 .columns
5957 .iter()
5958 .any(|column| column.id == *column_id)
5959 {
5960 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5961 }
5962 }
5963 if let Some(crate::query::Rerank::ExactVector {
5964 embedding_column,
5965 query,
5966 candidate_limit,
5967 weight,
5968 ..
5969 }) = &request.rerank
5970 {
5971 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5972 {
5973 return Err(MongrelError::InvalidArgument(format!(
5974 "rerank candidate_limit must be between search limit and {}",
5975 crate::query::MAX_RETRIEVER_K
5976 )));
5977 }
5978 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5979 return Err(MongrelError::InvalidArgument(format!(
5980 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5981 )));
5982 }
5983 let column = self
5984 .schema
5985 .columns
5986 .iter()
5987 .find(|column| column.id == *embedding_column)
5988 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5989 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5990 return Err(MongrelError::InvalidArgument(format!(
5991 "rerank column {embedding_column} is not an embedding"
5992 )));
5993 };
5994 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5995 return Err(MongrelError::InvalidArgument(format!(
5996 "rerank query must contain {dim} finite values"
5997 )));
5998 }
5999 }
6000
6001 let hard_filter_started = std::time::Instant::now();
6002 let hard_filter = if request.must.is_empty() {
6003 None
6004 } else {
6005 let mut sets = Vec::with_capacity(request.must.len());
6006 for condition in &request.must {
6007 if let Some(context) = context {
6008 context.checkpoint()?;
6009 }
6010 sets.push(self.resolve_condition(condition, snapshot)?);
6011 }
6012 Some(RowIdSet::intersect_many(sets))
6013 };
6014 crate::trace::QueryTrace::record(|trace| {
6015 trace.hard_filter_nanos = trace
6016 .hard_filter_nanos
6017 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6018 });
6019 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6020 return Ok(Vec::new());
6021 }
6022
6023 let constant = match request.fusion {
6024 Fusion::ReciprocalRank { constant } => constant,
6025 };
6026 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6027 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6028 let mut fusion_nanos = 0u64;
6029 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6030 std::collections::HashMap::new();
6031 for named in retrievers {
6032 if named.weight == 0.0 {
6033 continue;
6034 }
6035 if let Some(context) = context {
6036 context.checkpoint()?;
6037 }
6038 let hits = self.retrieve_filtered(
6039 &named.retriever,
6040 snapshot,
6041 hard_filter.as_ref(),
6042 authorized,
6043 candidate_authorization,
6044 context,
6045 )?;
6046 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6047 let fusion_started = std::time::Instant::now();
6048 for hit in hits {
6049 if let Some(context) = context {
6050 context.consume(1)?;
6051 }
6052 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6053 if !contribution.is_finite() {
6054 return Err(MongrelError::InvalidArgument(
6055 "retriever contribution must be finite".into(),
6056 ));
6057 }
6058 let max_fused_candidates = context.map_or(
6059 crate::query::MAX_FUSED_CANDIDATES,
6060 crate::query::AiExecutionContext::max_fused_candidates,
6061 );
6062 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6063 return Err(MongrelError::WorkBudgetExceeded);
6064 }
6065 let entry = fused.entry(hit.row_id).or_default();
6066 entry.0 += contribution;
6067 if !entry.0.is_finite() {
6068 return Err(MongrelError::InvalidArgument(
6069 "fused score must be finite".into(),
6070 ));
6071 }
6072 entry.1.push(ComponentScore {
6073 retriever_name: retriever_name.clone(),
6074 rank: hit.rank,
6075 raw_score: hit.score,
6076 contribution,
6077 });
6078 }
6079 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6080 }
6081 let union_size = fused.len();
6082 let mut ranked: Vec<_> = fused
6083 .into_iter()
6084 .map(|(row_id, (fused_score, components))| {
6085 (row_id, fused_score, components, None, fused_score)
6086 })
6087 .collect();
6088 let order = |(a_row, _, _, _, a_score): &(
6089 RowId,
6090 f64,
6091 Vec<ComponentScore>,
6092 Option<f32>,
6093 f64,
6094 ),
6095 (b_row, _, _, _, b_score): &(
6096 RowId,
6097 f64,
6098 Vec<ComponentScore>,
6099 Option<f32>,
6100 f64,
6101 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6102 if let Some(crate::query::Rerank::ExactVector {
6103 embedding_column,
6104 query,
6105 metric,
6106 candidate_limit,
6107 weight,
6108 }) = &request.rerank
6109 {
6110 let fused_order = |(a_row, a_score, ..): &(
6111 RowId,
6112 f64,
6113 Vec<ComponentScore>,
6114 Option<f32>,
6115 f64,
6116 ),
6117 (b_row, b_score, ..): &(
6118 RowId,
6119 f64,
6120 Vec<ComponentScore>,
6121 Option<f32>,
6122 f64,
6123 )| {
6124 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6125 };
6126 let selection_started = std::time::Instant::now();
6127 if let Some(context) = context {
6128 context.consume(ranked.len())?;
6129 }
6130 if ranked.len() > *candidate_limit {
6131 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6132 ranked.truncate(*candidate_limit);
6133 }
6134 ranked.sort_by(fused_order);
6135 fusion_nanos =
6136 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6137 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6138 if let Some(context) = context {
6139 context.consume(row_ids.len())?;
6140 }
6141 let query_now =
6142 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6143 let gather_started = std::time::Instant::now();
6144 let vectors = self.values_for_rids_batch_at_with_context(
6145 &row_ids,
6146 *embedding_column,
6147 snapshot,
6148 query_now,
6149 context,
6150 )?;
6151 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6152 let vector_work =
6153 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6154 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6155 if let Some(context) = context {
6156 context.consume(vector_work)?;
6157 }
6158 query
6159 .iter()
6160 .map(|value| f64::from(*value).powi(2))
6161 .sum::<f64>()
6162 .sqrt()
6163 } else {
6164 0.0
6165 };
6166 let score_started = std::time::Instant::now();
6167 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6168 for (row_id, value) in vectors {
6169 let Some(vector) = value.as_embedding() else {
6170 continue;
6171 };
6172 let score = match metric {
6173 crate::query::VectorMetric::DotProduct => {
6174 if let Some(context) = context {
6175 context.consume(vector_work)?;
6176 }
6177 query
6178 .iter()
6179 .zip(vector)
6180 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6181 .sum::<f64>()
6182 }
6183 crate::query::VectorMetric::Cosine => {
6184 if let Some(context) = context {
6185 context.consume(vector_work.saturating_mul(2))?;
6186 }
6187 let dot = query
6188 .iter()
6189 .zip(vector)
6190 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6191 .sum::<f64>();
6192 let norm = vector
6193 .iter()
6194 .map(|value| f64::from(*value).powi(2))
6195 .sum::<f64>()
6196 .sqrt();
6197 if query_norm == 0.0 || norm == 0.0 {
6198 0.0
6199 } else {
6200 dot / (query_norm * norm)
6201 }
6202 }
6203 crate::query::VectorMetric::Euclidean => {
6204 if let Some(context) = context {
6205 context.consume(vector_work)?;
6206 }
6207 query
6208 .iter()
6209 .zip(vector)
6210 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6211 .sum::<f64>()
6212 .sqrt()
6213 }
6214 };
6215 if !score.is_finite() {
6216 return Err(MongrelError::InvalidArgument(
6217 "exact rerank score must be finite".into(),
6218 ));
6219 }
6220 scores.insert(row_id, score as f32);
6221 }
6222 let mut reranked = Vec::with_capacity(ranked.len());
6223 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6224 let Some(score) = scores.get(&row_id).copied() else {
6225 continue;
6226 };
6227 let ordering_score = match metric {
6228 crate::query::VectorMetric::Euclidean => -f64::from(score),
6229 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6230 f64::from(score)
6231 }
6232 };
6233 let final_score = fused_score + *weight * ordering_score;
6234 if !final_score.is_finite() {
6235 return Err(MongrelError::InvalidArgument(
6236 "final rerank score must be finite".into(),
6237 ));
6238 }
6239 reranked.push((row_id, fused_score, components, Some(score), final_score));
6240 }
6241 ranked = reranked;
6242 ranked.sort_by(order);
6243 crate::trace::QueryTrace::record(|trace| {
6244 trace.exact_vector_gather_nanos =
6245 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6246 trace.exact_vector_score_nanos = trace
6247 .exact_vector_score_nanos
6248 .saturating_add(score_started.elapsed().as_nanos() as u64);
6249 });
6250 }
6251 if let Some(after) = after {
6252 ranked.retain(|(row_id, _, _, _, final_score)| {
6253 final_score.total_cmp(&after.final_score).is_lt()
6254 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6255 });
6256 }
6257 let projection_started = std::time::Instant::now();
6258 let sentinel = projection
6259 .first()
6260 .copied()
6261 .or_else(|| self.schema.columns.first().map(|column| column.id));
6262 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6263 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6264 let mut projection_rows = 0usize;
6265 let mut projection_cells = 0usize;
6266 while out.len() < request.limit && !ranked.is_empty() {
6267 if let Some(context) = context {
6268 context.checkpoint()?;
6269 context.consume(ranked.len())?;
6270 }
6271 let needed = request.limit - out.len();
6272 let window_size = ranked
6273 .len()
6274 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6275 let selection_started = std::time::Instant::now();
6276 let mut remainder = if ranked.len() > window_size {
6277 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6278 ranked.split_off(window_size)
6279 } else {
6280 Vec::new()
6281 };
6282 ranked.sort_by(order);
6283 fusion_nanos =
6284 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6285 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6286 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6287 if let Some(context) = context {
6288 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6289 }
6290 projection_rows = projection_rows.saturating_add(row_ids.len());
6291 projection_cells =
6292 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6293 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6294 std::collections::HashMap::new();
6295 if let Some(column_id) = sentinel {
6296 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6297 &row_ids, column_id, snapshot, query_now, context,
6298 )? {
6299 cells.entry(row_id).or_default().insert(column_id, value);
6300 }
6301 }
6302 for &column_id in &projection {
6303 if Some(column_id) == sentinel {
6304 continue;
6305 }
6306 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6307 &row_ids, column_id, snapshot, query_now, context,
6308 )? {
6309 cells.entry(row_id).or_default().insert(column_id, value);
6310 }
6311 }
6312 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6313 ranked.drain(..)
6314 {
6315 let Some(row_cells) = cells.remove(&row_id) else {
6316 continue;
6317 };
6318 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6319 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6320 out.push(SearchHit {
6321 row_id,
6322 cells: projection
6323 .iter()
6324 .filter_map(|column_id| {
6325 row_cells
6326 .get(column_id)
6327 .cloned()
6328 .map(|value| (*column_id, value))
6329 })
6330 .collect(),
6331 components,
6332 fused_score,
6333 exact_rerank_score,
6334 final_score,
6335 final_rank,
6336 });
6337 if out.len() == request.limit {
6338 break;
6339 }
6340 }
6341 ranked.append(&mut remainder);
6342 }
6343 crate::trace::QueryTrace::record(|trace| {
6344 trace.union_size = union_size;
6345 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6346 trace.projection_nanos = trace
6347 .projection_nanos
6348 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6349 trace.total_nanos = trace
6350 .total_nanos
6351 .saturating_add(total_started.elapsed().as_nanos() as u64);
6352 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6353 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6354 if let Some(context) = context {
6355 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6356 }
6357 });
6358 Ok(out)
6359 }
6360
6361 pub fn set_similarity(
6364 &mut self,
6365 request: &crate::query::SetSimilarityRequest,
6366 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6367 self.set_similarity_with_allowed(request, None)
6368 }
6369
6370 pub fn set_similarity_at(
6371 &mut self,
6372 request: &crate::query::SetSimilarityRequest,
6373 snapshot: Snapshot,
6374 allowed: Option<&std::collections::HashSet<RowId>>,
6375 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6376 self.set_similarity_explained_at(request, snapshot, allowed)
6377 .map(|(hits, _)| hits)
6378 }
6379
6380 pub fn ann_rerank(
6382 &mut self,
6383 request: &crate::query::AnnRerankRequest,
6384 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6385 self.ann_rerank_with_allowed(request, None)
6386 }
6387
6388 pub fn ann_rerank_with_allowed(
6389 &mut self,
6390 request: &crate::query::AnnRerankRequest,
6391 allowed: Option<&std::collections::HashSet<RowId>>,
6392 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6393 self.ann_rerank_at(request, self.snapshot(), allowed)
6394 }
6395
6396 pub fn ann_rerank_at(
6397 &mut self,
6398 request: &crate::query::AnnRerankRequest,
6399 snapshot: Snapshot,
6400 allowed: Option<&std::collections::HashSet<RowId>>,
6401 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6402 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6403 }
6404
6405 pub fn ann_rerank_at_with_context(
6406 &mut self,
6407 request: &crate::query::AnnRerankRequest,
6408 snapshot: Snapshot,
6409 allowed: Option<&std::collections::HashSet<RowId>>,
6410 context: Option<&crate::query::AiExecutionContext>,
6411 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6412 self.ensure_indexes_complete()?;
6413 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6414 }
6415
6416 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6417 &mut self,
6418 request: &crate::query::AnnRerankRequest,
6419 snapshot: Snapshot,
6420 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6421 context: Option<&crate::query::AiExecutionContext>,
6422 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6423 self.ensure_indexes_complete()?;
6424 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6425 }
6426
6427 #[doc(hidden)]
6428 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6429 &self,
6430 request: &crate::query::AnnRerankRequest,
6431 snapshot: Snapshot,
6432 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6433 context: Option<&crate::query::AiExecutionContext>,
6434 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6435 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6436 }
6437
6438 fn ann_rerank_at_with_filters_and_context(
6439 &self,
6440 request: &crate::query::AnnRerankRequest,
6441 snapshot: Snapshot,
6442 allowed: Option<&std::collections::HashSet<RowId>>,
6443 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6444 context: Option<&crate::query::AiExecutionContext>,
6445 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6446 use crate::query::{
6447 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6448 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6449 };
6450 if request.candidate_k == 0 || request.limit == 0 {
6451 return Err(MongrelError::InvalidArgument(
6452 "candidate_k and limit must be > 0".into(),
6453 ));
6454 }
6455 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6456 return Err(MongrelError::InvalidArgument(format!(
6457 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6458 )));
6459 }
6460 let retriever = Retriever::Ann {
6461 column_id: request.column_id,
6462 query: request.query.clone(),
6463 k: request.candidate_k,
6464 };
6465 self.require_select()?;
6466 self.validate_retriever(&retriever)?;
6467 let hits = self.retrieve_filtered(
6468 &retriever,
6469 snapshot,
6470 None,
6471 allowed,
6472 candidate_authorization,
6473 context,
6474 )?;
6475 let distances: std::collections::HashMap<_, _> = hits
6476 .iter()
6477 .filter_map(|hit| match hit.score {
6478 RetrieverScore::AnnHammingDistance(distance) => {
6479 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
6480 }
6481 RetrieverScore::AnnCosineDistance(distance) => {
6482 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
6483 }
6484 _ => None,
6485 })
6486 .collect();
6487 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6488 if let Some(context) = context {
6489 context.consume(row_ids.len())?;
6490 }
6491 let gather_started = std::time::Instant::now();
6492 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6493 let values = self.values_for_rids_batch_at_with_context(
6494 &row_ids,
6495 request.column_id,
6496 snapshot,
6497 query_now,
6498 context,
6499 )?;
6500 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6501 let score_started = std::time::Instant::now();
6502 let vector_work =
6503 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6504 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6505 if let Some(context) = context {
6506 context.consume(vector_work)?;
6507 }
6508 request
6509 .query
6510 .iter()
6511 .map(|value| f64::from(*value).powi(2))
6512 .sum::<f64>()
6513 .sqrt()
6514 } else {
6515 0.0
6516 };
6517 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6518 for (row_id, value) in values {
6519 let Some(vector) = value.as_embedding() else {
6520 continue;
6521 };
6522 let exact_score = match request.metric {
6523 VectorMetric::DotProduct => {
6524 if let Some(context) = context {
6525 context.consume(vector_work)?;
6526 }
6527 request
6528 .query
6529 .iter()
6530 .zip(vector)
6531 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6532 .sum::<f64>()
6533 }
6534 VectorMetric::Cosine => {
6535 if let Some(context) = context {
6536 context.consume(vector_work.saturating_mul(2))?;
6537 }
6538 let dot = request
6539 .query
6540 .iter()
6541 .zip(vector)
6542 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6543 .sum::<f64>();
6544 let norm = vector
6545 .iter()
6546 .map(|value| f64::from(*value).powi(2))
6547 .sum::<f64>()
6548 .sqrt();
6549 if query_norm == 0.0 || norm == 0.0 {
6550 0.0
6551 } else {
6552 dot / (query_norm * norm)
6553 }
6554 }
6555 VectorMetric::Euclidean => {
6556 if let Some(context) = context {
6557 context.consume(vector_work)?;
6558 }
6559 request
6560 .query
6561 .iter()
6562 .zip(vector)
6563 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6564 .sum::<f64>()
6565 .sqrt()
6566 }
6567 };
6568 let exact_score = exact_score as f32;
6569 if !exact_score.is_finite() {
6570 return Err(MongrelError::InvalidArgument(
6571 "exact ANN score must be finite".into(),
6572 ));
6573 }
6574 let Some(candidate_distance) = distances.get(&row_id).copied() else {
6575 continue;
6576 };
6577 reranked.push(AnnRerankHit {
6578 row_id,
6579 candidate_distance,
6580 exact_score,
6581 });
6582 }
6583 reranked.sort_by(|left, right| {
6584 let score = match request.metric {
6585 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6586 VectorMetric::Cosine | VectorMetric::DotProduct => {
6587 right.exact_score.total_cmp(&left.exact_score)
6588 }
6589 };
6590 score.then_with(|| left.row_id.cmp(&right.row_id))
6591 });
6592 reranked.truncate(request.limit);
6593 crate::trace::QueryTrace::record(|trace| {
6594 trace.exact_vector_gather_nanos =
6595 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6596 trace.exact_vector_score_nanos = trace
6597 .exact_vector_score_nanos
6598 .saturating_add(score_started.elapsed().as_nanos() as u64);
6599 });
6600 Ok(reranked)
6601 }
6602
6603 pub fn set_similarity_with_allowed(
6604 &mut self,
6605 request: &crate::query::SetSimilarityRequest,
6606 allowed: Option<&std::collections::HashSet<RowId>>,
6607 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6608 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6609 .map(|(hits, _)| hits)
6610 }
6611
6612 pub fn set_similarity_explained(
6613 &mut self,
6614 request: &crate::query::SetSimilarityRequest,
6615 ) -> Result<(
6616 Vec<crate::query::SetSimilarityHit>,
6617 crate::query::SetSimilarityTrace,
6618 )> {
6619 self.set_similarity_explained_at(request, self.snapshot(), None)
6620 }
6621
6622 fn set_similarity_explained_at(
6623 &mut self,
6624 request: &crate::query::SetSimilarityRequest,
6625 snapshot: Snapshot,
6626 allowed: Option<&std::collections::HashSet<RowId>>,
6627 ) -> Result<(
6628 Vec<crate::query::SetSimilarityHit>,
6629 crate::query::SetSimilarityTrace,
6630 )> {
6631 self.ensure_indexes_complete()?;
6632 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6633 }
6634
6635 pub fn set_similarity_at_with_context(
6636 &mut self,
6637 request: &crate::query::SetSimilarityRequest,
6638 snapshot: Snapshot,
6639 allowed: Option<&std::collections::HashSet<RowId>>,
6640 context: Option<&crate::query::AiExecutionContext>,
6641 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6642 self.ensure_indexes_complete()?;
6643 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6644 .map(|(hits, _)| hits)
6645 }
6646
6647 pub fn set_similarity_at_with_candidate_authorization_and_context(
6648 &mut self,
6649 request: &crate::query::SetSimilarityRequest,
6650 snapshot: Snapshot,
6651 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6652 context: Option<&crate::query::AiExecutionContext>,
6653 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6654 self.ensure_indexes_complete()?;
6655 self.set_similarity_explained_at_with_context(
6656 request,
6657 snapshot,
6658 None,
6659 authorization,
6660 context,
6661 )
6662 .map(|(hits, _)| hits)
6663 }
6664
6665 #[doc(hidden)]
6666 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6667 &self,
6668 request: &crate::query::SetSimilarityRequest,
6669 snapshot: Snapshot,
6670 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6671 context: Option<&crate::query::AiExecutionContext>,
6672 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6673 self.set_similarity_explained_at_with_context(
6674 request,
6675 snapshot,
6676 None,
6677 authorization,
6678 context,
6679 )
6680 .map(|(hits, _)| hits)
6681 }
6682
6683 fn set_similarity_explained_at_with_context(
6684 &self,
6685 request: &crate::query::SetSimilarityRequest,
6686 snapshot: Snapshot,
6687 allowed: Option<&std::collections::HashSet<RowId>>,
6688 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6689 context: Option<&crate::query::AiExecutionContext>,
6690 ) -> Result<(
6691 Vec<crate::query::SetSimilarityHit>,
6692 crate::query::SetSimilarityTrace,
6693 )> {
6694 use crate::query::{
6695 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6696 MAX_SET_MEMBERS,
6697 };
6698 let mut trace = crate::query::SetSimilarityTrace::default();
6699 if request.members.is_empty() {
6700 return Ok((Vec::new(), trace));
6701 }
6702 if request.candidate_k == 0 || request.limit == 0 {
6703 return Err(MongrelError::InvalidArgument(
6704 "candidate_k and limit must be > 0".into(),
6705 ));
6706 }
6707 if request.candidate_k > MAX_RETRIEVER_K
6708 || request.limit > MAX_FINAL_LIMIT
6709 || request.members.len() > MAX_SET_MEMBERS
6710 {
6711 return Err(MongrelError::InvalidArgument(format!(
6712 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6713 )));
6714 }
6715 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6716 return Err(MongrelError::InvalidArgument(
6717 "min_jaccard must be finite and between 0 and 1".into(),
6718 ));
6719 }
6720 let started = std::time::Instant::now();
6721 let retriever = Retriever::MinHash {
6722 column_id: request.column_id,
6723 members: request.members.clone(),
6724 k: request.candidate_k,
6725 };
6726 self.require_select()?;
6727 self.validate_retriever(&retriever)?;
6728 let hits = self.retrieve_filtered(
6729 &retriever,
6730 snapshot,
6731 None,
6732 allowed,
6733 candidate_authorization,
6734 context,
6735 )?;
6736 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6737 trace.candidate_count = hits.len();
6738 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6739 if let Some(context) = context {
6740 context.consume(row_ids.len())?;
6741 }
6742 let started = std::time::Instant::now();
6743 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6744 let values = self.values_for_rids_batch_at_with_context(
6745 &row_ids,
6746 request.column_id,
6747 snapshot,
6748 query_now,
6749 context,
6750 )?;
6751 trace.gather_us = started.elapsed().as_micros() as u64;
6752 if let Some(context) = context {
6753 context.consume(request.members.len())?;
6754 }
6755 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6756 let estimates: std::collections::HashMap<_, _> = hits
6757 .into_iter()
6758 .filter_map(|hit| match hit.score {
6759 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6760 _ => None,
6761 })
6762 .collect();
6763 let started = std::time::Instant::now();
6764 let mut parsed = Vec::with_capacity(values.len());
6765 for (row_id, value) in values {
6766 let Value::Bytes(bytes) = value else {
6767 continue;
6768 };
6769 if let Some(context) = context {
6770 context.consume(crate::query::work_units(
6771 bytes.len(),
6772 crate::query::PARSE_WORK_QUANTUM,
6773 ))?;
6774 }
6775 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6776 continue;
6777 };
6778 if let Some(context) = context {
6779 context.consume(members.len())?;
6780 }
6781 let stored = members
6782 .into_iter()
6783 .filter_map(|member| match member {
6784 serde_json::Value::String(value) => {
6785 Some(crate::query::SetMember::String(value))
6786 }
6787 serde_json::Value::Number(value) => {
6788 Some(crate::query::SetMember::Number(value))
6789 }
6790 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6791 _ => None,
6792 })
6793 .collect::<std::collections::HashSet<_>>();
6794 parsed.push((row_id, stored));
6795 }
6796 trace.parse_us = started.elapsed().as_micros() as u64;
6797 trace.verified_count = parsed.len();
6798 let started = std::time::Instant::now();
6799 let mut exact = Vec::new();
6800 for (row_id, stored) in parsed {
6801 if let Some(context) = context {
6802 context.consume(query.len().saturating_add(stored.len()))?;
6803 }
6804 let union = query.union(&stored).count();
6805 let score = if union == 0 {
6806 1.0
6807 } else {
6808 query.intersection(&stored).count() as f32 / union as f32
6809 };
6810 if score >= request.min_jaccard {
6811 exact.push(SetSimilarityHit {
6812 row_id,
6813 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6814 exact_jaccard: score,
6815 });
6816 }
6817 }
6818 exact.sort_by(|a, b| {
6819 b.exact_jaccard
6820 .total_cmp(&a.exact_jaccard)
6821 .then_with(|| a.row_id.cmp(&b.row_id))
6822 });
6823 exact.truncate(request.limit);
6824 trace.score_us = started.elapsed().as_micros() as u64;
6825 crate::trace::QueryTrace::record(|query_trace| {
6826 query_trace.exact_set_gather_nanos = query_trace
6827 .exact_set_gather_nanos
6828 .saturating_add(trace.gather_us.saturating_mul(1_000));
6829 query_trace.exact_set_parse_nanos = query_trace
6830 .exact_set_parse_nanos
6831 .saturating_add(trace.parse_us.saturating_mul(1_000));
6832 query_trace.exact_set_score_nanos = query_trace
6833 .exact_set_score_nanos
6834 .saturating_add(trace.score_us.saturating_mul(1_000));
6835 });
6836 Ok((exact, trace))
6837 }
6838
6839 fn values_for_rids_batch_at(
6841 &self,
6842 row_ids: &[u64],
6843 column_id: u16,
6844 snapshot: Snapshot,
6845 now: i64,
6846 ) -> Result<Vec<(RowId, Value)>> {
6847 if self.ttl.is_none()
6848 && self.memtable.is_empty()
6849 && self.mutable_run.is_empty()
6850 && self.run_refs.len() == 1
6851 {
6852 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6853 if row_ids.len().saturating_mul(24) < reader.row_count() {
6858 let mut values = Vec::with_capacity(row_ids.len());
6859 for &raw_row_id in row_ids {
6860 let row_id = RowId(raw_row_id);
6861 if let Some((_, false, Some(value))) =
6862 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6863 {
6864 values.push((row_id, value));
6865 }
6866 }
6867 return Ok(values);
6868 }
6869 let (positions, visible_row_ids) =
6870 reader.visible_positions_with_rids(snapshot.epoch)?;
6871 let requested: Vec<(RowId, usize)> = row_ids
6872 .iter()
6873 .filter_map(|raw| {
6874 visible_row_ids
6875 .binary_search(&(*raw as i64))
6876 .ok()
6877 .map(|index| (RowId(*raw), positions[index]))
6878 })
6879 .collect();
6880 let values = reader.gather_column(
6881 column_id,
6882 &requested
6883 .iter()
6884 .map(|(_, position)| *position)
6885 .collect::<Vec<_>>(),
6886 )?;
6887 return Ok(requested
6888 .into_iter()
6889 .zip(values)
6890 .map(|((row_id, _), value)| (row_id, value))
6891 .collect());
6892 }
6893 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6894 }
6895
6896 fn values_for_rids_batch_at_with_context(
6897 &self,
6898 row_ids: &[u64],
6899 column_id: u16,
6900 snapshot: Snapshot,
6901 now: i64,
6902 context: Option<&crate::query::AiExecutionContext>,
6903 ) -> Result<Vec<(RowId, Value)>> {
6904 let Some(context) = context else {
6905 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6906 };
6907 let mut values = Vec::with_capacity(row_ids.len());
6908 for chunk in row_ids.chunks(256) {
6909 context.checkpoint()?;
6910 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6911 }
6912 Ok(values)
6913 }
6914
6915 fn values_for_rids_at(
6917 &self,
6918 row_ids: &[u64],
6919 column_id: u16,
6920 snapshot: Snapshot,
6921 now: i64,
6922 ) -> Result<Vec<(RowId, Value)>> {
6923 let mut readers: Vec<_> = self
6924 .run_refs
6925 .iter()
6926 .map(|run| self.open_reader(run.run_id))
6927 .collect::<Result<_>>()?;
6928 let mut out = Vec::with_capacity(row_ids.len());
6929 for &raw_row_id in row_ids {
6930 let row_id = RowId(raw_row_id);
6931 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6932 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6933 let overlay = match (mem, mutable) {
6934 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6935 (a_epoch, a)
6936 } else {
6937 (b_epoch, b)
6938 }),
6939 (Some(value), None) | (None, Some(value)) => Some(value),
6940 (None, None) => None,
6941 };
6942 if let Some((_, row)) = overlay {
6943 if !row.deleted && !self.row_expired_at(&row, now) {
6944 if let Some(value) = row.columns.get(&column_id) {
6945 out.push((row_id, value.clone()));
6946 }
6947 }
6948 continue;
6949 }
6950
6951 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6952 for (index, reader) in readers.iter_mut().enumerate() {
6953 if let Some((epoch, deleted, value)) =
6954 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6955 {
6956 if best
6957 .as_ref()
6958 .map(|(best_epoch, ..)| epoch > *best_epoch)
6959 .unwrap_or(true)
6960 {
6961 best = Some((epoch, deleted, value, index));
6962 }
6963 }
6964 }
6965 let Some((_, false, Some(value), reader_index)) = best else {
6966 continue;
6967 };
6968 if let Some(ttl) = self.ttl {
6969 if ttl.column_id != column_id {
6970 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6971 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6972 {
6973 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6974 continue;
6975 }
6976 }
6977 } else if let Value::Int64(timestamp) = value {
6978 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6979 continue;
6980 }
6981 }
6982 }
6983 out.push((row_id, value));
6984 }
6985 Ok(out)
6986 }
6987
6988 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6993 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6994 }
6995
6996 pub fn rows_for_rids_with_context(
6997 &self,
6998 rids: &[u64],
6999 snapshot: Snapshot,
7000 context: &crate::query::AiExecutionContext,
7001 ) -> Result<Vec<Row>> {
7002 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
7003 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
7004 }
7005
7006 fn rows_for_rids_at_time(
7007 &self,
7008 rids: &[u64],
7009 snapshot: Snapshot,
7010 ttl_now: i64,
7011 control: Option<&crate::ExecutionControl>,
7012 ) -> Result<Vec<Row>> {
7013 use std::collections::HashMap;
7014 let mut rows = Vec::with_capacity(rids.len());
7015 let tier_size = self.memtable.len() + self.mutable_run.len();
7032 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7033 if rids.len().saturating_mul(24) < tier_size {
7034 for &rid in rids {
7035 if overlay.len() & 255 == 0 {
7036 control
7037 .map(crate::ExecutionControl::checkpoint)
7038 .transpose()?;
7039 }
7040 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
7041 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
7042 let newest = match (mem, mrun) {
7043 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
7044 (Some((_, mr)), None) => Some(mr),
7045 (None, Some((_, rr))) => Some(rr),
7046 (None, None) => None,
7047 };
7048 if let Some(row) = newest {
7049 overlay.insert(rid, row);
7050 }
7051 }
7052 } else {
7053 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7054 overlay
7055 .entry(row.row_id.0)
7056 .and_modify(|e| {
7057 if row.committed_epoch > e.committed_epoch {
7058 *e = row.clone();
7059 }
7060 })
7061 .or_insert(row);
7062 };
7063 for (index, row) in self
7064 .memtable
7065 .visible_versions(snapshot.epoch)
7066 .into_iter()
7067 .enumerate()
7068 {
7069 if index & 255 == 0 {
7070 control
7071 .map(crate::ExecutionControl::checkpoint)
7072 .transpose()?;
7073 }
7074 fold_newest(row, &mut overlay);
7075 }
7076 for (index, row) in self
7077 .mutable_run
7078 .visible_versions(snapshot.epoch)
7079 .into_iter()
7080 .enumerate()
7081 {
7082 if index & 255 == 0 {
7083 control
7084 .map(crate::ExecutionControl::checkpoint)
7085 .transpose()?;
7086 }
7087 fold_newest(row, &mut overlay);
7088 }
7089 }
7090 if self.run_refs.len() == 1 {
7091 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7092 if rids.len().saturating_mul(24) < reader.row_count() {
7100 for (index, &rid) in rids.iter().enumerate() {
7101 if index & 255 == 0 {
7102 control
7103 .map(crate::ExecutionControl::checkpoint)
7104 .transpose()?;
7105 }
7106 if let Some(r) = overlay.get(&rid) {
7107 if !r.deleted {
7108 rows.push(r.clone());
7109 }
7110 continue;
7111 }
7112 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7113 if !row.deleted {
7114 rows.push(row);
7115 }
7116 }
7117 }
7118 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7119 return Ok(rows);
7120 }
7121 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7130 enum Src {
7133 Overlay,
7134 Run,
7135 }
7136 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7137 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7138 for (index, rid) in rids.iter().enumerate() {
7139 if index & 255 == 0 {
7140 control
7141 .map(crate::ExecutionControl::checkpoint)
7142 .transpose()?;
7143 }
7144 if overlay.contains_key(rid) {
7145 plan.push(Src::Overlay);
7146 continue;
7147 }
7148 match vis_rids.binary_search(&(*rid as i64)) {
7149 Ok(i) => {
7150 plan.push(Src::Run);
7151 fetch.push(positions[i]);
7152 }
7153 Err(_) => { }
7154 }
7155 }
7156 let fetched = reader.materialize_batch(&fetch)?;
7157 let mut fetched_iter = fetched.into_iter();
7158 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7159 if index & 255 == 0 {
7160 control
7161 .map(crate::ExecutionControl::checkpoint)
7162 .transpose()?;
7163 }
7164 match src {
7165 Src::Overlay => {
7166 if let Some(r) = overlay.get(rid) {
7167 if !r.deleted {
7168 rows.push(r.clone());
7169 }
7170 }
7171 }
7172 Src::Run => {
7173 if let Some(row) = fetched_iter.next() {
7174 if !row.deleted {
7175 rows.push(row);
7176 }
7177 }
7178 }
7179 }
7180 }
7181 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7182 return Ok(rows);
7183 }
7184 let mut readers: Vec<_> = self
7188 .run_refs
7189 .iter()
7190 .map(|rr| self.open_reader(rr.run_id))
7191 .collect::<Result<Vec<_>>>()?;
7192 for (index, rid) in rids.iter().enumerate() {
7193 if index & 255 == 0 {
7194 control
7195 .map(crate::ExecutionControl::checkpoint)
7196 .transpose()?;
7197 }
7198 if let Some(r) = overlay.get(rid) {
7199 if !r.deleted {
7200 rows.push(r.clone());
7201 }
7202 continue;
7203 }
7204 let mut best: Option<(Epoch, Row)> = None;
7205 for reader in readers.iter_mut() {
7206 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7207 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7208 best = Some((epoch, row));
7209 }
7210 }
7211 }
7212 if let Some((_, r)) = best {
7213 if !r.deleted {
7214 rows.push(r);
7215 }
7216 }
7217 }
7218 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7219 Ok(rows)
7220 }
7221
7222 pub fn indexes_complete(&self) -> bool {
7232 self.indexes_complete
7233 }
7234
7235 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7237 self.index_build_policy
7238 }
7239
7240 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7244 self.index_build_policy = policy;
7245 }
7246
7247 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7252 if !self.indexes_complete {
7256 return None;
7257 }
7258 let b = self.bitmap.get(&column_id)?;
7259 let result: Vec<Vec<u8>> = b
7260 .keys()
7261 .into_iter()
7262 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7263 .collect();
7264 Some(result)
7265 }
7266
7267 pub fn fk_join_row_ids(
7268 &self,
7269 fk_column_id: u16,
7270 pk_values: &[Vec<u8>],
7271 fk_conditions: &[crate::query::Condition],
7272 snapshot: Snapshot,
7273 ) -> Result<Vec<u64>> {
7274 let Some(b) = self.bitmap.get(&fk_column_id) else {
7275 return Ok(Vec::new());
7276 };
7277 let mut join_set = {
7278 let mut acc = roaring::RoaringBitmap::new();
7279 for v in pk_values {
7280 acc |= b.get(v);
7281 }
7282 RowIdSet::from_roaring(acc)
7283 };
7284 if !fk_conditions.is_empty() {
7285 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7286 sets.push(join_set);
7287 for c in fk_conditions {
7288 sets.push(self.resolve_condition(c, snapshot)?);
7289 }
7290 join_set = RowIdSet::intersect_many(sets);
7291 }
7292 Ok(join_set.into_sorted_vec())
7293 }
7294
7295 pub fn fk_join_count(
7301 &self,
7302 fk_column_id: u16,
7303 pk_values: &[Vec<u8>],
7304 fk_conditions: &[crate::query::Condition],
7305 snapshot: Snapshot,
7306 ) -> Result<u64> {
7307 let Some(b) = self.bitmap.get(&fk_column_id) else {
7308 return Ok(0);
7309 };
7310 let mut acc = roaring::RoaringBitmap::new();
7311 for v in pk_values {
7312 acc |= b.get(v);
7313 }
7314 if fk_conditions.is_empty() {
7315 return Ok(acc.len());
7316 }
7317 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7318 sets.push(RowIdSet::from_roaring(acc));
7319 for c in fk_conditions {
7320 sets.push(self.resolve_condition(c, snapshot)?);
7321 }
7322 Ok(RowIdSet::intersect_many(sets).len() as u64)
7323 }
7324
7325 fn resolve_condition(
7330 &self,
7331 c: &crate::query::Condition,
7332 snapshot: Snapshot,
7333 ) -> Result<RowIdSet> {
7334 self.resolve_condition_with_allowed(c, snapshot, None)
7335 }
7336
7337 fn resolve_condition_with_allowed(
7338 &self,
7339 c: &crate::query::Condition,
7340 snapshot: Snapshot,
7341 allowed: Option<&std::collections::HashSet<RowId>>,
7342 ) -> Result<RowIdSet> {
7343 use crate::query::Condition;
7344 self.validate_condition(c)?;
7345 Ok(match c {
7346 Condition::Pk(key) => {
7347 let lookup = self
7348 .schema
7349 .primary_key()
7350 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7351 .unwrap_or_else(|| key.clone());
7352 self.hot
7353 .get(&lookup)
7354 .map(|r| RowIdSet::one(r.0))
7355 .unwrap_or_else(RowIdSet::empty)
7356 }
7357 Condition::BitmapEq { column_id, value } => {
7358 let lookup = self.index_lookup_key_bytes(*column_id, value);
7359 self.bitmap
7360 .get(column_id)
7361 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7362 .unwrap_or_else(RowIdSet::empty)
7363 }
7364 Condition::BitmapIn { column_id, values } => {
7365 let bm = self.bitmap.get(column_id);
7366 let mut acc = roaring::RoaringBitmap::new();
7367 if let Some(b) = bm {
7368 for v in values {
7369 let lookup = self.index_lookup_key_bytes(*column_id, v);
7370 acc |= b.get(&lookup);
7371 }
7372 }
7373 RowIdSet::from_roaring(acc)
7374 }
7375 Condition::BytesPrefix { column_id, prefix } => {
7376 if let Some(b) = self.bitmap.get(column_id) {
7381 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7382 let mut acc = roaring::RoaringBitmap::new();
7383 for key in b.keys() {
7384 if key.starts_with(&lookup_prefix) {
7385 acc |= b.get(&key);
7386 }
7387 }
7388 RowIdSet::from_roaring(acc)
7389 } else {
7390 RowIdSet::empty()
7391 }
7392 }
7393 Condition::FmContains { column_id, pattern } => self
7394 .fm
7395 .get(column_id)
7396 .map(|f| {
7397 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7398 })
7399 .unwrap_or_else(RowIdSet::empty),
7400 Condition::FmContainsAll {
7401 column_id,
7402 patterns,
7403 } => {
7404 if let Some(f) = self.fm.get(column_id) {
7407 let sets: Vec<RowIdSet> = patterns
7408 .iter()
7409 .map(|pat| {
7410 RowIdSet::from_unsorted(
7411 f.locate(pat).into_iter().map(|r| r.0).collect(),
7412 )
7413 })
7414 .collect();
7415 RowIdSet::intersect_many(sets)
7416 } else {
7417 RowIdSet::empty()
7418 }
7419 }
7420 Condition::Ann {
7421 column_id,
7422 query,
7423 k,
7424 } => RowIdSet::from_unsorted(
7425 self.retrieve_filtered(
7426 &crate::query::Retriever::Ann {
7427 column_id: *column_id,
7428 query: query.clone(),
7429 k: *k,
7430 },
7431 snapshot,
7432 None,
7433 allowed,
7434 None,
7435 None,
7436 )?
7437 .into_iter()
7438 .map(|hit| hit.row_id.0)
7439 .collect(),
7440 ),
7441 Condition::SparseMatch {
7442 column_id,
7443 query,
7444 k,
7445 } => RowIdSet::from_unsorted(
7446 self.retrieve_filtered(
7447 &crate::query::Retriever::Sparse {
7448 column_id: *column_id,
7449 query: query.clone(),
7450 k: *k,
7451 },
7452 snapshot,
7453 None,
7454 allowed,
7455 None,
7456 None,
7457 )?
7458 .into_iter()
7459 .map(|hit| hit.row_id.0)
7460 .collect(),
7461 ),
7462 Condition::MinHashSimilar {
7463 column_id,
7464 query,
7465 k,
7466 } => match self.minhash.get(column_id) {
7467 Some(index) => {
7468 let candidates = index.candidate_row_ids(query);
7469 let eligible =
7470 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7471 RowIdSet::from_unsorted(
7472 index
7473 .search_filtered(query, *k, |row_id| {
7474 eligible.contains(&row_id)
7475 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7476 })
7477 .into_iter()
7478 .map(|(row_id, _)| row_id.0)
7479 .collect(),
7480 )
7481 }
7482 None => RowIdSet::empty(),
7483 },
7484 Condition::Range { column_id, lo, hi } => {
7485 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7494 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7495 } else if self.run_refs.len() == 1 {
7496 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7497 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7498 } else {
7499 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7500 };
7501 set.remove_many(self.overlay_rid_set(snapshot));
7502 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7503 set
7504 }
7505 Condition::RangeF64 {
7506 column_id,
7507 lo,
7508 lo_inclusive,
7509 hi,
7510 hi_inclusive,
7511 } => {
7512 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7515 RowIdSet::from_unsorted(
7516 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7517 .into_iter()
7518 .collect(),
7519 )
7520 } else if self.run_refs.len() == 1 {
7521 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7522 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7523 } else {
7524 return self.range_scan_f64(
7525 *column_id,
7526 *lo,
7527 *lo_inclusive,
7528 *hi,
7529 *hi_inclusive,
7530 snapshot,
7531 );
7532 };
7533 set.remove_many(self.overlay_rid_set(snapshot));
7534 self.range_scan_overlay_f64(
7535 &mut set,
7536 *column_id,
7537 *lo,
7538 *lo_inclusive,
7539 *hi,
7540 *hi_inclusive,
7541 snapshot,
7542 );
7543 set
7544 }
7545 Condition::IsNull { column_id } => {
7546 let mut set = if self.run_refs.len() == 1 {
7547 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7548 r.null_row_id_set(*column_id, true)?
7549 } else {
7550 return self.null_scan(*column_id, true, snapshot);
7551 };
7552 set.remove_many(self.overlay_rid_set(snapshot));
7553 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7554 set
7555 }
7556 Condition::IsNotNull { column_id } => {
7557 let mut set = if self.run_refs.len() == 1 {
7558 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7559 r.null_row_id_set(*column_id, false)?
7560 } else {
7561 return self.null_scan(*column_id, false, snapshot);
7562 };
7563 set.remove_many(self.overlay_rid_set(snapshot));
7564 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7565 set
7566 }
7567 })
7568 }
7569
7570 fn range_scan_i64(
7578 &self,
7579 column_id: u16,
7580 lo: i64,
7581 hi: i64,
7582 snapshot: Snapshot,
7583 ) -> Result<RowIdSet> {
7584 let mut row_ids = Vec::new();
7585 let overlay_rids = self.overlay_rid_set(snapshot);
7586 for rr in &self.run_refs {
7587 let mut reader = self.open_reader(rr.run_id)?;
7588 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7589 for rid in matched {
7590 if !overlay_rids.contains(&rid) {
7591 row_ids.push(rid);
7592 }
7593 }
7594 }
7595 let mut s = RowIdSet::from_unsorted(row_ids);
7596 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7597 Ok(s)
7598 }
7599
7600 fn range_scan_f64(
7603 &self,
7604 column_id: u16,
7605 lo: f64,
7606 lo_inclusive: bool,
7607 hi: f64,
7608 hi_inclusive: bool,
7609 snapshot: Snapshot,
7610 ) -> Result<RowIdSet> {
7611 let mut row_ids = Vec::new();
7612 let overlay_rids = self.overlay_rid_set(snapshot);
7613 for rr in &self.run_refs {
7614 let mut reader = self.open_reader(rr.run_id)?;
7615 let matched = reader.range_row_ids_visible_f64(
7616 column_id,
7617 lo,
7618 lo_inclusive,
7619 hi,
7620 hi_inclusive,
7621 snapshot.epoch,
7622 )?;
7623 for rid in matched {
7624 if !overlay_rids.contains(&rid) {
7625 row_ids.push(rid);
7626 }
7627 }
7628 }
7629 let mut s = RowIdSet::from_unsorted(row_ids);
7630 self.range_scan_overlay_f64(
7631 &mut s,
7632 column_id,
7633 lo,
7634 lo_inclusive,
7635 hi,
7636 hi_inclusive,
7637 snapshot,
7638 );
7639 Ok(s)
7640 }
7641
7642 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7644 let mut s = HashSet::new();
7645 for row in self.memtable.visible_versions(snapshot.epoch) {
7646 s.insert(row.row_id.0);
7647 }
7648 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7649 s.insert(row.row_id.0);
7650 }
7651 s
7652 }
7653
7654 fn range_scan_overlay_i64(
7655 &self,
7656 s: &mut RowIdSet,
7657 column_id: u16,
7658 lo: i64,
7659 hi: i64,
7660 snapshot: Snapshot,
7661 ) {
7662 let mut newest: HashMap<u64, &Row> = HashMap::new();
7667 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7668 let memtable = self.memtable.visible_versions(snapshot.epoch);
7669 for r in &mutable {
7670 newest.entry(r.row_id.0).or_insert(r);
7671 }
7672 for r in &memtable {
7673 newest.insert(r.row_id.0, r);
7674 }
7675 for row in newest.values() {
7676 if !row.deleted {
7677 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7678 if *v >= lo && *v <= hi {
7679 s.insert(row.row_id.0);
7680 }
7681 }
7682 }
7683 }
7684 }
7685
7686 #[allow(clippy::too_many_arguments)]
7687 fn range_scan_overlay_f64(
7688 &self,
7689 s: &mut RowIdSet,
7690 column_id: u16,
7691 lo: f64,
7692 lo_inclusive: bool,
7693 hi: f64,
7694 hi_inclusive: bool,
7695 snapshot: Snapshot,
7696 ) {
7697 let mut newest: HashMap<u64, &Row> = HashMap::new();
7700 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7701 let memtable = self.memtable.visible_versions(snapshot.epoch);
7702 for r in &mutable {
7703 newest.entry(r.row_id.0).or_insert(r);
7704 }
7705 for r in &memtable {
7706 newest.insert(r.row_id.0, r);
7707 }
7708 for row in newest.values() {
7709 if !row.deleted {
7710 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7711 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7712 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7713 if ok_lo && ok_hi {
7714 s.insert(row.row_id.0);
7715 }
7716 }
7717 }
7718 }
7719 }
7720
7721 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7724 let mut row_ids = Vec::new();
7725 let overlay_rids = self.overlay_rid_set(snapshot);
7726 for rr in &self.run_refs {
7727 let mut reader = self.open_reader(rr.run_id)?;
7728 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7729 for rid in matched {
7730 if !overlay_rids.contains(&rid) {
7731 row_ids.push(rid);
7732 }
7733 }
7734 }
7735 let mut s = RowIdSet::from_unsorted(row_ids);
7736 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7737 Ok(s)
7738 }
7739
7740 fn null_scan_overlay(
7744 &self,
7745 s: &mut RowIdSet,
7746 column_id: u16,
7747 want_nulls: bool,
7748 snapshot: Snapshot,
7749 ) {
7750 let mut newest: HashMap<u64, &Row> = HashMap::new();
7751 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7752 let memtable = self.memtable.visible_versions(snapshot.epoch);
7753 for r in &mutable {
7754 newest.entry(r.row_id.0).or_insert(r);
7755 }
7756 for r in &memtable {
7757 newest.insert(r.row_id.0, r);
7758 }
7759 for row in newest.values() {
7760 if row.deleted {
7761 continue;
7762 }
7763 let is_null = !row.columns.contains_key(&column_id)
7764 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7765 if is_null == want_nulls {
7766 s.insert(row.row_id.0);
7767 }
7768 }
7769 }
7770
7771 pub fn snapshot(&self) -> Snapshot {
7772 Snapshot::at(self.epoch.visible())
7773 }
7774
7775 pub fn data_generation(&self) -> u64 {
7777 self.data_generation
7778 }
7779
7780 pub(crate) fn bump_data_generation(&mut self) {
7781 self.data_generation = self.data_generation.wrapping_add(1);
7782 }
7783
7784 pub fn table_id(&self) -> u64 {
7786 self.table_id
7787 }
7788
7789 fn seal_generations(&mut self) {
7795 self.memtable.seal();
7796 self.mutable_run.seal();
7797 self.hot.seal();
7798 for index in self.bitmap.values_mut() {
7799 index.seal();
7800 }
7801 for index in self.ann.values_mut() {
7802 index.seal();
7803 }
7804 for index in self.fm.values_mut() {
7805 index.seal();
7806 }
7807 for index in self.sparse.values_mut() {
7808 index.seal();
7809 }
7810 for index in self.minhash.values_mut() {
7811 index.seal();
7812 }
7813 self.pk_by_row.seal();
7814 }
7815
7816 fn capture_read_generation(&self) -> ReadGeneration {
7821 let visible_through = self.current_epoch();
7822 ReadGeneration {
7823 schema: Arc::new(self.schema.clone()),
7824 base_runs: Arc::new(self.run_refs.clone()),
7825 deltas: TableDeltas {
7826 memtable: self.memtable.clone(),
7827 mutable_run: self.mutable_run.clone(),
7828 hot: self.hot.clone(),
7829 pk_by_row: self.pk_by_row.clone(),
7830 },
7831 indexes: Arc::new(IndexGeneration::capture(
7832 &self.bitmap,
7833 &self.learned_range,
7834 &self.fm,
7835 &self.ann,
7836 &self.sparse,
7837 &self.minhash,
7838 visible_through,
7839 )),
7840 visible_through,
7841 }
7842 }
7843
7844 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
7849 self.ensure_indexes_complete()?;
7850 self.seal_generations();
7851 let view = Arc::new(self.capture_read_generation());
7852 self.published.store(Arc::clone(&view));
7853 Ok(view)
7854 }
7855
7856 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
7862 self.published.load_full()
7863 }
7864
7865 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
7867 &self.pins
7868 }
7869
7870 pub fn version_gc_floor(&self) -> Epoch {
7875 self.min_active_snapshot()
7876 .unwrap_or_else(|| self.current_epoch())
7877 }
7878
7879 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
7886 let mut report = self.pins.report();
7887 let transaction_floor = [
7888 self.pinned.keys().next().copied(),
7889 self.snapshots.min_pinned(),
7890 ]
7891 .into_iter()
7892 .flatten()
7893 .min();
7894 if let Some(epoch) = transaction_floor {
7895 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
7896 }
7897 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
7898 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
7899 }
7900 report
7901 }
7902
7903 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7904 self.publish_read_generation()?;
7905 let mut generation = self.clone();
7906 generation.read_only = true;
7907 generation.wal = WalSink::ReadOnly;
7908 generation.pending_delete_rids.clear();
7909 generation.pending_put_cols.clear();
7910 generation.pending_rows.clear();
7911 generation.pending_rows_auto_inc.clear();
7912 generation.pending_dels.clear();
7913 generation.pending_truncate = None;
7914 generation.agg_cache = Arc::new(HashMap::new());
7915 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
7918 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
7921 crate::retention::PinSource::ReadGeneration,
7922 self.current_epoch(),
7923 )));
7924 Ok(generation)
7925 }
7926
7927 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7928 (std::mem::size_of::<Self>() as u64)
7929 .saturating_add(self.memtable.approx_bytes())
7930 .saturating_add(self.mutable_run.approx_bytes())
7931 .saturating_add(self.live_count.saturating_mul(64))
7932 }
7933
7934 pub fn pin_snapshot(&mut self) -> Snapshot {
7937 let e = self.epoch.visible();
7938 *self.pinned.entry(e).or_insert(0) += 1;
7939 Snapshot::at(e)
7940 }
7941
7942 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7944 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7945 *count -= 1;
7946 if *count == 0 {
7947 self.pinned.remove(&snap.epoch);
7948 }
7949 }
7950 }
7951
7952 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7964 let local = self.pinned.keys().next().copied();
7965 let global = self.snapshots.min_pinned();
7966 let history = self.snapshots.history_floor(self.current_epoch());
7967 let pinned = self.pins.oldest_pinned();
7968 [local, global, history, pinned].into_iter().flatten().min()
7969 }
7970
7971 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7975 self.ensure_writable()?;
7976 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7977 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7978 }
7979
7980 pub fn clear_ttl(&mut self) -> Result<()> {
7981 self.ensure_writable()?;
7982 self.apply_ttl_policy_at(None, self.current_epoch())
7983 }
7984
7985 pub fn ttl(&self) -> Option<TtlPolicy> {
7986 self.ttl
7987 }
7988
7989 pub(crate) fn prepare_ttl_policy(
7990 &self,
7991 column_name: &str,
7992 duration_nanos: u64,
7993 ) -> Result<TtlPolicy> {
7994 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7995 return Err(MongrelError::InvalidArgument(
7996 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7997 ));
7998 }
7999 let column = self
8000 .schema
8001 .columns
8002 .iter()
8003 .find(|column| column.name == column_name)
8004 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
8005 if column.ty != TypeId::TimestampNanos {
8006 return Err(MongrelError::Schema(format!(
8007 "TTL column {column_name} must be TimestampNanos, is {:?}",
8008 column.ty
8009 )));
8010 }
8011 Ok(TtlPolicy {
8012 column_id: column.id,
8013 duration_nanos,
8014 })
8015 }
8016
8017 pub(crate) fn apply_ttl_policy_at(
8018 &mut self,
8019 policy: Option<TtlPolicy>,
8020 epoch: Epoch,
8021 ) -> Result<()> {
8022 if let Some(policy) = policy {
8023 let column = self
8024 .schema
8025 .columns
8026 .iter()
8027 .find(|column| column.id == policy.column_id)
8028 .ok_or_else(|| {
8029 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8030 })?;
8031 if column.ty != TypeId::TimestampNanos
8032 || policy.duration_nanos == 0
8033 || policy.duration_nanos > i64::MAX as u64
8034 {
8035 return Err(MongrelError::Schema("invalid TTL policy".into()));
8036 }
8037 }
8038 self.ttl = policy;
8039 self.agg_cache = Arc::new(HashMap::new());
8040 self.clear_result_cache();
8041 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8042 self.persist_manifest(epoch)
8043 }
8044
8045 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8046 let Some(policy) = self.ttl else {
8047 return false;
8048 };
8049 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8050 return false;
8051 };
8052 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8053 }
8054
8055 pub fn current_epoch(&self) -> Epoch {
8056 self.epoch.visible()
8057 }
8058
8059 pub fn memtable_len(&self) -> usize {
8060 self.memtable.len()
8061 }
8062
8063 pub fn count(&self) -> u64 {
8066 if self.ttl.is_none()
8067 && self.pending_put_cols.is_empty()
8068 && self.pending_delete_rids.is_empty()
8069 && self.pending_rows.is_empty()
8070 && self.pending_dels.is_empty()
8071 && self.pending_truncate.is_none()
8072 {
8073 self.live_count
8074 } else {
8075 self.visible_rows(self.snapshot())
8076 .map(|rows| rows.len() as u64)
8077 .unwrap_or(self.live_count)
8078 }
8079 }
8080
8081 pub fn count_conditions(
8085 &mut self,
8086 conditions: &[crate::query::Condition],
8087 snapshot: Snapshot,
8088 ) -> Result<Option<u64>> {
8089 use crate::query::Condition;
8090 if self.ttl.is_some() {
8091 if conditions.is_empty() {
8092 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8093 }
8094 let mut sets = Vec::with_capacity(conditions.len());
8095 for condition in conditions {
8096 sets.push(self.resolve_condition(condition, snapshot)?);
8097 }
8098 let survivors = RowIdSet::intersect_many(sets);
8099 let rows = self.visible_rows(snapshot)?;
8100 return Ok(Some(
8101 rows.into_iter()
8102 .filter(|row| survivors.contains(row.row_id.0))
8103 .count() as u64,
8104 ));
8105 }
8106 if conditions.is_empty() {
8107 return Ok(Some(self.count()));
8108 }
8109 let served = |c: &Condition| {
8110 matches!(
8111 c,
8112 Condition::Pk(_)
8113 | Condition::BitmapEq { .. }
8114 | Condition::BitmapIn { .. }
8115 | Condition::BytesPrefix { .. }
8116 | Condition::FmContains { .. }
8117 | Condition::FmContainsAll { .. }
8118 | Condition::Ann { .. }
8119 | Condition::Range { .. }
8120 | Condition::RangeF64 { .. }
8121 | Condition::SparseMatch { .. }
8122 | Condition::MinHashSimilar { .. }
8123 | Condition::IsNull { .. }
8124 | Condition::IsNotNull { .. }
8125 )
8126 };
8127 if !conditions.iter().all(served) {
8128 return Ok(None);
8129 }
8130 self.ensure_indexes_complete()?;
8131 if !self.pending_put_cols.is_empty()
8132 || !self.pending_delete_rids.is_empty()
8133 || !self.pending_rows.is_empty()
8134 || !self.pending_dels.is_empty()
8135 || self.pending_truncate.is_some()
8136 {
8137 let mut sets = Vec::with_capacity(conditions.len());
8138 for condition in conditions {
8139 sets.push(self.resolve_condition(condition, snapshot)?);
8140 }
8141 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8142 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8143 }
8144 let mut sets = Vec::with_capacity(conditions.len());
8145 for condition in conditions {
8146 sets.push(self.resolve_condition(condition, snapshot)?);
8147 }
8148 let mut rids = RowIdSet::intersect_many(sets);
8149 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8159 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8160 }
8161 let count = rids.len() as u64;
8162 crate::trace::QueryTrace::record(|t| {
8163 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8164 t.survivor_count = Some(count as usize);
8165 t.conditions_pushed = conditions.len();
8166 });
8167 Ok(Some(count))
8168 }
8169
8170 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8175 let mut out = Vec::new();
8176 for row in self.memtable.visible_versions(snapshot.epoch) {
8177 if row.deleted {
8178 out.push(row.row_id.0);
8179 }
8180 }
8181 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8182 if row.deleted {
8183 out.push(row.row_id.0);
8184 }
8185 }
8186 out
8187 }
8188
8189 pub fn bulk_load_columns(
8198 &mut self,
8199 user_columns: Vec<(u16, columnar::NativeColumn)>,
8200 ) -> Result<Epoch> {
8201 self.bulk_load_columns_with(user_columns, 3, false, true)
8202 }
8203
8204 pub fn bulk_load_fast(
8211 &mut self,
8212 user_columns: Vec<(u16, columnar::NativeColumn)>,
8213 ) -> Result<Epoch> {
8214 self.bulk_load_columns_with(user_columns, -1, true, false)
8215 }
8216
8217 fn bulk_load_columns_with(
8218 &mut self,
8219 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8220 zstd_level: i32,
8221 force_plain: bool,
8222 lz4: bool,
8223 ) -> Result<Epoch> {
8224 self.ensure_writable()?;
8225 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8226 if n == 0 {
8227 return Ok(self.current_epoch());
8228 }
8229 let epoch = self.commit_new_epoch()?;
8230 let live_before = self.live_count;
8231 self.spill_mutable_run(epoch)?;
8233 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8234 && self.indexes_complete
8235 && self.run_refs.is_empty()
8236 && self.memtable.is_empty()
8237 && self.mutable_run.is_empty();
8238 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8241 self.validate_columns_not_null(&user_columns, n)?;
8242 let winner_idx = self
8243 .bulk_pk_winner_indices(&user_columns, n)
8244 .filter(|idx| idx.len() != n);
8245 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8246 match winner_idx.as_deref() {
8247 Some(idx) => {
8248 let compacted = user_columns
8249 .iter()
8250 .map(|(id, c)| (*id, c.gather(idx)))
8251 .collect();
8252 (compacted, idx.len())
8253 }
8254 None => (user_columns, n),
8255 };
8256 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8257 let first = self.allocator.alloc_range(write_n as u64)?.0;
8258 for rid in first..first + write_n as u64 {
8259 self.reservoir.offer(rid);
8260 }
8261 let run_id = self.alloc_run_id()?;
8262 let path = self.run_path(run_id);
8263 let mut writer =
8264 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8265 if force_plain {
8266 writer = writer.with_plain();
8267 } else if lz4 {
8268 writer = writer.with_lz4();
8271 } else {
8272 writer = writer.with_zstd_level(zstd_level);
8273 }
8274 if let Some(kek) = &self.kek {
8275 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8276 }
8277 let header = match self.create_run_file(run_id)? {
8278 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8279 None => writer.write_native(&path, &write_columns, write_n, first)?,
8280 };
8281 self.run_refs.push(RunRef {
8282 run_id: run_id as u128,
8283 level: 0,
8284 epoch_created: epoch.0,
8285 row_count: header.row_count,
8286 });
8287 self.live_count = self.live_count.saturating_add(write_n as u64);
8288 if eager_index_build {
8289 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8290 self.index_columns_bulk(&write_columns, &row_ids);
8291 self.indexes_complete = true;
8292 self.build_learned_ranges()?;
8293 } else {
8294 self.indexes_complete = false;
8298 }
8299 self.mark_flushed(epoch)?;
8300 self.persist_manifest(epoch)?;
8301 if eager_index_build {
8302 self.checkpoint_indexes(epoch);
8303 }
8304 self.clear_result_cache();
8305 self.data_generation = self.data_generation.wrapping_add(1);
8306 Ok(epoch)
8307 }
8308
8309 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8327 let n = row_ids.len();
8328 if n == 0 {
8329 return;
8330 }
8331 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8332 columns.iter().map(|(id, c)| (*id, c)).collect();
8333 let ty_of: std::collections::HashMap<u16, TypeId> = self
8334 .schema
8335 .columns
8336 .iter()
8337 .map(|c| (c.id, c.ty.clone()))
8338 .collect();
8339 let pk_id = self.schema.primary_key().map(|c| c.id);
8340
8341 for (i, &rid) in row_ids.iter().enumerate() {
8342 let row_id = RowId(rid);
8343 if let Some(pid) = pk_id {
8344 if let Some(col) = by_id.get(&pid) {
8345 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8346 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8347 self.insert_hot_pk(key, row_id);
8348 }
8349 }
8350 }
8351 for idef in &self.schema.indexes {
8352 let Some(col) = by_id.get(&idef.column_id) else {
8353 continue;
8354 };
8355 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8356 match idef.kind {
8357 IndexKind::Bitmap => {
8358 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8359 if let Some(key) =
8360 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8361 {
8362 b.insert(key, row_id);
8363 }
8364 }
8365 }
8366 IndexKind::FmIndex => {
8367 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8368 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8369 f.insert(bytes.to_vec(), row_id);
8370 }
8371 }
8372 }
8373 IndexKind::Sparse => {
8374 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8375 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8376 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8377 s.insert(&terms, row_id);
8378 }
8379 }
8380 }
8381 }
8382 IndexKind::MinHash => {
8383 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8384 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8385 let tokens = crate::index::token_hashes_from_bytes(bytes);
8386 mh.insert(&tokens, row_id);
8387 }
8388 }
8389 }
8390 _ => {}
8391 }
8392 }
8393 }
8394 }
8395
8396 pub fn visible_columns_native(
8401 &self,
8402 snapshot: Snapshot,
8403 projection: Option<&[u16]>,
8404 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8405 self.visible_columns_native_inner(snapshot, projection, None)
8406 }
8407
8408 pub fn visible_columns_native_with_control(
8409 &self,
8410 snapshot: Snapshot,
8411 projection: Option<&[u16]>,
8412 control: &crate::ExecutionControl,
8413 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8414 self.visible_columns_native_inner(snapshot, projection, Some(control))
8415 }
8416
8417 fn visible_columns_native_inner(
8418 &self,
8419 snapshot: Snapshot,
8420 projection: Option<&[u16]>,
8421 control: Option<&crate::ExecutionControl>,
8422 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8423 execution_checkpoint(control, 0)?;
8424 let wanted: Vec<u16> = match projection {
8425 Some(p) => p.to_vec(),
8426 None => self.schema.columns.iter().map(|c| c.id).collect(),
8427 };
8428 if self.ttl.is_none()
8429 && self.memtable.is_empty()
8430 && self.mutable_run.is_empty()
8431 && self.run_refs.len() == 1
8432 {
8433 let rr = self.run_refs[0].clone();
8434 let mut reader = self.open_reader(rr.run_id)?;
8435 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8436 execution_checkpoint(control, 0)?;
8437 let all_visible = idxs.len() == reader.row_count();
8438 if reader.has_mmap() && control.is_none() {
8444 use rayon::prelude::*;
8445 let valid: Vec<u16> = wanted
8448 .iter()
8449 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8450 .copied()
8451 .collect();
8452 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8454 .par_iter()
8455 .filter_map(|cid| {
8456 reader
8457 .column_native_shared(*cid)
8458 .ok()
8459 .map(|col| (*cid, col))
8460 })
8461 .collect();
8462 let cols = decoded
8463 .into_iter()
8464 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8465 .collect();
8466 return Ok(cols);
8467 }
8468 let mut cols = Vec::with_capacity(wanted.len());
8469 for (index, cid) in wanted.iter().enumerate() {
8470 execution_checkpoint(control, index)?;
8471 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8472 Some(c) => c,
8473 None => continue,
8474 };
8475 let col = reader.column_native(cdef.id)?;
8476 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8477 }
8478 return Ok(cols);
8479 }
8480 let vcols = self.visible_columns(snapshot)?;
8481 execution_checkpoint(control, 0)?;
8482 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8483 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8484 .into_iter()
8485 .filter(|(id, _)| want_set.contains(id))
8486 .map(|(id, vals)| {
8487 let ty = self
8488 .schema
8489 .columns
8490 .iter()
8491 .find(|c| c.id == id)
8492 .map(|c| c.ty.clone())
8493 .unwrap_or(TypeId::Bytes);
8494 (id, columnar::values_to_native(ty, &vals))
8495 })
8496 .collect();
8497 Ok(out)
8498 }
8499
8500 pub fn run_count(&self) -> usize {
8501 self.run_refs.len()
8502 }
8503
8504 pub fn memtable_is_empty(&self) -> bool {
8506 self.memtable.is_empty()
8507 }
8508
8509 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8513 self.page_cache.stats()
8514 }
8515
8516 pub fn reset_page_cache_stats(&self) {
8518 self.page_cache.reset_stats();
8519 }
8520
8521 pub fn run_ids(&self) -> Vec<u128> {
8524 self.run_refs.iter().map(|r| r.run_id).collect()
8525 }
8526
8527 pub fn single_run_is_clean(&self) -> bool {
8531 if self.ttl.is_some() || self.run_refs.len() != 1 {
8532 return false;
8533 }
8534 self.open_reader(self.run_refs[0].run_id)
8535 .map(|r| r.is_clean())
8536 .unwrap_or(false)
8537 }
8538
8539 fn resolve_footprint(
8546 &self,
8547 conditions: &[crate::query::Condition],
8548 snapshot: Snapshot,
8549 ) -> roaring::RoaringBitmap {
8550 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8551 return roaring::RoaringBitmap::new();
8552 }
8553 if self.run_refs.is_empty() {
8554 return roaring::RoaringBitmap::new();
8555 }
8556 if self.run_refs.len() == 1 {
8558 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8559 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8560 return rids.to_roaring_lossy();
8561 }
8562 }
8563 }
8564 roaring::RoaringBitmap::new()
8565 }
8566
8567 pub fn query_columns_native_cached(
8578 &mut self,
8579 conditions: &[crate::query::Condition],
8580 projection: Option<&[u16]>,
8581 snapshot: Snapshot,
8582 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8583 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8584 }
8585
8586 pub fn query_columns_native_cached_with_control(
8587 &mut self,
8588 conditions: &[crate::query::Condition],
8589 projection: Option<&[u16]>,
8590 snapshot: Snapshot,
8591 control: &crate::ExecutionControl,
8592 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8593 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8594 }
8595
8596 fn query_columns_native_cached_inner(
8597 &mut self,
8598 conditions: &[crate::query::Condition],
8599 projection: Option<&[u16]>,
8600 snapshot: Snapshot,
8601 control: Option<&crate::ExecutionControl>,
8602 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8603 execution_checkpoint(control, 0)?;
8604 if self.ttl.is_some() {
8607 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8608 }
8609 if conditions.is_empty() {
8610 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8611 }
8612 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8616 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8617 crate::trace::QueryTrace::record(|t| {
8618 t.result_cache_hit = true;
8619 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8620 });
8621 return Ok(Some((*hit).clone()));
8622 }
8623 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8624 execution_checkpoint(control, 0)?;
8625 if let Some(cols) = &res {
8626 let footprint = self.resolve_footprint(conditions, snapshot);
8627 let condition_cols = crate::query::condition_columns(conditions);
8628 execution_checkpoint(control, 0)?;
8629 self.result_cache.lock().insert(
8630 key,
8631 CachedEntry {
8632 data: CachedData::Columns(Arc::new(cols.clone())),
8633 footprint,
8634 condition_cols,
8635 },
8636 );
8637 }
8638 Ok(res)
8639 }
8640
8641 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8646 if self.ttl.is_some() {
8647 return self.query(q);
8648 }
8649 if q.conditions.is_empty() {
8650 return self.query(q);
8651 }
8652 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8653 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8654 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8655 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8656 crate::trace::QueryTrace::record(|t| {
8657 t.result_cache_hit = true;
8658 t.scan_mode = crate::trace::ScanMode::Materialized;
8659 });
8660 return Ok((*hit).clone());
8661 }
8662 let rows = self.query(q)?;
8663 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8664 let condition_cols = crate::query::condition_columns(&q.conditions);
8665 self.result_cache.lock().insert(
8666 key,
8667 CachedEntry {
8668 data: CachedData::Rows(Arc::new(rows.clone())),
8669 footprint,
8670 condition_cols,
8671 },
8672 );
8673 Ok(rows)
8674 }
8675
8676 #[allow(clippy::type_complexity)]
8691 pub fn query_columns_native_traced(
8692 &mut self,
8693 conditions: &[crate::query::Condition],
8694 projection: Option<&[u16]>,
8695 snapshot: Snapshot,
8696 ) -> Result<(
8697 Option<Vec<(u16, columnar::NativeColumn)>>,
8698 crate::trace::QueryTrace,
8699 )> {
8700 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8701 self.query_columns_native(conditions, projection, snapshot)
8702 });
8703 Ok((result?, trace))
8704 }
8705
8706 #[allow(clippy::type_complexity)]
8709 pub fn query_columns_native_cached_traced(
8710 &mut self,
8711 conditions: &[crate::query::Condition],
8712 projection: Option<&[u16]>,
8713 snapshot: Snapshot,
8714 ) -> Result<(
8715 Option<Vec<(u16, columnar::NativeColumn)>>,
8716 crate::trace::QueryTrace,
8717 )> {
8718 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8719 self.query_columns_native_cached(conditions, projection, snapshot)
8720 });
8721 Ok((result?, trace))
8722 }
8723
8724 pub fn native_page_cursor_traced(
8726 &self,
8727 snapshot: Snapshot,
8728 projection: Vec<(u16, TypeId)>,
8729 conditions: &[crate::query::Condition],
8730 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8731 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8732 self.native_page_cursor(snapshot, projection, conditions)
8733 });
8734 Ok((result?, trace))
8735 }
8736
8737 pub fn native_multi_run_cursor_traced(
8739 &self,
8740 snapshot: Snapshot,
8741 projection: Vec<(u16, TypeId)>,
8742 conditions: &[crate::query::Condition],
8743 ) -> Result<(
8744 Option<crate::cursor::MultiRunCursor>,
8745 crate::trace::QueryTrace,
8746 )> {
8747 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8748 self.native_multi_run_cursor(snapshot, projection, conditions)
8749 });
8750 Ok((result?, trace))
8751 }
8752
8753 pub fn count_conditions_traced(
8755 &mut self,
8756 conditions: &[crate::query::Condition],
8757 snapshot: Snapshot,
8758 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8759 let (result, trace) =
8760 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8761 Ok((result?, trace))
8762 }
8763
8764 pub fn query_traced(
8766 &mut self,
8767 q: &crate::query::Query,
8768 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8769 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8770 Ok((result?, trace))
8771 }
8772
8773 pub fn query_columns_native(
8778 &mut self,
8779 conditions: &[crate::query::Condition],
8780 projection: Option<&[u16]>,
8781 snapshot: Snapshot,
8782 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8783 self.query_columns_native_inner(conditions, projection, snapshot, None)
8784 }
8785
8786 pub fn query_columns_native_with_control(
8787 &mut self,
8788 conditions: &[crate::query::Condition],
8789 projection: Option<&[u16]>,
8790 snapshot: Snapshot,
8791 control: &crate::ExecutionControl,
8792 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8793 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8794 }
8795
8796 fn query_columns_native_inner(
8797 &mut self,
8798 conditions: &[crate::query::Condition],
8799 projection: Option<&[u16]>,
8800 snapshot: Snapshot,
8801 control: Option<&crate::ExecutionControl>,
8802 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8803 use crate::query::Condition;
8804 execution_checkpoint(control, 0)?;
8805 if self.ttl.is_some() {
8808 return Ok(None);
8809 }
8810 if conditions.is_empty() {
8811 return Ok(None);
8812 }
8813 self.ensure_indexes_complete()?;
8814
8815 let served = |c: &Condition| {
8820 matches!(
8821 c,
8822 Condition::Pk(_)
8823 | Condition::BitmapEq { .. }
8824 | Condition::BitmapIn { .. }
8825 | Condition::BytesPrefix { .. }
8826 | Condition::FmContains { .. }
8827 | Condition::FmContainsAll { .. }
8828 | Condition::Ann { .. }
8829 | Condition::Range { .. }
8830 | Condition::RangeF64 { .. }
8831 | Condition::SparseMatch { .. }
8832 | Condition::MinHashSimilar { .. }
8833 | Condition::IsNull { .. }
8834 | Condition::IsNotNull { .. }
8835 )
8836 };
8837 if !conditions.iter().all(served) {
8838 return Ok(None);
8839 }
8840 let fast_path =
8841 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8842 crate::trace::QueryTrace::record(|t| {
8843 t.run_count = self.run_refs.len();
8844 t.memtable_rows = self.memtable.len();
8845 t.mutable_run_rows = self.mutable_run.len();
8846 t.conditions_pushed = conditions.len();
8847 t.learned_range_used = conditions.iter().any(|c| match c {
8848 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8849 self.learned_range.contains_key(column_id)
8850 }
8851 _ => false,
8852 });
8853 });
8854 let col_ids: Vec<u16> = projection
8856 .map(|p| p.to_vec())
8857 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8858 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8859 .iter()
8860 .map(|&cid| {
8861 let ty = self
8862 .schema
8863 .columns
8864 .iter()
8865 .find(|c| c.id == cid)
8866 .map(|c| c.ty.clone())
8867 .unwrap_or(TypeId::Bytes);
8868 (cid, ty)
8869 })
8870 .collect();
8871
8872 if fast_path {
8878 let needs_column = conditions.iter().any(|c| match c {
8881 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8882 Condition::RangeF64 { column_id, .. } => {
8883 !self.learned_range.contains_key(column_id)
8884 }
8885 _ => false,
8886 });
8887 let mut reader_opt: Option<RunReader> = if needs_column {
8888 Some(self.open_reader(self.run_refs[0].run_id)?)
8889 } else {
8890 None
8891 };
8892 let mut sets: Vec<RowIdSet> = Vec::new();
8893 for (index, c) in conditions.iter().enumerate() {
8894 execution_checkpoint(control, index)?;
8895 let s = match c {
8896 Condition::Range { column_id, lo, hi }
8897 if !self.learned_range.contains_key(column_id) =>
8898 {
8899 if reader_opt.is_none() {
8900 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8901 }
8902 reader_opt
8903 .as_mut()
8904 .expect("reader opened for range")
8905 .range_row_id_set_i64(*column_id, *lo, *hi)?
8906 }
8907 Condition::RangeF64 {
8908 column_id,
8909 lo,
8910 lo_inclusive,
8911 hi,
8912 hi_inclusive,
8913 } if !self.learned_range.contains_key(column_id) => {
8914 if reader_opt.is_none() {
8915 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8916 }
8917 reader_opt
8918 .as_mut()
8919 .expect("reader opened for range")
8920 .range_row_id_set_f64(
8921 *column_id,
8922 *lo,
8923 *lo_inclusive,
8924 *hi,
8925 *hi_inclusive,
8926 )?
8927 }
8928 _ => self.resolve_condition(c, snapshot)?,
8929 };
8930 sets.push(s);
8931 }
8932 let candidates = RowIdSet::intersect_many(sets);
8933 crate::trace::QueryTrace::record(|t| {
8934 t.survivor_count = Some(candidates.len());
8935 });
8936 if candidates.is_empty() {
8937 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8938 .iter()
8939 .map(|&id| {
8940 (
8941 id,
8942 columnar::null_native(
8943 proj_pairs
8944 .iter()
8945 .find(|(c, _)| c == &id)
8946 .map(|(_, t)| t.clone())
8947 .unwrap_or(TypeId::Bytes),
8948 0,
8949 ),
8950 )
8951 })
8952 .collect();
8953 return Ok(Some(cols));
8954 }
8955 let mut reader = match reader_opt.take() {
8956 Some(r) => r,
8957 None => self.open_reader(self.run_refs[0].run_id)?,
8958 };
8959 let candidate_ids = candidates.into_sorted_vec();
8960 let (positions, fast_rid) = if let Some(positions) =
8961 reader.positions_for_row_ids_fast(&candidate_ids)
8962 {
8963 (positions, true)
8964 } else {
8965 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8966 match col {
8967 columnar::NativeColumn::Int64 { data, .. } => {
8968 let mut p = Vec::with_capacity(candidate_ids.len());
8969 for (index, rid) in candidate_ids.iter().enumerate() {
8970 execution_checkpoint(control, index)?;
8971 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8972 p.push(position);
8973 }
8974 }
8975 p.sort_unstable();
8976 (p, false)
8977 }
8978 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8979 }
8980 };
8981 crate::trace::QueryTrace::record(|t| {
8982 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8983 t.fast_row_id_map = fast_rid;
8984 });
8985 let mut cols = Vec::with_capacity(col_ids.len());
8986 for (index, cid) in col_ids.iter().enumerate() {
8987 execution_checkpoint(control, index)?;
8988 let col = reader.column_native(*cid)?;
8989 cols.push((*cid, col.gather(&positions)));
8990 }
8991 return Ok(Some(cols));
8992 }
8993
8994 if !self.run_refs.is_empty() {
9007 use crate::cursor::{
9008 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
9009 };
9010 let remaining: usize;
9011 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9012 let c = self
9013 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9014 .expect("single-run cursor should build when run_refs.len() == 1");
9015 remaining = c.remaining_rows();
9016 Box::new(c)
9017 } else {
9018 let c = self
9019 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9020 .expect("multi-run cursor should build when run_refs.len() >= 1");
9021 remaining = c.remaining_rows();
9022 Box::new(c)
9023 };
9024 crate::trace::QueryTrace::record(|t| {
9025 if t.survivor_count.is_none() {
9026 t.survivor_count = Some(remaining);
9027 }
9028 });
9029 let cols = match control {
9030 Some(control) => {
9031 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9032 }
9033 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9034 };
9035 return Ok(Some(cols));
9036 }
9037
9038 crate::trace::QueryTrace::record(|t| {
9043 t.scan_mode = crate::trace::ScanMode::Materialized;
9044 t.row_materialized = true;
9045 });
9046 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9047 for (index, c) in conditions.iter().enumerate() {
9048 execution_checkpoint(control, index)?;
9049 sets.push(self.resolve_condition(c, snapshot)?);
9050 }
9051 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9052 let rows = self.rows_for_rids(&rids, snapshot)?;
9053 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9054 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9055 execution_checkpoint(control, index)?;
9056 let vals: Vec<Value> = rows
9057 .iter()
9058 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9059 .collect();
9060 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9061 }
9062 Ok(Some(cols))
9063 }
9064
9065 pub fn native_page_cursor(
9080 &self,
9081 snapshot: Snapshot,
9082 projection: Vec<(u16, TypeId)>,
9083 conditions: &[crate::query::Condition],
9084 ) -> Result<Option<NativePageCursor>> {
9085 use crate::cursor::build_page_plans;
9086 if self.ttl.is_some() {
9087 return Ok(None);
9088 }
9089 if !conditions.is_empty() && !self.indexes_complete {
9092 return Ok(None);
9093 }
9094 if self.run_refs.len() != 1 {
9095 return Ok(None);
9096 }
9097 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9098 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9099
9100 let overlay_rids: HashSet<u64> = {
9103 let mut s = HashSet::new();
9104 for row in self.memtable.visible_versions(snapshot.epoch) {
9105 s.insert(row.row_id.0);
9106 }
9107 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9108 s.insert(row.row_id.0);
9109 }
9110 s
9111 };
9112
9113 let survivors = if conditions.is_empty() {
9117 None
9118 } else {
9119 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9120 };
9121
9122 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9129 survivors.clone()
9130 } else if let Some(s) = &survivors {
9131 let mut run_set = s.clone();
9132 run_set.remove_many(overlay_rids.iter().copied());
9133 Some(run_set)
9134 } else {
9135 Some(RowIdSet::from_unsorted(
9136 rids.iter()
9137 .map(|&r| r as u64)
9138 .filter(|r| !overlay_rids.contains(r))
9139 .collect(),
9140 ))
9141 };
9142
9143 let overlay_rows = if overlay_rids.is_empty() {
9144 Vec::new()
9145 } else {
9146 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9147 self.overlay_visible_rows(snapshot, bound)
9148 };
9149
9150 let plans = if positions.is_empty() {
9152 Vec::new()
9153 } else {
9154 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9155 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9156 };
9157
9158 let overlay = if overlay_rows.is_empty() {
9160 None
9161 } else {
9162 let filtered =
9163 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9164 if filtered.is_empty() {
9165 None
9166 } else {
9167 Some(self.materialize_overlay(&filtered, &projection))
9168 }
9169 };
9170
9171 let overlay_row_count = overlay
9172 .as_ref()
9173 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9174 .unwrap_or(0);
9175 crate::trace::QueryTrace::record(|t| {
9176 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9177 t.run_count = self.run_refs.len();
9178 t.memtable_rows = self.memtable.len();
9179 t.mutable_run_rows = self.mutable_run.len();
9180 t.overlay_rows = overlay_row_count;
9181 t.conditions_pushed = conditions.len();
9182 t.pages_decoded = plans
9183 .iter()
9184 .map(|p| p.positions.len())
9185 .sum::<usize>()
9186 .min(1);
9187 });
9188
9189 Ok(Some(NativePageCursor::new_with_overlay(
9190 reader, projection, plans, overlay,
9191 )))
9192 }
9193 #[allow(clippy::type_complexity)]
9203 pub fn native_multi_run_cursor(
9204 &self,
9205 snapshot: Snapshot,
9206 projection: Vec<(u16, TypeId)>,
9207 conditions: &[crate::query::Condition],
9208 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9209 use crate::cursor::{MultiRunCursor, RunStream};
9210 use crate::sorted_run::SYS_ROW_ID;
9211 use std::collections::{BinaryHeap, HashMap, HashSet};
9212 if self.ttl.is_some() {
9213 return Ok(None);
9214 }
9215 if !conditions.is_empty() && !self.indexes_complete {
9218 return Ok(None);
9219 }
9220 if self.run_refs.is_empty() {
9221 return Ok(None);
9222 }
9223
9224 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9226 Vec::with_capacity(self.run_refs.len());
9227 for rr in &self.run_refs {
9228 let mut reader = self.open_reader(rr.run_id)?;
9229 let (rids, eps, del) = reader.system_columns_native()?;
9230 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9231 run_meta.push((reader, rids, eps, del, page_rows));
9232 }
9233
9234 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9238 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9239 for i in 0..rids.len() {
9240 let rid = rids[i] as u64;
9241 let e = eps[i] as u64;
9242 if e > snapshot.epoch.0 {
9243 continue;
9244 }
9245 let is_del = del[i] != 0;
9246 best.entry(rid)
9247 .and_modify(|cur| {
9248 if e > cur.0 {
9249 *cur = (e, run_idx, i, is_del);
9250 }
9251 })
9252 .or_insert((e, run_idx, i, is_del));
9253 }
9254 }
9255
9256 let overlay_rids: HashSet<u64> = {
9258 let mut s = HashSet::new();
9259 for row in self.memtable.visible_versions(snapshot.epoch) {
9260 s.insert(row.row_id.0);
9261 }
9262 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9263 s.insert(row.row_id.0);
9264 }
9265 s
9266 };
9267
9268 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9270 None
9271 } else {
9272 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9273 for c in conditions {
9274 sets.push(self.resolve_condition(c, snapshot)?);
9275 }
9276 Some(RowIdSet::intersect_many(sets))
9277 };
9278
9279 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9283 for (rid, (_, run_idx, pos, deleted)) in &best {
9284 if *deleted {
9285 continue;
9286 }
9287 if overlay_rids.contains(rid) {
9288 continue;
9289 }
9290 if let Some(s) = &survivors {
9291 if !s.contains(*rid) {
9292 continue;
9293 }
9294 }
9295 per_run[*run_idx].push((*rid, *pos));
9296 }
9297 for v in per_run.iter_mut() {
9298 v.sort_unstable_by_key(|&(rid, _)| rid);
9299 }
9300
9301 let mut streams = Vec::with_capacity(run_meta.len());
9303 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9304 let mut total = 0usize;
9305 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9306 let mut starts = Vec::with_capacity(page_rows.len());
9307 let mut acc = 0usize;
9308 for &r in &page_rows {
9309 starts.push(acc);
9310 acc += r;
9311 }
9312 let mut survivors_vec: Vec<(u64, usize, usize)> =
9313 Vec::with_capacity(per_run[run_idx].len());
9314 for &(rid, pos) in &per_run[run_idx] {
9315 let page_seq = match starts.partition_point(|&s| s <= pos) {
9316 0 => continue,
9317 p => p - 1,
9318 };
9319 let within = pos - starts[page_seq];
9320 survivors_vec.push((rid, page_seq, within));
9321 }
9322 total += survivors_vec.len();
9323 if let Some(&(rid, _, _)) = survivors_vec.first() {
9324 heap.push(std::cmp::Reverse((rid, run_idx)));
9325 }
9326 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9327 }
9328
9329 let overlay_rows = if overlay_rids.is_empty() {
9331 Vec::new()
9332 } else {
9333 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9334 self.overlay_visible_rows(snapshot, bound)
9335 };
9336 let overlay = if overlay_rows.is_empty() {
9337 None
9338 } else {
9339 let filtered =
9340 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9341 if filtered.is_empty() {
9342 None
9343 } else {
9344 Some(self.materialize_overlay(&filtered, &projection))
9345 }
9346 };
9347
9348 let overlay_row_count = overlay
9349 .as_ref()
9350 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9351 .unwrap_or(0);
9352 crate::trace::QueryTrace::record(|t| {
9353 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9354 t.run_count = self.run_refs.len();
9355 t.memtable_rows = self.memtable.len();
9356 t.mutable_run_rows = self.mutable_run.len();
9357 t.overlay_rows = overlay_row_count;
9358 t.conditions_pushed = conditions.len();
9359 t.survivor_count = Some(total);
9360 });
9361
9362 Ok(Some(MultiRunCursor::new(
9363 streams, projection, heap, total, overlay,
9364 )))
9365 }
9366
9367 fn overlay_materialization_bound<'a>(
9379 conditions: &[crate::query::Condition],
9380 survivors: &'a Option<RowIdSet>,
9381 ) -> Option<&'a RowIdSet> {
9382 use crate::query::Condition;
9383 let has_range = conditions
9384 .iter()
9385 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9386 if has_range {
9387 None
9388 } else {
9389 survivors.as_ref()
9390 }
9391 }
9392
9393 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9405 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9406 let mut fold = |row: Row| {
9407 if let Some(b) = bound {
9408 if !b.contains(row.row_id.0) {
9409 return;
9410 }
9411 }
9412 best.entry(row.row_id.0)
9413 .and_modify(|(be, br)| {
9414 if row.committed_epoch > *be {
9415 *be = row.committed_epoch;
9416 *br = row.clone();
9417 }
9418 })
9419 .or_insert_with(|| (row.committed_epoch, row));
9420 };
9421 for row in self.memtable.visible_versions(snapshot.epoch) {
9422 fold(row);
9423 }
9424 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9425 fold(row);
9426 }
9427 let mut out: Vec<Row> = best
9428 .into_values()
9429 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9430 .collect();
9431 out.sort_by_key(|r| r.row_id);
9432 out
9433 }
9434
9435 fn filter_overlay_rows(
9443 &self,
9444 rows: Vec<Row>,
9445 conditions: &[crate::query::Condition],
9446 survivors: Option<&RowIdSet>,
9447 snapshot: Snapshot,
9448 ) -> Result<Vec<Row>> {
9449 if conditions.is_empty() {
9450 return Ok(rows);
9451 }
9452 use crate::query::Condition;
9453 let all_index_served = !conditions
9457 .iter()
9458 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9459 if all_index_served {
9460 return Ok(rows
9461 .into_iter()
9462 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9463 .collect());
9464 }
9465 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9468 for c in conditions {
9469 let s = match c {
9470 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9471 _ => self.resolve_condition(c, snapshot)?,
9472 };
9473 per_cond_sets.push(s);
9474 }
9475 Ok(rows
9476 .into_iter()
9477 .filter(|row| {
9478 conditions.iter().enumerate().all(|(i, c)| match c {
9479 Condition::Range { column_id, lo, hi } => {
9480 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9481 }
9482 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9483 match row.columns.get(column_id) {
9484 Some(Value::Float64(v)) => {
9485 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9486 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9487 lo_ok && hi_ok
9488 }
9489 _ => false,
9490 }
9491 }
9492 _ => per_cond_sets[i].contains(row.row_id.0),
9493 })
9494 })
9495 .collect())
9496 }
9497
9498 fn materialize_overlay(
9501 &self,
9502 rows: &[Row],
9503 projection: &[(u16, TypeId)],
9504 ) -> Vec<columnar::NativeColumn> {
9505 if projection.is_empty() {
9506 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9507 }
9508 let mut cols = Vec::with_capacity(projection.len());
9509 for (cid, ty) in projection {
9510 let vals: Vec<Value> = rows
9511 .iter()
9512 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9513 .collect();
9514 cols.push(columnar::values_to_native(ty.clone(), &vals));
9515 }
9516 cols
9517 }
9518
9519 fn resolve_survivor_rids(
9524 &self,
9525 conditions: &[crate::query::Condition],
9526 reader: &mut RunReader,
9527 snapshot: Snapshot,
9528 ) -> Result<RowIdSet> {
9529 use crate::query::Condition;
9530 let mut sets: Vec<RowIdSet> = Vec::new();
9531 for c in conditions {
9532 self.validate_condition(c)?;
9533 let s: RowIdSet = match c {
9534 Condition::Pk(key) => {
9535 let lookup = self
9536 .schema
9537 .primary_key()
9538 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9539 .unwrap_or_else(|| key.clone());
9540 self.hot
9541 .get(&lookup)
9542 .map(|r| RowIdSet::one(r.0))
9543 .unwrap_or_else(RowIdSet::empty)
9544 }
9545 Condition::BitmapEq { column_id, value } => {
9546 let lookup = self.index_lookup_key_bytes(*column_id, value);
9547 self.bitmap
9548 .get(column_id)
9549 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9550 .unwrap_or_else(RowIdSet::empty)
9551 }
9552 Condition::BitmapIn { column_id, values } => {
9553 let bm = self.bitmap.get(column_id);
9554 let mut acc = roaring::RoaringBitmap::new();
9555 if let Some(b) = bm {
9556 for v in values {
9557 let lookup = self.index_lookup_key_bytes(*column_id, v);
9558 acc |= b.get(&lookup);
9559 }
9560 }
9561 RowIdSet::from_roaring(acc)
9562 }
9563 Condition::BytesPrefix { column_id, prefix } => {
9564 if let Some(b) = self.bitmap.get(column_id) {
9565 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9566 let mut acc = roaring::RoaringBitmap::new();
9567 for key in b.keys() {
9568 if key.starts_with(&lookup_prefix) {
9569 acc |= b.get(&key);
9570 }
9571 }
9572 RowIdSet::from_roaring(acc)
9573 } else {
9574 RowIdSet::empty()
9575 }
9576 }
9577 Condition::FmContains { column_id, pattern } => self
9578 .fm
9579 .get(column_id)
9580 .map(|f| {
9581 RowIdSet::from_unsorted(
9582 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9583 )
9584 })
9585 .unwrap_or_else(RowIdSet::empty),
9586 Condition::FmContainsAll {
9587 column_id,
9588 patterns,
9589 } => {
9590 if let Some(f) = self.fm.get(column_id) {
9591 let sets: Vec<RowIdSet> = patterns
9592 .iter()
9593 .map(|pat| {
9594 RowIdSet::from_unsorted(
9595 f.locate(pat).into_iter().map(|r| r.0).collect(),
9596 )
9597 })
9598 .collect();
9599 RowIdSet::intersect_many(sets)
9600 } else {
9601 RowIdSet::empty()
9602 }
9603 }
9604 Condition::Ann {
9605 column_id,
9606 query,
9607 k,
9608 } => RowIdSet::from_unsorted(
9609 self.retrieve_filtered(
9610 &crate::query::Retriever::Ann {
9611 column_id: *column_id,
9612 query: query.clone(),
9613 k: *k,
9614 },
9615 snapshot,
9616 None,
9617 None,
9618 None,
9619 None,
9620 )?
9621 .into_iter()
9622 .map(|hit| hit.row_id.0)
9623 .collect(),
9624 ),
9625 Condition::SparseMatch {
9626 column_id,
9627 query,
9628 k,
9629 } => RowIdSet::from_unsorted(
9630 self.retrieve_filtered(
9631 &crate::query::Retriever::Sparse {
9632 column_id: *column_id,
9633 query: query.clone(),
9634 k: *k,
9635 },
9636 snapshot,
9637 None,
9638 None,
9639 None,
9640 None,
9641 )?
9642 .into_iter()
9643 .map(|hit| hit.row_id.0)
9644 .collect(),
9645 ),
9646 Condition::MinHashSimilar {
9647 column_id,
9648 query,
9649 k,
9650 } => match self.minhash.get(column_id) {
9651 Some(index) => {
9652 let candidates = index.candidate_row_ids(query);
9653 let eligible =
9654 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9655 RowIdSet::from_unsorted(
9656 index
9657 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9658 .into_iter()
9659 .map(|(row_id, _)| row_id.0)
9660 .collect(),
9661 )
9662 }
9663 None => RowIdSet::empty(),
9664 },
9665 Condition::Range { column_id, lo, hi } => {
9666 if let Some(li) = self.learned_range.get(column_id) {
9667 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9668 } else {
9669 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9670 }
9671 }
9672 Condition::RangeF64 {
9673 column_id,
9674 lo,
9675 lo_inclusive,
9676 hi,
9677 hi_inclusive,
9678 } => {
9679 if let Some(li) = self.learned_range.get(column_id) {
9680 RowIdSet::from_unsorted(
9681 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9682 .into_iter()
9683 .collect(),
9684 )
9685 } else {
9686 reader.range_row_id_set_f64(
9687 *column_id,
9688 *lo,
9689 *lo_inclusive,
9690 *hi,
9691 *hi_inclusive,
9692 )?
9693 }
9694 }
9695 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9696 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9697 };
9698 sets.push(s);
9699 }
9700 Ok(RowIdSet::intersect_many(sets))
9701 }
9702
9703 pub fn scan_cursor(
9724 &self,
9725 snapshot: Snapshot,
9726 projection: Vec<(u16, TypeId)>,
9727 conditions: &[crate::query::Condition],
9728 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9729 if self.ttl.is_some() {
9730 return Ok(None);
9731 }
9732 if !conditions.is_empty() && !self.indexes_complete {
9738 return Ok(None);
9739 }
9740 if self.run_refs.len() == 1 {
9741 Ok(self
9742 .native_page_cursor(snapshot, projection, conditions)?
9743 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9744 } else {
9745 Ok(self
9746 .native_multi_run_cursor(snapshot, projection, conditions)?
9747 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9748 }
9749 }
9750
9751 pub fn aggregate_native(
9765 &self,
9766 snapshot: Snapshot,
9767 column: Option<u16>,
9768 conditions: &[crate::query::Condition],
9769 agg: NativeAgg,
9770 ) -> Result<Option<NativeAggResult>> {
9771 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9772 }
9773
9774 pub fn aggregate_native_with_control(
9775 &self,
9776 snapshot: Snapshot,
9777 column: Option<u16>,
9778 conditions: &[crate::query::Condition],
9779 agg: NativeAgg,
9780 control: &crate::ExecutionControl,
9781 ) -> Result<Option<NativeAggResult>> {
9782 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9783 }
9784
9785 fn aggregate_native_inner(
9786 &self,
9787 snapshot: Snapshot,
9788 column: Option<u16>,
9789 conditions: &[crate::query::Condition],
9790 agg: NativeAgg,
9791 control: Option<&crate::ExecutionControl>,
9792 ) -> Result<Option<NativeAggResult>> {
9793 execution_checkpoint(control, 0)?;
9794 if self.ttl.is_some() {
9795 return Ok(None);
9796 }
9797 if self.run_refs.len() == 1 && conditions.is_empty() {
9799 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9800 return Ok(Some(res));
9801 }
9802 }
9803 if matches!(agg, NativeAgg::Count) && column.is_none() {
9807 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
9808 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
9809 }
9810 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9811 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
9812 }
9813 let cid = match column {
9816 Some(c) => c,
9817 None => return Ok(None),
9818 };
9819 let ty = self.column_type(cid);
9820 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
9821 execution_checkpoint(control, 0)?;
9822 return match ty {
9823 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9824 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9825 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9826 }
9827 TypeId::Float64 => {
9828 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9829 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9830 }
9831 _ => Ok(None),
9832 };
9833 }
9834 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9836 execution_checkpoint(control, 0)?;
9837 match ty {
9838 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9839 let mut count = 0u64;
9840 let mut sum = 0i128;
9841 let mut mn = i64::MAX;
9842 let mut mx = i64::MIN;
9843 for row in &rows {
9844 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
9845 count += 1;
9846 sum += i128::from(*v);
9847 mn = mn.min(*v);
9848 mx = mx.max(*v);
9849 }
9850 }
9851 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9852 }
9853 TypeId::Float64 => {
9854 let mut count = 0u64;
9855 let mut sum = 0.0f64;
9856 let mut mn = f64::INFINITY;
9857 let mut mx = f64::NEG_INFINITY;
9858 for row in &rows {
9859 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
9860 count += 1;
9861 sum += *v;
9862 mn = mn.min(*v);
9863 mx = mx.max(*v);
9864 }
9865 }
9866 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9867 }
9868 _ => Ok(None),
9869 }
9870 }
9871
9872 fn visible_rows_filtered(
9874 &self,
9875 snapshot: Snapshot,
9876 conditions: &[crate::query::Condition],
9877 control: Option<&crate::ExecutionControl>,
9878 ) -> Result<Vec<Row>> {
9879 let rows = if let Some(control) = control {
9880 self.visible_rows_controlled(snapshot, control)?
9881 } else {
9882 self.visible_rows(snapshot)?
9883 };
9884 if conditions.is_empty() {
9885 return Ok(rows);
9886 }
9887 Ok(rows
9888 .into_iter()
9889 .filter(|row| {
9890 conditions
9891 .iter()
9892 .all(|cond| condition_matches_row(cond, row, &self.schema))
9893 })
9894 .collect())
9895 }
9896
9897 fn aggregate_from_stats(
9905 &self,
9906 snapshot: Snapshot,
9907 column: Option<u16>,
9908 agg: NativeAgg,
9909 ) -> Result<Option<NativeAggResult>> {
9910 let cid = match (agg, column) {
9911 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9912 _ => return Ok(None), };
9914 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9915 return Ok(None);
9916 };
9917 let Some(cs) = stats.get(&cid) else {
9918 return Ok(None);
9919 };
9920 match agg {
9921 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9923 self.live_count.saturating_sub(cs.null_count),
9924 ))),
9925 NativeAgg::Min | NativeAgg::Max => {
9926 let bound = if agg == NativeAgg::Min {
9927 &cs.min
9928 } else {
9929 &cs.max
9930 };
9931 match bound {
9932 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9933 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9934 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9939 None => Ok(None),
9940 }
9941 }
9942 _ => Ok(None),
9943 }
9944 }
9945
9946 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9955 if self.ttl.is_some() {
9956 return Ok(None);
9957 }
9958 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9959 return Ok(None);
9960 }
9961 self.ensure_indexes_complete()?;
9964 let reader = self.open_reader(self.run_refs[0].run_id)?;
9965 if self.live_count != reader.row_count() as u64 {
9966 return Ok(None);
9967 }
9968 let Some(bm) = self.bitmap.get(&column_id) else {
9969 return Ok(None); };
9971 let mut distinct = bm.value_count() as u64;
9972 if !bm.get(&Value::Null.encode_key()).is_empty() {
9975 distinct = distinct.saturating_sub(1);
9976 }
9977 Ok(Some(distinct))
9978 }
9979
9980 pub fn aggregate_incremental(
9992 &mut self,
9993 cache_key: u64,
9994 conditions: &[crate::query::Condition],
9995 column: Option<u16>,
9996 agg: NativeAgg,
9997 ) -> Result<IncrementalAggResult> {
9998 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9999 }
10000
10001 pub fn aggregate_incremental_with_control(
10002 &mut self,
10003 cache_key: u64,
10004 conditions: &[crate::query::Condition],
10005 column: Option<u16>,
10006 agg: NativeAgg,
10007 control: &crate::ExecutionControl,
10008 ) -> Result<IncrementalAggResult> {
10009 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
10010 }
10011
10012 fn aggregate_incremental_inner(
10013 &mut self,
10014 cache_key: u64,
10015 conditions: &[crate::query::Condition],
10016 column: Option<u16>,
10017 agg: NativeAgg,
10018 control: Option<&crate::ExecutionControl>,
10019 ) -> Result<IncrementalAggResult> {
10020 execution_checkpoint(control, 0)?;
10021 let snap = self.snapshot();
10022 let cur_wm = self.allocator.current().0;
10023 let cur_epoch = snap.epoch.0;
10024 let incremental_ok = self.ttl.is_none()
10031 && !self.had_deletes
10032 && self.memtable.is_empty()
10033 && self.mutable_run.is_empty();
10034
10035 if incremental_ok {
10038 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10039 if cached.epoch == cur_epoch {
10040 return Ok(IncrementalAggResult {
10041 state: cached.state,
10042 incremental: true,
10043 delta_rows: 0,
10044 });
10045 }
10046 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10047 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10048 let mut delta_rids = Vec::with_capacity(delta_len);
10049 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10050 execution_checkpoint(control, index)?;
10051 delta_rids.push(row_id);
10052 }
10053 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10054 execution_checkpoint(control, 0)?;
10055 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10056 let delta_state = agg_state_from_rows(
10057 &delta_rows,
10058 conditions,
10059 &index_sets,
10060 column,
10061 agg,
10062 &self.schema,
10063 control,
10064 )?;
10065 let merged = cached.state.merge(delta_state);
10066 let delta_n = delta_rids.len() as u64;
10067 Arc::make_mut(&mut self.agg_cache).insert(
10068 cache_key,
10069 CachedAgg {
10070 state: merged.clone(),
10071 watermark: cur_wm,
10072 epoch: cur_epoch,
10073 },
10074 );
10075 return Ok(IncrementalAggResult {
10076 state: merged,
10077 incremental: true,
10078 delta_rows: delta_n,
10079 });
10080 }
10081 }
10082 }
10083
10084 let cursor_ok =
10089 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10090 let state = if cursor_ok && agg != NativeAgg::Avg {
10091 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10092 Some(result) => {
10093 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10094 }
10095 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10096 }
10097 } else {
10098 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10099 };
10100 if incremental_ok {
10102 Arc::make_mut(&mut self.agg_cache).insert(
10103 cache_key,
10104 CachedAgg {
10105 state: state.clone(),
10106 watermark: cur_wm,
10107 epoch: cur_epoch,
10108 },
10109 );
10110 }
10111 Ok(IncrementalAggResult {
10112 state,
10113 incremental: false,
10114 delta_rows: 0,
10115 })
10116 }
10117
10118 fn agg_state_full_scan(
10121 &self,
10122 conditions: &[crate::query::Condition],
10123 column: Option<u16>,
10124 agg: NativeAgg,
10125 snap: Snapshot,
10126 control: Option<&crate::ExecutionControl>,
10127 ) -> Result<AggState> {
10128 execution_checkpoint(control, 0)?;
10129 let rows = self.visible_rows(snap)?;
10130 execution_checkpoint(control, 0)?;
10131 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10132 agg_state_from_rows(
10133 &rows,
10134 conditions,
10135 &index_sets,
10136 column,
10137 agg,
10138 &self.schema,
10139 control,
10140 )
10141 }
10142
10143 fn resolve_index_conditions(
10146 &self,
10147 conditions: &[crate::query::Condition],
10148 snapshot: Snapshot,
10149 ) -> Result<Vec<RowIdSet>> {
10150 use crate::query::Condition;
10151 let mut sets = Vec::new();
10152 for c in conditions {
10153 if matches!(
10154 c,
10155 Condition::Ann { .. }
10156 | Condition::SparseMatch { .. }
10157 | Condition::MinHashSimilar { .. }
10158 ) {
10159 sets.push(self.resolve_condition(c, snapshot)?);
10160 }
10161 }
10162 Ok(sets)
10163 }
10164
10165 fn column_type(&self, cid: u16) -> TypeId {
10166 self.schema
10167 .columns
10168 .iter()
10169 .find(|c| c.id == cid)
10170 .map(|c| c.ty.clone())
10171 .unwrap_or(TypeId::Bytes)
10172 }
10173
10174 pub fn approx_aggregate(
10183 &mut self,
10184 conditions: &[crate::query::Condition],
10185 column: Option<u16>,
10186 agg: ApproxAgg,
10187 z: f64,
10188 ) -> Result<Option<ApproxResult>> {
10189 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10190 }
10191
10192 pub fn approx_aggregate_with_candidate_authorization(
10195 &mut self,
10196 conditions: &[crate::query::Condition],
10197 column: Option<u16>,
10198 agg: ApproxAgg,
10199 z: f64,
10200 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10201 ) -> Result<Option<ApproxResult>> {
10202 use crate::query::Condition;
10203 self.ensure_reservoir_complete()?;
10204 let snapshot = self.snapshot();
10205 let n_pop = self.count();
10206 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10207 if sample_rids.is_empty() {
10208 return Ok(None);
10209 }
10210 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10212 let s = live_sample.len();
10213 if s == 0 {
10214 return Ok(None);
10215 }
10216 let authorized = authorization
10217 .map(|authorization| {
10218 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10219 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10220 })
10221 .transpose()?;
10222
10223 let mut index_sets: Vec<RowIdSet> = Vec::new();
10226 for c in conditions {
10227 if matches!(
10228 c,
10229 Condition::Ann { .. }
10230 | Condition::SparseMatch { .. }
10231 | Condition::MinHashSimilar { .. }
10232 ) {
10233 index_sets.push(self.resolve_condition(c, snapshot)?);
10234 }
10235 }
10236
10237 let cid = match (agg, column) {
10239 (ApproxAgg::Count, _) => None,
10240 (_, Some(c)) => Some(c),
10241 _ => return Ok(None),
10242 };
10243 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10244 for r in &live_sample {
10245 if authorized
10246 .as_ref()
10247 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10248 {
10249 continue;
10250 }
10251 if !conditions
10253 .iter()
10254 .all(|c| condition_matches_row(c, r, &self.schema))
10255 {
10256 continue;
10257 }
10258 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10260 continue;
10261 }
10262 if let Some(cid) = cid {
10263 let mut cells = r
10264 .columns
10265 .get(&cid)
10266 .cloned()
10267 .map(|value| vec![(cid, value)])
10268 .unwrap_or_default();
10269 if let Some(authorization) = authorization {
10270 authorization.security.apply_masks_to_cells(
10271 authorization.table,
10272 &mut cells,
10273 authorization.principal,
10274 );
10275 }
10276 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10277 passing_vals.push(v);
10278 } } else {
10280 passing_vals.push(0.0); }
10282 }
10283 let m = passing_vals.len();
10284
10285 let (point, half) = match agg {
10286 ApproxAgg::Count => {
10287 let p = m as f64 / s as f64;
10289 let point = n_pop as f64 * p;
10290 let var = if s > 1 {
10291 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10292 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10293 } else {
10294 0.0
10295 };
10296 (point, z * var.sqrt())
10297 }
10298 ApproxAgg::Sum => {
10299 let y: Vec<f64> = live_sample
10301 .iter()
10302 .map(|r| {
10303 let passes_row = authorized
10304 .as_ref()
10305 .is_none_or(|authorized| authorized.contains(&r.row_id))
10306 && conditions
10307 .iter()
10308 .all(|c| condition_matches_row(c, r, &self.schema))
10309 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10310 if passes_row {
10311 cid.and_then(|cid| {
10312 let mut cells = r
10313 .columns
10314 .get(&cid)
10315 .cloned()
10316 .map(|value| vec![(cid, value)])
10317 .unwrap_or_default();
10318 if let Some(authorization) = authorization {
10319 authorization.security.apply_masks_to_cells(
10320 authorization.table,
10321 &mut cells,
10322 authorization.principal,
10323 );
10324 }
10325 as_f64(cells.first().map(|(_, value)| value))
10326 })
10327 .unwrap_or(0.0)
10328 } else {
10329 0.0
10330 }
10331 })
10332 .collect();
10333 let mean_y = y.iter().sum::<f64>() / s as f64;
10334 let point = n_pop as f64 * mean_y;
10335 let var = if s > 1 {
10336 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10337 let var_y = ss / (s - 1) as f64;
10338 n_pop as f64 * n_pop as f64 * var_y / s as f64
10339 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10340 } else {
10341 0.0
10342 };
10343 (point, z * var.sqrt())
10344 }
10345 ApproxAgg::Avg => {
10346 if m == 0 {
10347 return Ok(Some(ApproxResult {
10348 point: 0.0,
10349 ci_low: 0.0,
10350 ci_high: 0.0,
10351 n_population: n_pop,
10352 n_sample_live: s,
10353 n_passing: 0,
10354 }));
10355 }
10356 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10357 let half = if m > 1 {
10358 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10359 let sd = (ss / (m - 1) as f64).sqrt();
10360 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10361 z * sd / (m as f64).sqrt() * fpc.sqrt()
10362 } else {
10363 0.0
10364 };
10365 (mean, half)
10366 }
10367 };
10368
10369 Ok(Some(ApproxResult {
10370 point,
10371 ci_low: point - half,
10372 ci_high: point + half,
10373 n_population: n_pop,
10374 n_sample_live: s,
10375 n_passing: m,
10376 }))
10377 }
10378
10379 pub fn exact_column_stats(
10387 &self,
10388 _snapshot: Snapshot,
10389 projection: &[u16],
10390 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10391 if self.ttl.is_some()
10392 || !(self.memtable.is_empty()
10393 && self.mutable_run.is_empty()
10394 && self.run_refs.len() == 1)
10395 {
10396 return Ok(None);
10397 }
10398 let reader = self.open_reader(self.run_refs[0].run_id)?;
10399 if self.live_count != reader.row_count() as u64 {
10400 return Ok(None);
10401 }
10402 let mut out = HashMap::new();
10403 for &cid in projection {
10404 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10405 Some(c) => c,
10406 None => continue,
10407 };
10408 let Some(stats) = reader.column_page_stats(cid) else {
10410 out.insert(
10411 cid,
10412 ColumnStat {
10413 min: None,
10414 max: None,
10415 null_count: self.live_count,
10416 },
10417 );
10418 continue;
10419 };
10420 let stat = match cdef.ty {
10421 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10422 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10423 min: mn.map(Value::Int64),
10424 max: mx.map(Value::Int64),
10425 null_count: n,
10426 })
10427 }
10428 TypeId::Float64 => {
10429 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10430 min: mn.map(Value::Float64),
10431 max: mx.map(Value::Float64),
10432 null_count: n,
10433 })
10434 }
10435 _ => None,
10436 };
10437 if let Some(s) = stat {
10438 out.insert(cid, s);
10439 }
10440 }
10441 Ok(Some(out))
10442 }
10443
10444 pub fn dir(&self) -> &Path {
10445 &self.dir
10446 }
10447
10448 pub fn schema(&self) -> &Schema {
10449 &self.schema
10450 }
10451
10452 pub(crate) fn set_catalog_name(&mut self, name: String) {
10453 self.name = name;
10454 }
10455
10456 pub(crate) fn prepare_alter_column(
10457 &mut self,
10458 column_name: &str,
10459 change: &AlterColumn,
10460 ) -> Result<(ColumnDef, Option<Schema>)> {
10461 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10462 return Err(MongrelError::InvalidArgument(
10463 "ALTER COLUMN requires committing staged writes first".into(),
10464 ));
10465 }
10466 let old = self
10467 .schema
10468 .columns
10469 .iter()
10470 .find(|c| c.name == column_name)
10471 .cloned()
10472 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10473 let mut next = old.clone();
10474
10475 if let Some(name) = &change.name {
10476 let trimmed = name.trim();
10477 if trimmed.is_empty() {
10478 return Err(MongrelError::InvalidArgument(
10479 "ALTER COLUMN name must not be empty".into(),
10480 ));
10481 }
10482 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10483 return Err(MongrelError::Schema(format!(
10484 "column {trimmed} already exists"
10485 )));
10486 }
10487 next.name = trimmed.to_string();
10488 }
10489
10490 if let Some(ty) = &change.ty {
10491 next.ty = ty.clone();
10492 }
10493 if let Some(flags) = change.flags {
10494 validate_alter_column_flags(old.flags, flags)?;
10495 next.flags = flags;
10496 }
10497
10498 if let Some(default_change) = &change.default_value {
10499 next.default_value = default_change.clone();
10500 }
10501 if let Some(source_change) = &change.embedding_source {
10502 next.embedding_source = source_change.clone();
10503 }
10504
10505 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10506 if old.flags.contains(ColumnFlags::NULLABLE)
10507 && !next.flags.contains(ColumnFlags::NULLABLE)
10508 && self.column_has_nulls(old.id)?
10509 {
10510 return Err(MongrelError::InvalidArgument(format!(
10511 "column '{}' contains NULL values",
10512 old.name
10513 )));
10514 }
10515 if next == old {
10516 return Ok((next, None));
10517 }
10518 let mut schema = self.schema.clone();
10519 let index = schema
10520 .columns
10521 .iter()
10522 .position(|column| column.id == next.id)
10523 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10524 schema.columns[index] = next.clone();
10525 schema.schema_id = schema
10526 .schema_id
10527 .checked_add(1)
10528 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10529 schema.validate_auto_increment()?;
10530 schema.validate_defaults()?;
10531 Ok((next, Some(schema)))
10532 }
10533
10534 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10535 self.schema = schema;
10536 self.auto_inc = resolve_auto_inc(&self.schema);
10537 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10538 self.clear_result_cache();
10539 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10540 }
10541
10542 pub(crate) fn publish_index_schema_change(
10546 &mut self,
10547 schema: Schema,
10548 artifact: SecondaryIndexArtifact,
10549 ) -> Result<()> {
10550 self.apply_altered_schema_prepared(schema);
10551 self.prune_index_maps_to_schema();
10552 match artifact {
10553 SecondaryIndexArtifact::Bitmap(column_id, index) => {
10554 self.bitmap.insert(column_id, index);
10555 }
10556 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
10557 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
10558 }
10559 SecondaryIndexArtifact::Fm(column_id, index) => {
10560 self.fm.insert(column_id, *index);
10561 }
10562 SecondaryIndexArtifact::Ann(column_id, index) => {
10563 self.ann.insert(column_id, index);
10564 }
10565 SecondaryIndexArtifact::Sparse(column_id, index) => {
10566 self.sparse.insert(column_id, index);
10567 }
10568 SecondaryIndexArtifact::MinHash(column_id, index) => {
10569 self.minhash.insert(column_id, index);
10570 }
10571 }
10572 self.indexes_complete = true;
10573 self.seal_generations();
10574 let view = Arc::new(self.capture_read_generation());
10575 self.published.store(Arc::clone(&view));
10576 checkpoint_current_schema(self)?;
10577 self.invalidate_index_checkpoint();
10580 Ok(())
10581 }
10582
10583 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
10589 self.apply_altered_schema_prepared(schema);
10590 self.prune_index_maps_to_schema();
10591 self.indexes_complete = true;
10592 self.seal_generations();
10593 let view = Arc::new(self.capture_read_generation());
10594 self.published.store(Arc::clone(&view));
10595 checkpoint_current_schema(self)?;
10596 self.invalidate_index_checkpoint();
10598 Ok(())
10599 }
10600
10601 fn prune_index_maps_to_schema(&mut self) {
10602 let keep_bitmap: std::collections::HashSet<u16> = self
10603 .schema
10604 .indexes
10605 .iter()
10606 .filter(|index| index.kind == IndexKind::Bitmap)
10607 .map(|index| index.column_id)
10608 .collect();
10609 let keep_ann: std::collections::HashSet<u16> = self
10610 .schema
10611 .indexes
10612 .iter()
10613 .filter(|index| index.kind == IndexKind::Ann)
10614 .map(|index| index.column_id)
10615 .collect();
10616 let keep_fm: std::collections::HashSet<u16> = self
10617 .schema
10618 .indexes
10619 .iter()
10620 .filter(|index| index.kind == IndexKind::FmIndex)
10621 .map(|index| index.column_id)
10622 .collect();
10623 let keep_sparse: std::collections::HashSet<u16> = self
10624 .schema
10625 .indexes
10626 .iter()
10627 .filter(|index| index.kind == IndexKind::Sparse)
10628 .map(|index| index.column_id)
10629 .collect();
10630 let keep_minhash: std::collections::HashSet<u16> = self
10631 .schema
10632 .indexes
10633 .iter()
10634 .filter(|index| index.kind == IndexKind::MinHash)
10635 .map(|index| index.column_id)
10636 .collect();
10637 let keep_learned: std::collections::HashSet<u16> = self
10638 .schema
10639 .indexes
10640 .iter()
10641 .filter(|index| index.kind == IndexKind::LearnedRange)
10642 .map(|index| index.column_id)
10643 .collect();
10644 self.bitmap
10645 .retain(|column_id, _| keep_bitmap.contains(column_id));
10646 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
10647 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
10648 self.sparse
10649 .retain(|column_id, _| keep_sparse.contains(column_id));
10650 self.minhash
10651 .retain(|column_id, _| keep_minhash.contains(column_id));
10652 {
10653 let learned = Arc::make_mut(&mut self.learned_range);
10654 learned.retain(|column_id, _| keep_learned.contains(column_id));
10655 }
10656 }
10657
10658 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10659 checkpoint_current_schema(self)
10660 }
10661
10662 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10663 self.ensure_writable()?;
10664 let previous_schema = self.schema.clone();
10665 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10666 if let Some(schema) = schema {
10667 self.apply_altered_schema_prepared(schema);
10668 self.checkpoint_standalone_schema_change(previous_schema)?;
10669 }
10670 Ok(column)
10671 }
10672
10673 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10674 if self.live_count == 0 {
10675 return Ok(false);
10676 }
10677 let snap = self.snapshot();
10678 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10679 Ok(columns
10680 .first()
10681 .map(|(_, col)| col.null_count(col.len()) != 0)
10682 .unwrap_or(true))
10683 }
10684
10685 fn has_stored_versions(&self) -> bool {
10686 !self.memtable.is_empty()
10687 || !self.mutable_run.is_empty()
10688 || self.run_refs.iter().any(|r| r.row_count > 0)
10689 || !self.retiring.is_empty()
10690 }
10691
10692 pub fn add_column(
10697 &mut self,
10698 name: &str,
10699 ty: TypeId,
10700 flags: ColumnFlags,
10701 default_value: Option<crate::schema::DefaultExpr>,
10702 ) -> Result<u16> {
10703 self.add_column_with_id(name, ty, flags, default_value, None)
10704 }
10705
10706 pub fn add_column_with_id(
10707 &mut self,
10708 name: &str,
10709 ty: TypeId,
10710 flags: ColumnFlags,
10711 default_value: Option<crate::schema::DefaultExpr>,
10712 requested_id: Option<u16>,
10713 ) -> Result<u16> {
10714 self.ensure_writable()?;
10715 let previous_schema = self.schema.clone();
10716 let (column, schema) =
10717 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
10718 self.apply_altered_schema_prepared(schema);
10719 self.checkpoint_standalone_schema_change(previous_schema)?;
10720 Ok(column.id)
10721 }
10722
10723 pub(crate) fn prepare_add_column(
10724 &mut self,
10725 name: &str,
10726 ty: TypeId,
10727 flags: ColumnFlags,
10728 default_value: Option<crate::schema::DefaultExpr>,
10729 requested_id: Option<u16>,
10730 ) -> Result<(ColumnDef, Schema)> {
10731 self.ensure_writable()?;
10732 if self.schema.columns.iter().any(|c| c.name == name) {
10733 return Err(MongrelError::Schema(format!(
10734 "column {name} already exists"
10735 )));
10736 }
10737 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10738 if self.schema.columns.iter().any(|c| c.id == id) {
10739 return Err(MongrelError::Schema(format!(
10740 "column id {id} already exists"
10741 )));
10742 }
10743 id
10744 } else {
10745 self.schema
10746 .columns
10747 .iter()
10748 .map(|c| c.id)
10749 .max()
10750 .unwrap_or(0)
10751 .checked_add(1)
10752 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10753 };
10754 let column = ColumnDef {
10755 id,
10756 name: name.to_string(),
10757 ty,
10758 flags,
10759 default_value,
10760 embedding_source: None,
10761 };
10762 let mut next_schema = self.schema.clone();
10763 next_schema.columns.push(column.clone());
10764 next_schema.schema_id = next_schema
10765 .schema_id
10766 .checked_add(1)
10767 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10768 next_schema.validate_auto_increment()?;
10769 next_schema.validate_defaults()?;
10770 Ok((column, next_schema))
10771 }
10772
10773 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10782 self.ensure_writable()?;
10783 let cid = self
10784 .schema
10785 .columns
10786 .iter()
10787 .find(|c| c.name == column_name)
10788 .map(|c| c.id)
10789 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10790 let ty = self
10791 .schema
10792 .columns
10793 .iter()
10794 .find(|c| c.id == cid)
10795 .map(|c| c.ty.clone())
10796 .unwrap_or(TypeId::Int64);
10797 if !matches!(
10798 ty,
10799 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10800 ) {
10801 return Err(MongrelError::Schema(format!(
10802 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10803 )));
10804 }
10805 if self
10806 .schema
10807 .indexes
10808 .iter()
10809 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10810 {
10811 return Ok(()); }
10813 let previous_schema = self.schema.clone();
10814 let previous_learned_range = Arc::clone(&self.learned_range);
10815 let mut next_schema = previous_schema.clone();
10816 next_schema.indexes.push(IndexDef {
10817 name: format!("{}_learned_range", column_name),
10818 column_id: cid,
10819 kind: IndexKind::LearnedRange,
10820 predicate: None,
10821 options: Default::default(),
10822 });
10823 next_schema.schema_id = next_schema
10824 .schema_id
10825 .checked_add(1)
10826 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10827 self.apply_altered_schema_prepared(next_schema);
10828 if let Err(error) = self.build_learned_ranges() {
10829 self.apply_altered_schema_prepared(previous_schema);
10830 self.learned_range = previous_learned_range;
10831 return Err(error);
10832 }
10833 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10834 if !matches!(
10835 &error,
10836 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10837 ) {
10838 self.learned_range = previous_learned_range;
10839 }
10840 return Err(error);
10841 }
10842 Ok(())
10843 }
10844
10845 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10846 let mut schema_published = false;
10847 let schema_result = match self._root_guard.as_deref() {
10848 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10849 schema_published = true;
10850 }),
10851 None => write_schema_with_after(&self.dir, &self.schema, || {
10852 schema_published = true;
10853 }),
10854 };
10855 if schema_result.is_err() && !schema_published {
10856 self.apply_altered_schema_prepared(previous_schema);
10857 return schema_result;
10858 }
10859
10860 let manifest_result = self.persist_manifest(self.current_epoch());
10861 match (schema_result, manifest_result) {
10862 (_, Ok(())) => Ok(()),
10863 (Ok(()), Err(error)) => {
10864 self.poison_after_maintenance_publish_failure();
10865 Err(MongrelError::DurableCommit {
10866 epoch: self.current_epoch().0,
10867 message: format!(
10868 "schema is durable but matching manifest publication failed: {error}"
10869 ),
10870 })
10871 }
10872 (Err(schema_error), Err(manifest_error)) => {
10873 self.poison_after_maintenance_publish_failure();
10874 Err(MongrelError::CommitOutcomeUnknown {
10875 epoch: self.current_epoch().0,
10876 message: format!(
10877 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10878 ),
10879 })
10880 }
10881 }
10882 }
10883
10884 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10887 self.sync_byte_threshold = threshold;
10888 if let WalSink::Private(w) = &mut self.wal {
10889 w.set_sync_byte_threshold(threshold);
10890 }
10891 }
10892
10893 pub fn page_cache_flush(&self) {
10897 self.page_cache.flush_to_disk();
10898 }
10899
10900 pub fn page_cache_len(&self) -> usize {
10902 self.page_cache.len()
10903 }
10904
10905 pub fn decoded_cache_len(&self) -> usize {
10908 self.decoded_cache.len()
10909 }
10910
10911 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10914 self.memtable.drain_sorted()
10915 }
10916
10917 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10918 self.runs_dir().join(format!("r-{run_id}.sr"))
10919 }
10920
10921 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10922 match self.runs_root.as_deref() {
10923 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10924 None => Ok(None),
10925 }
10926 }
10927
10928 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10929 match self.runs_root.as_deref() {
10930 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10931 None => Ok(None),
10932 }
10933 }
10934
10935 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10936 match self.runs_root.as_deref() {
10937 Some(root) => match root.remove_file(name) {
10938 Ok(()) => Ok(()),
10939 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10940 Err(error) => Err(error.into()),
10941 },
10942 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10943 Ok(()) => Ok(()),
10944 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10945 Err(error) => Err(error.into()),
10946 },
10947 }
10948 }
10949
10950 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10951 match self.runs_root.as_deref() {
10952 Some(root) => root
10953 .rename_file_new(source, destination)
10954 .map_err(Into::into),
10955 None => crate::durable_file::rename(
10956 &self.runs_dir().join(source),
10957 &self.runs_dir().join(destination),
10958 )
10959 .map_err(Into::into),
10960 }
10961 }
10962
10963 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10964 self.run_refs.iter().map(|run| run.run_id)
10965 }
10966
10967 pub(crate) fn table_dir(&self) -> &Path {
10968 &self.dir
10969 }
10970
10971 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10972 &self.schema
10973 }
10974
10975 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10976 let id = self.next_run_id;
10977 self.next_run_id = self
10978 .next_run_id
10979 .checked_add(1)
10980 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10981 Ok(id)
10982 }
10983
10984 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10985 self.run_refs.push(run_ref);
10986 }
10987
10988 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10998 self.retiring.push(crate::manifest::RetiredRun {
10999 run_id,
11000 retire_epoch,
11001 });
11002 }
11003
11004 pub(crate) fn reap_retiring(
11008 &mut self,
11009 min_active: Epoch,
11010 backup_pinned: &std::collections::HashSet<u128>,
11011 ) -> Result<usize> {
11012 if self.retiring.is_empty() {
11013 return Ok(0);
11014 }
11015 let mut reaped = 0;
11016 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11017 for r in std::mem::take(&mut self.retiring) {
11023 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11024 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11025 reaped += 1;
11026 } else {
11027 kept.push(r);
11028 }
11029 }
11030 self.retiring = kept;
11031 if reaped > 0 {
11032 self.persist_manifest(self.current_epoch())?;
11033 }
11034 Ok(reaped)
11035 }
11036
11037 pub(crate) fn has_reapable_retiring(
11038 &self,
11039 min_active: Epoch,
11040 backup_pinned: &std::collections::HashSet<u128>,
11041 ) -> bool {
11042 self.retiring
11043 .iter()
11044 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11045 }
11046
11047 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11048 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11049 return false;
11050 }
11051 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11052 self.run_refs.push(run_ref);
11053 self.indexes_complete = false;
11054 true
11055 }
11056
11057 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11058 self.kek.as_ref()
11059 }
11060
11061 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11062 let mut reader = match self.runs_root.as_deref() {
11063 Some(root) => RunReader::open_file_with_cache(
11064 root.open_regular(format!("r-{run_id}.sr"))?,
11065 self.schema.clone(),
11066 self.kek.clone(),
11067 Some(self.page_cache.clone()),
11068 Some(self.decoded_cache.clone()),
11069 self.table_id,
11070 Some(&self.verified_runs),
11071 None,
11072 )?,
11073 None => RunReader::open_with_cache(
11074 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11075 self.schema.clone(),
11076 self.kek.clone(),
11077 Some(self.page_cache.clone()),
11078 Some(self.decoded_cache.clone()),
11079 self.table_id,
11080 Some(&self.verified_runs),
11081 )?,
11082 };
11083 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11087 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11088 }
11089 Ok(reader)
11090 }
11091
11092 pub(crate) fn run_refs(&self) -> &[RunRef] {
11093 &self.run_refs
11094 }
11095
11096 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11097 self.retiring.iter().map(|run| run.run_id)
11098 }
11099
11100 pub(crate) fn runs_dir(&self) -> PathBuf {
11101 self.runs_root
11102 .as_deref()
11103 .and_then(|root| root.io_path().ok())
11104 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11105 }
11106
11107 pub(crate) fn wal_dir(&self) -> PathBuf {
11108 self.dir.join(WAL_DIR)
11109 }
11110
11111 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11112 self.run_refs = refs;
11113 }
11114
11115 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11116 self.compaction_zstd_level
11117 }
11118
11119 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11120 self.kek.clone()
11121 }
11122
11123 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11127 self.kek.as_ref().map(|k| k.derive_idx_key())
11128 }
11129
11130 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11134 self.kek.as_ref().map(|k| *k.derive_meta_key())
11135 }
11136
11137 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11140 self.column_keys
11141 .iter()
11142 .map(|(&id, &(_, scheme))| (id, scheme))
11143 .collect()
11144 }
11145
11146 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11151 self.tokenize_value_enc(column_id, v)
11152 }
11153
11154 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11155 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11156 let (key, scheme) = self.column_keys.get(&column_id)?;
11157 let token: Vec<u8> = match (*scheme, v) {
11158 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11159 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11160 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11161 _ => hmac_token(key, &v.encode_key()).to_vec(),
11162 };
11163 Some(Value::Bytes(token))
11164 }
11165
11166 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11168 self.index_lookup_key_bytes(column_id, &v.encode_key())
11169 }
11170
11171 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11174 {
11175 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11176 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11177 if *scheme == SCHEME_HMAC_EQ {
11178 return hmac_token(key, encoded).to_vec();
11179 }
11180 }
11181 }
11182 let _ = column_id;
11183 encoded.to_vec()
11184 }
11185}
11186
11187fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11188 let columnar::NativeColumn::Int64 { data, validity } = col else {
11189 return false;
11190 };
11191 if data.len() < n || !columnar::all_non_null(validity, n) {
11192 return false;
11193 }
11194 data.iter()
11195 .take(n)
11196 .zip(data.iter().skip(1))
11197 .all(|(a, b)| a < b)
11198}
11199
11200#[derive(Debug, Clone)]
11204pub struct ColumnStat {
11205 pub min: Option<Value>,
11206 pub max: Option<Value>,
11207 pub null_count: u64,
11208}
11209
11210#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11212pub enum NativeAgg {
11213 Count,
11214 Sum,
11215 Min,
11216 Max,
11217 Avg,
11218}
11219
11220#[derive(Debug, Clone, PartialEq)]
11222pub enum NativeAggResult {
11223 Count(u64),
11224 Int(i64),
11225 Float(f64),
11226 Null,
11228}
11229
11230#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11232pub enum ApproxAgg {
11233 Count,
11234 Sum,
11235 Avg,
11236}
11237
11238#[derive(Debug, Clone)]
11242pub struct ApproxResult {
11243 pub point: f64,
11245 pub ci_low: f64,
11247 pub ci_high: f64,
11249 pub n_population: u64,
11251 pub n_sample_live: usize,
11253 pub n_passing: usize,
11255}
11256
11257#[derive(Debug, Clone, PartialEq)]
11262pub enum AggState {
11263 Count(u64),
11265 SumI {
11267 sum: i128,
11268 count: u64,
11269 },
11270 SumF {
11272 sum: f64,
11273 count: u64,
11274 },
11275 AvgI {
11277 sum: i128,
11278 count: u64,
11279 },
11280 AvgF {
11282 sum: f64,
11283 count: u64,
11284 },
11285 MinI(i64),
11287 MaxI(i64),
11288 MinF(f64),
11290 MaxF(f64),
11291 Empty,
11293}
11294
11295impl AggState {
11296 pub fn merge(self, other: AggState) -> AggState {
11298 use AggState::*;
11299 match (self, other) {
11300 (Empty, x) | (x, Empty) => x,
11301 (Count(a), Count(b)) => Count(a + b),
11302 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11303 sum: sa + sb,
11304 count: ca + cb,
11305 },
11306 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11307 sum: sa + sb,
11308 count: ca + cb,
11309 },
11310 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11311 sum: sa + sb,
11312 count: ca + cb,
11313 },
11314 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11315 sum: sa + sb,
11316 count: ca + cb,
11317 },
11318 (MinI(a), MinI(b)) => MinI(a.min(b)),
11319 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11320 (MinF(a), MinF(b)) => MinF(a.min(b)),
11321 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11322 _ => Empty, }
11324 }
11325
11326 pub fn point(&self) -> Option<f64> {
11328 match self {
11329 AggState::Count(n) => Some(*n as f64),
11330 AggState::SumI { sum, .. } => Some(*sum as f64),
11331 AggState::SumF { sum, .. } => Some(*sum),
11332 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11333 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11334 AggState::MinI(n) => Some(*n as f64),
11335 AggState::MaxI(n) => Some(*n as f64),
11336 AggState::MinF(n) => Some(*n),
11337 AggState::MaxF(n) => Some(*n),
11338 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11339 }
11340 }
11341
11342 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11346 let is_float = matches!(ty, Some(TypeId::Float64));
11347 match (agg, result) {
11348 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11349 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11350 sum: x as i128,
11351 count: 1, },
11353 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11354 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11355 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11356 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11357 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11358 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11359 (NativeAgg::Count, _) => AggState::Empty,
11360 (_, NativeAggResult::Null) => AggState::Empty,
11361 _ => {
11362 let _ = is_float;
11363 AggState::Empty
11364 }
11365 }
11366 }
11367}
11368
11369#[derive(Debug, Clone)]
11372pub struct CachedAgg {
11373 pub state: AggState,
11374 pub watermark: u64,
11375 pub epoch: u64,
11376}
11377
11378#[derive(Debug, Clone)]
11380pub struct IncrementalAggResult {
11381 pub state: AggState,
11383 pub incremental: bool,
11386 pub delta_rows: u64,
11388}
11389
11390fn agg_state_from_rows(
11394 rows: &[Row],
11395 conditions: &[crate::query::Condition],
11396 index_sets: &[RowIdSet],
11397 column: Option<u16>,
11398 agg: NativeAgg,
11399 schema: &Schema,
11400 control: Option<&crate::ExecutionControl>,
11401) -> Result<AggState> {
11402 let mut count: u64 = 0;
11403 let mut sum_i: i128 = 0;
11404 let mut sum_f: f64 = 0.0;
11405 let mut mn_i: i64 = i64::MAX;
11406 let mut mx_i: i64 = i64::MIN;
11407 let mut mn_f: f64 = f64::INFINITY;
11408 let mut mx_f: f64 = f64::NEG_INFINITY;
11409 let mut saw_int = false;
11410 let mut saw_float = false;
11411 for (index, r) in rows.iter().enumerate() {
11412 execution_checkpoint(control, index)?;
11413 if !conditions
11414 .iter()
11415 .all(|c| condition_matches_row(c, r, schema))
11416 {
11417 continue;
11418 }
11419 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11420 continue;
11421 }
11422 match agg {
11423 NativeAgg::Count => match column {
11424 None => count += 1,
11426 Some(cid) => match r.columns.get(&cid) {
11429 None | Some(Value::Null) => {}
11430 Some(_) => count += 1,
11431 },
11432 },
11433 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11434 Some(Value::Int64(n)) => {
11435 count += 1;
11436 sum_i += *n as i128;
11437 mn_i = mn_i.min(*n);
11438 mx_i = mx_i.max(*n);
11439 saw_int = true;
11440 }
11441 Some(Value::Float64(f)) => {
11442 count += 1;
11443 sum_f += f;
11444 mn_f = mn_f.min(*f);
11445 mx_f = mx_f.max(*f);
11446 saw_float = true;
11447 }
11448 _ => {}
11449 },
11450 }
11451 }
11452 Ok(match agg {
11453 NativeAgg::Count => {
11454 if count == 0 {
11455 AggState::Empty
11456 } else {
11457 AggState::Count(count)
11458 }
11459 }
11460 NativeAgg::Sum => {
11461 if count == 0 {
11462 AggState::Empty
11463 } else if saw_int {
11464 AggState::SumI { sum: sum_i, count }
11465 } else {
11466 AggState::SumF { sum: sum_f, count }
11467 }
11468 }
11469 NativeAgg::Avg => {
11470 if count == 0 {
11471 AggState::Empty
11472 } else if saw_int {
11473 AggState::AvgI { sum: sum_i, count }
11474 } else {
11475 AggState::AvgF { sum: sum_f, count }
11476 }
11477 }
11478 NativeAgg::Min => {
11479 if !saw_int && !saw_float {
11480 AggState::Empty
11481 } else if saw_int {
11482 AggState::MinI(mn_i)
11483 } else {
11484 AggState::MinF(mn_f)
11485 }
11486 }
11487 NativeAgg::Max => {
11488 if !saw_int && !saw_float {
11489 AggState::Empty
11490 } else if saw_int {
11491 AggState::MaxI(mx_i)
11492 } else {
11493 AggState::MaxF(mx_f)
11494 }
11495 }
11496 })
11497}
11498
11499fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11503 use crate::query::Condition;
11504 match c {
11505 Condition::Pk(key) => match schema.primary_key() {
11506 Some(pk) => row
11507 .columns
11508 .get(&pk.id)
11509 .map(|v| v.encode_key() == *key)
11510 .unwrap_or(false),
11511 None => false,
11512 },
11513 Condition::BitmapEq { column_id, value } => row
11514 .columns
11515 .get(column_id)
11516 .map(|v| v.encode_key() == *value)
11517 .unwrap_or(false),
11518 Condition::BitmapIn { column_id, values } => {
11519 let key = row.columns.get(column_id).map(|v| v.encode_key());
11520 match key {
11521 Some(k) => values.contains(&k),
11522 None => false,
11523 }
11524 }
11525 Condition::BytesPrefix { column_id, prefix } => row
11526 .columns
11527 .get(column_id)
11528 .map(|v| v.encode_key().starts_with(prefix))
11529 .unwrap_or(false),
11530 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11531 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11532 _ => false,
11533 },
11534 Condition::RangeF64 {
11535 column_id,
11536 lo,
11537 lo_inclusive,
11538 hi,
11539 hi_inclusive,
11540 } => match row.columns.get(column_id) {
11541 Some(Value::Float64(n)) => {
11542 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11543 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11544 lo_ok && hi_ok
11545 }
11546 _ => false,
11547 },
11548 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11549 Some(Value::Bytes(b)) => {
11550 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11551 }
11552 _ => false,
11553 },
11554 Condition::FmContainsAll {
11555 column_id,
11556 patterns,
11557 } => match row.columns.get(column_id) {
11558 Some(Value::Bytes(b)) => patterns
11559 .iter()
11560 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11561 _ => false,
11562 },
11563 Condition::Ann { .. }
11564 | Condition::SparseMatch { .. }
11565 | Condition::MinHashSimilar { .. } => true,
11566 Condition::IsNull { column_id } => {
11567 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11568 }
11569 Condition::IsNotNull { column_id } => {
11570 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11571 }
11572 }
11573}
11574
11575fn as_f64(v: Option<&Value>) -> Option<f64> {
11577 match v {
11578 Some(Value::Int64(n)) => Some(*n as f64),
11579 Some(Value::Float64(f)) => Some(*f),
11580 _ => None,
11581 }
11582}
11583
11584fn accumulate_int(
11588 cursor: &mut dyn crate::cursor::Cursor,
11589 control: Option<&crate::ExecutionControl>,
11590) -> Result<(u64, i128, i64, i64)> {
11591 let mut count: u64 = 0;
11592 let mut sum: i128 = 0;
11593 let mut mn: i64 = i64::MAX;
11594 let mut mx: i64 = i64::MIN;
11595 while let Some(cols) = cursor.next_batch()? {
11596 execution_checkpoint(control, 0)?;
11597 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11598 if crate::columnar::all_non_null(validity, data.len()) {
11599 count += data.len() as u64;
11601 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11602 execution_checkpoint(control, chunk_index * 1024)?;
11603 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11604 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11605 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11606 }
11607 } else {
11608 for (i, &v) in data.iter().enumerate() {
11609 execution_checkpoint(control, i)?;
11610 if crate::columnar::validity_bit(validity, i) {
11611 count += 1;
11612 sum += v as i128;
11613 mn = mn.min(v);
11614 mx = mx.max(v);
11615 }
11616 }
11617 }
11618 }
11619 }
11620 Ok((count, sum, mn, mx))
11621}
11622
11623fn accumulate_float(
11625 cursor: &mut dyn crate::cursor::Cursor,
11626 control: Option<&crate::ExecutionControl>,
11627) -> Result<(u64, f64, f64, f64)> {
11628 let mut count: u64 = 0;
11629 let mut sum: f64 = 0.0;
11630 let mut mn: f64 = f64::INFINITY;
11631 let mut mx: f64 = f64::NEG_INFINITY;
11632 while let Some(cols) = cursor.next_batch()? {
11633 execution_checkpoint(control, 0)?;
11634 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11635 if crate::columnar::all_non_null(validity, data.len()) {
11636 count += data.len() as u64;
11637 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11638 execution_checkpoint(control, chunk_index * 1024)?;
11639 sum += chunk.iter().sum::<f64>();
11640 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11641 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11642 }
11643 } else {
11644 for (i, &v) in data.iter().enumerate() {
11645 execution_checkpoint(control, i)?;
11646 if crate::columnar::validity_bit(validity, i) {
11647 count += 1;
11648 sum += v;
11649 mn = mn.min(v);
11650 mx = mx.max(v);
11651 }
11652 }
11653 }
11654 }
11655 }
11656 Ok((count, sum, mn, mx))
11657}
11658
11659#[inline]
11660fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11661 if index.is_multiple_of(256) {
11662 control
11663 .map(crate::ExecutionControl::checkpoint)
11664 .transpose()?;
11665 }
11666 Ok(())
11667}
11668
11669fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11670 if count == 0 && !matches!(agg, NativeAgg::Count) {
11671 return NativeAggResult::Null;
11672 }
11673 match agg {
11674 NativeAgg::Count => NativeAggResult::Count(count),
11675 NativeAgg::Sum => match sum.try_into() {
11678 Ok(v) => NativeAggResult::Int(v),
11679 Err(_) => NativeAggResult::Null,
11680 },
11681 NativeAgg::Min => NativeAggResult::Int(mn),
11682 NativeAgg::Max => NativeAggResult::Int(mx),
11683 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11684 }
11685}
11686
11687fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11688 if count == 0 && !matches!(agg, NativeAgg::Count) {
11689 return NativeAggResult::Null;
11690 }
11691 match agg {
11692 NativeAgg::Count => NativeAggResult::Count(count),
11693 NativeAgg::Sum => NativeAggResult::Float(sum),
11694 NativeAgg::Min => NativeAggResult::Float(mn),
11695 NativeAgg::Max => NativeAggResult::Float(mx),
11696 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11697 }
11698}
11699
11700fn agg_int(
11703 stats: &[crate::page::PageStat],
11704 decode: fn(Option<&[u8]>) -> Option<i64>,
11705) -> Option<(Option<i64>, Option<i64>, u64)> {
11706 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11707 let mut any = false;
11708 for s in stats {
11709 if let Some(v) = decode(s.min.as_deref()) {
11710 mn = mn.min(v);
11711 any = true;
11712 }
11713 if let Some(v) = decode(s.max.as_deref()) {
11714 mx = mx.max(v);
11715 any = true;
11716 }
11717 nulls += s.null_count;
11718 }
11719 any.then_some((Some(mn), Some(mx), nulls))
11720}
11721
11722fn agg_float(
11724 stats: &[crate::page::PageStat],
11725 decode: fn(Option<&[u8]>) -> Option<f64>,
11726) -> Option<(Option<f64>, Option<f64>, u64)> {
11727 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11728 let mut any = false;
11729 for s in stats {
11730 if let Some(v) = decode(s.min.as_deref()) {
11731 mn = mn.min(v);
11732 any = true;
11733 }
11734 if let Some(v) = decode(s.max.as_deref()) {
11735 mx = mx.max(v);
11736 any = true;
11737 }
11738 nulls += s.null_count;
11739 }
11740 any.then_some((Some(mn), Some(mx), nulls))
11741}
11742
11743type SecondaryIndexes = (
11745 HashMap<u16, BitmapIndex>,
11746 HashMap<u16, AnnIndex>,
11747 HashMap<u16, FmIndex>,
11748 HashMap<u16, SparseIndex>,
11749 HashMap<u16, MinHashIndex>,
11750);
11751
11752fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11753 let mut bitmap = HashMap::new();
11754 let mut ann = HashMap::new();
11755 let mut fm = HashMap::new();
11756 let mut sparse = HashMap::new();
11757 let mut minhash = HashMap::new();
11758 for idef in &schema.indexes {
11759 match idef.kind {
11760 IndexKind::Bitmap => {
11761 bitmap.insert(idef.column_id, BitmapIndex::new());
11762 }
11763 IndexKind::Ann => {
11764 let dim = schema
11765 .columns
11766 .iter()
11767 .find(|c| c.id == idef.column_id)
11768 .and_then(|c| match c.ty {
11769 TypeId::Embedding { dim } => Some(dim as usize),
11770 _ => None,
11771 })
11772 .unwrap_or(0);
11773 let options = idef.options.ann.clone().unwrap_or_default();
11774 ann.insert(
11775 idef.column_id,
11776 AnnIndex::with_full_options(
11777 dim,
11778 options.m,
11779 options.ef_construction,
11780 options.ef_search,
11781 &options,
11782 ),
11783 );
11784 }
11785 IndexKind::FmIndex => {
11786 fm.insert(idef.column_id, FmIndex::new());
11787 }
11788 IndexKind::Sparse => {
11789 sparse.insert(idef.column_id, SparseIndex::new());
11790 }
11791 IndexKind::MinHash => {
11792 let options = idef.options.minhash.clone().unwrap_or_default();
11793 minhash.insert(
11794 idef.column_id,
11795 MinHashIndex::with_options(options.permutations, options.bands),
11796 );
11797 }
11798 _ => {}
11799 }
11800 }
11801 (bitmap, ann, fm, sparse, minhash)
11802}
11803
11804const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11805 | ColumnFlags::AUTO_INCREMENT
11806 | ColumnFlags::ENCRYPTED
11807 | ColumnFlags::ENCRYPTED_INDEXABLE
11808 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11809
11810fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11811 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11812 return Err(MongrelError::Schema(
11813 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11814 ));
11815 }
11816 Ok(())
11817}
11818
11819fn validate_alter_column_type(
11820 schema: &Schema,
11821 old: &ColumnDef,
11822 next: &ColumnDef,
11823 has_stored_versions: bool,
11824) -> Result<()> {
11825 if old.ty == next.ty {
11826 return Ok(());
11827 }
11828 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11829 return Err(MongrelError::Schema(format!(
11830 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11831 old.name
11832 )));
11833 }
11834 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11835 return Ok(());
11836 }
11837 Err(MongrelError::Schema(format!(
11838 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11839 old.ty, next.ty
11840 )))
11841}
11842
11843fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11844 matches!(
11845 (old, new),
11846 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11847 )
11848}
11849
11850fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11856 let mut prev: Option<i64> = None;
11857 for r in rows {
11858 match r.columns.get(&pk_id) {
11859 Some(Value::Int64(v)) => {
11860 if prev.is_some_and(|p| p >= *v) {
11861 return false;
11862 }
11863 prev = Some(*v);
11864 }
11865 _ => return false,
11866 }
11867 }
11868 true
11869}
11870
11871#[allow(clippy::too_many_arguments)]
11872fn index_into(
11873 schema: &Schema,
11874 row: &Row,
11875 hot: &mut HotIndex,
11876 bitmap: &mut HashMap<u16, BitmapIndex>,
11877 ann: &mut HashMap<u16, AnnIndex>,
11878 fm: &mut HashMap<u16, FmIndex>,
11879 sparse: &mut HashMap<u16, SparseIndex>,
11880 minhash: &mut HashMap<u16, MinHashIndex>,
11881) {
11882 for idef in &schema.indexes {
11883 let Some(val) = row.columns.get(&idef.column_id) else {
11884 continue;
11885 };
11886 match idef.kind {
11887 IndexKind::Bitmap => {
11888 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11889 b.insert(val.encode_key(), row.row_id);
11890 }
11891 }
11892 IndexKind::Ann => {
11893 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
11894 a.insert_validated(v, row.row_id);
11895 }
11896 }
11897 IndexKind::FmIndex => {
11898 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11899 f.insert(b.clone(), row.row_id);
11900 }
11901 }
11902 IndexKind::Sparse => {
11903 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11904 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11907 s.insert(&terms, row.row_id);
11908 }
11909 }
11910 }
11911 IndexKind::MinHash => {
11912 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11913 let tokens = crate::index::token_hashes_from_bytes(b);
11916 mh.insert(&tokens, row.row_id);
11917 }
11918 }
11919 _ => {}
11920 }
11921 }
11922 if let Some(pk_col) = schema.primary_key() {
11923 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11924 hot.insert(pk_val.encode_key(), row.row_id);
11925 }
11926 }
11927}
11928
11929#[allow(clippy::too_many_arguments)]
11932fn index_into_single(
11933 idef: &IndexDef,
11934 _schema: &Schema,
11935 row: &Row,
11936 _hot: &mut HotIndex,
11937 bitmap: &mut HashMap<u16, BitmapIndex>,
11938 ann: &mut HashMap<u16, AnnIndex>,
11939 fm: &mut HashMap<u16, FmIndex>,
11940 sparse: &mut HashMap<u16, SparseIndex>,
11941 minhash: &mut HashMap<u16, MinHashIndex>,
11942) {
11943 let Some(val) = row.columns.get(&idef.column_id) else {
11944 return;
11945 };
11946 match idef.kind {
11947 IndexKind::Bitmap => {
11948 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11949 b.insert(val.encode_key(), row.row_id);
11950 }
11951 }
11952 IndexKind::Ann => {
11953 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
11954 a.insert_validated(v, row.row_id);
11955 }
11956 }
11957 IndexKind::FmIndex => {
11958 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11959 f.insert(b.clone(), row.row_id);
11960 }
11961 }
11962 IndexKind::Sparse => {
11963 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11964 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11965 s.insert(&terms, row.row_id);
11966 }
11967 }
11968 }
11969 IndexKind::MinHash => {
11970 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11971 let tokens = crate::index::token_hashes_from_bytes(b);
11972 mh.insert(&tokens, row.row_id);
11973 }
11974 }
11975 _ => {}
11976 }
11977}
11978
11979fn eval_partial_predicate(
11985 pred: &str,
11986 columns_map: &HashMap<u16, &Value>,
11987 name_to_id: &HashMap<&str, u16>,
11988) -> bool {
11989 let lower = pred.trim().to_ascii_lowercase();
11990 if let Some(rest) = lower.strip_suffix(" is not null") {
11992 let col_name = rest.trim();
11993 if let Some(col_id) = name_to_id.get(col_name) {
11994 return columns_map
11995 .get(col_id)
11996 .is_some_and(|v| !matches!(v, Value::Null));
11997 }
11998 }
11999 if let Some(rest) = lower.strip_suffix(" is null") {
12001 let col_name = rest.trim();
12002 if let Some(col_id) = name_to_id.get(col_name) {
12003 return columns_map
12004 .get(col_id)
12005 .is_none_or(|v| matches!(v, Value::Null));
12006 }
12007 }
12008 true
12011}
12012
12013#[allow(dead_code)]
12019fn bulk_index_key(
12020 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12021 column_id: u16,
12022 ty: TypeId,
12023 col: &columnar::NativeColumn,
12024 i: usize,
12025) -> Option<Vec<u8>> {
12026 let encoded = columnar::encode_key_native(ty, col, i)?;
12027 {
12028 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12029 if let Some((key, scheme)) = column_keys.get(&column_id) {
12030 return Some(match (*scheme, col) {
12031 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12032 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12033 ope_token_i64(key, data[i]).to_vec()
12034 }
12035 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12036 ope_token_f64(key, data[i]).to_vec()
12037 }
12038 _ => hmac_token(key, &encoded).to_vec(),
12039 });
12040 }
12041 }
12042 Some(encoded)
12043}
12044
12045pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12046 write_schema_with_after(dir, schema, || {})
12047}
12048
12049pub(crate) fn write_schema_durable(
12050 root: &crate::durable_file::DurableRoot,
12051 schema: &Schema,
12052) -> Result<()> {
12053 write_schema_durable_with_after(root, schema, || {})
12054}
12055
12056fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12057where
12058 F: FnOnce(),
12059{
12060 let json = serde_json::to_string_pretty(schema)
12061 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12062 crate::durable_file::write_atomic_with_after(
12063 &dir.join(SCHEMA_FILENAME),
12064 json.as_bytes(),
12065 after_publish,
12066 )?;
12067 Ok(())
12068}
12069
12070fn write_schema_durable_with_after<F>(
12071 root: &crate::durable_file::DurableRoot,
12072 schema: &Schema,
12073 after_publish: F,
12074) -> Result<()>
12075where
12076 F: FnOnce(),
12077{
12078 let json = serde_json::to_string_pretty(schema)
12079 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12080 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12081 Ok(())
12082}
12083
12084fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12085 let mut schema_published = false;
12086 let schema_result = match table._root_guard.as_deref() {
12087 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12088 schema_published = true;
12089 }),
12090 None => write_schema_with_after(&table.dir, &table.schema, || {
12091 schema_published = true;
12092 }),
12093 };
12094 if schema_result.is_err() && !schema_published {
12095 return schema_result;
12096 }
12097 match table.persist_manifest(table.current_epoch()) {
12098 Ok(()) => Ok(()),
12099 Err(manifest_error) => Err(match schema_result {
12100 Ok(()) => manifest_error,
12101 Err(schema_error) => MongrelError::Other(format!(
12102 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12103 )),
12104 }),
12105 }
12106}
12107
12108fn read_schema(dir: &Path) -> Result<Schema> {
12109 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12110 read_schema_file(file)
12111}
12112
12113fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12114 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12115 use std::io::Read;
12116
12117 let length = file.metadata()?.len();
12118 if length > MAX_SCHEMA_BYTES {
12119 return Err(MongrelError::ResourceLimitExceeded {
12120 resource: "schema bytes",
12121 requested: usize::try_from(length).unwrap_or(usize::MAX),
12122 limit: MAX_SCHEMA_BYTES as usize,
12123 });
12124 }
12125 let mut bytes = Vec::with_capacity(length as usize);
12126 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12127 if bytes.len() as u64 != length {
12128 return Err(MongrelError::Schema(
12129 "schema length changed while reading".into(),
12130 ));
12131 }
12132 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12133}
12134
12135fn preflight_standalone_open(
12136 dir: &Path,
12137 runs_root: Option<&crate::durable_file::DurableRoot>,
12138 idx_root: Option<&crate::durable_file::DurableRoot>,
12139 manifest: &Manifest,
12140 schema: &Schema,
12141 records: &[crate::wal::Record],
12142 kek: Option<Arc<Kek>>,
12143) -> Result<()> {
12144 crate::wal::validate_shared_transaction_framing(records)?;
12145 if manifest.schema_id > schema.schema_id
12146 || manifest.flushed_epoch > manifest.current_epoch
12147 || manifest.global_idx_epoch > manifest.current_epoch
12148 || manifest.next_row_id == u64::MAX
12149 || manifest.auto_inc_next < 0
12150 || manifest.auto_inc_next == i64::MAX
12151 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12152 {
12153 return Err(MongrelError::InvalidArgument(
12154 "manifest counters or schema identity are invalid".into(),
12155 ));
12156 }
12157 let mut run_ids = HashSet::new();
12158 let mut maximum_row_id = None::<u64>;
12159 for run in &manifest.runs {
12160 if run.run_id >= u64::MAX as u128
12161 || !run_ids.insert(run.run_id)
12162 || run.epoch_created > manifest.current_epoch
12163 {
12164 return Err(MongrelError::InvalidArgument(
12165 "manifest contains an invalid or duplicate active run".into(),
12166 ));
12167 }
12168 let mut reader = match runs_root {
12169 Some(root) => RunReader::open_file(
12170 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12171 schema.clone(),
12172 kek.clone(),
12173 )?,
12174 None => RunReader::open(
12175 dir.join(RUNS_DIR)
12176 .join(format!("r-{}.sr", run.run_id as u64)),
12177 schema.clone(),
12178 kek.clone(),
12179 )?,
12180 };
12181 let header = reader.header();
12182 if header.run_id != run.run_id
12183 || header.level != run.level
12184 || header.row_count != run.row_count
12185 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12186 || header.is_uniform_epoch() && header.epoch_created != 0
12187 || header.schema_id > schema.schema_id
12188 {
12189 return Err(MongrelError::InvalidArgument(format!(
12190 "run {} differs from its manifest",
12191 run.run_id
12192 )));
12193 }
12194 if header.row_count != 0 {
12195 maximum_row_id = Some(
12196 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12197 );
12198 }
12199 reader.validate_all_pages()?;
12200 }
12201 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12202 return Err(MongrelError::InvalidArgument(
12203 "manifest next_row_id does not advance beyond persisted rows".into(),
12204 ));
12205 }
12206 for run in &manifest.retiring {
12207 if run.run_id >= u64::MAX as u128
12208 || run.retire_epoch > manifest.current_epoch
12209 || !run_ids.insert(run.run_id)
12210 {
12211 return Err(MongrelError::InvalidArgument(
12212 "manifest contains an invalid or duplicate retired run".into(),
12213 ));
12214 }
12215 }
12216 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12217 match idx_root {
12218 Some(root) => {
12219 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12220 }
12221 None => {
12222 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12223 }
12224 }
12225
12226 let committed = records
12227 .iter()
12228 .filter_map(|record| match record.op {
12229 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12230 _ => None,
12231 })
12232 .collect::<HashMap<_, _>>();
12233 for record in records {
12234 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12235 continue;
12236 };
12237 match &record.op {
12238 Op::Put { table_id, rows } => {
12239 if *table_id != manifest.table_id {
12240 return Err(MongrelError::CorruptWal {
12241 offset: record.seq.0,
12242 reason: format!(
12243 "private WAL record references table {table_id}, expected {}",
12244 manifest.table_id
12245 ),
12246 });
12247 }
12248 let rows: Vec<Row> =
12249 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12250 offset: record.seq.0,
12251 reason: format!("committed Put payload could not be decoded: {error}"),
12252 })?;
12253 for row in rows {
12254 if row.deleted || row.row_id.0 == u64::MAX {
12255 return Err(MongrelError::CorruptWal {
12256 offset: record.seq.0,
12257 reason: "committed Put contains an invalid row identity".into(),
12258 });
12259 }
12260 let cells = row.columns.into_iter().collect::<Vec<_>>();
12261 schema
12262 .validate_values(&cells)
12263 .map_err(|error| MongrelError::CorruptWal {
12264 offset: record.seq.0,
12265 reason: format!("committed Put violates table schema: {error}"),
12266 })?;
12267 if schema.auto_increment_column().is_some_and(|column| {
12268 matches!(
12269 cells.iter().find(|(id, _)| *id == column.id),
12270 Some((_, Value::Int64(value))) if *value == i64::MAX
12271 )
12272 }) {
12273 return Err(MongrelError::CorruptWal {
12274 offset: record.seq.0,
12275 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12276 });
12277 }
12278 }
12279 }
12280 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12281 if *table_id != manifest.table_id =>
12282 {
12283 return Err(MongrelError::CorruptWal {
12284 offset: record.seq.0,
12285 reason: format!(
12286 "private WAL record references table {table_id}, expected {}",
12287 manifest.table_id
12288 ),
12289 });
12290 }
12291 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12292 return Err(MongrelError::CorruptWal {
12293 offset: record.seq.0,
12294 reason: "private WAL contains shared spilled-run metadata".into(),
12295 });
12296 }
12297 _ => {}
12298 }
12299 }
12300 Ok(())
12301}
12302
12303fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12304 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12305}
12306
12307fn wal_segment_number(path: &Path) -> Option<u64> {
12308 path.file_stem()
12309 .and_then(|stem| stem.to_str())
12310 .and_then(|stem| stem.strip_prefix("seg-"))
12311 .and_then(|number| number.parse().ok())
12312}
12313
12314fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12315 let n = list_wal_numbers(wal_dir)?;
12316 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12317}
12318
12319fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12320 list_wal_numbers(wal_dir)?
12321 .map(|maximum| {
12322 maximum
12323 .checked_add(1)
12324 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12325 })
12326 .unwrap_or(Ok(0))
12327}
12328
12329fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12330 let mut max_n = None;
12331 let entries = match std::fs::read_dir(wal_dir) {
12332 Ok(entries) => entries,
12333 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12334 Err(error) => return Err(error.into()),
12335 };
12336 for entry in entries {
12337 let entry = entry?;
12338 let fname = entry.file_name();
12339 let Some(s) = fname.to_str() else {
12340 continue;
12341 };
12342 let Some(stripped) = s.strip_prefix("seg-") else {
12343 continue;
12344 };
12345 let Some(number) = stripped.strip_suffix(".wal") else {
12346 return Err(MongrelError::CorruptWal {
12347 offset: 0,
12348 reason: format!("malformed WAL segment name {s:?}"),
12349 });
12350 };
12351 let n = number
12352 .parse::<u32>()
12353 .map_err(|_| MongrelError::CorruptWal {
12354 offset: 0,
12355 reason: format!("malformed WAL segment name {s:?}"),
12356 })?;
12357 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12358 return Err(MongrelError::CorruptWal {
12359 offset: n as u64,
12360 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12361 });
12362 }
12363 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12364 }
12365 Ok(max_n)
12366}