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, Box<AnnIndex>),
601 Sparse(u16, SparseIndex),
602 MinHash(u16, MinHashIndex),
603}
604
605#[derive(Clone)]
611pub struct TableDeltas {
612 memtable: Memtable,
613 mutable_run: MutableRun,
614 hot: HotIndex,
615 pk_by_row: ReversePkMap,
616}
617
618impl ReadGeneration {
619 fn empty(schema: &Schema) -> Self {
622 Self {
623 schema: Arc::new(schema.clone()),
624 base_runs: Arc::new(Vec::new()),
625 deltas: TableDeltas {
626 memtable: Memtable::new(),
627 mutable_run: MutableRun::new(),
628 hot: HotIndex::new(),
629 pk_by_row: ReversePkMap::new(),
630 },
631 indexes: Arc::new(IndexGeneration::default()),
632 visible_through: Epoch(0),
633 }
634 }
635
636 pub fn schema(&self) -> &Arc<Schema> {
638 &self.schema
639 }
640
641 pub fn base_runs(&self) -> &[RunRef] {
643 &self.base_runs
644 }
645
646 pub fn indexes(&self) -> &Arc<IndexGeneration> {
648 &self.indexes
649 }
650
651 pub fn visible_through(&self) -> Epoch {
653 self.visible_through
654 }
655
656 pub fn deltas(&self) -> &TableDeltas {
660 &self.deltas
661 }
662}
663
664impl TableDeltas {
665 pub fn approx_bytes(&self) -> u64 {
668 self.memtable
669 .approx_bytes()
670 .saturating_add(self.mutable_run.approx_bytes())
671 }
672
673 pub fn memtable_len(&self) -> usize {
675 self.memtable.len()
676 }
677
678 pub fn mutable_run_len(&self) -> usize {
680 self.mutable_run.len()
681 }
682
683 pub fn hot_len(&self) -> usize {
685 self.hot.len()
686 }
687
688 pub fn reverse_pk_len(&self) -> usize {
690 self.pk_by_row.entries().len()
691 }
692}
693
694#[derive(Clone)]
696pub struct Table {
697 dir: PathBuf,
698 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
699 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
700 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
701 table_id: u64,
702 name: String,
706 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
711 read_only: bool,
714 durable_commit_failed: bool,
718 wal: WalSink,
719 memtable: Memtable,
720 mutable_run: MutableRun,
725 mutable_run_spill_bytes: u64,
727 compaction_zstd_level: i32,
730 allocator: RowIdAllocator,
731 epoch: Arc<EpochAuthority>,
732 data_generation: u64,
735 schema: Schema,
736 hot: HotIndex,
737 kek: Option<Arc<Kek>>,
740 column_keys: HashMap<u16, ([u8; 32], u8)>,
744 run_refs: Vec<RunRef>,
745 retiring: Vec<crate::manifest::RetiredRun>,
748 next_run_id: u64,
749 sync_byte_threshold: u64,
750 current_txn_id: u64,
755 pending_private_mutations: bool,
759 bitmap: HashMap<u16, BitmapIndex>,
760 ann: HashMap<u16, AnnIndex>,
761 fm: HashMap<u16, FmIndex>,
762 sparse: HashMap<u16, SparseIndex>,
763 minhash: HashMap<u16, MinHashIndex>,
764 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
767 pk_by_row: ReversePkMap,
769 pinned: BTreeMap<Epoch, usize>,
772 pub(crate) live_count: u64,
775 reservoir: crate::reservoir::Reservoir,
778 reservoir_complete: bool,
786 had_deletes: bool,
790 agg_cache: Arc<HashMap<u64, CachedAgg>>,
794 global_idx_epoch: u64,
798 indexes_complete: bool,
803 index_build_policy: IndexBuildPolicy,
805 pk_by_row_complete: bool,
812 flushed_epoch: u64,
815 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
818 snapshots: Arc<crate::retention::SnapshotRegistry>,
821 commit_lock: Arc<parking_lot::Mutex<()>>,
823 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
826 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
836 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
845 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
847 pending_delete_rids: roaring::RoaringBitmap,
850 pending_put_cols: std::collections::HashSet<u16>,
853 pending_rows: Vec<Row>,
859 pending_rows_auto_inc: Vec<bool>,
860 pending_dels: Vec<RowId>,
863 pending_truncate: Option<Epoch>,
867 auto_inc: Option<AutoIncState>,
870 ttl: Option<TtlPolicy>,
873 pins: Arc<crate::retention::PinRegistry>,
880 published: Arc<ArcSwap<ReadGeneration>>,
885 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
890}
891
892const _: () = {
899 const fn assert_sync<T: ?Sized + Sync>() {}
900 assert_sync::<Table>();
901};
902
903enum CachedData {
909 Rows(Arc<Vec<Row>>),
910 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
911}
912
913impl CachedData {
914 fn approx_bytes(&self) -> u64 {
915 match self {
916 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
917 CachedData::Columns(c) => c
918 .iter()
919 .map(|(_, c)| c.approx_bytes())
920 .sum::<u64>()
921 .saturating_add(c.len() as u64 * 16),
922 }
923 }
924}
925
926struct CachedEntry {
930 data: CachedData,
931 footprint: roaring::RoaringBitmap,
932 condition_cols: Vec<u16>,
933}
934
935struct ResultCache {
946 entries: std::collections::HashMap<u64, CachedEntry>,
947 order: 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 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
1454}
1455
1456enum WalSink {
1459 Private(Wal),
1460 Shared(SharedWalCtx),
1461 ReadOnly,
1462}
1463
1464impl Clone for WalSink {
1465 fn clone(&self) -> Self {
1466 match self {
1467 Self::Shared(shared) => Self::Shared(shared.clone()),
1468 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1469 }
1470 }
1471}
1472
1473impl SharedCtx {
1474 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1478 let n_shards = if cache_dir.is_some() {
1482 1
1483 } else {
1484 crate::cache::CACHE_SHARDS
1485 };
1486 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1487 let page_cache = if let Some(d) = cache_dir {
1488 Arc::new(crate::cache::Sharded::new(1, || {
1489 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1490 }))
1491 } else {
1492 Arc::new(crate::cache::Sharded::new(n_shards, || {
1493 crate::cache::PageCache::new(per_shard)
1494 }))
1495 };
1496 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1497 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1498 crate::cache::CACHE_SHARDS,
1499 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1500 ));
1501 Self {
1502 root_guard: None,
1503 epoch: Arc::new(EpochAuthority::new(0)),
1504 page_cache,
1505 decoded_cache,
1506 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1507 kek,
1508 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1509 shared: None,
1510 table_name: None,
1511 auth: None,
1512 read_only: false,
1513 }
1514 }
1515}
1516
1517fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1521 use crate::query::Condition;
1522 match c {
1523 Condition::Pk(_)
1525 | Condition::BitmapEq { .. }
1526 | Condition::BitmapIn { .. }
1527 | Condition::BytesPrefix { .. }
1528 | Condition::IsNull { .. }
1529 | Condition::IsNotNull { .. } => 0,
1530 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1532 1
1533 }
1534 Condition::FmContains { .. }
1536 | Condition::FmContainsAll { .. }
1537 | Condition::Ann { .. }
1538 | Condition::SparseMatch { .. } => 2,
1539 }
1540}
1541
1542impl Table {
1543 pub(crate) fn build_secondary_index_artifact<F>(
1549 &self,
1550 definition: &IndexDef,
1551 rows: &[Row],
1552 mut checkpoint: F,
1553 ) -> Result<SecondaryIndexArtifact>
1554 where
1555 F: FnMut(usize, usize) -> Result<()>,
1556 {
1557 let mut build_schema = self.schema.clone();
1558 build_schema.indexes = vec![definition.clone()];
1559 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
1560 let name_to_id: HashMap<&str, u16> = self
1561 .schema
1562 .columns
1563 .iter()
1564 .map(|column| (column.name.as_str(), column.id))
1565 .collect();
1566 let mut hot = HotIndex::new();
1567 let mut accepted = Vec::with_capacity(rows.len());
1568
1569 for (offset, row) in rows.iter().enumerate() {
1570 checkpoint(offset, rows.len())?;
1571 if row.deleted {
1572 continue;
1573 }
1574 if let Some(predicate) = &definition.predicate {
1575 let columns: HashMap<u16, &Value> =
1576 row.columns.iter().map(|(id, value)| (*id, value)).collect();
1577 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
1578 continue;
1579 }
1580 }
1581 accepted.push(row);
1582 if definition.kind == IndexKind::LearnedRange {
1583 continue;
1584 }
1585 let effective = if self.column_keys.is_empty() {
1586 row.clone()
1587 } else {
1588 self.tokenized_for_indexes(row)
1589 };
1590 if definition.kind == IndexKind::Ann {
1591 if let (Some(index), Some(vector)) = (
1592 ann.get_mut(&definition.column_id),
1593 effective
1594 .columns
1595 .get(&definition.column_id)
1596 .and_then(Value::as_embedding),
1597 ) {
1598 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
1599 checkpoint(offset, rows.len())
1600 })?;
1601 }
1602 } else {
1603 index_into_single(
1604 definition,
1605 &build_schema,
1606 &effective,
1607 &mut hot,
1608 &mut bitmap,
1609 &mut ann,
1610 &mut fm,
1611 &mut sparse,
1612 &mut minhash,
1613 );
1614 }
1615 }
1616 checkpoint(rows.len(), rows.len())?;
1617
1618 if let Some(index) = ann.get_mut(&definition.column_id) {
1619 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
1620 }
1621
1622 let column_id = definition.column_id;
1623 match definition.kind {
1624 IndexKind::Bitmap => bitmap
1625 .remove(&column_id)
1626 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
1627 IndexKind::Ann => ann
1628 .remove(&column_id)
1629 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
1630 IndexKind::FmIndex => fm
1631 .remove(&column_id)
1632 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
1633 IndexKind::Sparse => sparse
1634 .remove(&column_id)
1635 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
1636 IndexKind::MinHash => minhash
1637 .remove(&column_id)
1638 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
1639 IndexKind::LearnedRange => {
1640 let epsilon = definition
1641 .options
1642 .learned_range
1643 .as_ref()
1644 .map(|options| options.epsilon)
1645 .unwrap_or(16);
1646 let column = self
1647 .schema
1648 .columns
1649 .iter()
1650 .find(|column| column.id == column_id)
1651 .ok_or_else(|| {
1652 MongrelError::Schema(format!(
1653 "index {} references unknown column {column_id}",
1654 definition.name
1655 ))
1656 })?;
1657 let index = match column.ty {
1658 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1659 let pairs: Vec<_> = accepted
1660 .iter()
1661 .filter_map(|row| match row.columns.get(&column_id) {
1662 Some(Value::Int64(value)) if !row.deleted => {
1663 Some((*value, row.row_id.0))
1664 }
1665 _ => None,
1666 })
1667 .collect();
1668 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
1669 }
1670 TypeId::Float32 | TypeId::Float64 => {
1671 let pairs: Vec<_> = accepted
1672 .iter()
1673 .filter_map(|row| match row.columns.get(&column_id) {
1674 Some(Value::Float64(value)) if !row.deleted => {
1675 Some((*value, row.row_id.0))
1676 }
1677 _ => None,
1678 })
1679 .collect();
1680 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
1681 }
1682 ref ty => {
1683 return Err(MongrelError::Schema(format!(
1684 "LearnedRange index {} does not support {ty:?}",
1685 definition.name
1686 )));
1687 }
1688 };
1689 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
1690 }
1691 }
1692 .ok_or_else(|| {
1693 MongrelError::Other(format!(
1694 "failed to construct hidden index {}",
1695 definition.name
1696 ))
1697 })
1698 }
1699
1700 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1701 let dir = dir.as_ref().to_path_buf();
1702 crate::durable_file::create_directory_all(&dir)?;
1703 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1704 let pinned = root.io_path()?;
1705 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1706 ctx.root_guard = Some(root);
1707 Self::create_in(&pinned, schema, table_id, ctx)
1708 }
1709
1710 pub fn create_encrypted(
1721 dir: impl AsRef<Path>,
1722 schema: Schema,
1723 table_id: u64,
1724 passphrase: &str,
1725 ) -> Result<Self> {
1726 let dir = dir.as_ref().to_path_buf();
1727 crate::durable_file::create_directory_all(&dir)?;
1728 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1729 root.create_directory_all(META_DIR)?;
1730 let salt = crate::encryption::random_salt()?;
1731 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1732 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1733 let pinned = root.io_path()?;
1734 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1735 ctx.root_guard = Some(root);
1736 Self::create_in(&pinned, schema, table_id, ctx)
1737 }
1738
1739 pub fn create_with_key(
1744 dir: impl AsRef<Path>,
1745 schema: Schema,
1746 table_id: u64,
1747 key: &[u8],
1748 ) -> Result<Self> {
1749 let dir = dir.as_ref().to_path_buf();
1750 crate::durable_file::create_directory_all(&dir)?;
1751 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1752 root.create_directory_all(META_DIR)?;
1753 let salt = crate::encryption::random_salt()?;
1754 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1755 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1756 let pinned = root.io_path()?;
1757 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1758 ctx.root_guard = Some(root);
1759 Self::create_in(&pinned, schema, table_id, ctx)
1760 }
1761
1762 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1764 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1765 let salt = read_table_encryption_salt_root(&root)?;
1766 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1767 let pinned = root.io_path()?;
1768 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1769 ctx.root_guard = Some(root);
1770 Self::open_in(&pinned, ctx)
1771 }
1772
1773 pub(crate) fn create_in(
1774 dir: impl AsRef<Path>,
1775 schema: Schema,
1776 table_id: u64,
1777 ctx: SharedCtx,
1778 ) -> Result<Self> {
1779 schema.validate_auto_increment()?;
1780 schema.validate_defaults()?;
1781 schema.validate_ai()?;
1782 for index in &schema.indexes {
1783 index.validate_options()?;
1784 }
1785 let dir = dir.as_ref().to_path_buf();
1786 let runs_root = match ctx.root_guard.as_ref() {
1787 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1788 None => {
1789 crate::durable_file::create_directory_all(&dir)?;
1790 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1791 None
1792 }
1793 };
1794 match ctx.root_guard.as_deref() {
1795 Some(root) => write_schema_durable(root, &schema)?,
1796 None => write_schema(&dir, &schema)?,
1797 }
1798 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1799 let (wal, current_txn_id) = match ctx.shared.clone() {
1802 Some(s) => (WalSink::Shared(s), 0),
1803 None => {
1804 let pinned_wal_root = match ctx.root_guard.as_deref() {
1805 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1806 None => None,
1807 };
1808 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1809 root.io_path()?
1810 } else {
1811 let wal_dir = dir.join(WAL_DIR);
1812 crate::durable_file::create_directory_all(&wal_dir)?;
1813 wal_dir
1814 };
1815 let mut w = if let Some(ref dk) = wal_dek {
1816 Wal::create_with_cipher(
1817 wal_dir.join("seg-000000.wal"),
1818 Epoch(0),
1819 Some(make_cipher(dk)),
1820 0,
1821 )?
1822 } else {
1823 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1824 };
1825 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1826 (WalSink::Private(w), 1)
1827 }
1828 };
1829 let mut manifest = Manifest::new(table_id, schema.schema_id);
1830 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1835 match ctx.root_guard.as_deref() {
1836 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1837 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1838 }
1839 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1840 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1841 let auto_inc = resolve_auto_inc(&schema);
1842 let rcache_dir = dir.join(RCACHE_DIR);
1843 let initial_view = ReadGeneration::empty(&schema);
1844 Ok(Self {
1845 dir,
1846 _root_guard: ctx.root_guard,
1847 runs_root,
1848 idx_root: None,
1849 table_id,
1850 name: ctx.table_name.unwrap_or_default(),
1851 auth: ctx.auth,
1852 read_only: ctx.read_only,
1853 durable_commit_failed: false,
1854 wal,
1855 memtable: Memtable::new(),
1856 mutable_run: MutableRun::new(),
1857 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1858 compaction_zstd_level: 3,
1859 allocator: RowIdAllocator::new(0),
1860 epoch: ctx.epoch,
1861 data_generation: 0,
1862 schema,
1863 hot: HotIndex::new(),
1864 kek: ctx.kek,
1865 column_keys,
1866 run_refs: Vec::new(),
1867 retiring: Vec::new(),
1868 next_run_id: 1,
1869 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1870 current_txn_id,
1871 pending_private_mutations: false,
1872 bitmap,
1873 ann,
1874 fm,
1875 sparse,
1876 minhash,
1877 learned_range: Arc::new(HashMap::new()),
1878 pk_by_row: ReversePkMap::new(),
1879 pinned: BTreeMap::new(),
1880 live_count: 0,
1881 reservoir: crate::reservoir::Reservoir::default(),
1882 reservoir_complete: true,
1883 had_deletes: false,
1884 agg_cache: Arc::new(HashMap::new()),
1885 global_idx_epoch: 0,
1886 indexes_complete: true,
1887 index_build_policy: IndexBuildPolicy::default(),
1888 pk_by_row_complete: false,
1889 flushed_epoch: 0,
1890 page_cache: ctx.page_cache,
1891 decoded_cache: ctx.decoded_cache,
1892 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1893 snapshots: ctx.snapshots,
1894 commit_lock: ctx.commit_lock,
1895 result_cache: Arc::new(parking_lot::Mutex::new(
1896 ResultCache::new()
1897 .with_dir(rcache_dir)
1898 .with_cache_dek(cache_dek.clone()),
1899 )),
1900 pending_delete_rids: roaring::RoaringBitmap::new(),
1901 pending_put_cols: std::collections::HashSet::new(),
1902 pending_rows: Vec::new(),
1903 pending_rows_auto_inc: Vec::new(),
1904 pending_dels: Vec::new(),
1905 pending_truncate: None,
1906 wal_dek,
1907 auto_inc,
1908 ttl: None,
1909 pins: Arc::new(crate::retention::PinRegistry::new()),
1910 published: Arc::new(ArcSwap::from_pointee(initial_view)),
1911 read_generation_pin: None,
1912 })
1913 }
1914
1915 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1919 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1920 let pinned = root.io_path()?;
1921 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1922 ctx.root_guard = Some(root);
1923 Self::open_in(&pinned, ctx)
1924 }
1925
1926 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1929 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1930 let salt = read_table_encryption_salt_root(&root)?;
1931 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1932 let pinned = root.io_path()?;
1933 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1934 ctx.root_guard = Some(root);
1935 let t = Self::open_in(&pinned, ctx)?;
1936 Ok(t)
1937 }
1938
1939 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1940 let dir = dir.as_ref().to_path_buf();
1941 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1942 let mut manifest = match ctx.root_guard.as_ref() {
1943 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1944 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1945 };
1946 let schema: Schema = match ctx.root_guard.as_ref() {
1947 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1948 None => read_schema(&dir)?,
1949 };
1950 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1956 let runs_root = match ctx.root_guard.as_ref() {
1957 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1958 None => None,
1959 };
1960 let idx_root = match ctx.root_guard.as_ref() {
1961 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1962 Ok(root) => Some(Arc::new(root)),
1963 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1964 Err(error) => return Err(error.into()),
1965 },
1966 None => None,
1967 };
1968 schema.validate_auto_increment()?;
1969 schema.validate_defaults()?;
1970 schema.validate_ai()?;
1971 for index in &schema.indexes {
1972 index.validate_options()?;
1973 }
1974 let replay_epoch = Epoch(manifest.current_epoch);
1975 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1976 let private_replayed = if ctx.shared.is_none() {
1977 match latest_wal_segment(&dir.join(WAL_DIR))? {
1978 Some(path) => {
1979 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1980 crate::wal::replay_with_cipher(path, cipher)?
1981 }
1982 None => Vec::new(),
1983 }
1984 } else {
1985 Vec::new()
1986 };
1987 if ctx.shared.is_none() {
1988 preflight_standalone_open(
1989 &dir,
1990 runs_root.as_deref(),
1991 idx_root.as_deref(),
1992 &manifest,
1993 &schema,
1994 &private_replayed,
1995 ctx.kek.clone(),
1996 )?;
1997 }
1998 let next_run_id = derive_next_run_id(
1999 &dir,
2000 runs_root.as_deref(),
2001 &manifest.runs,
2002 &manifest.retiring,
2003 )?;
2004 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2008 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2009 None => {
2010 let replayed = private_replayed;
2011 let wal_dir = dir.join(WAL_DIR);
2017 crate::durable_file::create_directory_all(&wal_dir)?;
2018 let segment = next_wal_segment(&wal_dir)?;
2019 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2020 let temporary = wal_dir.join(format!(
2021 ".recovery-{}-{}-{segment_no:06}.tmp",
2022 std::process::id(),
2023 std::time::SystemTime::now()
2024 .duration_since(std::time::UNIX_EPOCH)
2025 .unwrap_or_default()
2026 .as_nanos()
2027 ));
2028 let mut w = Wal::create_with_cipher(
2029 &temporary,
2030 replay_epoch,
2031 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2032 segment_no,
2033 )?;
2034 for record in &replayed {
2035 w.append_txn(record.txn_id, record.op.clone())?;
2036 }
2037 let mut w = w.publish_as(segment)?;
2038 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2039 let next_txn_id = replayed
2040 .iter()
2041 .map(|record| record.txn_id)
2042 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2043 .max()
2044 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2045 .unwrap_or(1);
2046 (WalSink::Private(w), replayed, next_txn_id)
2047 }
2048 };
2049
2050 let mut memtable = Memtable::new();
2051 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2052 let persisted_epoch = manifest.current_epoch;
2053 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2060 s.next = manifest.auto_inc_next;
2061 s.seeded = manifest.auto_inc_next > 0;
2062 s
2063 });
2064
2065 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2072 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2073 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2074 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2075 std::collections::BTreeMap::new();
2076 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2077 let mut recovered_epoch = manifest.current_epoch;
2078 let mut recovered_manifest_dirty = schema_manifest_repair;
2079 let mut saw_delete = false;
2080 for record in replayed {
2081 let txn_id = record.txn_id;
2082 match record.op {
2083 Op::Put { rows, .. } => {
2084 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2085 for row in &rows {
2086 allocator.advance_to(row.row_id)?;
2087 if let Some(ai) = auto_inc.as_mut() {
2088 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2089 let next = n.checked_add(1).ok_or_else(|| {
2090 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2091 })?;
2092 if next > ai.next {
2093 ai.next = next;
2094 }
2095 }
2096 }
2097 }
2098 staged_puts.entry(txn_id).or_default().extend(rows);
2099 }
2100 Op::Delete { row_ids, .. } => {
2101 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2102 }
2103 Op::TxnCommit { epoch, .. } => {
2104 let commit_epoch = Epoch(epoch);
2105 recovered_epoch = recovered_epoch.max(epoch);
2106 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2107 memtable = Memtable::new();
2108 replayed_puts.clear();
2109 replayed_deletes.clear();
2110 manifest.runs.clear();
2111 manifest.retiring.clear();
2112 manifest.live_count = 0;
2113 manifest.global_idx_epoch = 0;
2114 manifest.current_epoch = manifest.current_epoch.max(epoch);
2115 recovered_manifest_dirty = true;
2116 saw_delete = true;
2117 }
2118 if let Some(puts) = staged_puts.remove(&txn_id) {
2119 if commit_epoch.0 > manifest.flushed_epoch {
2120 for row in &puts {
2121 memtable.upsert(row.clone());
2122 }
2123 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2124 }
2125 }
2126 if let Some(dels) = staged_deletes.remove(&txn_id) {
2127 saw_delete = true;
2128 if commit_epoch.0 > manifest.flushed_epoch {
2129 for rid in dels {
2130 memtable.tombstone(rid, commit_epoch);
2131 replayed_deletes.push((rid, commit_epoch));
2132 }
2133 }
2134 }
2135 }
2136 Op::TxnAbort => {
2137 staged_puts.remove(&txn_id);
2138 staged_deletes.remove(&txn_id);
2139 staged_truncates.remove(&txn_id);
2140 }
2141 Op::TruncateTable { .. } => {
2142 staged_puts.remove(&txn_id);
2143 staged_deletes.remove(&txn_id);
2144 staged_truncates.insert(txn_id);
2145 }
2146 Op::ExternalTableState { .. }
2147 | Op::Flush { .. }
2148 | Op::Ddl(_)
2149 | Op::BeforeImage { .. }
2150 | Op::CommitTimestamp { .. }
2151 | Op::SpilledRows { .. } => {}
2152 }
2153 }
2154
2155 let rcache_dir = dir.join(RCACHE_DIR);
2156 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2157 let initial_view = ReadGeneration::empty(&schema);
2158 let mut db = Self {
2159 dir,
2160 _root_guard: ctx.root_guard,
2161 runs_root,
2162 idx_root,
2163 table_id: manifest.table_id,
2164 name: ctx.table_name.unwrap_or_default(),
2165 auth: ctx.auth,
2166 read_only: ctx.read_only,
2167 durable_commit_failed: false,
2168 wal,
2169 memtable,
2170 mutable_run: MutableRun::new(),
2171 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2172 compaction_zstd_level: 3,
2173 allocator,
2174 epoch: ctx.epoch,
2175 data_generation: persisted_epoch,
2176 schema,
2177 hot: HotIndex::new(),
2178 kek: ctx.kek,
2179 column_keys,
2180 run_refs: manifest.runs.clone(),
2181 retiring: manifest.retiring.clone(),
2182 next_run_id,
2183 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2184 current_txn_id,
2185 pending_private_mutations: false,
2186 bitmap: HashMap::new(),
2187 ann: HashMap::new(),
2188 fm: HashMap::new(),
2189 sparse: HashMap::new(),
2190 minhash: HashMap::new(),
2191 learned_range: Arc::new(HashMap::new()),
2192 pk_by_row: ReversePkMap::new(),
2193 pinned: BTreeMap::new(),
2194 live_count: manifest.live_count,
2195 reservoir: crate::reservoir::Reservoir::default(),
2196 reservoir_complete: false,
2197 had_deletes: saw_delete
2198 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2199 != manifest.live_count,
2200 agg_cache: Arc::new(HashMap::new()),
2201 global_idx_epoch: manifest.global_idx_epoch,
2202 indexes_complete: true,
2203 index_build_policy: IndexBuildPolicy::default(),
2204 pk_by_row_complete: false,
2205 flushed_epoch: manifest.flushed_epoch,
2206 page_cache: ctx.page_cache,
2207 decoded_cache: ctx.decoded_cache,
2208 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2209 snapshots: ctx.snapshots,
2210 commit_lock: ctx.commit_lock,
2211 result_cache: Arc::new(parking_lot::Mutex::new(
2212 ResultCache::new()
2213 .with_dir(rcache_dir)
2214 .with_cache_dek(cache_dek.clone()),
2215 )),
2216 pending_delete_rids: roaring::RoaringBitmap::new(),
2217 pending_put_cols: std::collections::HashSet::new(),
2218 pending_rows: Vec::new(),
2219 pending_rows_auto_inc: Vec::new(),
2220 pending_dels: Vec::new(),
2221 pending_truncate: None,
2222 wal_dek,
2223 auto_inc,
2224 ttl: manifest.ttl,
2225 pins: Arc::new(crate::retention::PinRegistry::new()),
2226 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2227 read_generation_pin: None,
2228 };
2229
2230 db.epoch.advance_recovered(Epoch(recovered_epoch));
2233
2234 let checkpoint = match db.idx_root.as_deref() {
2239 Some(root) => {
2240 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2241 }
2242 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2243 };
2244 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2245 c.epoch_built == manifest.global_idx_epoch
2246 && manifest.global_idx_epoch > 0
2247 && manifest
2248 .runs
2249 .iter()
2250 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2251 });
2252 if let Some(loaded) = checkpoint {
2253 if checkpoint_valid {
2254 db.hot = loaded.hot;
2255 db.bitmap = loaded.bitmap;
2256 db.ann = loaded.ann;
2257 db.fm = loaded.fm;
2258 db.sparse = loaded.sparse;
2259 db.minhash = loaded.minhash;
2260 db.learned_range = Arc::new(loaded.learned_range);
2261 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2266 for (cid, idx) in bitmap0 {
2267 db.bitmap.entry(cid).or_insert(idx);
2268 }
2269 for (cid, idx) in ann0 {
2270 db.ann.entry(cid).or_insert(idx);
2271 }
2272 for (cid, idx) in fm0 {
2273 db.fm.entry(cid).or_insert(idx);
2274 }
2275 for (cid, idx) in sparse0 {
2276 db.sparse.entry(cid).or_insert(idx);
2277 }
2278 for (cid, idx) in minhash0 {
2279 db.minhash.entry(cid).or_insert(idx);
2280 }
2281 }
2284 }
2285 if !checkpoint_valid {
2286 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2287 db.bitmap = bitmap;
2288 db.ann = ann;
2289 db.fm = fm;
2290 db.sparse = sparse;
2291 db.minhash = minhash;
2292 db.rebuild_indexes_from_runs()?;
2293 db.build_learned_ranges()?;
2294 }
2295
2296 for (epoch, group) in replayed_puts {
2301 let (losers, winner_pks) = db.partition_pk_winners(&group);
2302 for (key, &row_id) in &winner_pks {
2303 if let Some(old_rid) = db.hot.get(key) {
2304 if old_rid != row_id {
2305 db.tombstone_row(old_rid, epoch, None, false);
2306 }
2307 }
2308 }
2309 for &loser_rid in &losers {
2310 db.tombstone_row(loser_rid, epoch, None, false);
2311 }
2312 for (key, row_id) in winner_pks {
2313 db.insert_hot_pk(key, row_id);
2314 }
2315 if db.schema.primary_key().is_none() {
2316 for r in &group {
2317 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2318 }
2319 }
2320 for r in &group {
2321 if !losers.contains(&r.row_id) {
2322 db.index_row(r);
2323 }
2324 }
2325 }
2326 for (rid, epoch) in &replayed_deletes {
2330 db.remove_hot_for_row(*rid, *epoch);
2331 }
2332
2333 if recovered_manifest_dirty {
2334 let rows = db.visible_rows(Snapshot::unbounded())?;
2335 db.live_count = rows.len() as u64;
2336 db.persist_manifest(Epoch(recovered_epoch))?;
2337 }
2338
2339 db.result_cache.lock().load_persistent();
2346 Ok(db)
2347 }
2348
2349 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2355 if self.reservoir_complete {
2356 return Ok(());
2357 }
2358 self.rebuild_reservoir()?;
2359 self.reservoir_complete = true;
2360 Ok(())
2361 }
2362
2363 fn rebuild_reservoir(&mut self) -> Result<()> {
2366 let snap = self.snapshot();
2367 let rows = self.visible_rows(snap)?;
2368 self.reservoir.reset();
2369 for r in rows {
2370 self.reservoir.offer(r.row_id.0);
2371 }
2372 Ok(())
2373 }
2374
2375 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2376 self.rebuild_indexes_from_runs_inner(None)
2377 }
2378
2379 fn rebuild_indexes_from_runs_inner(
2380 &mut self,
2381 control: Option<&crate::ExecutionControl>,
2382 ) -> Result<()> {
2383 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2386 crate::retention::PinSource::OnlineIndexBuild,
2387 self.current_epoch(),
2388 );
2389 self.hot = HotIndex::new();
2390 self.pk_by_row.clear();
2391 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2392 self.bitmap = bitmap;
2393 self.ann = ann;
2394 self.fm = fm;
2395 self.sparse = sparse;
2396 self.minhash = minhash;
2397 let snapshot = Epoch(u64::MAX);
2398 let ttl_now = unix_nanos_now();
2399 let mut scanned = 0_usize;
2400 for rr in self.run_refs.clone() {
2401 if let Some(control) = control {
2402 control.checkpoint()?;
2403 }
2404 let mut reader = self.open_reader(rr.run_id)?;
2405 for row in reader.visible_rows(snapshot)? {
2406 if scanned.is_multiple_of(256) {
2407 if let Some(control) = control {
2408 control.checkpoint()?;
2409 }
2410 }
2411 scanned += 1;
2412 if self.row_expired_at(&row, ttl_now) {
2413 continue;
2414 }
2415 let tok_row = self.tokenized_for_indexes(&row);
2416 index_into(
2417 &self.schema,
2418 &tok_row,
2419 &mut self.hot,
2420 &mut self.bitmap,
2421 &mut self.ann,
2422 &mut self.fm,
2423 &mut self.sparse,
2424 &mut self.minhash,
2425 );
2426 }
2427 }
2428 for row in self.mutable_run.visible_versions(snapshot) {
2429 if scanned.is_multiple_of(256) {
2430 if let Some(control) = control {
2431 control.checkpoint()?;
2432 }
2433 }
2434 scanned += 1;
2435 if row.deleted {
2436 self.remove_hot_for_row(row.row_id, snapshot);
2437 } else if !self.row_expired_at(&row, ttl_now) {
2438 self.index_row(&row);
2439 }
2440 }
2441 for row in self.memtable.visible_versions(snapshot) {
2442 if scanned.is_multiple_of(256) {
2443 if let Some(control) = control {
2444 control.checkpoint()?;
2445 }
2446 }
2447 scanned += 1;
2448 if row.deleted {
2449 self.remove_hot_for_row(row.row_id, snapshot);
2450 } else if !self.row_expired_at(&row, ttl_now) {
2451 self.index_row(&row);
2452 }
2453 }
2454 self.refresh_pk_by_row_from_hot();
2455 Ok(())
2456 }
2457
2458 fn refresh_pk_by_row_from_hot(&mut self) {
2459 self.pk_by_row_complete = true;
2460 if self.schema.primary_key().is_none() {
2461 self.pk_by_row.clear();
2462 return;
2463 }
2464 self.pk_by_row = ReversePkMap::from_entries(
2470 self.hot
2471 .entries()
2472 .into_iter()
2473 .map(|(key, row_id)| (row_id, key)),
2474 );
2475 }
2476
2477 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2478 if self.schema.primary_key().is_some() {
2479 self.pk_by_row.insert(row_id, key.clone());
2480 }
2481 self.hot.insert(key, row_id);
2482 }
2483
2484 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2488 self.build_learned_ranges_inner(None)
2489 }
2490
2491 fn build_learned_ranges_inner(
2492 &mut self,
2493 control: Option<&crate::ExecutionControl>,
2494 ) -> Result<()> {
2495 self.learned_range = Arc::new(HashMap::new());
2496 if self.run_refs.len() != 1 {
2497 return Ok(());
2498 }
2499 let cols: Vec<(u16, usize)> = self
2500 .schema
2501 .indexes
2502 .iter()
2503 .filter(|i| i.kind == IndexKind::LearnedRange)
2504 .map(|i| {
2505 (
2506 i.column_id,
2507 i.options
2508 .learned_range
2509 .as_ref()
2510 .map(|options| options.epsilon)
2511 .unwrap_or(16),
2512 )
2513 })
2514 .collect();
2515 if cols.is_empty() {
2516 return Ok(());
2517 }
2518 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2519 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2520 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2521 _ => return Ok(()),
2522 };
2523 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2524 if column_index % 256 == 0 {
2525 if let Some(control) = control {
2526 control.checkpoint()?;
2527 }
2528 }
2529 let ty = self
2530 .schema
2531 .columns
2532 .iter()
2533 .find(|c| c.id == cid)
2534 .map(|c| c.ty.clone())
2535 .unwrap_or(TypeId::Int64);
2536 match ty {
2537 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2538 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2539 let pairs: Vec<(i64, u64)> = data
2540 .iter()
2541 .zip(row_ids.iter())
2542 .map(|(v, r)| (*v, *r))
2543 .collect();
2544 Arc::make_mut(&mut self.learned_range).insert(
2545 cid,
2546 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2547 );
2548 }
2549 }
2550 TypeId::Float64 => {
2551 if let columnar::NativeColumn::Float64 { data, .. } =
2552 reader.column_native(cid)?
2553 {
2554 let pairs: Vec<(f64, u64)> = data
2555 .iter()
2556 .zip(row_ids.iter())
2557 .map(|(v, r)| (*v, *r))
2558 .collect();
2559 Arc::make_mut(&mut self.learned_range).insert(
2560 cid,
2561 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2562 );
2563 }
2564 }
2565 _ => {}
2566 }
2567 }
2568 Ok(())
2569 }
2570
2571 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2578 if self.indexes_complete {
2579 crate::trace::QueryTrace::record(|t| {
2580 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2581 });
2582 return Ok(());
2583 }
2584 crate::trace::QueryTrace::record(|t| {
2585 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2586 });
2587 self.rebuild_indexes_from_runs()?;
2588 self.build_learned_ranges()?;
2589 self.indexes_complete = true;
2590 let epoch = self.current_epoch();
2591 self.checkpoint_indexes(epoch);
2592 Ok(())
2593 }
2594
2595 #[doc(hidden)]
2598 pub fn ensure_indexes_complete_controlled<F>(
2599 &mut self,
2600 control: &crate::ExecutionControl,
2601 before_publish: F,
2602 ) -> Result<bool>
2603 where
2604 F: FnOnce() -> bool,
2605 {
2606 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2607 .map(|(changed, _)| changed)
2608 }
2609
2610 #[doc(hidden)]
2613 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2614 &mut self,
2615 control: &crate::ExecutionControl,
2616 before_publish: F,
2617 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2618 where
2619 F: FnOnce() -> bool,
2620 {
2621 if self.indexes_complete {
2622 crate::trace::QueryTrace::record(|trace| {
2623 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2624 });
2625 return Ok((false, None));
2626 }
2627 crate::trace::QueryTrace::record(|trace| {
2628 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2629 });
2630 control.checkpoint()?;
2631 let maintenance_epoch = self.current_epoch();
2632 self.rebuild_indexes_from_runs_inner(Some(control))?;
2633 self.build_learned_ranges_inner(Some(control))?;
2634 control.checkpoint()?;
2635 if !before_publish() {
2636 return Err(MongrelError::Cancelled);
2637 }
2638 self.indexes_complete = true;
2639 self.checkpoint_indexes(maintenance_epoch);
2640 Ok((
2641 true,
2642 Some(MaintenanceReceipt {
2643 epoch: maintenance_epoch,
2644 }),
2645 ))
2646 }
2647
2648 fn pending_epoch(&self) -> Epoch {
2649 Epoch(self.epoch.visible().0 + 1)
2650 }
2651
2652 fn is_shared(&self) -> bool {
2655 matches!(self.wal, WalSink::Shared(_))
2656 }
2657
2658 fn ensure_txn_id(&mut self) -> Result<u64> {
2662 if self.current_txn_id == 0 {
2663 let id = match &self.wal {
2664 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2665 WalSink::Private(_) => {
2666 return Err(MongrelError::Full(
2667 "standalone transaction id namespace exhausted".into(),
2668 ))
2669 }
2670 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2671 };
2672 self.current_txn_id = id;
2673 }
2674 Ok(self.current_txn_id)
2675 }
2676
2677 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2680 self.ensure_writable()?;
2681 let txn_id = self.ensure_txn_id()?;
2682 let table_id = self.table_id;
2683 match &mut self.wal {
2684 WalSink::Private(w) => {
2685 w.append_txn(txn_id, op)?;
2686 self.pending_private_mutations = true;
2687 }
2688 WalSink::Shared(s) => {
2689 s.wal.lock().append(txn_id, table_id, op)?;
2690 }
2691 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2692 }
2693 Ok(())
2694 }
2695
2696 fn ensure_writable(&self) -> Result<()> {
2697 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2698 return Err(MongrelError::ReadOnlyReplica);
2699 }
2700 if self.durable_commit_failed {
2701 return Err(MongrelError::Other(
2702 "table poisoned by post-commit failure; reopen required".into(),
2703 ));
2704 }
2705 Ok(())
2706 }
2707
2708 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2719 match &self.auth {
2720 Some(checker) => checker.check(&self.name, perm),
2721 None => Ok(()),
2722 }
2723 }
2724 pub fn require_select(&self) -> Result<()> {
2729 self.require(crate::auth_state::RequiredPermission::Select)
2730 }
2731 fn require_insert(&self) -> Result<()> {
2732 self.require(crate::auth_state::RequiredPermission::Insert)
2733 }
2734 #[allow(dead_code)]
2738 fn require_update(&self) -> Result<()> {
2739 self.require(crate::auth_state::RequiredPermission::Update)
2740 }
2741 fn require_delete(&self) -> Result<()> {
2742 self.require(crate::auth_state::RequiredPermission::Delete)
2743 }
2744
2745 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2748 self.require_insert()?;
2749 Ok(self.put_returning(columns)?.0)
2750 }
2751
2752 pub fn put_returning(
2757 &mut self,
2758 mut columns: Vec<(u16, Value)>,
2759 ) -> Result<(RowId, Option<i64>)> {
2760 self.require_insert()?;
2761 let assigned = self.fill_auto_inc(&mut columns)?;
2762 self.apply_defaults(&mut columns)?;
2763 self.schema.validate_values(&columns)?;
2764 let row_id = if self.schema.clustered {
2769 self.derive_clustered_row_id(&columns)?
2770 } else {
2771 self.allocator.alloc()?
2772 };
2773 let epoch = self.pending_epoch();
2774 let mut row = Row::new(row_id, epoch);
2775 for (col_id, val) in columns {
2776 row.columns.insert(col_id, val);
2777 }
2778 self.commit_rows(vec![row], assigned.is_some())?;
2779 Ok((row_id, assigned))
2780 }
2781
2782 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2785 self.require_insert()?;
2786 Ok(self
2787 .put_batch_returning(batch)?
2788 .into_iter()
2789 .map(|(r, _)| r)
2790 .collect())
2791 }
2792
2793 pub fn put_batch_returning(
2796 &mut self,
2797 batch: Vec<Vec<(u16, Value)>>,
2798 ) -> Result<Vec<(RowId, Option<i64>)>> {
2799 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2800 for mut cols in batch {
2801 let assigned = self.fill_auto_inc(&mut cols)?;
2802 self.apply_defaults(&mut cols)?;
2803 filled.push((cols, assigned));
2804 }
2805 for (cols, _) in &filled {
2806 self.schema.validate_values(cols)?;
2807 }
2808 let epoch = self.pending_epoch();
2809 let mut rows = Vec::with_capacity(filled.len());
2810 let mut ids = Vec::with_capacity(filled.len());
2811 let first_row_id = if self.schema.clustered {
2812 None
2813 } else {
2814 let count = u64::try_from(filled.len())
2815 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2816 Some(self.allocator.alloc_range(count)?.0)
2817 };
2818 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2819 let row_id = match first_row_id {
2820 Some(first) => RowId(first + row_index as u64),
2821 None => self.derive_clustered_row_id(&cols)?,
2822 };
2823 let mut row = Row::new(row_id, epoch);
2824 for (c, v) in cols {
2825 row.columns.insert(c, v);
2826 }
2827 ids.push((row_id, assigned));
2828 rows.push(row);
2829 }
2830 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2831 self.commit_rows(rows, all_auto_generated)?;
2832 Ok(ids)
2833 }
2834
2835 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2841 self.ensure_writable()?;
2842 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2843 return Ok(None);
2844 };
2845 let pos = columns.iter().position(|(c, _)| *c == cid);
2846 let assigned = match pos {
2847 Some(i) => match &columns[i].1 {
2848 Value::Null => {
2849 let next = self.alloc_auto_inc_value()?;
2850 columns[i].1 = Value::Int64(next);
2851 Some(next)
2852 }
2853 Value::Int64(n) => {
2854 self.advance_auto_inc_past(*n)?;
2855 None
2856 }
2857 other => {
2858 return Err(MongrelError::InvalidArgument(format!(
2859 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2860 other
2861 )))
2862 }
2863 },
2864 None => {
2865 let next = self.alloc_auto_inc_value()?;
2866 columns.push((cid, Value::Int64(next)));
2867 Some(next)
2868 }
2869 };
2870 Ok(assigned)
2871 }
2872
2873 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2879 for col in &self.schema.columns {
2880 let Some(expr) = &col.default_value else {
2881 continue;
2882 };
2883 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2885 continue;
2886 }
2887 let pos = columns.iter().position(|(c, _)| *c == col.id);
2888 let needs_default = match pos {
2889 None => true,
2890 Some(i) => matches!(columns[i].1, Value::Null),
2891 };
2892 if !needs_default {
2893 continue;
2894 }
2895 let v = match expr {
2896 crate::schema::DefaultExpr::Static(v) => v.clone(),
2897 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2898 crate::schema::DefaultExpr::Uuid => {
2899 let mut buf = [0u8; 16];
2900 getrandom::getrandom(&mut buf)
2901 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2902 Value::Uuid(buf)
2903 }
2904 };
2905 match pos {
2906 None => columns.push((col.id, v)),
2907 Some(i) => columns[i].1 = v,
2908 }
2909 }
2910 Ok(())
2911 }
2912
2913 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2915 self.ensure_auto_inc_seeded()?;
2916 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2918 let v = ai.next;
2919 ai.next = ai
2920 .next
2921 .checked_add(1)
2922 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2923 Ok(v)
2924 }
2925
2926 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2929 self.ensure_auto_inc_seeded()?;
2930 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2931 let floor = used
2932 .checked_add(1)
2933 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2934 .max(1);
2935 if ai.next < floor {
2936 ai.next = floor;
2937 }
2938 Ok(())
2939 }
2940
2941 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2946 let needs_seed = match self.auto_inc {
2947 Some(ai) => !ai.seeded,
2948 None => return Ok(()),
2949 };
2950 if !needs_seed {
2951 return Ok(());
2952 }
2953 if self.seed_empty_auto_inc() {
2954 return Ok(());
2955 }
2956 let cid = self
2957 .auto_inc
2958 .as_ref()
2959 .expect("auto-inc column present")
2960 .column_id;
2961 let max = self.scan_max_int64(cid)?;
2962 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2963 let floor = max
2964 .checked_add(1)
2965 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2966 .max(1);
2967 if ai.next < floor {
2968 ai.next = floor;
2969 }
2970 ai.seeded = true;
2971 Ok(())
2972 }
2973
2974 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2975 if n == 0 || self.auto_inc.is_none() {
2976 return Ok(None);
2977 }
2978 self.ensure_auto_inc_seeded()?;
2979 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2980 let start = ai.next;
2981 let count = i64::try_from(n)
2982 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2983 ai.next = ai
2984 .next
2985 .checked_add(count)
2986 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2987 Ok(Some(start))
2988 }
2989
2990 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2994 let mut max: i64 = 0;
2995 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
2996 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2997 if *n > max {
2998 max = *n;
2999 }
3000 }
3001 }
3002 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3003 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3004 if *n > max {
3005 max = *n;
3006 }
3007 }
3008 }
3009 for rr in self.run_refs.clone() {
3010 let reader = self.open_reader(rr.run_id)?;
3011 if let Some(stats) = reader.column_page_stats(column_id) {
3012 for s in stats {
3013 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3014 if n > max {
3015 max = n;
3016 }
3017 }
3018 }
3019 } else if reader.has_column(column_id) {
3020 if let columnar::NativeColumn::Int64 { data, validity } =
3021 reader.column_native_shared(column_id)?
3022 {
3023 for (i, n) in data.iter().enumerate() {
3024 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3025 {
3026 max = *n;
3027 }
3028 }
3029 }
3030 }
3031 }
3032 Ok(max)
3033 }
3034
3035 fn seed_empty_auto_inc(&mut self) -> bool {
3036 let Some(ai) = self.auto_inc.as_mut() else {
3037 return false;
3038 };
3039 if ai.seeded || self.live_count != 0 {
3040 return false;
3041 }
3042 if ai.next < 1 {
3043 ai.next = 1;
3044 }
3045 ai.seeded = true;
3046 true
3047 }
3048
3049 fn advance_auto_inc_from_native_columns(
3050 &mut self,
3051 columns: &[(u16, columnar::NativeColumn)],
3052 n: usize,
3053 live_before: u64,
3054 ) -> Result<()> {
3055 let Some(ai) = self.auto_inc.as_mut() else {
3056 return Ok(());
3057 };
3058 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3059 return Ok(());
3060 };
3061 let columnar::NativeColumn::Int64 { data, validity } = col else {
3062 return Err(MongrelError::InvalidArgument(format!(
3063 "AUTO_INCREMENT column {} must be Int64",
3064 ai.column_id
3065 )));
3066 };
3067 let max = if native_int64_strictly_increasing(col, n) {
3068 data.get(n.saturating_sub(1)).copied()
3069 } else {
3070 data.iter()
3071 .take(n)
3072 .enumerate()
3073 .filter_map(|(i, v)| {
3074 if validity.is_empty() || columnar::validity_bit(validity, i) {
3075 Some(*v)
3076 } else {
3077 None
3078 }
3079 })
3080 .max()
3081 };
3082 if let Some(max) = max {
3083 let floor = max
3084 .checked_add(1)
3085 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3086 .max(1);
3087 if ai.next < floor {
3088 ai.next = floor;
3089 }
3090 if ai.seeded || live_before == 0 {
3091 ai.seeded = true;
3092 }
3093 }
3094 Ok(())
3095 }
3096
3097 fn fill_auto_inc_native_columns(
3098 &mut self,
3099 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3100 n: usize,
3101 ) -> Result<()> {
3102 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3103 return Ok(());
3104 };
3105 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3106 if let Some(start) = self.alloc_auto_inc_range(n)? {
3107 columns.push((
3108 cid,
3109 columnar::NativeColumn::Int64 {
3110 data: (start..start.saturating_add(n as i64)).collect(),
3111 validity: vec![0xFF; n.div_ceil(8)],
3112 },
3113 ));
3114 }
3115 return Ok(());
3116 };
3117
3118 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3119 return Err(MongrelError::InvalidArgument(format!(
3120 "AUTO_INCREMENT column {cid} must be Int64"
3121 )));
3122 };
3123 if data.len() < n {
3124 return Err(MongrelError::InvalidArgument(format!(
3125 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3126 data.len()
3127 )));
3128 }
3129 if columnar::all_non_null(validity, n) {
3130 return Ok(());
3131 }
3132 if validity.iter().all(|b| *b == 0) {
3133 if let Some(start) = self.alloc_auto_inc_range(n)? {
3134 for (i, slot) in data.iter_mut().take(n).enumerate() {
3135 *slot = start.saturating_add(i as i64);
3136 }
3137 *validity = vec![0xFF; n.div_ceil(8)];
3138 }
3139 return Ok(());
3140 }
3141
3142 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3143 for (i, slot) in data.iter_mut().enumerate().take(n) {
3144 if columnar::validity_bit(validity, i) {
3145 self.advance_auto_inc_past(*slot)?;
3146 } else {
3147 *slot = self.alloc_auto_inc_value()?;
3148 }
3149 }
3150 *validity = new_validity;
3151 Ok(())
3152 }
3153
3154 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3168 self.ensure_writable()?;
3169 if self.auto_inc.is_none() {
3170 return Ok(None);
3171 }
3172 Ok(Some(self.alloc_auto_inc_value()?))
3173 }
3174
3175 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3181 let payload = bincode::serialize(&rows)?;
3182 self.wal_append_data(Op::Put {
3183 table_id: self.table_id,
3184 rows: payload,
3185 })?;
3186 if self.is_shared() {
3187 self.pending_rows_auto_inc
3188 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3189 self.pending_rows.extend(rows);
3190 } else {
3191 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3192 }
3193 Ok(())
3194 }
3195
3196 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3199 self.ensure_indexes_complete()
3200 }
3201
3202 pub(crate) fn prepare_durable_publish_controlled(
3203 &mut self,
3204 control: &crate::ExecutionControl,
3205 ) -> Result<()> {
3206 self.ensure_indexes_complete_controlled(control, || true)?;
3207 Ok(())
3208 }
3209
3210 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3211 self.apply_put_rows_inner_prepared(rows, true);
3212 }
3213
3214 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3215 if check_existing_pk {
3216 self.ensure_indexes_complete()?;
3217 }
3218 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3219 Ok(())
3220 }
3221
3222 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3226 if rows.len() == 1 {
3230 let row = rows.into_iter().next().expect("len checked");
3231 self.apply_put_row_single(row, check_existing_pk);
3232 return;
3233 }
3234 let pk_id = self.schema.primary_key().map(|c| c.id);
3251 let probe = match pk_id {
3252 Some(pid) => {
3253 check_existing_pk
3254 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3255 }
3256 None => false,
3257 };
3258 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3261 for r in rows {
3262 for &cid in r.columns.keys() {
3263 self.pending_put_cols.insert(cid);
3264 }
3265 match pk_id {
3266 Some(pid) if probe || maintain_pk_by_row => {
3267 if let Some(pk_val) = r.columns.get(&pid) {
3268 let key = self.index_lookup_key(pid, pk_val);
3269 if probe {
3270 if let Some(old_rid) = self.hot.get(&key) {
3271 if old_rid != r.row_id {
3272 self.tombstone_row(
3273 old_rid,
3274 r.committed_epoch,
3275 r.commit_ts,
3276 true,
3277 );
3278 }
3279 }
3280 }
3281 if maintain_pk_by_row {
3282 self.pk_by_row.insert(r.row_id, key);
3283 }
3284 }
3285 }
3286 Some(_) => {}
3287 None => {
3288 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3289 }
3290 }
3291 self.index_row(&r);
3292 self.reservoir.offer(r.row_id.0);
3293 self.memtable.upsert(r);
3294 self.live_count = self.live_count.saturating_add(1);
3297 }
3298 self.data_generation = self.data_generation.wrapping_add(1);
3299 }
3300
3301 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3305 for &cid in row.columns.keys() {
3306 self.pending_put_cols.insert(cid);
3307 }
3308 let epoch = row.committed_epoch;
3309 if let Some(pk_col) = self.schema.primary_key() {
3310 let pk_id = pk_col.id;
3311 if let Some(pk_val) = row.columns.get(&pk_id) {
3312 let maintain_pk_by_row = self.pk_by_row_complete;
3316 if check_existing_pk || maintain_pk_by_row {
3317 let key = self.index_lookup_key(pk_id, pk_val);
3318 if check_existing_pk {
3319 if let Some(old_rid) = self.hot.get(&key) {
3320 if old_rid != row.row_id {
3321 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
3322 }
3323 }
3324 }
3325 if maintain_pk_by_row {
3326 self.pk_by_row.insert(row.row_id, key);
3327 }
3328 }
3329 }
3330 } else {
3331 self.hot
3332 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3333 }
3334 self.index_row(&row);
3335 self.reservoir.offer(row.row_id.0);
3336 self.memtable.upsert(row);
3337 self.live_count = self.live_count.saturating_add(1);
3338 self.data_generation = self.data_generation.wrapping_add(1);
3339 }
3340
3341 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3344 self.allocator.alloc()
3345 }
3346
3347 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3353 let pk = self.schema.primary_key().ok_or_else(|| {
3354 MongrelError::Schema("clustered table requires a single-column primary key".into())
3355 })?;
3356 let pk_val = columns
3357 .iter()
3358 .find(|(id, _)| *id == pk.id)
3359 .map(|(_, v)| v)
3360 .ok_or_else(|| {
3361 MongrelError::Schema(format!(
3362 "clustered table missing primary key column {} ({})",
3363 pk.id, pk.name
3364 ))
3365 })?;
3366 Ok(clustered_row_id(pk_val))
3367 }
3368
3369 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3377 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3378 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3379 }
3380 let n = rows.len();
3381 for r in rows {
3382 for &cid in r.columns.keys() {
3383 self.pending_put_cols.insert(cid);
3384 }
3385 }
3386 let (losers, winner_pks) = self.partition_pk_winners(rows);
3387 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3388 let group_ts = rows.first().and_then(|r| r.commit_ts);
3390 for (key, &row_id) in &winner_pks {
3391 if let Some(old_rid) = self.hot.get(key) {
3392 if old_rid != row_id {
3393 self.tombstone_row(old_rid, epoch, group_ts, true);
3394 }
3395 }
3396 }
3397 for &loser_rid in &losers {
3400 self.tombstone_row(loser_rid, epoch, group_ts, false);
3401 }
3402 for (key, row_id) in winner_pks {
3404 self.insert_hot_pk(key, row_id);
3405 }
3406 if self.schema.primary_key().is_none() {
3407 for r in rows {
3408 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3409 }
3410 }
3411 for r in rows {
3412 self.allocator.advance_to(r.row_id)?;
3413 if !losers.contains(&r.row_id) {
3414 self.index_row(r);
3415 }
3416 }
3417 for r in rows {
3418 if !losers.contains(&r.row_id) {
3419 self.reservoir.offer(r.row_id.0);
3420 }
3421 }
3422 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3423 self.data_generation = self.data_generation.wrapping_add(1);
3424 Ok(())
3425 }
3426
3427 pub(crate) fn recover_apply(
3432 &mut self,
3433 rows: Vec<Row>,
3434 deletes: Vec<(RowId, Epoch)>,
3435 ) -> Result<()> {
3436 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3440 std::collections::BTreeMap::new();
3441 for row in rows {
3442 if row.row_id.0 >= u64::MAX - 1 {
3443 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3444 }
3445 self.allocator.advance_to(row.row_id)?;
3446 if let Some(ai) = self.auto_inc.as_mut() {
3451 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3452 let next = n.checked_add(1).ok_or_else(|| {
3453 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3454 })?;
3455 if next > ai.next {
3456 ai.next = next;
3457 }
3458 }
3459 }
3460 by_epoch.entry(row.committed_epoch).or_default().push(row);
3461 }
3462 for (epoch, group) in by_epoch {
3463 let (losers, winner_pks) = self.partition_pk_winners(&group);
3464 let group_ts = group.first().and_then(|r| r.commit_ts);
3466 for (key, &row_id) in &winner_pks {
3467 if let Some(old_rid) = self.hot.get(key) {
3468 if old_rid != row_id {
3469 self.tombstone_row(old_rid, epoch, group_ts, false);
3470 }
3471 }
3472 }
3473 for (key, row_id) in winner_pks {
3474 self.insert_hot_pk(key, row_id);
3475 }
3476 if self.schema.primary_key().is_none() {
3477 for r in &group {
3478 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3479 }
3480 }
3481 for r in &group {
3482 if !losers.contains(&r.row_id) {
3483 self.memtable.upsert(r.clone());
3484 self.index_row(r);
3485 }
3486 }
3487 }
3488 for (rid, epoch) in deletes {
3489 self.memtable.tombstone(rid, epoch);
3490 self.remove_hot_for_row(rid, epoch);
3491 }
3492 self.reservoir_complete = false;
3495 Ok(())
3496 }
3497
3498 pub(crate) fn flushed_epoch(&self) -> u64 {
3500 self.flushed_epoch
3501 }
3502
3503 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3504 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3505 }
3506
3507 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3509 self.schema.validate_values(cells)
3510 }
3511
3512 fn validate_columns_not_null(
3516 &self,
3517 columns: &[(u16, columnar::NativeColumn)],
3518 n: usize,
3519 ) -> Result<()> {
3520 let by_id: HashMap<u16, &columnar::NativeColumn> =
3521 columns.iter().map(|(id, c)| (*id, c)).collect();
3522 for col in &self.schema.columns {
3523 if !col.flags.contains(ColumnFlags::NULLABLE) {
3524 match by_id.get(&col.id) {
3525 None => {
3526 return Err(MongrelError::InvalidArgument(format!(
3527 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3528 col.name, col.id
3529 )));
3530 }
3531 Some(c) => {
3532 if c.null_count(n) != 0 {
3533 return Err(MongrelError::InvalidArgument(format!(
3534 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3535 col.name, col.id
3536 )));
3537 }
3538 }
3539 }
3540 }
3541 if let TypeId::Enum { variants } = &col.ty {
3542 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3543 if by_id.contains_key(&col.id) {
3544 return Err(MongrelError::InvalidArgument(format!(
3545 "column '{}' ({}) enum requires a bytes column",
3546 col.name, col.id
3547 )));
3548 }
3549 continue;
3550 };
3551 for index in 0..n {
3552 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3553 continue;
3554 };
3555 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3556 return Err(MongrelError::InvalidArgument(format!(
3557 "column '{}' ({}) enum value {:?} is not one of {:?}",
3558 col.name,
3559 col.id,
3560 String::from_utf8_lossy(value),
3561 variants
3562 )));
3563 }
3564 }
3565 }
3566 }
3567 Ok(())
3568 }
3569
3570 fn bulk_pk_winner_indices(
3575 &self,
3576 columns: &[(u16, columnar::NativeColumn)],
3577 n: usize,
3578 ) -> Option<Vec<usize>> {
3579 let pk_col = self.schema.primary_key()?;
3580 let pk_id = pk_col.id;
3581 let pk_ty = pk_col.ty.clone();
3582 let by_id: HashMap<u16, &columnar::NativeColumn> =
3583 columns.iter().map(|(id, c)| (*id, c)).collect();
3584 let pk_native = by_id.get(&pk_id)?;
3585 if native_int64_strictly_increasing(pk_native, n) {
3586 return None;
3587 }
3588 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3590 let mut null_pk_rows: Vec<usize> = Vec::new();
3591 for i in 0..n {
3592 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3593 Some(key) => {
3594 last.insert(key, i);
3595 }
3596 None => null_pk_rows.push(i),
3597 }
3598 }
3599 let mut winners: HashSet<usize> = last.values().copied().collect();
3600 for i in null_pk_rows {
3601 winners.insert(i);
3602 }
3603 Some((0..n).filter(|i| winners.contains(i)).collect())
3604 }
3605
3606 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3608 self.require_delete()?;
3609 let epoch = self.pending_epoch();
3610 self.wal_append_data(Op::Delete {
3611 table_id: self.table_id,
3612 row_ids: vec![row_id],
3613 })?;
3614 if self.is_shared() {
3615 self.pending_dels.push(row_id);
3616 } else {
3617 self.apply_delete(row_id, epoch);
3618 }
3619 Ok(())
3620 }
3621
3622 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3623 let pre = self.get(row_id, self.snapshot());
3624 self.delete(row_id)?;
3625 Ok(pre.map(|row| {
3626 let mut columns: Vec<_> = row.columns.into_iter().collect();
3627 columns.sort_by_key(|(id, _)| *id);
3628 OwnedRow { columns }
3629 }))
3630 }
3631
3632 pub fn truncate(&mut self) -> Result<()> {
3634 self.require_delete()?;
3635 let epoch = self.pending_epoch();
3636 self.wal_append_data(Op::TruncateTable {
3637 table_id: self.table_id,
3638 })?;
3639 self.pending_rows.clear();
3640 self.pending_rows_auto_inc.clear();
3641 self.pending_dels.clear();
3642 self.pending_truncate = Some(epoch);
3643 Ok(())
3644 }
3645
3646 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3648 self.run_refs.clear();
3653 self.retiring.clear();
3654 self.memtable = Memtable::new();
3655 self.mutable_run = MutableRun::new();
3656 self.hot = HotIndex::new();
3657 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3658 self.bitmap = bitmap;
3659 self.ann = ann;
3660 self.fm = fm;
3661 self.sparse = sparse;
3662 self.minhash = minhash;
3663 self.learned_range = Arc::new(HashMap::new());
3664 self.pk_by_row.clear();
3665 self.pk_by_row_complete = false;
3666 self.live_count = 0;
3667 self.reservoir = crate::reservoir::Reservoir::default();
3668 self.reservoir_complete = true;
3669 self.had_deletes = true;
3670 self.agg_cache = Arc::new(HashMap::new());
3671 self.global_idx_epoch = 0;
3672 self.indexes_complete = true;
3673 self.pending_delete_rids.clear();
3674 self.pending_put_cols.clear();
3675 self.pending_rows.clear();
3676 self.pending_rows_auto_inc.clear();
3677 self.pending_dels.clear();
3678 self.clear_result_cache();
3679 self.invalidate_index_checkpoint();
3680 self.data_generation = self.data_generation.wrapping_add(1);
3681 }
3682
3683 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3686 self.apply_delete_at(row_id, epoch, None);
3687 }
3688
3689 pub(crate) fn apply_delete_at(
3691 &mut self,
3692 row_id: RowId,
3693 epoch: Epoch,
3694 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
3695 ) {
3696 self.remove_hot_for_row(row_id, epoch);
3697 self.tombstone_row(row_id, epoch, commit_ts, true);
3698 self.data_generation = self.data_generation.wrapping_add(1);
3699 }
3700
3701 fn tombstone_row(
3705 &mut self,
3706 row_id: RowId,
3707 epoch: Epoch,
3708 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
3709 adjust_live_count: bool,
3710 ) {
3711 let tombstone = Row {
3712 row_id,
3713 committed_epoch: epoch,
3714 commit_ts,
3715 columns: std::collections::HashMap::new(),
3716 deleted: true,
3717 };
3718 self.memtable.upsert(tombstone);
3719 self.pk_by_row.remove(&row_id);
3720 if adjust_live_count {
3721 self.live_count = self.live_count.saturating_sub(1);
3722 }
3723 self.pending_delete_rids.insert(row_id.0 as u32);
3725 self.had_deletes = true;
3728 self.agg_cache = Arc::new(HashMap::new());
3729 }
3730
3731 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3735 let Some(pk_col) = self.schema.primary_key() else {
3736 return;
3737 };
3738 if self.pk_by_row_complete {
3741 if let Some(key) = self.pk_by_row.remove(&row_id) {
3742 if self.hot.get(&key) == Some(row_id) {
3743 self.hot.remove(&key);
3744 }
3745 }
3746 return;
3747 }
3748 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3767 if self.indexes_complete {
3768 let pk_val = self
3769 .memtable
3770 .get_version(row_id, lookup_epoch)
3771 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3772 .or_else(|| {
3773 self.mutable_run
3774 .get_version(row_id, lookup_epoch)
3775 .filter(|(_, r)| !r.deleted)
3776 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3777 })
3778 .or_else(|| {
3779 self.run_refs.iter().find_map(|rr| {
3780 let mut reader = self.open_reader(rr.run_id).ok()?;
3781 let (_, deleted, val) = reader
3782 .get_version_column(row_id, lookup_epoch, pk_col.id)
3783 .ok()??;
3784 if deleted {
3785 return None;
3786 }
3787 val
3788 })
3789 });
3790 if let Some(pk_val) = pk_val {
3791 let key = self.index_lookup_key(pk_col.id, &pk_val);
3792 if self.hot.get(&key) == Some(row_id) {
3793 self.hot.remove(&key);
3794 }
3795 return;
3796 }
3797 }
3798 self.refresh_pk_by_row_from_hot();
3803 if let Some(key) = self.pk_by_row.remove(&row_id) {
3804 if self.hot.get(&key) == Some(row_id) {
3805 self.hot.remove(&key);
3806 }
3807 }
3808 }
3809
3810 fn partition_pk_winners(
3815 &self,
3816 rows: &[Row],
3817 ) -> (
3818 std::collections::HashSet<RowId>,
3819 std::collections::HashMap<Vec<u8>, RowId>,
3820 ) {
3821 let mut losers = std::collections::HashSet::new();
3822 let Some(pk_col) = self.schema.primary_key() else {
3823 return (losers, std::collections::HashMap::new());
3824 };
3825 let pk_id = pk_col.id;
3826 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3827 std::collections::HashMap::new();
3828 for r in rows {
3829 let Some(pk_val) = r.columns.get(&pk_id) else {
3830 continue;
3831 };
3832 let key = self.index_lookup_key(pk_id, pk_val);
3833 if let Some(&old_rid) = winners.get(&key) {
3834 losers.insert(old_rid);
3835 }
3836 winners.insert(key, r.row_id);
3837 }
3838 (losers, winners)
3839 }
3840
3841 fn index_row(&mut self, row: &Row) {
3842 if row.deleted {
3843 return;
3844 }
3845 let any_predicate = self
3853 .schema
3854 .indexes
3855 .iter()
3856 .any(|idx| idx.predicate.is_some());
3857 if any_predicate {
3858 let columns_map: HashMap<u16, &Value> =
3859 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3860 let name_to_id: HashMap<&str, u16> = self
3861 .schema
3862 .columns
3863 .iter()
3864 .map(|c| (c.name.as_str(), c.id))
3865 .collect();
3866 for idx in &self.schema.indexes {
3867 if let Some(pred) = &idx.predicate {
3868 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3869 continue; }
3871 }
3872 index_into_single(
3874 idx,
3875 &self.schema,
3876 row,
3877 &mut self.hot,
3878 &mut self.bitmap,
3879 &mut self.ann,
3880 &mut self.fm,
3881 &mut self.sparse,
3882 &mut self.minhash,
3883 );
3884 }
3885 return;
3886 }
3887 if self.column_keys.is_empty() {
3891 index_into(
3892 &self.schema,
3893 row,
3894 &mut self.hot,
3895 &mut self.bitmap,
3896 &mut self.ann,
3897 &mut self.fm,
3898 &mut self.sparse,
3899 &mut self.minhash,
3900 );
3901 return;
3902 }
3903 let effective_row = self.tokenized_for_indexes(row);
3904 index_into(
3905 &self.schema,
3906 &effective_row,
3907 &mut self.hot,
3908 &mut self.bitmap,
3909 &mut self.ann,
3910 &mut self.fm,
3911 &mut self.sparse,
3912 &mut self.minhash,
3913 );
3914 }
3915
3916 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3922 if self.column_keys.is_empty() {
3923 return row.clone();
3924 }
3925 {
3926 use crate::encryption::SCHEME_HMAC_EQ;
3927 let mut tok = row.clone();
3928 for (&cid, &(_, scheme)) in &self.column_keys {
3929 if scheme != SCHEME_HMAC_EQ {
3930 continue;
3931 }
3932 if let Some(v) = tok.columns.get(&cid).cloned() {
3933 if let Some(t) = self.tokenize_value(cid, &v) {
3934 tok.columns.insert(cid, t);
3935 }
3936 }
3937 }
3938 tok
3939 }
3940 }
3941
3942 pub fn commit(&mut self) -> Result<Epoch> {
3947 self.commit_inner(None)
3948 }
3949
3950 #[doc(hidden)]
3953 pub fn commit_controlled<F>(
3954 &mut self,
3955 control: &crate::ExecutionControl,
3956 mut before_commit: F,
3957 ) -> Result<Epoch>
3958 where
3959 F: FnMut() -> Result<()>,
3960 {
3961 self.commit_inner(Some((control, &mut before_commit)))
3962 }
3963
3964 fn commit_inner(
3965 &mut self,
3966 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3967 ) -> Result<Epoch> {
3968 self.ensure_writable()?;
3969 if !self.has_pending_mutations() {
3970 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3971 return Err(MongrelError::Full(
3972 "standalone transaction id namespace exhausted".into(),
3973 ));
3974 }
3975 return Ok(self.epoch.visible());
3976 }
3977 self.commit_new_epoch_inner(controlled)
3978 }
3979
3980 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3984 self.commit_new_epoch_inner(None)
3985 }
3986
3987 fn commit_new_epoch_inner(
3988 &mut self,
3989 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3990 ) -> Result<Epoch> {
3991 self.ensure_writable()?;
3992 if self.is_shared() {
3993 self.commit_shared(controlled)
3994 } else {
3995 self.commit_private(controlled)
3996 }
3997 }
3998
3999 fn commit_private(
4001 &mut self,
4002 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4003 ) -> Result<Epoch> {
4004 let commit_lock = Arc::clone(&self.commit_lock);
4008 let _g = commit_lock.lock();
4009 let txn_id = self.ensure_txn_id()?;
4012 if let Some((control, before_commit)) = controlled {
4013 control.checkpoint()?;
4014 before_commit()?;
4015 }
4016 let new_epoch = self.epoch.bump_assigned();
4017 let epoch_authority = Arc::clone(&self.epoch);
4018 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4019 let wal_result = match &mut self.wal {
4023 WalSink::Private(w) => w
4024 .append_txn(
4025 txn_id,
4026 Op::TxnCommit {
4027 epoch: new_epoch.0,
4028 added_runs: Vec::new(),
4029 },
4030 )
4031 .and_then(|_| w.sync()),
4032 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4033 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4034 };
4035 if let Err(error) = wal_result {
4036 self.durable_commit_failed = true;
4037 return Err(MongrelError::CommitOutcomeUnknown {
4038 epoch: new_epoch.0,
4039 message: error.to_string(),
4040 });
4041 }
4042 if let Some(epoch) = self.pending_truncate.take() {
4045 self.apply_truncate(epoch);
4046 }
4047 self.invalidate_pending_cache();
4048 let publish_result = self.persist_manifest(new_epoch);
4049 self.epoch.publish_in_order(new_epoch);
4053 epoch_guard.disarm();
4054 if let Err(error) = publish_result {
4055 self.durable_commit_failed = true;
4056 return Err(MongrelError::DurableCommit {
4057 epoch: new_epoch.0,
4058 message: error.to_string(),
4059 });
4060 }
4061 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4062 self.pending_private_mutations = false;
4063 self.data_generation = self.data_generation.wrapping_add(1);
4064 Ok(new_epoch)
4065 }
4066
4067 fn commit_shared(
4073 &mut self,
4074 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4075 ) -> Result<Epoch> {
4076 use std::sync::atomic::Ordering;
4077 let s = match &self.wal {
4078 WalSink::Shared(s) => s.clone(),
4079 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4080 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4081 };
4082 if s.poisoned.load(Ordering::Relaxed) {
4083 return Err(MongrelError::Other(
4084 "database poisoned by fsync error".into(),
4085 ));
4086 }
4087 let commit_lock = Arc::clone(&self.commit_lock);
4094 let _g = commit_lock.lock();
4095 if !self.pending_rows.is_empty() {
4096 match controlled.as_ref() {
4097 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4098 None => self.prepare_durable_publish()?,
4099 }
4100 }
4101 let txn_id = self.ensure_txn_id()?;
4104 let mut wal = s.wal.lock();
4105 if let Some((control, before_commit)) = controlled {
4106 control.checkpoint()?;
4107 before_commit()?;
4108 }
4109 let new_epoch = self.epoch.bump_assigned();
4110 let epoch_authority = Arc::clone(&self.epoch);
4111 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4112 let commit_ts = s.hlc.now().map_err(|skew| {
4115 MongrelError::Other(format!(
4116 "clock skew rejected commit timestamp allocation: {skew}"
4117 ))
4118 })?;
4119 let commit_seq = match wal.append_commit_at(
4120 txn_id,
4121 new_epoch,
4122 &[],
4123 commit_ts.physical_micros.saturating_mul(1_000),
4124 ) {
4125 Ok(commit_seq) => commit_seq,
4126 Err(error) => {
4127 s.poisoned.store(true, Ordering::Relaxed);
4128 s.lifecycle.poison();
4129 return Err(MongrelError::CommitOutcomeUnknown {
4130 epoch: new_epoch.0,
4131 message: error.to_string(),
4132 });
4133 }
4134 };
4135 drop(wal);
4136 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4137 s.poisoned.store(true, Ordering::Relaxed);
4138 s.lifecycle.poison();
4139 return Err(MongrelError::CommitOutcomeUnknown {
4140 epoch: new_epoch.0,
4141 message: error.to_string(),
4142 });
4143 }
4144
4145 if self.pending_truncate.take().is_some() {
4148 self.apply_truncate(new_epoch);
4149 }
4150 let mut rows = std::mem::take(&mut self.pending_rows);
4151 if !rows.is_empty() {
4152 for r in &mut rows {
4153 r.committed_epoch = new_epoch;
4154 r.commit_ts = Some(commit_ts);
4155 }
4156 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4157 let all_auto_generated =
4158 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4159 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4160 } else {
4161 self.pending_rows_auto_inc.clear();
4162 }
4163 let dels = std::mem::take(&mut self.pending_dels);
4164 for rid in dels {
4165 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
4166 }
4167
4168 self.invalidate_pending_cache();
4169 let publish_result = self.persist_manifest(new_epoch);
4170 self.epoch.publish_in_order(new_epoch);
4171 epoch_guard.disarm();
4172 let _ = s.change_wake.send(());
4173 if let Err(error) = publish_result {
4174 self.durable_commit_failed = true;
4175 s.poisoned.store(true, Ordering::Relaxed);
4176 s.lifecycle.poison();
4177 return Err(MongrelError::DurableCommit {
4178 epoch: new_epoch.0,
4179 message: error.to_string(),
4180 });
4181 }
4182 self.current_txn_id = 0;
4184 self.data_generation = self.data_generation.wrapping_add(1);
4185 Ok(new_epoch)
4186 }
4187
4188 pub fn flush(&mut self) -> Result<Epoch> {
4196 self.flush_with_outcome().map(|(epoch, _)| epoch)
4197 }
4198
4199 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4201 self.flush_with_outcome_inner(None)
4202 }
4203
4204 #[doc(hidden)]
4207 pub fn flush_with_outcome_controlled<F>(
4208 &mut self,
4209 control: &crate::ExecutionControl,
4210 mut before_commit: F,
4211 ) -> Result<(Epoch, bool)>
4212 where
4213 F: FnMut() -> Result<()>,
4214 {
4215 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4216 }
4217
4218 fn flush_with_outcome_inner(
4219 &mut self,
4220 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4221 ) -> Result<(Epoch, bool)> {
4222 match controlled.as_ref() {
4223 Some((control, _)) => {
4224 self.ensure_indexes_complete_controlled(control, || true)?;
4225 }
4226 None => self.ensure_indexes_complete()?,
4227 }
4228 let committed = self.has_pending_mutations();
4229 let epoch = self.commit_inner(controlled)?;
4230 let finish: Result<(Epoch, bool)> = (|| {
4231 let rows = self.memtable.drain_sorted();
4232 if !rows.is_empty() {
4233 self.mutable_run.insert_many(rows);
4234 }
4235 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4236 self.spill_mutable_run(epoch)?;
4237 self.mark_flushed(epoch)?;
4241 self.persist_manifest(epoch)?;
4242 self.build_learned_ranges()?;
4243 self.checkpoint_indexes(epoch);
4246 }
4247 Ok((epoch, committed))
4250 })();
4251 let outcome = match finish {
4252 Err(error) if committed => Err(MongrelError::DurableCommit {
4253 epoch: epoch.0,
4254 message: error.to_string(),
4255 }),
4256 result => result,
4257 };
4258 if outcome.is_ok() {
4259 let _ = self.publish_read_generation();
4265 }
4266 outcome
4267 }
4268
4269 fn has_pending_mutations(&self) -> bool {
4270 self.pending_private_mutations
4271 || !self.pending_rows.is_empty()
4272 || !self.pending_dels.is_empty()
4273 || self.pending_truncate.is_some()
4274 }
4275
4276 pub fn has_pending_writes(&self) -> bool {
4277 self.has_pending_mutations()
4278 }
4279
4280 pub fn force_flush(&mut self) -> Result<Epoch> {
4289 let saved = self.mutable_run_spill_bytes;
4290 self.mutable_run_spill_bytes = 1;
4291 let result = self.flush();
4292 self.mutable_run_spill_bytes = saved;
4293 result
4294 }
4295
4296 pub fn close(&mut self) -> Result<()> {
4303 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4304 self.force_flush()?;
4305 }
4306 Ok(())
4307 }
4308
4309 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4316 let op = Op::Flush {
4317 table_id: self.table_id,
4318 flushed_epoch: epoch.0,
4319 };
4320 match &mut self.wal {
4321 WalSink::Private(w) => {
4322 w.append_system(op)?;
4323 w.sync()?;
4324 }
4325 WalSink::Shared(s) => {
4326 s.wal.lock().append_system(op)?;
4331 }
4332 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4333 }
4334 self.flushed_epoch = epoch.0;
4335 if matches!(self.wal, WalSink::Private(_)) {
4336 self.rotate_wal(epoch)?;
4337 }
4338 Ok(())
4339 }
4340
4341 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4345 if self.mutable_run.is_empty() {
4346 return Ok(());
4347 }
4348 let run_id = self.alloc_run_id()?;
4349 let rows = self.mutable_run.drain_sorted();
4350 if rows.is_empty() {
4351 return Ok(());
4352 }
4353 let path = self.run_path(run_id);
4354 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4355 if let Some(kek) = &self.kek {
4356 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4357 }
4358 let header = match self.create_run_file(run_id)? {
4359 Some(file) => writer.write_file(file, &rows)?,
4360 None => writer.write(&path, &rows)?,
4361 };
4362 self.run_refs.push(RunRef {
4363 run_id: run_id as u128,
4364 level: 0,
4365 epoch_created: epoch.0,
4366 row_count: header.row_count,
4367 });
4368 Ok(())
4369 }
4370
4371 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4375 self.mutable_run_spill_bytes = bytes.max(1);
4376 }
4377
4378 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4382 self.compaction_zstd_level = level;
4383 }
4384
4385 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4389 self.result_cache.lock().set_max_bytes(max_bytes);
4390 }
4391
4392 pub(crate) fn clear_result_cache(&mut self) {
4396 self.result_cache.lock().clear();
4397 }
4398
4399 pub fn mutable_run_len(&self) -> usize {
4401 self.mutable_run.len()
4402 }
4403
4404 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4407 self.mutable_run.drain_sorted()
4408 }
4409
4410 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4412 let mut snapshot = self.mutable_run.clone();
4413 snapshot.drain_sorted()
4414 }
4415
4416 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4421 self.ensure_writable()?;
4422 let n = batch.len();
4423 if n == 0 {
4424 return Ok(self.current_epoch());
4425 }
4426 for row in &batch {
4427 self.schema.validate_values(row)?;
4428 }
4429 let epoch = self.commit_new_epoch()?;
4430 let live_before = self.live_count;
4431 self.spill_mutable_run(epoch)?;
4435 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4436 && self.indexes_complete
4437 && self.run_refs.is_empty()
4438 && self.memtable.is_empty()
4439 && self.mutable_run.is_empty();
4440 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4446 use rayon::prelude::*;
4447 self.schema
4448 .columns
4449 .par_iter()
4450 .map(|cdef| {
4451 (
4452 cdef.id,
4453 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4454 )
4455 })
4456 .collect::<Vec<_>>()
4457 };
4458 drop(batch);
4459 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4464 self.validate_columns_not_null(&user_columns, n)?;
4465 let winner_idx = self
4466 .bulk_pk_winner_indices(&user_columns, n)
4467 .filter(|idx| idx.len() != n);
4468 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4469 match winner_idx.as_deref() {
4470 Some(idx) => {
4471 let compacted = user_columns
4472 .iter()
4473 .map(|(id, c)| (*id, c.gather(idx)))
4474 .collect();
4475 (compacted, idx.len())
4476 }
4477 None => (std::mem::take(&mut user_columns), n),
4478 };
4479 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4480 let first = self.allocator.alloc_range(write_n as u64)?.0;
4481 for rid in first..first + write_n as u64 {
4482 self.reservoir.offer(rid);
4483 }
4484 let run_id = self.alloc_run_id()?;
4485 let path = self.run_path(run_id);
4486 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4487 .clean(true)
4488 .with_lz4()
4489 .with_native_endian();
4490 if let Some(kek) = &self.kek {
4491 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4492 }
4493 let header = match self.create_run_file(run_id)? {
4494 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4495 None => writer.write_native(&path, &write_columns, write_n, first)?,
4496 };
4497 self.run_refs.push(RunRef {
4498 run_id: run_id as u128,
4499 level: 0,
4500 epoch_created: epoch.0,
4501 row_count: header.row_count,
4502 });
4503 self.live_count = self.live_count.saturating_add(write_n as u64);
4504 if eager_index_build {
4505 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4506 self.index_columns_bulk(&write_columns, &row_ids);
4507 self.indexes_complete = true;
4508 self.build_learned_ranges()?;
4509 } else {
4510 self.indexes_complete = false;
4511 }
4512 self.mark_flushed(epoch)?;
4513 self.persist_manifest(epoch)?;
4514 if eager_index_build {
4515 self.checkpoint_indexes(epoch);
4516 }
4517 self.clear_result_cache();
4518 Ok(epoch)
4519 }
4520
4521 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4524 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4525 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4526 let segment_no = segment
4529 .file_stem()
4530 .and_then(|s| s.to_str())
4531 .and_then(|s| s.strip_prefix("seg-"))
4532 .and_then(|s| s.parse::<u64>().ok())
4533 .unwrap_or(0);
4534 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4535 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4536 wal.sync()?;
4537 self.wal = WalSink::Private(wal);
4538 Ok(())
4539 }
4540
4541 pub(crate) fn invalidate_pending_cache(&mut self) {
4546 self.result_cache
4547 .lock()
4548 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4549 self.pending_delete_rids.clear();
4550 self.pending_put_cols.clear();
4551 }
4552
4553 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4554 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4555 m.current_epoch = epoch.0;
4556 m.next_row_id = self.allocator.current().0;
4557 m.runs = self.run_refs.clone();
4558 m.live_count = self.live_count;
4559 m.global_idx_epoch = self.global_idx_epoch;
4560 m.flushed_epoch = self.flushed_epoch;
4561 m.retiring = self.retiring.clone();
4562 m.auto_inc_next = match self.auto_inc {
4566 Some(ai) if ai.seeded => ai.next,
4567 _ => 0,
4568 };
4569 m.ttl = self.ttl;
4570 let meta_dek = self.manifest_meta_dek();
4571 match self._root_guard.as_deref() {
4572 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4573 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4574 }
4575 Ok(())
4576 }
4577
4578 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4579 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
4583 let live_count = u64::try_from(rows.len())
4584 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4585 let auto_inc = match self.auto_inc {
4586 Some(mut state) => {
4587 let maximum = self.scan_max_int64(state.column_id)?;
4588 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4589 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4590 })?;
4591 state.next = state.next.max(after_maximum).max(1);
4592 state.seeded = true;
4593 Some(state)
4594 }
4595 None => None,
4596 };
4597 Ok(RecoveryMetadataPlan {
4598 live_count,
4599 auto_inc,
4600 changed: live_count != self.live_count
4601 || auto_inc.is_some_and(|planned| {
4602 self.auto_inc.is_none_or(|current| {
4603 current.next != planned.next || current.seeded != planned.seeded
4604 })
4605 }),
4606 })
4607 }
4608
4609 pub(crate) fn apply_recovered_metadata(
4610 &mut self,
4611 plan: RecoveryMetadataPlan,
4612 epoch: Epoch,
4613 ) -> Result<()> {
4614 if !plan.changed {
4615 return Ok(());
4616 }
4617 self.live_count = plan.live_count;
4618 self.auto_inc = plan.auto_inc;
4619 self.persist_manifest(epoch)
4620 }
4621
4622 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4628 if !self.indexes_complete {
4631 return;
4632 }
4633 if crate::catalog::inject_hook("index.publish.before").is_err() {
4636 return;
4637 }
4638 if self.idx_root.is_none() {
4639 if let Some(root) = self._root_guard.as_ref() {
4640 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4641 return;
4642 };
4643 self.idx_root = Some(Arc::new(idx_root));
4644 }
4645 }
4646 let snap = global_idx::IndexSnapshot {
4647 hot: &self.hot,
4648 bitmap: &self.bitmap,
4649 ann: &self.ann,
4650 fm: &self.fm,
4651 sparse: &self.sparse,
4652 minhash: &self.minhash,
4653 learned_range: &self.learned_range,
4654 };
4655 let idx_dek = self.idx_dek();
4657 let written = match self.idx_root.as_deref() {
4658 Some(root) => global_idx::write_atomic_root(
4659 root,
4660 self.table_id,
4661 epoch.0,
4662 snap,
4663 idx_dek.as_deref(),
4664 ),
4665 None => global_idx::write_atomic(
4666 &self.dir,
4667 self.table_id,
4668 epoch.0,
4669 snap,
4670 idx_dek.as_deref(),
4671 ),
4672 };
4673 if written.is_ok() {
4674 self.global_idx_epoch = epoch.0;
4675 let _ = self.persist_manifest(epoch);
4676 let _ = crate::catalog::inject_hook("index.publish.after");
4678 }
4679 }
4680
4681 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4684 self.global_idx_epoch = 0;
4685 if let Some(root) = self.idx_root.as_deref() {
4686 let _ = root.remove_file(global_idx::IDX_FILENAME);
4687 } else {
4688 global_idx::remove(&self.dir);
4689 }
4690 let _ = self.persist_manifest(self.epoch.visible());
4691 }
4692
4693 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4698 self.indexes_complete = false;
4699 self.global_idx_epoch = 0;
4700 if let Some(root) = self.idx_root.as_deref() {
4701 let _ = root.remove_file(global_idx::IDX_FILENAME);
4702 } else {
4703 global_idx::remove(&self.dir);
4704 }
4705 }
4706
4707 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4708 self.indexes_complete = true;
4709 }
4710
4711 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4717 self.durable_commit_failed = true;
4718 if let WalSink::Shared(shared) = &self.wal {
4719 shared
4720 .poisoned
4721 .store(true, std::sync::atomic::Ordering::Relaxed);
4722 }
4723 }
4724
4725 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4729 self.durable_commit_failed = true;
4730 }
4731
4732 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4741 let mut best: Option<Row> = None;
4742 let mut consider = |row: Row| {
4743 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
4744 return;
4745 }
4746 if best.as_ref().is_none_or(|current| {
4747 Snapshot::version_is_newer(
4748 row.committed_epoch,
4749 row.commit_ts,
4750 current.committed_epoch,
4751 current.commit_ts,
4752 )
4753 }) {
4754 best = Some(row);
4755 }
4756 };
4757 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
4758 consider(row);
4759 }
4760 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
4761 consider(row);
4762 }
4763 for rr in &self.run_refs {
4764 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4765 continue;
4766 };
4767 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
4770 continue;
4771 };
4772 consider(row);
4773 }
4774 let now_nanos = unix_nanos_now();
4775 match best {
4776 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4777 Some(r) => Some(r),
4778 None => None,
4779 }
4780 }
4781
4782 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4786 self.visible_rows_at_time(snapshot, unix_nanos_now())
4787 }
4788
4789 #[doc(hidden)]
4792 pub fn visible_rows_controlled(
4793 &self,
4794 snapshot: Snapshot,
4795 control: &crate::ExecutionControl,
4796 ) -> Result<Vec<Row>> {
4797 let mut out = Vec::new();
4798 self.for_each_visible_row_controlled(snapshot, control, |row| {
4799 out.push(row);
4800 Ok(())
4801 })?;
4802 Ok(out)
4803 }
4804
4805 #[doc(hidden)]
4808 pub fn for_each_visible_row_controlled<F>(
4809 &self,
4810 snapshot: Snapshot,
4811 control: &crate::ExecutionControl,
4812 mut visit: F,
4813 ) -> Result<()>
4814 where
4815 F: FnMut(Row) -> Result<()>,
4816 {
4817 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4818 control.checkpoint()?;
4819 let memtable = self.memtable.visible_versions_at(snapshot);
4821 if !memtable.is_empty() {
4822 sources.push(ControlledVisibleSource::memory(memtable));
4823 }
4824 control.checkpoint()?;
4825 let mutable = self.mutable_run.visible_versions_at(snapshot);
4828 if !mutable.is_empty() {
4829 sources.push(ControlledVisibleSource::memory(mutable));
4830 }
4831 for run in &self.run_refs {
4832 control.checkpoint()?;
4833 let reader = self.open_reader(run.run_id)?;
4834 sources.push(ControlledVisibleSource::run(
4836 reader.into_visible_version_cursor(snapshot.epoch)?,
4837 ));
4838 }
4839 let now_nanos = unix_nanos_now();
4840 merge_controlled_visible_sources(
4841 &mut sources,
4842 control,
4843 |row| self.row_expired_at(row, now_nanos),
4844 |row| {
4845 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
4848 return Ok(());
4849 }
4850 visit(row)
4851 },
4852 )
4853 }
4854
4855 #[doc(hidden)]
4856 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4857 let mut best: HashMap<u64, Row> = HashMap::new();
4858 let mut fold = |row: Row| {
4859 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
4860 return;
4861 }
4862 best.entry(row.row_id.0)
4863 .and_modify(|existing| {
4864 if Snapshot::version_is_newer(
4865 row.committed_epoch,
4866 row.commit_ts,
4867 existing.committed_epoch,
4868 existing.commit_ts,
4869 ) {
4870 *existing = row.clone();
4871 }
4872 })
4873 .or_insert(row);
4874 };
4875 for row in self.memtable.visible_versions_at(snapshot) {
4876 fold(row);
4877 }
4878 for row in self.mutable_run.visible_versions_at(snapshot) {
4879 fold(row);
4880 }
4881 for rr in &self.run_refs {
4882 let mut reader = self.open_reader(rr.run_id)?;
4883 for row in reader.visible_versions(snapshot.epoch)? {
4885 fold(row);
4886 }
4887 }
4888 let mut out: Vec<Row> = best
4889 .into_values()
4890 .filter_map(|r| {
4891 if r.deleted || self.row_expired_at(&r, now_nanos) {
4892 None
4893 } else {
4894 Some(r)
4895 }
4896 })
4897 .collect();
4898 out.sort_by_key(|r| r.row_id);
4899 Ok(out)
4900 }
4901
4902 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4909 if self.ttl.is_none()
4910 && self.memtable.is_empty()
4911 && self.mutable_run.is_empty()
4912 && self.run_refs.len() == 1
4913 {
4914 let rr = self.run_refs[0].clone();
4915 let mut reader = self.open_reader(rr.run_id)?;
4916 let idxs = reader.visible_indices(snapshot.epoch)?;
4917 let mut cols = Vec::with_capacity(self.schema.columns.len());
4918 for cdef in &self.schema.columns {
4919 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4920 }
4921 return Ok(cols);
4922 }
4923 let rows = self.visible_rows(snapshot)?;
4925 let mut cols: Vec<(u16, Vec<Value>)> = self
4926 .schema
4927 .columns
4928 .iter()
4929 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4930 .collect();
4931 for r in &rows {
4932 for (cid, vec) in cols.iter_mut() {
4933 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4934 }
4935 }
4936 Ok(cols)
4937 }
4938
4939 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4941 let row_id = self.hot.get(key)?;
4942 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
4943 Some(row_id)
4944 } else {
4945 None
4946 }
4947 }
4948
4949 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4954 self.query_at_with_allowed(q, self.snapshot(), None)
4955 }
4956
4957 pub fn query_controlled(
4960 &mut self,
4961 q: &crate::query::Query,
4962 control: &crate::ExecutionControl,
4963 ) -> Result<Vec<Row>> {
4964 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4965 }
4966
4967 pub fn query_at_with_allowed(
4970 &mut self,
4971 q: &crate::query::Query,
4972 snapshot: Snapshot,
4973 allowed: Option<&std::collections::HashSet<RowId>>,
4974 ) -> Result<Vec<Row>> {
4975 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4976 }
4977
4978 #[doc(hidden)]
4979 pub fn query_at_with_allowed_controlled(
4980 &mut self,
4981 q: &crate::query::Query,
4982 snapshot: Snapshot,
4983 allowed: Option<&std::collections::HashSet<RowId>>,
4984 control: &crate::ExecutionControl,
4985 ) -> Result<Vec<Row>> {
4986 self.require_select()?;
4987 self.ensure_indexes_complete_controlled(control, || true)?;
4988 self.validate_native_query(q)?;
4989 self.query_conditions_at(
4990 &q.conditions,
4991 snapshot,
4992 allowed,
4993 q.limit,
4994 q.offset,
4995 None,
4996 unix_nanos_now(),
4997 Some(control),
4998 )
4999 }
5000
5001 #[doc(hidden)]
5002 pub fn query_at_with_allowed_after(
5003 &mut self,
5004 q: &crate::query::Query,
5005 snapshot: Snapshot,
5006 allowed: Option<&std::collections::HashSet<RowId>>,
5007 after_row_id: Option<RowId>,
5008 ) -> Result<Vec<Row>> {
5009 self.query_at_with_allowed_after_at_time(
5010 q,
5011 snapshot,
5012 allowed,
5013 after_row_id,
5014 unix_nanos_now(),
5015 )
5016 }
5017
5018 #[doc(hidden)]
5019 pub fn query_at_with_allowed_after_at_time(
5020 &mut self,
5021 q: &crate::query::Query,
5022 snapshot: Snapshot,
5023 allowed: Option<&std::collections::HashSet<RowId>>,
5024 after_row_id: Option<RowId>,
5025 query_time_nanos: i64,
5026 ) -> Result<Vec<Row>> {
5027 self.require_select()?;
5028 self.ensure_indexes_complete()?;
5029 self.validate_native_query(q)?;
5030 self.query_conditions_at(
5031 &q.conditions,
5032 snapshot,
5033 allowed,
5034 q.limit,
5035 q.offset,
5036 after_row_id,
5037 query_time_nanos,
5038 None,
5039 )
5040 }
5041
5042 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
5043 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5044 return Err(MongrelError::InvalidArgument(format!(
5045 "query exceeds {} conditions",
5046 crate::query::MAX_HARD_CONDITIONS
5047 )));
5048 }
5049 if let Some(limit) = q.limit {
5050 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
5051 return Err(MongrelError::InvalidArgument(format!(
5052 "query limit must be between 1 and {}",
5053 crate::query::MAX_FINAL_LIMIT
5054 )));
5055 }
5056 }
5057 if q.offset > crate::query::MAX_QUERY_OFFSET {
5058 return Err(MongrelError::InvalidArgument(format!(
5059 "query offset exceeds {}",
5060 crate::query::MAX_QUERY_OFFSET
5061 )));
5062 }
5063 Ok(())
5064 }
5065
5066 #[doc(hidden)]
5069 pub fn query_all_at(
5070 &mut self,
5071 conditions: &[crate::query::Condition],
5072 snapshot: Snapshot,
5073 ) -> Result<Vec<Row>> {
5074 self.require_select()?;
5075 self.ensure_indexes_complete()?;
5076 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
5077 return Err(MongrelError::InvalidArgument(format!(
5078 "query exceeds {} conditions",
5079 crate::query::MAX_HARD_CONDITIONS
5080 )));
5081 }
5082 self.query_conditions_at(
5083 conditions,
5084 snapshot,
5085 None,
5086 None,
5087 0,
5088 None,
5089 unix_nanos_now(),
5090 None,
5091 )
5092 }
5093
5094 #[allow(clippy::too_many_arguments)]
5095 fn query_conditions_at(
5096 &self,
5097 conditions: &[crate::query::Condition],
5098 snapshot: Snapshot,
5099 allowed: Option<&std::collections::HashSet<RowId>>,
5100 limit: Option<usize>,
5101 offset: usize,
5102 after_row_id: Option<RowId>,
5103 query_time_nanos: i64,
5104 control: Option<&crate::ExecutionControl>,
5105 ) -> Result<Vec<Row>> {
5106 control
5107 .map(crate::ExecutionControl::checkpoint)
5108 .transpose()?;
5109 crate::trace::QueryTrace::record(|t| {
5110 t.run_count = self.run_refs.len();
5111 t.memtable_rows = self.memtable.len();
5112 t.mutable_run_rows = self.mutable_run.len();
5113 });
5114 if conditions.is_empty() {
5118 crate::trace::QueryTrace::record(|t| {
5119 t.scan_mode = crate::trace::ScanMode::Materialized;
5120 t.row_materialized = true;
5121 });
5122 let mut rows = match control {
5123 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5124 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5125 };
5126 if let Some(allowed) = allowed {
5127 let mut filtered = Vec::with_capacity(rows.len());
5128 for (index, row) in rows.into_iter().enumerate() {
5129 if index & 255 == 0 {
5130 control
5131 .map(crate::ExecutionControl::checkpoint)
5132 .transpose()?;
5133 }
5134 if allowed.contains(&row.row_id) {
5135 filtered.push(row);
5136 }
5137 }
5138 rows = filtered;
5139 }
5140 if let Some(after_row_id) = after_row_id {
5141 rows.retain(|row| row.row_id > after_row_id);
5142 }
5143 rows.drain(..offset.min(rows.len()));
5144 if let Some(limit) = limit {
5145 rows.truncate(limit);
5146 }
5147 return Ok(rows);
5148 }
5149 crate::trace::QueryTrace::record(|t| {
5150 t.conditions_pushed = conditions.len();
5151 t.scan_mode = crate::trace::ScanMode::Materialized;
5152 t.row_materialized = true;
5153 });
5154 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5161 ordered.sort_by_key(|c| condition_cost_rank(c));
5162 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5163 for c in &ordered {
5164 control
5165 .map(crate::ExecutionControl::checkpoint)
5166 .transpose()?;
5167 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5168 let empty = s.is_empty();
5169 sets.push(s);
5170 if empty {
5171 break;
5172 }
5173 }
5174 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5175 if let Some(allowed) = allowed {
5176 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5177 }
5178 if let Some(after_row_id) = after_row_id {
5179 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5180 rids.drain(..first);
5181 }
5182 rids.drain(..offset.min(rids.len()));
5183 if let Some(limit) = limit {
5184 rids.truncate(limit);
5185 }
5186 control
5187 .map(crate::ExecutionControl::checkpoint)
5188 .transpose()?;
5189 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5190 }
5191
5192 pub fn retrieve(
5194 &mut self,
5195 retriever: &crate::query::Retriever,
5196 ) -> Result<Vec<crate::query::RetrieverHit>> {
5197 self.retrieve_with_allowed(retriever, None)
5198 }
5199
5200 pub fn retrieve_at(
5201 &mut self,
5202 retriever: &crate::query::Retriever,
5203 snapshot: Snapshot,
5204 allowed: Option<&std::collections::HashSet<RowId>>,
5205 ) -> Result<Vec<crate::query::RetrieverHit>> {
5206 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5207 }
5208
5209 pub fn retrieve_with_allowed(
5212 &mut self,
5213 retriever: &crate::query::Retriever,
5214 allowed: Option<&std::collections::HashSet<RowId>>,
5215 ) -> Result<Vec<crate::query::RetrieverHit>> {
5216 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5217 }
5218
5219 pub fn retrieve_at_with_allowed(
5220 &mut self,
5221 retriever: &crate::query::Retriever,
5222 snapshot: Snapshot,
5223 allowed: Option<&std::collections::HashSet<RowId>>,
5224 ) -> Result<Vec<crate::query::RetrieverHit>> {
5225 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5226 }
5227
5228 pub fn retrieve_at_with_allowed_and_context(
5229 &mut self,
5230 retriever: &crate::query::Retriever,
5231 snapshot: Snapshot,
5232 allowed: Option<&std::collections::HashSet<RowId>>,
5233 context: Option<&crate::query::AiExecutionContext>,
5234 ) -> Result<Vec<crate::query::RetrieverHit>> {
5235 self.require_select()?;
5236 self.ensure_indexes_complete()?;
5237 self.validate_retriever(retriever)?;
5238 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5239 }
5240
5241 pub fn retrieve_at_with_candidate_authorization_and_context(
5242 &mut self,
5243 retriever: &crate::query::Retriever,
5244 snapshot: Snapshot,
5245 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5246 context: Option<&crate::query::AiExecutionContext>,
5247 ) -> Result<Vec<crate::query::RetrieverHit>> {
5248 self.require_select()?;
5249 self.ensure_indexes_complete()?;
5250 self.retrieve_at_with_candidate_authorization_on_generation(
5251 retriever,
5252 snapshot,
5253 authorization,
5254 context,
5255 )
5256 }
5257
5258 #[doc(hidden)]
5259 pub fn retrieve_at_with_candidate_authorization_on_generation(
5260 &self,
5261 retriever: &crate::query::Retriever,
5262 snapshot: Snapshot,
5263 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5264 context: Option<&crate::query::AiExecutionContext>,
5265 ) -> Result<Vec<crate::query::RetrieverHit>> {
5266 self.require_select()?;
5267 self.validate_retriever(retriever)?;
5268 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5269 }
5270
5271 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5272 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5273 let (column_id, k) = match retriever {
5274 Retriever::Ann {
5275 column_id,
5276 query,
5277 k,
5278 } => {
5279 let index = self.ann.get(column_id).ok_or_else(|| {
5280 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5281 })?;
5282 if query.len() != index.dim() {
5283 return Err(MongrelError::InvalidArgument(format!(
5284 "ANN query dimension must be {}, got {}",
5285 index.dim(),
5286 query.len()
5287 )));
5288 }
5289 if query.iter().any(|value| !value.is_finite()) {
5290 return Err(MongrelError::InvalidArgument(
5291 "ANN query values must be finite".into(),
5292 ));
5293 }
5294 (*column_id, *k)
5295 }
5296 Retriever::Sparse {
5297 column_id,
5298 query,
5299 k,
5300 } => {
5301 if !self.sparse.contains_key(column_id) {
5302 return Err(MongrelError::InvalidArgument(format!(
5303 "column {column_id} has no Sparse index"
5304 )));
5305 }
5306 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5307 return Err(MongrelError::InvalidArgument(
5308 "Sparse query must be non-empty with finite weights".into(),
5309 ));
5310 }
5311 if query.len() > MAX_SPARSE_TERMS {
5312 return Err(MongrelError::InvalidArgument(format!(
5313 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5314 )));
5315 }
5316 (*column_id, *k)
5317 }
5318 Retriever::MinHash {
5319 column_id,
5320 members,
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 members.is_empty() {
5329 return Err(MongrelError::InvalidArgument(
5330 "MinHash members must not be empty".into(),
5331 ));
5332 }
5333 if members.len() > MAX_SET_MEMBERS {
5334 return Err(MongrelError::InvalidArgument(format!(
5335 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5336 )));
5337 }
5338 let mut total_bytes = 0usize;
5339 for member in members {
5340 let bytes = member.encoded_len();
5341 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5342 return Err(MongrelError::InvalidArgument(format!(
5343 "MinHash member exceeds {} bytes",
5344 crate::query::MAX_SET_MEMBER_BYTES
5345 )));
5346 }
5347 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5348 MongrelError::InvalidArgument("MinHash input size overflow".into())
5349 })?;
5350 }
5351 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5352 return Err(MongrelError::InvalidArgument(format!(
5353 "MinHash input exceeds {} bytes",
5354 crate::query::MAX_SET_INPUT_BYTES
5355 )));
5356 }
5357 (*column_id, *k)
5358 }
5359 };
5360 if k == 0 {
5361 return Err(MongrelError::InvalidArgument(
5362 "retriever k must be > 0".into(),
5363 ));
5364 }
5365 if k > MAX_RETRIEVER_K {
5366 return Err(MongrelError::InvalidArgument(format!(
5367 "retriever k exceeds {MAX_RETRIEVER_K}"
5368 )));
5369 }
5370 debug_assert!(self
5371 .schema
5372 .columns
5373 .iter()
5374 .any(|column| column.id == column_id));
5375 Ok(())
5376 }
5377
5378 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5379 use crate::query::Condition;
5380 match condition {
5381 Condition::Ann {
5382 column_id,
5383 query,
5384 k,
5385 } => self.validate_retriever(&crate::query::Retriever::Ann {
5386 column_id: *column_id,
5387 query: query.clone(),
5388 k: *k,
5389 }),
5390 Condition::SparseMatch {
5391 column_id,
5392 query,
5393 k,
5394 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5395 column_id: *column_id,
5396 query: query.clone(),
5397 k: *k,
5398 }),
5399 Condition::MinHashSimilar {
5400 column_id,
5401 query,
5402 k,
5403 } => {
5404 if !self.minhash.contains_key(column_id) {
5405 return Err(MongrelError::InvalidArgument(format!(
5406 "column {column_id} has no MinHash index"
5407 )));
5408 }
5409 if query.is_empty() || *k == 0 {
5410 return Err(MongrelError::InvalidArgument(
5411 "MinHash query must be non-empty and k must be > 0".into(),
5412 ));
5413 }
5414 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5415 {
5416 return Err(MongrelError::InvalidArgument(format!(
5417 "MinHash query must have <= {} members and k <= {}",
5418 crate::query::MAX_SET_MEMBERS,
5419 crate::query::MAX_RETRIEVER_K
5420 )));
5421 }
5422 Ok(())
5423 }
5424 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5425 Err(MongrelError::InvalidArgument(format!(
5426 "bitmap IN exceeds {} values",
5427 crate::query::MAX_SET_MEMBERS
5428 )))
5429 }
5430 Condition::FmContainsAll { patterns, .. }
5431 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5432 {
5433 Err(MongrelError::InvalidArgument(format!(
5434 "FM query exceeds {} patterns",
5435 crate::query::MAX_HARD_CONDITIONS
5436 )))
5437 }
5438 _ => Ok(()),
5439 }
5440 }
5441
5442 fn retrieve_filtered(
5443 &self,
5444 retriever: &crate::query::Retriever,
5445 snapshot: Snapshot,
5446 hard_filter: Option<&RowIdSet>,
5447 allowed: Option<&std::collections::HashSet<RowId>>,
5448 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5449 context: Option<&crate::query::AiExecutionContext>,
5450 ) -> Result<Vec<crate::query::RetrieverHit>> {
5451 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5452 let started = std::time::Instant::now();
5453 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5454 Retriever::Ann {
5455 column_id,
5456 query,
5457 k,
5458 } => {
5459 let Some(index) = self.ann.get(column_id) else {
5460 return Ok(Vec::new());
5461 };
5462 let cap = ann_candidate_cap(index.len(), context);
5463 if cap == 0 {
5464 return Ok(Vec::new());
5465 }
5466 let mut breadth = (*k).max(1).min(cap);
5467 let mut eligibility = std::collections::HashMap::new();
5468 let mut filtered = loop {
5469 let mut seen = std::collections::HashSet::new();
5470 if let Some(context) = context {
5471 context.checkpoint()?;
5472 }
5473 let raw = index.search_with_context(query, breadth, context)?;
5474 let unchecked: Vec<_> = raw
5475 .iter()
5476 .map(|(row_id, _)| *row_id)
5477 .filter(|row_id| !eligibility.contains_key(row_id))
5478 .filter(|row_id| {
5479 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5480 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5481 })
5482 .collect();
5483 let eligible = self.eligible_and_authorized_candidate_ids(
5484 &unchecked,
5485 *column_id,
5486 snapshot,
5487 candidate_authorization,
5488 context,
5489 )?;
5490 for row_id in unchecked {
5491 eligibility.insert(row_id, eligible.contains(&row_id));
5492 }
5493 let filtered: Vec<_> = raw
5494 .into_iter()
5495 .filter(|(row_id, _)| {
5496 seen.insert(*row_id)
5497 && eligibility.get(row_id).copied().unwrap_or(false)
5498 })
5499 .map(|(row_id, score)| {
5500 let score = match score {
5501 crate::index::AnnDistance::Hamming(d) => {
5502 RetrieverScore::AnnHammingDistance(d)
5503 }
5504 crate::index::AnnDistance::Cosine(d) => {
5505 RetrieverScore::AnnCosineDistance(d)
5506 }
5507 };
5508 (row_id, score)
5509 })
5510 .collect();
5511 if filtered.len() >= *k || breadth >= cap {
5512 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5513 crate::trace::QueryTrace::record(|trace| {
5514 trace.ann_candidate_cap_hit = true;
5515 });
5516 }
5517 break filtered;
5518 }
5519 breadth = breadth.saturating_mul(2).min(cap);
5520 };
5521 filtered.truncate(*k);
5522 filtered
5523 }
5524 Retriever::Sparse {
5525 column_id,
5526 query,
5527 k,
5528 } => self
5529 .sparse
5530 .get(column_id)
5531 .map(|index| -> Result<Vec<_>> {
5532 let mut breadth = (*k).max(1);
5533 let mut eligibility = std::collections::HashMap::new();
5534 loop {
5535 if let Some(context) = context {
5536 context.checkpoint()?;
5537 }
5538 let raw = index.search_with_context(query, breadth, context)?;
5539 let unchecked: Vec<_> = raw
5540 .iter()
5541 .map(|(row_id, _)| *row_id)
5542 .filter(|row_id| !eligibility.contains_key(row_id))
5543 .filter(|row_id| {
5544 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5545 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5546 })
5547 .collect();
5548 let eligible = self.eligible_and_authorized_candidate_ids(
5549 &unchecked,
5550 *column_id,
5551 snapshot,
5552 candidate_authorization,
5553 context,
5554 )?;
5555 for row_id in unchecked {
5556 eligibility.insert(row_id, eligible.contains(&row_id));
5557 }
5558 let filtered: Vec<_> = raw
5559 .iter()
5560 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5561 .take(*k)
5562 .map(|(row_id, score)| {
5563 (*row_id, RetrieverScore::SparseDotProduct(*score))
5564 })
5565 .collect();
5566 if filtered.len() >= *k || raw.len() < breadth {
5567 break Ok(filtered);
5568 }
5569 let next = breadth.saturating_mul(2);
5570 if next == breadth {
5571 break Ok(filtered);
5572 }
5573 breadth = next;
5574 }
5575 })
5576 .transpose()?
5577 .unwrap_or_default(),
5578 Retriever::MinHash {
5579 column_id,
5580 members,
5581 k,
5582 } => self
5583 .minhash
5584 .get(column_id)
5585 .map(|index| -> Result<Vec<_>> {
5586 let mut hashes = Vec::with_capacity(members.len());
5587 for member in members {
5588 if let Some(context) = context {
5589 context.consume(crate::query::work_units(
5590 member.encoded_len(),
5591 crate::query::PARSE_WORK_QUANTUM,
5592 ))?;
5593 }
5594 hashes.push(member.hash_v1());
5595 }
5596 let mut breadth = (*k).max(1);
5597 let mut eligibility = std::collections::HashMap::new();
5598 loop {
5599 if let Some(context) = context {
5600 context.checkpoint()?;
5601 }
5602 let raw = index.search_with_context(&hashes, breadth, context)?;
5603 let unchecked: Vec<_> = raw
5604 .iter()
5605 .map(|(row_id, _)| *row_id)
5606 .filter(|row_id| !eligibility.contains_key(row_id))
5607 .filter(|row_id| {
5608 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5609 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5610 })
5611 .collect();
5612 let eligible = self.eligible_and_authorized_candidate_ids(
5613 &unchecked,
5614 *column_id,
5615 snapshot,
5616 candidate_authorization,
5617 context,
5618 )?;
5619 for row_id in unchecked {
5620 eligibility.insert(row_id, eligible.contains(&row_id));
5621 }
5622 let filtered: Vec<_> = raw
5623 .iter()
5624 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5625 .take(*k)
5626 .map(|(row_id, score)| {
5627 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5628 })
5629 .collect();
5630 if filtered.len() >= *k || raw.len() < breadth {
5631 break Ok(filtered);
5632 }
5633 let next = breadth.saturating_mul(2);
5634 if next == breadth {
5635 break Ok(filtered);
5636 }
5637 breadth = next;
5638 }
5639 })
5640 .transpose()?
5641 .unwrap_or_default(),
5642 };
5643 let elapsed = started.elapsed().as_nanos() as u64;
5644 crate::trace::QueryTrace::record(|trace| {
5645 match retriever {
5646 Retriever::Ann { .. } => {
5647 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
5648 if let Retriever::Ann { column_id, .. } = retriever {
5649 if let Some(index) = self.ann.get(column_id) {
5650 trace.ann_algorithm = Some(index.algorithm());
5651 trace.ann_quantization = Some(index.quantization());
5652 trace.ann_backend = Some(index.backend_name());
5653 }
5654 }
5655 }
5656 Retriever::Sparse { .. } => {
5657 trace.sparse_candidate_nanos =
5658 trace.sparse_candidate_nanos.saturating_add(elapsed)
5659 }
5660 Retriever::MinHash { .. } => {
5661 trace.minhash_candidate_nanos =
5662 trace.minhash_candidate_nanos.saturating_add(elapsed)
5663 }
5664 }
5665 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5666 });
5667 Ok(scored
5668 .into_iter()
5669 .enumerate()
5670 .map(|(rank, (row_id, score))| RetrieverHit {
5671 row_id,
5672 rank: rank + 1,
5673 score,
5674 })
5675 .collect())
5676 }
5677
5678 fn eligible_candidate_ids(
5679 &self,
5680 candidates: &[RowId],
5681 _column_id: u16,
5682 snapshot: Snapshot,
5683 context: Option<&crate::query::AiExecutionContext>,
5684 ) -> Result<std::collections::HashSet<RowId>> {
5685 if !self.had_deletes
5686 && self.ttl.is_none()
5687 && self.pending_put_cols.is_empty()
5688 && snapshot.epoch == self.snapshot().epoch
5689 {
5690 return Ok(candidates.iter().copied().collect());
5691 }
5692 let mut readers: Vec<_> = self
5693 .run_refs
5694 .iter()
5695 .map(|run| self.open_reader(run.run_id))
5696 .collect::<Result<_>>()?;
5697 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5698 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5699 for &row_id in candidates {
5700 if let Some(context) = context {
5701 context.consume(1)?;
5702 }
5703 let mem = self.memtable.get_version_at(row_id, snapshot);
5704 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
5705 let overlay = match (mem, mutable) {
5706 (Some(left), Some(right)) => Some(
5707 if Snapshot::version_is_newer(
5708 left.1.committed_epoch,
5709 left.1.commit_ts,
5710 right.1.committed_epoch,
5711 right.1.commit_ts,
5712 ) {
5713 left
5714 } else {
5715 right
5716 },
5717 ),
5718 (Some(value), None) | (None, Some(value)) => Some(value),
5719 (None, None) => None,
5720 };
5721 if let Some((_, row)) = overlay {
5722 if !row.deleted && !self.row_expired_at(&row, now) {
5723 eligible.insert(row_id);
5724 }
5725 continue;
5726 }
5727 let mut best: Option<(Epoch, bool, usize)> = None;
5728 for (index, reader) in readers.iter_mut().enumerate() {
5729 if let Some((epoch, deleted)) =
5730 reader.get_version_visibility(row_id, snapshot.epoch)?
5731 {
5732 if best
5733 .as_ref()
5734 .map(|(best_epoch, ..)| epoch > *best_epoch)
5735 .unwrap_or(true)
5736 {
5737 best = Some((epoch, deleted, index));
5738 }
5739 }
5740 }
5741 let Some((_, false, reader_index)) = best else {
5742 continue;
5743 };
5744 if let Some(ttl) = self.ttl {
5745 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5746 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5747 {
5748 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5749 continue;
5750 }
5751 }
5752 }
5753 eligible.insert(row_id);
5754 }
5755 Ok(eligible)
5756 }
5757
5758 fn eligible_and_authorized_candidate_ids(
5759 &self,
5760 candidates: &[RowId],
5761 column_id: u16,
5762 snapshot: Snapshot,
5763 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5764 context: Option<&crate::query::AiExecutionContext>,
5765 ) -> Result<std::collections::HashSet<RowId>> {
5766 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5767 let Some(authorization) = authorization else {
5768 return Ok(eligible);
5769 };
5770 let candidates: Vec<_> = eligible.into_iter().collect();
5771 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5772 }
5773
5774 fn policy_allowed_candidate_ids(
5775 &self,
5776 candidates: &[RowId],
5777 snapshot: Snapshot,
5778 authorization: &crate::security::CandidateAuthorization<'_>,
5779 context: Option<&crate::query::AiExecutionContext>,
5780 ) -> Result<std::collections::HashSet<RowId>> {
5781 let started = std::time::Instant::now();
5782 if candidates.is_empty()
5783 || authorization.principal.is_admin
5784 || !authorization.security.rls_enabled(authorization.table)
5785 {
5786 return Ok(candidates.iter().copied().collect());
5787 }
5788 if let Some(context) = context {
5789 context.checkpoint()?;
5790 }
5791 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5792 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5793 .iter()
5794 .map(|row_id| {
5795 (
5796 *row_id,
5797 Row {
5798 row_id: *row_id,
5799 committed_epoch: snapshot.epoch,
5800 commit_ts: None,
5801 columns: std::collections::HashMap::new(),
5802 deleted: false,
5803 },
5804 )
5805 })
5806 .collect();
5807 let columns = authorization
5808 .security
5809 .select_policy_columns(authorization.table, authorization.principal);
5810 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5811 let mut decoded = 0usize;
5812 for column_id in &columns {
5813 if let Some(context) = context {
5814 context.checkpoint()?;
5815 }
5816 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5817 &row_ids, *column_id, snapshot, query_now, context,
5818 )? {
5819 if let Some(row) = rows.get_mut(&row_id) {
5820 row.columns.insert(*column_id, value);
5821 decoded = decoded.saturating_add(1);
5822 }
5823 }
5824 }
5825 if let Some(context) = context {
5826 context.consume(candidates.len().saturating_add(decoded))?;
5827 }
5828 let allowed = rows
5829 .into_values()
5830 .filter_map(|row| {
5831 authorization
5832 .security
5833 .row_allowed(
5834 authorization.table,
5835 crate::security::PolicyCommand::Select,
5836 &row,
5837 authorization.principal,
5838 false,
5839 )
5840 .then_some(row.row_id)
5841 })
5842 .collect();
5843 crate::trace::QueryTrace::record(|trace| {
5844 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5845 trace.rls_policy_columns_decoded =
5846 trace.rls_policy_columns_decoded.saturating_add(decoded);
5847 trace.authorization_nanos = trace
5848 .authorization_nanos
5849 .saturating_add(started.elapsed().as_nanos() as u64);
5850 });
5851 Ok(allowed)
5852 }
5853
5854 pub fn search(
5856 &mut self,
5857 request: &crate::query::SearchRequest,
5858 ) -> Result<Vec<crate::query::SearchHit>> {
5859 self.search_with_allowed(request, None)
5860 }
5861
5862 pub fn search_at(
5863 &mut self,
5864 request: &crate::query::SearchRequest,
5865 snapshot: Snapshot,
5866 authorized: Option<&std::collections::HashSet<RowId>>,
5867 ) -> Result<Vec<crate::query::SearchHit>> {
5868 self.search_at_with_allowed(request, snapshot, authorized)
5869 }
5870
5871 pub fn search_with_allowed(
5872 &mut self,
5873 request: &crate::query::SearchRequest,
5874 authorized: Option<&std::collections::HashSet<RowId>>,
5875 ) -> Result<Vec<crate::query::SearchHit>> {
5876 self.search_at_with_allowed(request, self.snapshot(), authorized)
5877 }
5878
5879 pub fn search_at_with_allowed(
5880 &mut self,
5881 request: &crate::query::SearchRequest,
5882 snapshot: Snapshot,
5883 authorized: Option<&std::collections::HashSet<RowId>>,
5884 ) -> Result<Vec<crate::query::SearchHit>> {
5885 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5886 }
5887
5888 pub fn search_at_with_allowed_and_context(
5889 &mut self,
5890 request: &crate::query::SearchRequest,
5891 snapshot: Snapshot,
5892 authorized: Option<&std::collections::HashSet<RowId>>,
5893 context: Option<&crate::query::AiExecutionContext>,
5894 ) -> Result<Vec<crate::query::SearchHit>> {
5895 self.ensure_indexes_complete()?;
5896 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5897 }
5898
5899 pub fn search_at_with_candidate_authorization_and_context(
5900 &mut self,
5901 request: &crate::query::SearchRequest,
5902 snapshot: Snapshot,
5903 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5904 context: Option<&crate::query::AiExecutionContext>,
5905 ) -> Result<Vec<crate::query::SearchHit>> {
5906 self.ensure_indexes_complete()?;
5907 self.search_at_with_filters_and_context(
5908 request,
5909 snapshot,
5910 None,
5911 authorization,
5912 context,
5913 None,
5914 )
5915 }
5916
5917 #[doc(hidden)]
5918 pub fn search_at_with_candidate_authorization_on_generation(
5919 &self,
5920 request: &crate::query::SearchRequest,
5921 snapshot: Snapshot,
5922 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5923 context: Option<&crate::query::AiExecutionContext>,
5924 ) -> Result<Vec<crate::query::SearchHit>> {
5925 self.search_at_with_filters_and_context(
5926 request,
5927 snapshot,
5928 None,
5929 authorization,
5930 context,
5931 None,
5932 )
5933 }
5934
5935 #[doc(hidden)]
5936 pub fn search_at_with_candidate_authorization_on_generation_after(
5937 &self,
5938 request: &crate::query::SearchRequest,
5939 snapshot: Snapshot,
5940 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5941 context: Option<&crate::query::AiExecutionContext>,
5942 after: Option<crate::query::SearchAfter>,
5943 ) -> Result<Vec<crate::query::SearchHit>> {
5944 self.search_at_with_filters_and_context(
5945 request,
5946 snapshot,
5947 None,
5948 authorization,
5949 context,
5950 after,
5951 )
5952 }
5953
5954 fn search_at_with_filters_and_context(
5955 &self,
5956 request: &crate::query::SearchRequest,
5957 snapshot: Snapshot,
5958 authorized: Option<&std::collections::HashSet<RowId>>,
5959 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5960 context: Option<&crate::query::AiExecutionContext>,
5961 after: Option<crate::query::SearchAfter>,
5962 ) -> Result<Vec<crate::query::SearchHit>> {
5963 use crate::query::{
5964 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5965 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5966 };
5967 let total_started = std::time::Instant::now();
5968 let rank_offset = after.map_or(0, |after| after.returned_count);
5969 self.require_select()?;
5970 if request.limit == 0 {
5971 return Err(MongrelError::InvalidArgument(
5972 "search limit must be > 0".into(),
5973 ));
5974 }
5975 if request.limit > MAX_FINAL_LIMIT {
5976 return Err(MongrelError::InvalidArgument(format!(
5977 "search limit exceeds {MAX_FINAL_LIMIT}"
5978 )));
5979 }
5980 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5981 return Err(MongrelError::InvalidArgument(
5982 "search-after score must be finite".into(),
5983 ));
5984 }
5985 if request.retrievers.is_empty() {
5986 return Err(MongrelError::InvalidArgument(
5987 "search requires at least one retriever".into(),
5988 ));
5989 }
5990 if request.retrievers.len() > MAX_RETRIEVERS {
5991 return Err(MongrelError::InvalidArgument(format!(
5992 "search exceeds {MAX_RETRIEVERS} retrievers"
5993 )));
5994 }
5995 if request.must.len() > MAX_HARD_CONDITIONS {
5996 return Err(MongrelError::InvalidArgument(format!(
5997 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5998 )));
5999 }
6000 for condition in &request.must {
6001 self.validate_condition(condition)?;
6002 }
6003 if request.must.iter().any(|condition| {
6004 matches!(
6005 condition,
6006 Condition::Ann { .. }
6007 | Condition::SparseMatch { .. }
6008 | Condition::MinHashSimilar { .. }
6009 )
6010 }) {
6011 return Err(MongrelError::InvalidArgument(
6012 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
6013 .into(),
6014 ));
6015 }
6016 let mut names = std::collections::HashSet::new();
6017 for named in &request.retrievers {
6018 if named.name.is_empty()
6019 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
6020 || !names.insert(named.name.as_str())
6021 {
6022 return Err(MongrelError::InvalidArgument(format!(
6023 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
6024 crate::query::MAX_RETRIEVER_NAME_BYTES
6025 )));
6026 }
6027 if !named.weight.is_finite()
6028 || named.weight < 0.0
6029 || named.weight > MAX_RETRIEVER_WEIGHT
6030 {
6031 return Err(MongrelError::InvalidArgument(format!(
6032 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6033 )));
6034 }
6035 self.validate_retriever(&named.retriever)?;
6036 }
6037 let projection = request
6038 .projection
6039 .clone()
6040 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
6041 if projection.len() > MAX_PROJECTION_COLUMNS {
6042 return Err(MongrelError::InvalidArgument(format!(
6043 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
6044 )));
6045 }
6046 for column_id in &projection {
6047 if !self
6048 .schema
6049 .columns
6050 .iter()
6051 .any(|column| column.id == *column_id)
6052 {
6053 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
6054 }
6055 }
6056 if let Some(crate::query::Rerank::ExactVector {
6057 embedding_column,
6058 query,
6059 candidate_limit,
6060 weight,
6061 ..
6062 }) = &request.rerank
6063 {
6064 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
6065 {
6066 return Err(MongrelError::InvalidArgument(format!(
6067 "rerank candidate_limit must be between search limit and {}",
6068 crate::query::MAX_RETRIEVER_K
6069 )));
6070 }
6071 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
6072 return Err(MongrelError::InvalidArgument(format!(
6073 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
6074 )));
6075 }
6076 let column = self
6077 .schema
6078 .columns
6079 .iter()
6080 .find(|column| column.id == *embedding_column)
6081 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
6082 let crate::schema::TypeId::Embedding { dim } = column.ty else {
6083 return Err(MongrelError::InvalidArgument(format!(
6084 "rerank column {embedding_column} is not an embedding"
6085 )));
6086 };
6087 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
6088 return Err(MongrelError::InvalidArgument(format!(
6089 "rerank query must contain {dim} finite values"
6090 )));
6091 }
6092 }
6093
6094 let hard_filter_started = std::time::Instant::now();
6095 let hard_filter = if request.must.is_empty() {
6096 None
6097 } else {
6098 let mut sets = Vec::with_capacity(request.must.len());
6099 for condition in &request.must {
6100 if let Some(context) = context {
6101 context.checkpoint()?;
6102 }
6103 sets.push(self.resolve_condition(condition, snapshot)?);
6104 }
6105 Some(RowIdSet::intersect_many(sets))
6106 };
6107 crate::trace::QueryTrace::record(|trace| {
6108 trace.hard_filter_nanos = trace
6109 .hard_filter_nanos
6110 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6111 });
6112 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6113 return Ok(Vec::new());
6114 }
6115
6116 let constant = match request.fusion {
6117 Fusion::ReciprocalRank { constant } => constant,
6118 };
6119 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6120 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6121 let mut fusion_nanos = 0u64;
6122 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6123 std::collections::HashMap::new();
6124 for named in retrievers {
6125 if named.weight == 0.0 {
6126 continue;
6127 }
6128 if let Some(context) = context {
6129 context.checkpoint()?;
6130 }
6131 let hits = self.retrieve_filtered(
6132 &named.retriever,
6133 snapshot,
6134 hard_filter.as_ref(),
6135 authorized,
6136 candidate_authorization,
6137 context,
6138 )?;
6139 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6140 let fusion_started = std::time::Instant::now();
6141 for hit in hits {
6142 if let Some(context) = context {
6143 context.consume(1)?;
6144 }
6145 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6146 if !contribution.is_finite() {
6147 return Err(MongrelError::InvalidArgument(
6148 "retriever contribution must be finite".into(),
6149 ));
6150 }
6151 let max_fused_candidates = context.map_or(
6152 crate::query::MAX_FUSED_CANDIDATES,
6153 crate::query::AiExecutionContext::max_fused_candidates,
6154 );
6155 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6156 return Err(MongrelError::WorkBudgetExceeded);
6157 }
6158 let entry = fused.entry(hit.row_id).or_default();
6159 entry.0 += contribution;
6160 if !entry.0.is_finite() {
6161 return Err(MongrelError::InvalidArgument(
6162 "fused score must be finite".into(),
6163 ));
6164 }
6165 entry.1.push(ComponentScore {
6166 retriever_name: retriever_name.clone(),
6167 rank: hit.rank,
6168 raw_score: hit.score,
6169 contribution,
6170 });
6171 }
6172 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6173 }
6174 let union_size = fused.len();
6175 let mut ranked: Vec<_> = fused
6176 .into_iter()
6177 .map(|(row_id, (fused_score, components))| {
6178 (row_id, fused_score, components, None, fused_score)
6179 })
6180 .collect();
6181 let order = |(a_row, _, _, _, a_score): &(
6182 RowId,
6183 f64,
6184 Vec<ComponentScore>,
6185 Option<f32>,
6186 f64,
6187 ),
6188 (b_row, _, _, _, b_score): &(
6189 RowId,
6190 f64,
6191 Vec<ComponentScore>,
6192 Option<f32>,
6193 f64,
6194 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6195 if let Some(crate::query::Rerank::ExactVector {
6196 embedding_column,
6197 query,
6198 metric,
6199 candidate_limit,
6200 weight,
6201 }) = &request.rerank
6202 {
6203 let fused_order = |(a_row, a_score, ..): &(
6204 RowId,
6205 f64,
6206 Vec<ComponentScore>,
6207 Option<f32>,
6208 f64,
6209 ),
6210 (b_row, b_score, ..): &(
6211 RowId,
6212 f64,
6213 Vec<ComponentScore>,
6214 Option<f32>,
6215 f64,
6216 )| {
6217 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6218 };
6219 let selection_started = std::time::Instant::now();
6220 if let Some(context) = context {
6221 context.consume(ranked.len())?;
6222 }
6223 if ranked.len() > *candidate_limit {
6224 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6225 ranked.truncate(*candidate_limit);
6226 }
6227 ranked.sort_by(fused_order);
6228 fusion_nanos =
6229 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6230 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6231 if let Some(context) = context {
6232 context.consume(row_ids.len())?;
6233 }
6234 let query_now =
6235 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6236 let gather_started = std::time::Instant::now();
6237 let vectors = self.values_for_rids_batch_at_with_context(
6238 &row_ids,
6239 *embedding_column,
6240 snapshot,
6241 query_now,
6242 context,
6243 )?;
6244 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6245 let vector_work =
6246 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6247 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6248 if let Some(context) = context {
6249 context.consume(vector_work)?;
6250 }
6251 query
6252 .iter()
6253 .map(|value| f64::from(*value).powi(2))
6254 .sum::<f64>()
6255 .sqrt()
6256 } else {
6257 0.0
6258 };
6259 let score_started = std::time::Instant::now();
6260 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6261 for (row_id, value) in vectors {
6262 if let Some(meta) = value.generated_embedding_metadata() {
6263 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6264 continue;
6265 }
6266 }
6267 let Some(vector) = value.as_embedding() else {
6268 continue;
6269 };
6270 let score = match metric {
6271 crate::query::VectorMetric::DotProduct => {
6272 if let Some(context) = context {
6273 context.consume(vector_work)?;
6274 }
6275 query
6276 .iter()
6277 .zip(vector)
6278 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6279 .sum::<f64>()
6280 }
6281 crate::query::VectorMetric::Cosine => {
6282 if let Some(context) = context {
6283 context.consume(vector_work.saturating_mul(2))?;
6284 }
6285 let dot = query
6286 .iter()
6287 .zip(vector)
6288 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6289 .sum::<f64>();
6290 let norm = vector
6291 .iter()
6292 .map(|value| f64::from(*value).powi(2))
6293 .sum::<f64>()
6294 .sqrt();
6295 if query_norm == 0.0 || norm == 0.0 {
6296 0.0
6297 } else {
6298 dot / (query_norm * norm)
6299 }
6300 }
6301 crate::query::VectorMetric::Euclidean => {
6302 if let Some(context) = context {
6303 context.consume(vector_work)?;
6304 }
6305 query
6306 .iter()
6307 .zip(vector)
6308 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6309 .sum::<f64>()
6310 .sqrt()
6311 }
6312 };
6313 if !score.is_finite() {
6314 return Err(MongrelError::InvalidArgument(
6315 "exact rerank score must be finite".into(),
6316 ));
6317 }
6318 scores.insert(row_id, score as f32);
6319 }
6320 let mut reranked = Vec::with_capacity(ranked.len());
6321 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6322 let Some(score) = scores.get(&row_id).copied() else {
6323 continue;
6324 };
6325 let ordering_score = match metric {
6326 crate::query::VectorMetric::Euclidean => -f64::from(score),
6327 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6328 f64::from(score)
6329 }
6330 };
6331 let final_score = fused_score + *weight * ordering_score;
6332 if !final_score.is_finite() {
6333 return Err(MongrelError::InvalidArgument(
6334 "final rerank score must be finite".into(),
6335 ));
6336 }
6337 reranked.push((row_id, fused_score, components, Some(score), final_score));
6338 }
6339 ranked = reranked;
6340 ranked.sort_by(order);
6341 crate::trace::QueryTrace::record(|trace| {
6342 trace.exact_vector_gather_nanos =
6343 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6344 trace.exact_vector_score_nanos = trace
6345 .exact_vector_score_nanos
6346 .saturating_add(score_started.elapsed().as_nanos() as u64);
6347 });
6348 }
6349 if let Some(after) = after {
6350 ranked.retain(|(row_id, _, _, _, final_score)| {
6351 final_score.total_cmp(&after.final_score).is_lt()
6352 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6353 });
6354 }
6355 let projection_started = std::time::Instant::now();
6356 let sentinel = projection
6357 .first()
6358 .copied()
6359 .or_else(|| self.schema.columns.first().map(|column| column.id));
6360 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6361 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6362 let mut projection_rows = 0usize;
6363 let mut projection_cells = 0usize;
6364 while out.len() < request.limit && !ranked.is_empty() {
6365 if let Some(context) = context {
6366 context.checkpoint()?;
6367 context.consume(ranked.len())?;
6368 }
6369 let needed = request.limit - out.len();
6370 let window_size = ranked
6371 .len()
6372 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6373 let selection_started = std::time::Instant::now();
6374 let mut remainder = if ranked.len() > window_size {
6375 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6376 ranked.split_off(window_size)
6377 } else {
6378 Vec::new()
6379 };
6380 ranked.sort_by(order);
6381 fusion_nanos =
6382 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6383 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6384 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6385 if let Some(context) = context {
6386 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6387 }
6388 projection_rows = projection_rows.saturating_add(row_ids.len());
6389 projection_cells =
6390 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6391 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6392 std::collections::HashMap::new();
6393 if let Some(column_id) = sentinel {
6394 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6395 &row_ids, column_id, snapshot, query_now, context,
6396 )? {
6397 cells.entry(row_id).or_default().insert(column_id, value);
6398 }
6399 }
6400 for &column_id in &projection {
6401 if Some(column_id) == sentinel {
6402 continue;
6403 }
6404 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6405 &row_ids, column_id, snapshot, query_now, context,
6406 )? {
6407 cells.entry(row_id).or_default().insert(column_id, value);
6408 }
6409 }
6410 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6411 ranked.drain(..)
6412 {
6413 let Some(row_cells) = cells.remove(&row_id) else {
6414 continue;
6415 };
6416 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6417 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6418 out.push(SearchHit {
6419 row_id,
6420 cells: projection
6421 .iter()
6422 .filter_map(|column_id| {
6423 row_cells
6424 .get(column_id)
6425 .cloned()
6426 .map(|value| (*column_id, value))
6427 })
6428 .collect(),
6429 components,
6430 fused_score,
6431 exact_rerank_score,
6432 final_score,
6433 final_rank,
6434 });
6435 if out.len() == request.limit {
6436 break;
6437 }
6438 }
6439 ranked.append(&mut remainder);
6440 }
6441 crate::trace::QueryTrace::record(|trace| {
6442 trace.union_size = union_size;
6443 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6444 trace.projection_nanos = trace
6445 .projection_nanos
6446 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6447 trace.total_nanos = trace
6448 .total_nanos
6449 .saturating_add(total_started.elapsed().as_nanos() as u64);
6450 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6451 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6452 if let Some(context) = context {
6453 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6454 }
6455 });
6456 Ok(out)
6457 }
6458
6459 pub fn set_similarity(
6462 &mut self,
6463 request: &crate::query::SetSimilarityRequest,
6464 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6465 self.set_similarity_with_allowed(request, None)
6466 }
6467
6468 pub fn set_similarity_at(
6469 &mut self,
6470 request: &crate::query::SetSimilarityRequest,
6471 snapshot: Snapshot,
6472 allowed: Option<&std::collections::HashSet<RowId>>,
6473 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6474 self.set_similarity_explained_at(request, snapshot, allowed)
6475 .map(|(hits, _)| hits)
6476 }
6477
6478 pub fn ann_rerank(
6480 &mut self,
6481 request: &crate::query::AnnRerankRequest,
6482 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6483 self.ann_rerank_with_allowed(request, None)
6484 }
6485
6486 pub fn ann_rerank_with_allowed(
6487 &mut self,
6488 request: &crate::query::AnnRerankRequest,
6489 allowed: Option<&std::collections::HashSet<RowId>>,
6490 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6491 self.ann_rerank_at(request, self.snapshot(), allowed)
6492 }
6493
6494 pub fn ann_rerank_at(
6495 &mut self,
6496 request: &crate::query::AnnRerankRequest,
6497 snapshot: Snapshot,
6498 allowed: Option<&std::collections::HashSet<RowId>>,
6499 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6500 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6501 }
6502
6503 pub fn ann_rerank_at_with_context(
6504 &mut self,
6505 request: &crate::query::AnnRerankRequest,
6506 snapshot: Snapshot,
6507 allowed: Option<&std::collections::HashSet<RowId>>,
6508 context: Option<&crate::query::AiExecutionContext>,
6509 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6510 self.ensure_indexes_complete()?;
6511 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6512 }
6513
6514 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6515 &mut self,
6516 request: &crate::query::AnnRerankRequest,
6517 snapshot: Snapshot,
6518 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6519 context: Option<&crate::query::AiExecutionContext>,
6520 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6521 self.ensure_indexes_complete()?;
6522 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6523 }
6524
6525 #[doc(hidden)]
6526 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6527 &self,
6528 request: &crate::query::AnnRerankRequest,
6529 snapshot: Snapshot,
6530 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6531 context: Option<&crate::query::AiExecutionContext>,
6532 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6533 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6534 }
6535
6536 fn ann_rerank_at_with_filters_and_context(
6537 &self,
6538 request: &crate::query::AnnRerankRequest,
6539 snapshot: Snapshot,
6540 allowed: Option<&std::collections::HashSet<RowId>>,
6541 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6542 context: Option<&crate::query::AiExecutionContext>,
6543 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6544 use crate::query::{
6545 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6546 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6547 };
6548 if request.candidate_k == 0 || request.limit == 0 {
6549 return Err(MongrelError::InvalidArgument(
6550 "candidate_k and limit must be > 0".into(),
6551 ));
6552 }
6553 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6554 return Err(MongrelError::InvalidArgument(format!(
6555 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6556 )));
6557 }
6558 let retriever = Retriever::Ann {
6559 column_id: request.column_id,
6560 query: request.query.clone(),
6561 k: request.candidate_k,
6562 };
6563 self.require_select()?;
6564 self.validate_retriever(&retriever)?;
6565 let hits = self.retrieve_filtered(
6566 &retriever,
6567 snapshot,
6568 None,
6569 allowed,
6570 candidate_authorization,
6571 context,
6572 )?;
6573 let distances: std::collections::HashMap<_, _> = hits
6574 .iter()
6575 .filter_map(|hit| match hit.score {
6576 RetrieverScore::AnnHammingDistance(distance) => {
6577 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
6578 }
6579 RetrieverScore::AnnCosineDistance(distance) => {
6580 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
6581 }
6582 _ => None,
6583 })
6584 .collect();
6585 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6586 if let Some(context) = context {
6587 context.consume(row_ids.len())?;
6588 }
6589 let gather_started = std::time::Instant::now();
6590 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6591 let values = self.values_for_rids_batch_at_with_context(
6592 &row_ids,
6593 request.column_id,
6594 snapshot,
6595 query_now,
6596 context,
6597 )?;
6598 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6599 let score_started = std::time::Instant::now();
6600 let vector_work =
6601 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6602 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6603 if let Some(context) = context {
6604 context.consume(vector_work)?;
6605 }
6606 request
6607 .query
6608 .iter()
6609 .map(|value| f64::from(*value).powi(2))
6610 .sum::<f64>()
6611 .sqrt()
6612 } else {
6613 0.0
6614 };
6615 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6616 for (row_id, value) in values {
6617 if let Some(meta) = value.generated_embedding_metadata() {
6618 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
6619 continue;
6620 }
6621 }
6622 let Some(vector) = value.as_embedding() else {
6623 continue;
6624 };
6625 let exact_score = match request.metric {
6626 VectorMetric::DotProduct => {
6627 if let Some(context) = context {
6628 context.consume(vector_work)?;
6629 }
6630 request
6631 .query
6632 .iter()
6633 .zip(vector)
6634 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6635 .sum::<f64>()
6636 }
6637 VectorMetric::Cosine => {
6638 if let Some(context) = context {
6639 context.consume(vector_work.saturating_mul(2))?;
6640 }
6641 let dot = request
6642 .query
6643 .iter()
6644 .zip(vector)
6645 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6646 .sum::<f64>();
6647 let norm = vector
6648 .iter()
6649 .map(|value| f64::from(*value).powi(2))
6650 .sum::<f64>()
6651 .sqrt();
6652 if query_norm == 0.0 || norm == 0.0 {
6653 0.0
6654 } else {
6655 dot / (query_norm * norm)
6656 }
6657 }
6658 VectorMetric::Euclidean => {
6659 if let Some(context) = context {
6660 context.consume(vector_work)?;
6661 }
6662 request
6663 .query
6664 .iter()
6665 .zip(vector)
6666 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6667 .sum::<f64>()
6668 .sqrt()
6669 }
6670 };
6671 let exact_score = exact_score as f32;
6672 if !exact_score.is_finite() {
6673 return Err(MongrelError::InvalidArgument(
6674 "exact ANN score must be finite".into(),
6675 ));
6676 }
6677 let Some(candidate_distance) = distances.get(&row_id).copied() else {
6678 continue;
6679 };
6680 reranked.push(AnnRerankHit {
6681 row_id,
6682 candidate_distance,
6683 exact_score,
6684 });
6685 }
6686 reranked.sort_by(|left, right| {
6687 let score = match request.metric {
6688 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6689 VectorMetric::Cosine | VectorMetric::DotProduct => {
6690 right.exact_score.total_cmp(&left.exact_score)
6691 }
6692 };
6693 score.then_with(|| left.row_id.cmp(&right.row_id))
6694 });
6695 reranked.truncate(request.limit);
6696 crate::trace::QueryTrace::record(|trace| {
6697 trace.exact_vector_gather_nanos =
6698 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6699 trace.exact_vector_score_nanos = trace
6700 .exact_vector_score_nanos
6701 .saturating_add(score_started.elapsed().as_nanos() as u64);
6702 });
6703 Ok(reranked)
6704 }
6705
6706 pub fn set_similarity_with_allowed(
6707 &mut self,
6708 request: &crate::query::SetSimilarityRequest,
6709 allowed: Option<&std::collections::HashSet<RowId>>,
6710 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6711 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6712 .map(|(hits, _)| hits)
6713 }
6714
6715 pub fn set_similarity_explained(
6716 &mut self,
6717 request: &crate::query::SetSimilarityRequest,
6718 ) -> Result<(
6719 Vec<crate::query::SetSimilarityHit>,
6720 crate::query::SetSimilarityTrace,
6721 )> {
6722 self.set_similarity_explained_at(request, self.snapshot(), None)
6723 }
6724
6725 fn set_similarity_explained_at(
6726 &mut self,
6727 request: &crate::query::SetSimilarityRequest,
6728 snapshot: Snapshot,
6729 allowed: Option<&std::collections::HashSet<RowId>>,
6730 ) -> Result<(
6731 Vec<crate::query::SetSimilarityHit>,
6732 crate::query::SetSimilarityTrace,
6733 )> {
6734 self.ensure_indexes_complete()?;
6735 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6736 }
6737
6738 pub fn set_similarity_at_with_context(
6739 &mut self,
6740 request: &crate::query::SetSimilarityRequest,
6741 snapshot: Snapshot,
6742 allowed: Option<&std::collections::HashSet<RowId>>,
6743 context: Option<&crate::query::AiExecutionContext>,
6744 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6745 self.ensure_indexes_complete()?;
6746 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6747 .map(|(hits, _)| hits)
6748 }
6749
6750 pub fn set_similarity_at_with_candidate_authorization_and_context(
6751 &mut self,
6752 request: &crate::query::SetSimilarityRequest,
6753 snapshot: Snapshot,
6754 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6755 context: Option<&crate::query::AiExecutionContext>,
6756 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6757 self.ensure_indexes_complete()?;
6758 self.set_similarity_explained_at_with_context(
6759 request,
6760 snapshot,
6761 None,
6762 authorization,
6763 context,
6764 )
6765 .map(|(hits, _)| hits)
6766 }
6767
6768 #[doc(hidden)]
6769 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6770 &self,
6771 request: &crate::query::SetSimilarityRequest,
6772 snapshot: Snapshot,
6773 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6774 context: Option<&crate::query::AiExecutionContext>,
6775 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6776 self.set_similarity_explained_at_with_context(
6777 request,
6778 snapshot,
6779 None,
6780 authorization,
6781 context,
6782 )
6783 .map(|(hits, _)| hits)
6784 }
6785
6786 fn set_similarity_explained_at_with_context(
6787 &self,
6788 request: &crate::query::SetSimilarityRequest,
6789 snapshot: Snapshot,
6790 allowed: Option<&std::collections::HashSet<RowId>>,
6791 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6792 context: Option<&crate::query::AiExecutionContext>,
6793 ) -> Result<(
6794 Vec<crate::query::SetSimilarityHit>,
6795 crate::query::SetSimilarityTrace,
6796 )> {
6797 use crate::query::{
6798 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6799 MAX_SET_MEMBERS,
6800 };
6801 let mut trace = crate::query::SetSimilarityTrace::default();
6802 if request.members.is_empty() {
6803 return Ok((Vec::new(), trace));
6804 }
6805 if request.candidate_k == 0 || request.limit == 0 {
6806 return Err(MongrelError::InvalidArgument(
6807 "candidate_k and limit must be > 0".into(),
6808 ));
6809 }
6810 if request.candidate_k > MAX_RETRIEVER_K
6811 || request.limit > MAX_FINAL_LIMIT
6812 || request.members.len() > MAX_SET_MEMBERS
6813 {
6814 return Err(MongrelError::InvalidArgument(format!(
6815 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6816 )));
6817 }
6818 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6819 return Err(MongrelError::InvalidArgument(
6820 "min_jaccard must be finite and between 0 and 1".into(),
6821 ));
6822 }
6823 let started = std::time::Instant::now();
6824 let retriever = Retriever::MinHash {
6825 column_id: request.column_id,
6826 members: request.members.clone(),
6827 k: request.candidate_k,
6828 };
6829 self.require_select()?;
6830 self.validate_retriever(&retriever)?;
6831 let hits = self.retrieve_filtered(
6832 &retriever,
6833 snapshot,
6834 None,
6835 allowed,
6836 candidate_authorization,
6837 context,
6838 )?;
6839 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6840 trace.candidate_count = hits.len();
6841 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6842 if let Some(context) = context {
6843 context.consume(row_ids.len())?;
6844 }
6845 let started = std::time::Instant::now();
6846 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6847 let values = self.values_for_rids_batch_at_with_context(
6848 &row_ids,
6849 request.column_id,
6850 snapshot,
6851 query_now,
6852 context,
6853 )?;
6854 trace.gather_us = started.elapsed().as_micros() as u64;
6855 if let Some(context) = context {
6856 context.consume(request.members.len())?;
6857 }
6858 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6859 let estimates: std::collections::HashMap<_, _> = hits
6860 .into_iter()
6861 .filter_map(|hit| match hit.score {
6862 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6863 _ => None,
6864 })
6865 .collect();
6866 let started = std::time::Instant::now();
6867 let mut parsed = Vec::with_capacity(values.len());
6868 for (row_id, value) in values {
6869 let Value::Bytes(bytes) = value else {
6870 continue;
6871 };
6872 if let Some(context) = context {
6873 context.consume(crate::query::work_units(
6874 bytes.len(),
6875 crate::query::PARSE_WORK_QUANTUM,
6876 ))?;
6877 }
6878 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6879 continue;
6880 };
6881 if let Some(context) = context {
6882 context.consume(members.len())?;
6883 }
6884 let stored = members
6885 .into_iter()
6886 .filter_map(|member| match member {
6887 serde_json::Value::String(value) => {
6888 Some(crate::query::SetMember::String(value))
6889 }
6890 serde_json::Value::Number(value) => {
6891 Some(crate::query::SetMember::Number(value))
6892 }
6893 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6894 _ => None,
6895 })
6896 .collect::<std::collections::HashSet<_>>();
6897 parsed.push((row_id, stored));
6898 }
6899 trace.parse_us = started.elapsed().as_micros() as u64;
6900 trace.verified_count = parsed.len();
6901 let started = std::time::Instant::now();
6902 let mut exact = Vec::new();
6903 for (row_id, stored) in parsed {
6904 if let Some(context) = context {
6905 context.consume(query.len().saturating_add(stored.len()))?;
6906 }
6907 let union = query.union(&stored).count();
6908 let score = if union == 0 {
6909 1.0
6910 } else {
6911 query.intersection(&stored).count() as f32 / union as f32
6912 };
6913 if score >= request.min_jaccard {
6914 exact.push(SetSimilarityHit {
6915 row_id,
6916 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6917 exact_jaccard: score,
6918 });
6919 }
6920 }
6921 exact.sort_by(|a, b| {
6922 b.exact_jaccard
6923 .total_cmp(&a.exact_jaccard)
6924 .then_with(|| a.row_id.cmp(&b.row_id))
6925 });
6926 exact.truncate(request.limit);
6927 trace.score_us = started.elapsed().as_micros() as u64;
6928 crate::trace::QueryTrace::record(|query_trace| {
6929 query_trace.exact_set_gather_nanos = query_trace
6930 .exact_set_gather_nanos
6931 .saturating_add(trace.gather_us.saturating_mul(1_000));
6932 query_trace.exact_set_parse_nanos = query_trace
6933 .exact_set_parse_nanos
6934 .saturating_add(trace.parse_us.saturating_mul(1_000));
6935 query_trace.exact_set_score_nanos = query_trace
6936 .exact_set_score_nanos
6937 .saturating_add(trace.score_us.saturating_mul(1_000));
6938 });
6939 Ok((exact, trace))
6940 }
6941
6942 fn values_for_rids_batch_at(
6944 &self,
6945 row_ids: &[u64],
6946 column_id: u16,
6947 snapshot: Snapshot,
6948 now: i64,
6949 ) -> Result<Vec<(RowId, Value)>> {
6950 if self.ttl.is_none()
6951 && self.memtable.is_empty()
6952 && self.mutable_run.is_empty()
6953 && self.run_refs.len() == 1
6954 {
6955 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6956 if row_ids.len().saturating_mul(24) < reader.row_count() {
6961 let mut values = Vec::with_capacity(row_ids.len());
6962 for &raw_row_id in row_ids {
6963 let row_id = RowId(raw_row_id);
6964 if let Some((_, false, Some(value))) =
6965 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6966 {
6967 values.push((row_id, value));
6968 }
6969 }
6970 return Ok(values);
6971 }
6972 let (positions, visible_row_ids) =
6973 reader.visible_positions_with_rids(snapshot.epoch)?;
6974 let requested: Vec<(RowId, usize)> = row_ids
6975 .iter()
6976 .filter_map(|raw| {
6977 visible_row_ids
6978 .binary_search(&(*raw as i64))
6979 .ok()
6980 .map(|index| (RowId(*raw), positions[index]))
6981 })
6982 .collect();
6983 let values = reader.gather_column(
6984 column_id,
6985 &requested
6986 .iter()
6987 .map(|(_, position)| *position)
6988 .collect::<Vec<_>>(),
6989 )?;
6990 return Ok(requested
6991 .into_iter()
6992 .zip(values)
6993 .map(|((row_id, _), value)| (row_id, value))
6994 .collect());
6995 }
6996 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6997 }
6998
6999 fn values_for_rids_batch_at_with_context(
7000 &self,
7001 row_ids: &[u64],
7002 column_id: u16,
7003 snapshot: Snapshot,
7004 now: i64,
7005 context: Option<&crate::query::AiExecutionContext>,
7006 ) -> Result<Vec<(RowId, Value)>> {
7007 let Some(context) = context else {
7008 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
7009 };
7010 let mut values = Vec::with_capacity(row_ids.len());
7011 for chunk in row_ids.chunks(256) {
7012 context.checkpoint()?;
7013 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
7014 }
7015 Ok(values)
7016 }
7017
7018 fn values_for_rids_at(
7020 &self,
7021 row_ids: &[u64],
7022 column_id: u16,
7023 snapshot: Snapshot,
7024 now: i64,
7025 ) -> Result<Vec<(RowId, Value)>> {
7026 let mut readers: Vec<_> = self
7027 .run_refs
7028 .iter()
7029 .map(|run| self.open_reader(run.run_id))
7030 .collect::<Result<_>>()?;
7031 let mut out = Vec::with_capacity(row_ids.len());
7032 for &raw_row_id in row_ids {
7033 let row_id = RowId(raw_row_id);
7034 let mem = self.memtable.get_version_at(row_id, snapshot);
7035 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
7036 let overlay = match (mem, mutable) {
7037 (Some((_, a)), Some((_, b))) => Some(
7038 if Snapshot::version_is_newer(
7039 a.committed_epoch,
7040 a.commit_ts,
7041 b.committed_epoch,
7042 b.commit_ts,
7043 ) {
7044 a
7045 } else {
7046 b
7047 },
7048 ),
7049 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
7050 (None, None) => None,
7051 };
7052 if let Some(row) = overlay {
7053 if !row.deleted && !self.row_expired_at(&row, now) {
7054 if let Some(value) = row.columns.get(&column_id) {
7055 out.push((row_id, value.clone()));
7056 }
7057 }
7058 continue;
7059 }
7060
7061 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
7062 for (index, reader) in readers.iter_mut().enumerate() {
7063 if let Some((epoch, deleted, value)) =
7064 reader.get_version_column(row_id, snapshot.epoch, column_id)?
7065 {
7066 if best
7067 .as_ref()
7068 .map(|(best_epoch, ..)| epoch > *best_epoch)
7069 .unwrap_or(true)
7070 {
7071 best = Some((epoch, deleted, value, index));
7072 }
7073 }
7074 }
7075 let Some((_, false, Some(value), reader_index)) = best else {
7076 continue;
7077 };
7078 if let Some(ttl) = self.ttl {
7079 if ttl.column_id != column_id {
7080 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7081 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
7082 {
7083 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7084 continue;
7085 }
7086 }
7087 } else if let Value::Int64(timestamp) = value {
7088 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7089 continue;
7090 }
7091 }
7092 }
7093 out.push((row_id, value));
7094 }
7095 Ok(out)
7096 }
7097
7098 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
7103 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
7104 }
7105
7106 pub fn rows_for_rids_with_context(
7107 &self,
7108 rids: &[u64],
7109 snapshot: Snapshot,
7110 context: &crate::query::AiExecutionContext,
7111 ) -> Result<Vec<Row>> {
7112 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
7113 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
7114 }
7115
7116 fn rows_for_rids_at_time(
7117 &self,
7118 rids: &[u64],
7119 snapshot: Snapshot,
7120 ttl_now: i64,
7121 control: Option<&crate::ExecutionControl>,
7122 ) -> Result<Vec<Row>> {
7123 use std::collections::HashMap;
7124 let mut rows = Vec::with_capacity(rids.len());
7125 let tier_size = self.memtable.len() + self.mutable_run.len();
7140 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7141 if rids.len().saturating_mul(24) < tier_size {
7142 for &rid in rids {
7143 if overlay.len() & 255 == 0 {
7144 control
7145 .map(crate::ExecutionControl::checkpoint)
7146 .transpose()?;
7147 }
7148 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
7149 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
7150 let newest = match (mem, mrun) {
7151 (Some((_, mr)), Some((_, rr))) => Some(
7152 if Snapshot::version_is_newer(
7153 mr.committed_epoch,
7154 mr.commit_ts,
7155 rr.committed_epoch,
7156 rr.commit_ts,
7157 ) {
7158 mr
7159 } else {
7160 rr
7161 },
7162 ),
7163 (Some((_, mr)), None) => Some(mr),
7164 (None, Some((_, rr))) => Some(rr),
7165 (None, None) => None,
7166 };
7167 if let Some(row) = newest {
7168 overlay.insert(rid, row);
7169 }
7170 }
7171 } else {
7172 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7173 overlay
7174 .entry(row.row_id.0)
7175 .and_modify(|e| {
7176 if Snapshot::version_is_newer(
7177 row.committed_epoch,
7178 row.commit_ts,
7179 e.committed_epoch,
7180 e.commit_ts,
7181 ) {
7182 *e = row.clone();
7183 }
7184 })
7185 .or_insert(row);
7186 };
7187 for (index, row) in self
7188 .memtable
7189 .visible_versions_at(snapshot)
7190 .into_iter()
7191 .enumerate()
7192 {
7193 if index & 255 == 0 {
7194 control
7195 .map(crate::ExecutionControl::checkpoint)
7196 .transpose()?;
7197 }
7198 fold_newest(row, &mut overlay);
7199 }
7200 for (index, row) in self
7201 .mutable_run
7202 .visible_versions_at(snapshot)
7203 .into_iter()
7204 .enumerate()
7205 {
7206 if index & 255 == 0 {
7207 control
7208 .map(crate::ExecutionControl::checkpoint)
7209 .transpose()?;
7210 }
7211 fold_newest(row, &mut overlay);
7212 }
7213 }
7214 if self.run_refs.len() == 1 {
7215 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7216 if rids.len().saturating_mul(24) < reader.row_count() {
7224 for (index, &rid) in rids.iter().enumerate() {
7225 if index & 255 == 0 {
7226 control
7227 .map(crate::ExecutionControl::checkpoint)
7228 .transpose()?;
7229 }
7230 if let Some(r) = overlay.get(&rid) {
7231 if !r.deleted {
7232 rows.push(r.clone());
7233 }
7234 continue;
7235 }
7236 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7237 if !row.deleted {
7238 rows.push(row);
7239 }
7240 }
7241 }
7242 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7243 return Ok(rows);
7244 }
7245 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7254 enum Src {
7257 Overlay,
7258 Run,
7259 }
7260 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7261 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7262 for (index, rid) in rids.iter().enumerate() {
7263 if index & 255 == 0 {
7264 control
7265 .map(crate::ExecutionControl::checkpoint)
7266 .transpose()?;
7267 }
7268 if overlay.contains_key(rid) {
7269 plan.push(Src::Overlay);
7270 continue;
7271 }
7272 match vis_rids.binary_search(&(*rid as i64)) {
7273 Ok(i) => {
7274 plan.push(Src::Run);
7275 fetch.push(positions[i]);
7276 }
7277 Err(_) => { }
7278 }
7279 }
7280 let fetched = reader.materialize_batch(&fetch)?;
7281 let mut fetched_iter = fetched.into_iter();
7282 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7283 if index & 255 == 0 {
7284 control
7285 .map(crate::ExecutionControl::checkpoint)
7286 .transpose()?;
7287 }
7288 match src {
7289 Src::Overlay => {
7290 if let Some(r) = overlay.get(rid) {
7291 if !r.deleted {
7292 rows.push(r.clone());
7293 }
7294 }
7295 }
7296 Src::Run => {
7297 if let Some(row) = fetched_iter.next() {
7298 if !row.deleted {
7299 rows.push(row);
7300 }
7301 }
7302 }
7303 }
7304 }
7305 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7306 return Ok(rows);
7307 }
7308 let mut readers: Vec<_> = self
7312 .run_refs
7313 .iter()
7314 .map(|rr| self.open_reader(rr.run_id))
7315 .collect::<Result<Vec<_>>>()?;
7316 for (index, rid) in rids.iter().enumerate() {
7317 if index & 255 == 0 {
7318 control
7319 .map(crate::ExecutionControl::checkpoint)
7320 .transpose()?;
7321 }
7322 if let Some(r) = overlay.get(rid) {
7323 if !r.deleted {
7324 rows.push(r.clone());
7325 }
7326 continue;
7327 }
7328 let mut best: Option<(Epoch, Row)> = None;
7329 for reader in readers.iter_mut() {
7330 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7331 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7332 best = Some((epoch, row));
7333 }
7334 }
7335 }
7336 if let Some((_, r)) = best {
7337 if !r.deleted {
7338 rows.push(r);
7339 }
7340 }
7341 }
7342 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7343 Ok(rows)
7344 }
7345
7346 pub fn indexes_complete(&self) -> bool {
7356 self.indexes_complete
7357 }
7358
7359 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7361 self.index_build_policy
7362 }
7363
7364 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7368 self.index_build_policy = policy;
7369 }
7370
7371 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7376 if !self.indexes_complete {
7380 return None;
7381 }
7382 let b = self.bitmap.get(&column_id)?;
7383 let result: Vec<Vec<u8>> = b
7384 .keys()
7385 .into_iter()
7386 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7387 .collect();
7388 Some(result)
7389 }
7390
7391 pub fn fk_join_row_ids(
7392 &self,
7393 fk_column_id: u16,
7394 pk_values: &[Vec<u8>],
7395 fk_conditions: &[crate::query::Condition],
7396 snapshot: Snapshot,
7397 ) -> Result<Vec<u64>> {
7398 let Some(b) = self.bitmap.get(&fk_column_id) else {
7399 return Ok(Vec::new());
7400 };
7401 let mut join_set = {
7402 let mut acc = roaring::RoaringBitmap::new();
7403 for v in pk_values {
7404 acc |= b.get(v);
7405 }
7406 RowIdSet::from_roaring(acc)
7407 };
7408 if !fk_conditions.is_empty() {
7409 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7410 sets.push(join_set);
7411 for c in fk_conditions {
7412 sets.push(self.resolve_condition(c, snapshot)?);
7413 }
7414 join_set = RowIdSet::intersect_many(sets);
7415 }
7416 Ok(join_set.into_sorted_vec())
7417 }
7418
7419 pub fn fk_join_count(
7425 &self,
7426 fk_column_id: u16,
7427 pk_values: &[Vec<u8>],
7428 fk_conditions: &[crate::query::Condition],
7429 snapshot: Snapshot,
7430 ) -> Result<u64> {
7431 let Some(b) = self.bitmap.get(&fk_column_id) else {
7432 return Ok(0);
7433 };
7434 let mut acc = roaring::RoaringBitmap::new();
7435 for v in pk_values {
7436 acc |= b.get(v);
7437 }
7438 if fk_conditions.is_empty() {
7439 return Ok(acc.len());
7440 }
7441 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7442 sets.push(RowIdSet::from_roaring(acc));
7443 for c in fk_conditions {
7444 sets.push(self.resolve_condition(c, snapshot)?);
7445 }
7446 Ok(RowIdSet::intersect_many(sets).len() as u64)
7447 }
7448
7449 fn resolve_condition(
7454 &self,
7455 c: &crate::query::Condition,
7456 snapshot: Snapshot,
7457 ) -> Result<RowIdSet> {
7458 self.resolve_condition_with_allowed(c, snapshot, None)
7459 }
7460
7461 fn resolve_condition_with_allowed(
7462 &self,
7463 c: &crate::query::Condition,
7464 snapshot: Snapshot,
7465 allowed: Option<&std::collections::HashSet<RowId>>,
7466 ) -> Result<RowIdSet> {
7467 use crate::query::Condition;
7468 self.validate_condition(c)?;
7469 Ok(match c {
7470 Condition::Pk(key) => {
7471 let lookup = self
7472 .schema
7473 .primary_key()
7474 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7475 .unwrap_or_else(|| key.clone());
7476 self.hot
7477 .get(&lookup)
7478 .map(|r| RowIdSet::one(r.0))
7479 .unwrap_or_else(RowIdSet::empty)
7480 }
7481 Condition::BitmapEq { column_id, value } => {
7482 let lookup = self.index_lookup_key_bytes(*column_id, value);
7483 self.bitmap
7484 .get(column_id)
7485 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7486 .unwrap_or_else(RowIdSet::empty)
7487 }
7488 Condition::BitmapIn { column_id, values } => {
7489 let bm = self.bitmap.get(column_id);
7490 let mut acc = roaring::RoaringBitmap::new();
7491 if let Some(b) = bm {
7492 for v in values {
7493 let lookup = self.index_lookup_key_bytes(*column_id, v);
7494 acc |= b.get(&lookup);
7495 }
7496 }
7497 RowIdSet::from_roaring(acc)
7498 }
7499 Condition::BytesPrefix { column_id, prefix } => {
7500 if let Some(b) = self.bitmap.get(column_id) {
7505 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7506 let mut acc = roaring::RoaringBitmap::new();
7507 for key in b.keys() {
7508 if key.starts_with(&lookup_prefix) {
7509 acc |= b.get(&key);
7510 }
7511 }
7512 RowIdSet::from_roaring(acc)
7513 } else {
7514 RowIdSet::empty()
7515 }
7516 }
7517 Condition::FmContains { column_id, pattern } => self
7518 .fm
7519 .get(column_id)
7520 .map(|f| {
7521 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7522 })
7523 .unwrap_or_else(RowIdSet::empty),
7524 Condition::FmContainsAll {
7525 column_id,
7526 patterns,
7527 } => {
7528 if let Some(f) = self.fm.get(column_id) {
7531 let sets: Vec<RowIdSet> = patterns
7532 .iter()
7533 .map(|pat| {
7534 RowIdSet::from_unsorted(
7535 f.locate(pat).into_iter().map(|r| r.0).collect(),
7536 )
7537 })
7538 .collect();
7539 RowIdSet::intersect_many(sets)
7540 } else {
7541 RowIdSet::empty()
7542 }
7543 }
7544 Condition::Ann {
7545 column_id,
7546 query,
7547 k,
7548 } => RowIdSet::from_unsorted(
7549 self.retrieve_filtered(
7550 &crate::query::Retriever::Ann {
7551 column_id: *column_id,
7552 query: query.clone(),
7553 k: *k,
7554 },
7555 snapshot,
7556 None,
7557 allowed,
7558 None,
7559 None,
7560 )?
7561 .into_iter()
7562 .map(|hit| hit.row_id.0)
7563 .collect(),
7564 ),
7565 Condition::SparseMatch {
7566 column_id,
7567 query,
7568 k,
7569 } => RowIdSet::from_unsorted(
7570 self.retrieve_filtered(
7571 &crate::query::Retriever::Sparse {
7572 column_id: *column_id,
7573 query: query.clone(),
7574 k: *k,
7575 },
7576 snapshot,
7577 None,
7578 allowed,
7579 None,
7580 None,
7581 )?
7582 .into_iter()
7583 .map(|hit| hit.row_id.0)
7584 .collect(),
7585 ),
7586 Condition::MinHashSimilar {
7587 column_id,
7588 query,
7589 k,
7590 } => match self.minhash.get(column_id) {
7591 Some(index) => {
7592 let candidates = index.candidate_row_ids(query);
7593 let eligible =
7594 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7595 RowIdSet::from_unsorted(
7596 index
7597 .search_filtered(query, *k, |row_id| {
7598 eligible.contains(&row_id)
7599 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7600 })
7601 .into_iter()
7602 .map(|(row_id, _)| row_id.0)
7603 .collect(),
7604 )
7605 }
7606 None => RowIdSet::empty(),
7607 },
7608 Condition::Range { column_id, lo, hi } => {
7609 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7618 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7619 } else if self.run_refs.len() == 1 {
7620 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7621 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7622 } else {
7623 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7624 };
7625 set.remove_many(self.overlay_rid_set(snapshot));
7626 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7627 set
7628 }
7629 Condition::RangeF64 {
7630 column_id,
7631 lo,
7632 lo_inclusive,
7633 hi,
7634 hi_inclusive,
7635 } => {
7636 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7639 RowIdSet::from_unsorted(
7640 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7641 .into_iter()
7642 .collect(),
7643 )
7644 } else if self.run_refs.len() == 1 {
7645 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7646 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7647 } else {
7648 return self.range_scan_f64(
7649 *column_id,
7650 *lo,
7651 *lo_inclusive,
7652 *hi,
7653 *hi_inclusive,
7654 snapshot,
7655 );
7656 };
7657 set.remove_many(self.overlay_rid_set(snapshot));
7658 self.range_scan_overlay_f64(
7659 &mut set,
7660 *column_id,
7661 *lo,
7662 *lo_inclusive,
7663 *hi,
7664 *hi_inclusive,
7665 snapshot,
7666 );
7667 set
7668 }
7669 Condition::IsNull { column_id } => {
7670 let mut set = if self.run_refs.len() == 1 {
7671 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7672 r.null_row_id_set(*column_id, true)?
7673 } else {
7674 return self.null_scan(*column_id, true, snapshot);
7675 };
7676 set.remove_many(self.overlay_rid_set(snapshot));
7677 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7678 set
7679 }
7680 Condition::IsNotNull { column_id } => {
7681 let mut set = if self.run_refs.len() == 1 {
7682 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7683 r.null_row_id_set(*column_id, false)?
7684 } else {
7685 return self.null_scan(*column_id, false, snapshot);
7686 };
7687 set.remove_many(self.overlay_rid_set(snapshot));
7688 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7689 set
7690 }
7691 })
7692 }
7693
7694 fn range_scan_i64(
7702 &self,
7703 column_id: u16,
7704 lo: i64,
7705 hi: i64,
7706 snapshot: Snapshot,
7707 ) -> Result<RowIdSet> {
7708 let mut row_ids = Vec::new();
7709 let overlay_rids = self.overlay_rid_set(snapshot);
7710 for rr in &self.run_refs {
7711 let mut reader = self.open_reader(rr.run_id)?;
7712 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7713 for rid in matched {
7714 if !overlay_rids.contains(&rid) {
7715 row_ids.push(rid);
7716 }
7717 }
7718 }
7719 let mut s = RowIdSet::from_unsorted(row_ids);
7720 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7721 Ok(s)
7722 }
7723
7724 fn range_scan_f64(
7727 &self,
7728 column_id: u16,
7729 lo: f64,
7730 lo_inclusive: bool,
7731 hi: f64,
7732 hi_inclusive: bool,
7733 snapshot: Snapshot,
7734 ) -> Result<RowIdSet> {
7735 let mut row_ids = Vec::new();
7736 let overlay_rids = self.overlay_rid_set(snapshot);
7737 for rr in &self.run_refs {
7738 let mut reader = self.open_reader(rr.run_id)?;
7739 let matched = reader.range_row_ids_visible_f64(
7740 column_id,
7741 lo,
7742 lo_inclusive,
7743 hi,
7744 hi_inclusive,
7745 snapshot.epoch,
7746 )?;
7747 for rid in matched {
7748 if !overlay_rids.contains(&rid) {
7749 row_ids.push(rid);
7750 }
7751 }
7752 }
7753 let mut s = RowIdSet::from_unsorted(row_ids);
7754 self.range_scan_overlay_f64(
7755 &mut s,
7756 column_id,
7757 lo,
7758 lo_inclusive,
7759 hi,
7760 hi_inclusive,
7761 snapshot,
7762 );
7763 Ok(s)
7764 }
7765
7766 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7768 let mut s = HashSet::new();
7769 for row in self.memtable.visible_versions_at(snapshot) {
7770 s.insert(row.row_id.0);
7771 }
7772 for row in self.mutable_run.visible_versions_at(snapshot) {
7773 s.insert(row.row_id.0);
7774 }
7775 s
7776 }
7777
7778 fn range_scan_overlay_i64(
7779 &self,
7780 s: &mut RowIdSet,
7781 column_id: u16,
7782 lo: i64,
7783 hi: i64,
7784 snapshot: Snapshot,
7785 ) {
7786 let mut newest: HashMap<u64, Row> = HashMap::new();
7793 for r in self.mutable_run.visible_versions_at(snapshot) {
7794 newest.insert(r.row_id.0, r);
7795 }
7796 for r in self.memtable.visible_versions_at(snapshot) {
7797 newest
7798 .entry(r.row_id.0)
7799 .and_modify(|cur| {
7800 if Snapshot::version_is_newer(
7801 r.committed_epoch,
7802 r.commit_ts,
7803 cur.committed_epoch,
7804 cur.commit_ts,
7805 ) {
7806 *cur = r.clone();
7807 }
7808 })
7809 .or_insert(r);
7810 }
7811 for row in newest.values() {
7812 if !row.deleted {
7813 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7814 if *v >= lo && *v <= hi {
7815 s.insert(row.row_id.0);
7816 }
7817 }
7818 }
7819 }
7820 }
7821
7822 #[allow(clippy::too_many_arguments)]
7823 fn range_scan_overlay_f64(
7824 &self,
7825 s: &mut RowIdSet,
7826 column_id: u16,
7827 lo: f64,
7828 lo_inclusive: bool,
7829 hi: f64,
7830 hi_inclusive: bool,
7831 snapshot: Snapshot,
7832 ) {
7833 let mut newest: HashMap<u64, Row> = HashMap::new();
7836 for r in self.mutable_run.visible_versions_at(snapshot) {
7837 newest.insert(r.row_id.0, r);
7838 }
7839 for r in self.memtable.visible_versions_at(snapshot) {
7840 newest
7841 .entry(r.row_id.0)
7842 .and_modify(|cur| {
7843 if Snapshot::version_is_newer(
7844 r.committed_epoch,
7845 r.commit_ts,
7846 cur.committed_epoch,
7847 cur.commit_ts,
7848 ) {
7849 *cur = r.clone();
7850 }
7851 })
7852 .or_insert(r);
7853 }
7854 for row in newest.values() {
7855 if !row.deleted {
7856 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7857 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7858 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7859 if ok_lo && ok_hi {
7860 s.insert(row.row_id.0);
7861 }
7862 }
7863 }
7864 }
7865 }
7866
7867 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7870 let mut row_ids = Vec::new();
7871 let overlay_rids = self.overlay_rid_set(snapshot);
7872 for rr in &self.run_refs {
7873 let mut reader = self.open_reader(rr.run_id)?;
7874 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7875 for rid in matched {
7876 if !overlay_rids.contains(&rid) {
7877 row_ids.push(rid);
7878 }
7879 }
7880 }
7881 let mut s = RowIdSet::from_unsorted(row_ids);
7882 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7883 Ok(s)
7884 }
7885
7886 fn null_scan_overlay(
7890 &self,
7891 s: &mut RowIdSet,
7892 column_id: u16,
7893 want_nulls: bool,
7894 snapshot: Snapshot,
7895 ) {
7896 let mut newest: HashMap<u64, Row> = HashMap::new();
7897 for r in self.mutable_run.visible_versions_at(snapshot) {
7898 newest.insert(r.row_id.0, r);
7899 }
7900 for r in self.memtable.visible_versions_at(snapshot) {
7901 newest
7902 .entry(r.row_id.0)
7903 .and_modify(|cur| {
7904 if Snapshot::version_is_newer(
7905 r.committed_epoch,
7906 r.commit_ts,
7907 cur.committed_epoch,
7908 cur.commit_ts,
7909 ) {
7910 *cur = r.clone();
7911 }
7912 })
7913 .or_insert(r);
7914 }
7915 for row in newest.values() {
7916 if row.deleted {
7917 continue;
7918 }
7919 let is_null = !row.columns.contains_key(&column_id)
7920 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7921 if is_null == want_nulls {
7922 s.insert(row.row_id.0);
7923 }
7924 }
7925 }
7926
7927 pub fn snapshot(&self) -> Snapshot {
7928 let epoch = self.epoch.visible();
7929 match &self.wal {
7933 WalSink::Shared(shared) => match shared.hlc.now() {
7934 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
7935 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
7937 },
7938 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
7939 }
7940 }
7941
7942 pub fn data_generation(&self) -> u64 {
7944 self.data_generation
7945 }
7946
7947 pub(crate) fn bump_data_generation(&mut self) {
7948 self.data_generation = self.data_generation.wrapping_add(1);
7949 }
7950
7951 pub fn table_id(&self) -> u64 {
7953 self.table_id
7954 }
7955
7956 fn seal_generations(&mut self) {
7962 self.memtable.seal();
7963 self.mutable_run.seal();
7964 self.hot.seal();
7965 for index in self.bitmap.values_mut() {
7966 index.seal();
7967 }
7968 for index in self.ann.values_mut() {
7969 index.seal();
7970 }
7971 for index in self.fm.values_mut() {
7972 index.seal();
7973 }
7974 for index in self.sparse.values_mut() {
7975 index.seal();
7976 }
7977 for index in self.minhash.values_mut() {
7978 index.seal();
7979 }
7980 self.pk_by_row.seal();
7981 }
7982
7983 fn capture_read_generation(&self) -> ReadGeneration {
7988 let visible_through = self.current_epoch();
7989 ReadGeneration {
7990 schema: Arc::new(self.schema.clone()),
7991 base_runs: Arc::new(self.run_refs.clone()),
7992 deltas: TableDeltas {
7993 memtable: self.memtable.clone(),
7994 mutable_run: self.mutable_run.clone(),
7995 hot: self.hot.clone(),
7996 pk_by_row: self.pk_by_row.clone(),
7997 },
7998 indexes: Arc::new(IndexGeneration::capture(
7999 &self.bitmap,
8000 &self.learned_range,
8001 &self.fm,
8002 &self.ann,
8003 &self.sparse,
8004 &self.minhash,
8005 visible_through,
8006 match &self.wal {
8008 WalSink::Shared(shared) => shared
8009 .hlc
8010 .now()
8011 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
8012 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
8013 },
8014 )),
8015 visible_through,
8016 }
8017 }
8018
8019 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
8024 self.ensure_indexes_complete()?;
8025 self.seal_generations();
8026 let view = Arc::new(self.capture_read_generation());
8027 self.published.store(Arc::clone(&view));
8028 Ok(view)
8029 }
8030
8031 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
8037 self.published.load_full()
8038 }
8039
8040 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
8042 &self.pins
8043 }
8044
8045 pub fn version_gc_floor(&self) -> Epoch {
8050 self.min_active_snapshot()
8051 .unwrap_or_else(|| self.current_epoch())
8052 }
8053
8054 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
8061 let mut report = self.pins.report();
8062 let transaction_floor = [
8063 self.pinned.keys().next().copied(),
8064 self.snapshots.min_pinned(),
8065 ]
8066 .into_iter()
8067 .flatten()
8068 .min();
8069 if let Some(epoch) = transaction_floor {
8070 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
8071 }
8072 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
8073 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
8074 }
8075 report
8076 }
8077
8078 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
8079 self.publish_read_generation()?;
8080 let mut generation = self.clone();
8081 generation.read_only = true;
8082 generation.wal = WalSink::ReadOnly;
8083 generation.pending_delete_rids.clear();
8084 generation.pending_put_cols.clear();
8085 generation.pending_rows.clear();
8086 generation.pending_rows_auto_inc.clear();
8087 generation.pending_dels.clear();
8088 generation.pending_truncate = None;
8089 generation.agg_cache = Arc::new(HashMap::new());
8090 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
8093 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
8096 crate::retention::PinSource::ReadGeneration,
8097 self.current_epoch(),
8098 )));
8099 Ok(generation)
8100 }
8101
8102 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
8103 (std::mem::size_of::<Self>() as u64)
8104 .saturating_add(self.memtable.approx_bytes())
8105 .saturating_add(self.mutable_run.approx_bytes())
8106 .saturating_add(self.live_count.saturating_mul(64))
8107 }
8108
8109 pub fn pin_snapshot(&mut self) -> Snapshot {
8116 let snap = self.snapshot();
8117 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
8118 snap
8119 }
8120
8121 pub fn hlc_gc_floor(
8131 &self,
8132 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
8133 ) -> crate::epoch::GcFloor {
8134 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
8135 epoch
8136 .and_then(&mut project_epoch)
8137 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
8138 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
8139 };
8140 let transaction = [
8141 self.pinned.keys().next().copied(),
8142 self.snapshots.min_pinned(),
8143 ]
8144 .into_iter()
8145 .flatten()
8146 .min();
8147 let history = self.snapshots.history_floor(self.current_epoch());
8148 crate::epoch::GcFloor {
8149 transaction_snapshot: project(transaction),
8150 history_retention: project(history),
8151 backup_pitr: project(
8152 self.pins
8153 .oldest_for(crate::retention::PinSource::BackupPitr),
8154 ),
8155 replication: project(
8156 self.pins
8157 .oldest_for(crate::retention::PinSource::Replication),
8158 ),
8159 read_generation: project(
8160 self.pins
8161 .oldest_for(crate::retention::PinSource::ReadGeneration),
8162 ),
8163 online_index_build: project(
8164 self.pins
8165 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
8166 ),
8167 }
8168 }
8169
8170 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
8172 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
8173 *count -= 1;
8174 if *count == 0 {
8175 self.pinned.remove(&snap.epoch);
8176 }
8177 }
8178 }
8179
8180 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
8192 let local = self.pinned.keys().next().copied();
8193 let global = self.snapshots.min_pinned();
8194 let history = self.snapshots.history_floor(self.current_epoch());
8195 let pinned = self.pins.oldest_pinned();
8196 [local, global, history, pinned].into_iter().flatten().min()
8197 }
8198
8199 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
8203 self.ensure_writable()?;
8204 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
8205 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
8206 }
8207
8208 pub fn clear_ttl(&mut self) -> Result<()> {
8209 self.ensure_writable()?;
8210 self.apply_ttl_policy_at(None, self.current_epoch())
8211 }
8212
8213 pub fn ttl(&self) -> Option<TtlPolicy> {
8214 self.ttl
8215 }
8216
8217 pub(crate) fn prepare_ttl_policy(
8218 &self,
8219 column_name: &str,
8220 duration_nanos: u64,
8221 ) -> Result<TtlPolicy> {
8222 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
8223 return Err(MongrelError::InvalidArgument(
8224 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
8225 ));
8226 }
8227 let column = self
8228 .schema
8229 .columns
8230 .iter()
8231 .find(|column| column.name == column_name)
8232 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
8233 if column.ty != TypeId::TimestampNanos {
8234 return Err(MongrelError::Schema(format!(
8235 "TTL column {column_name} must be TimestampNanos, is {:?}",
8236 column.ty
8237 )));
8238 }
8239 Ok(TtlPolicy {
8240 column_id: column.id,
8241 duration_nanos,
8242 })
8243 }
8244
8245 pub(crate) fn apply_ttl_policy_at(
8246 &mut self,
8247 policy: Option<TtlPolicy>,
8248 epoch: Epoch,
8249 ) -> Result<()> {
8250 if let Some(policy) = policy {
8251 let column = self
8252 .schema
8253 .columns
8254 .iter()
8255 .find(|column| column.id == policy.column_id)
8256 .ok_or_else(|| {
8257 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8258 })?;
8259 if column.ty != TypeId::TimestampNanos
8260 || policy.duration_nanos == 0
8261 || policy.duration_nanos > i64::MAX as u64
8262 {
8263 return Err(MongrelError::Schema("invalid TTL policy".into()));
8264 }
8265 }
8266 self.ttl = policy;
8267 self.agg_cache = Arc::new(HashMap::new());
8268 self.clear_result_cache();
8269 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8270 self.persist_manifest(epoch)
8271 }
8272
8273 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8274 let Some(policy) = self.ttl else {
8275 return false;
8276 };
8277 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8278 return false;
8279 };
8280 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8281 }
8282
8283 pub fn current_epoch(&self) -> Epoch {
8284 self.epoch.visible()
8285 }
8286
8287 pub fn memtable_len(&self) -> usize {
8288 self.memtable.len()
8289 }
8290
8291 pub fn count(&self) -> u64 {
8294 if self.ttl.is_none()
8295 && self.pending_put_cols.is_empty()
8296 && self.pending_delete_rids.is_empty()
8297 && self.pending_rows.is_empty()
8298 && self.pending_dels.is_empty()
8299 && self.pending_truncate.is_none()
8300 {
8301 self.live_count
8302 } else {
8303 self.visible_rows(self.snapshot())
8304 .map(|rows| rows.len() as u64)
8305 .unwrap_or(self.live_count)
8306 }
8307 }
8308
8309 pub fn count_conditions(
8313 &mut self,
8314 conditions: &[crate::query::Condition],
8315 snapshot: Snapshot,
8316 ) -> Result<Option<u64>> {
8317 use crate::query::Condition;
8318 if self.ttl.is_some() {
8319 if conditions.is_empty() {
8320 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8321 }
8322 let mut sets = Vec::with_capacity(conditions.len());
8323 for condition in conditions {
8324 sets.push(self.resolve_condition(condition, snapshot)?);
8325 }
8326 let survivors = RowIdSet::intersect_many(sets);
8327 let rows = self.visible_rows(snapshot)?;
8328 return Ok(Some(
8329 rows.into_iter()
8330 .filter(|row| survivors.contains(row.row_id.0))
8331 .count() as u64,
8332 ));
8333 }
8334 if conditions.is_empty() {
8335 return Ok(Some(self.count()));
8336 }
8337 let served = |c: &Condition| {
8338 matches!(
8339 c,
8340 Condition::Pk(_)
8341 | Condition::BitmapEq { .. }
8342 | Condition::BitmapIn { .. }
8343 | Condition::BytesPrefix { .. }
8344 | Condition::FmContains { .. }
8345 | Condition::FmContainsAll { .. }
8346 | Condition::Ann { .. }
8347 | Condition::Range { .. }
8348 | Condition::RangeF64 { .. }
8349 | Condition::SparseMatch { .. }
8350 | Condition::MinHashSimilar { .. }
8351 | Condition::IsNull { .. }
8352 | Condition::IsNotNull { .. }
8353 )
8354 };
8355 if !conditions.iter().all(served) {
8356 return Ok(None);
8357 }
8358 self.ensure_indexes_complete()?;
8359 if !self.pending_put_cols.is_empty()
8360 || !self.pending_delete_rids.is_empty()
8361 || !self.pending_rows.is_empty()
8362 || !self.pending_dels.is_empty()
8363 || self.pending_truncate.is_some()
8364 {
8365 let mut sets = Vec::with_capacity(conditions.len());
8366 for condition in conditions {
8367 sets.push(self.resolve_condition(condition, snapshot)?);
8368 }
8369 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8370 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8371 }
8372 let mut sets = Vec::with_capacity(conditions.len());
8373 for condition in conditions {
8374 sets.push(self.resolve_condition(condition, snapshot)?);
8375 }
8376 let mut rids = RowIdSet::intersect_many(sets);
8377 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8387 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8388 }
8389 let count = rids.len() as u64;
8390 crate::trace::QueryTrace::record(|t| {
8391 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8392 t.survivor_count = Some(count as usize);
8393 t.conditions_pushed = conditions.len();
8394 });
8395 Ok(Some(count))
8396 }
8397
8398 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8403 let mut out = Vec::new();
8404 for row in self.memtable.visible_versions(snapshot.epoch) {
8405 if row.deleted {
8406 out.push(row.row_id.0);
8407 }
8408 }
8409 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8410 if row.deleted {
8411 out.push(row.row_id.0);
8412 }
8413 }
8414 out
8415 }
8416
8417 pub fn bulk_load_columns(
8426 &mut self,
8427 user_columns: Vec<(u16, columnar::NativeColumn)>,
8428 ) -> Result<Epoch> {
8429 self.bulk_load_columns_with(user_columns, 3, false, true)
8430 }
8431
8432 pub fn bulk_load_fast(
8439 &mut self,
8440 user_columns: Vec<(u16, columnar::NativeColumn)>,
8441 ) -> Result<Epoch> {
8442 self.bulk_load_columns_with(user_columns, -1, true, false)
8443 }
8444
8445 fn bulk_load_columns_with(
8446 &mut self,
8447 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8448 zstd_level: i32,
8449 force_plain: bool,
8450 lz4: bool,
8451 ) -> Result<Epoch> {
8452 self.ensure_writable()?;
8453 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8454 if n == 0 {
8455 return Ok(self.current_epoch());
8456 }
8457 let epoch = self.commit_new_epoch()?;
8458 let live_before = self.live_count;
8459 self.spill_mutable_run(epoch)?;
8461 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8462 && self.indexes_complete
8463 && self.run_refs.is_empty()
8464 && self.memtable.is_empty()
8465 && self.mutable_run.is_empty();
8466 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8469 self.validate_columns_not_null(&user_columns, n)?;
8470 let winner_idx = self
8471 .bulk_pk_winner_indices(&user_columns, n)
8472 .filter(|idx| idx.len() != n);
8473 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8474 match winner_idx.as_deref() {
8475 Some(idx) => {
8476 let compacted = user_columns
8477 .iter()
8478 .map(|(id, c)| (*id, c.gather(idx)))
8479 .collect();
8480 (compacted, idx.len())
8481 }
8482 None => (user_columns, n),
8483 };
8484 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8485 let first = self.allocator.alloc_range(write_n as u64)?.0;
8486 for rid in first..first + write_n as u64 {
8487 self.reservoir.offer(rid);
8488 }
8489 let run_id = self.alloc_run_id()?;
8490 let path = self.run_path(run_id);
8491 let mut writer =
8492 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8493 if force_plain {
8494 writer = writer.with_plain();
8495 } else if lz4 {
8496 writer = writer.with_lz4();
8499 } else {
8500 writer = writer.with_zstd_level(zstd_level);
8501 }
8502 if let Some(kek) = &self.kek {
8503 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8504 }
8505 let header = match self.create_run_file(run_id)? {
8506 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8507 None => writer.write_native(&path, &write_columns, write_n, first)?,
8508 };
8509 self.run_refs.push(RunRef {
8510 run_id: run_id as u128,
8511 level: 0,
8512 epoch_created: epoch.0,
8513 row_count: header.row_count,
8514 });
8515 self.live_count = self.live_count.saturating_add(write_n as u64);
8516 if eager_index_build {
8517 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8518 self.index_columns_bulk(&write_columns, &row_ids);
8519 self.indexes_complete = true;
8520 self.build_learned_ranges()?;
8521 } else {
8522 self.indexes_complete = false;
8526 }
8527 self.mark_flushed(epoch)?;
8528 self.persist_manifest(epoch)?;
8529 if eager_index_build {
8530 self.checkpoint_indexes(epoch);
8531 }
8532 self.clear_result_cache();
8533 self.data_generation = self.data_generation.wrapping_add(1);
8534 Ok(epoch)
8535 }
8536
8537 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8555 let n = row_ids.len();
8556 if n == 0 {
8557 return;
8558 }
8559 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8560 columns.iter().map(|(id, c)| (*id, c)).collect();
8561 let ty_of: std::collections::HashMap<u16, TypeId> = self
8562 .schema
8563 .columns
8564 .iter()
8565 .map(|c| (c.id, c.ty.clone()))
8566 .collect();
8567 let pk_id = self.schema.primary_key().map(|c| c.id);
8568
8569 for (i, &rid) in row_ids.iter().enumerate() {
8570 let row_id = RowId(rid);
8571 if let Some(pid) = pk_id {
8572 if let Some(col) = by_id.get(&pid) {
8573 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8574 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8575 self.insert_hot_pk(key, row_id);
8576 }
8577 }
8578 }
8579 for idef in &self.schema.indexes {
8580 let Some(col) = by_id.get(&idef.column_id) else {
8581 continue;
8582 };
8583 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8584 match idef.kind {
8585 IndexKind::Bitmap => {
8586 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8587 if let Some(key) =
8588 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8589 {
8590 b.insert(key, row_id);
8591 }
8592 }
8593 }
8594 IndexKind::FmIndex => {
8595 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8596 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8597 f.insert(bytes.to_vec(), row_id);
8598 }
8599 }
8600 }
8601 IndexKind::Sparse => {
8602 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8603 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8604 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8605 s.insert(&terms, row_id);
8606 }
8607 }
8608 }
8609 }
8610 IndexKind::MinHash => {
8611 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8612 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8613 let tokens = crate::index::token_hashes_from_bytes(bytes);
8614 mh.insert(&tokens, row_id);
8615 }
8616 }
8617 }
8618 _ => {}
8619 }
8620 }
8621 }
8622 }
8623
8624 pub fn visible_columns_native(
8629 &self,
8630 snapshot: Snapshot,
8631 projection: Option<&[u16]>,
8632 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8633 self.visible_columns_native_inner(snapshot, projection, None)
8634 }
8635
8636 pub fn visible_columns_native_with_control(
8637 &self,
8638 snapshot: Snapshot,
8639 projection: Option<&[u16]>,
8640 control: &crate::ExecutionControl,
8641 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8642 self.visible_columns_native_inner(snapshot, projection, Some(control))
8643 }
8644
8645 fn visible_columns_native_inner(
8646 &self,
8647 snapshot: Snapshot,
8648 projection: Option<&[u16]>,
8649 control: Option<&crate::ExecutionControl>,
8650 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8651 execution_checkpoint(control, 0)?;
8652 let wanted: Vec<u16> = match projection {
8653 Some(p) => p.to_vec(),
8654 None => self.schema.columns.iter().map(|c| c.id).collect(),
8655 };
8656 if self.ttl.is_none()
8657 && self.memtable.is_empty()
8658 && self.mutable_run.is_empty()
8659 && self.run_refs.len() == 1
8660 {
8661 let rr = self.run_refs[0].clone();
8662 let mut reader = self.open_reader(rr.run_id)?;
8663 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8664 execution_checkpoint(control, 0)?;
8665 let all_visible = idxs.len() == reader.row_count();
8666 if reader.has_mmap() && control.is_none() {
8672 use rayon::prelude::*;
8673 let valid: Vec<u16> = wanted
8676 .iter()
8677 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8678 .copied()
8679 .collect();
8680 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8682 .par_iter()
8683 .filter_map(|cid| {
8684 reader
8685 .column_native_shared(*cid)
8686 .ok()
8687 .map(|col| (*cid, col))
8688 })
8689 .collect();
8690 let cols = decoded
8691 .into_iter()
8692 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8693 .collect();
8694 return Ok(cols);
8695 }
8696 let mut cols = Vec::with_capacity(wanted.len());
8697 for (index, cid) in wanted.iter().enumerate() {
8698 execution_checkpoint(control, index)?;
8699 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8700 Some(c) => c,
8701 None => continue,
8702 };
8703 let col = reader.column_native(cdef.id)?;
8704 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8705 }
8706 return Ok(cols);
8707 }
8708 let vcols = self.visible_columns(snapshot)?;
8709 execution_checkpoint(control, 0)?;
8710 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8711 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8712 .into_iter()
8713 .filter(|(id, _)| want_set.contains(id))
8714 .map(|(id, vals)| {
8715 let ty = self
8716 .schema
8717 .columns
8718 .iter()
8719 .find(|c| c.id == id)
8720 .map(|c| c.ty.clone())
8721 .unwrap_or(TypeId::Bytes);
8722 (id, columnar::values_to_native(ty, &vals))
8723 })
8724 .collect();
8725 Ok(out)
8726 }
8727
8728 pub fn run_count(&self) -> usize {
8729 self.run_refs.len()
8730 }
8731
8732 pub fn memtable_is_empty(&self) -> bool {
8734 self.memtable.is_empty()
8735 }
8736
8737 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8741 self.page_cache.stats()
8742 }
8743
8744 pub fn reset_page_cache_stats(&self) {
8746 self.page_cache.reset_stats();
8747 }
8748
8749 pub fn run_ids(&self) -> Vec<u128> {
8752 self.run_refs.iter().map(|r| r.run_id).collect()
8753 }
8754
8755 pub fn single_run_is_clean(&self) -> bool {
8759 if self.ttl.is_some() || self.run_refs.len() != 1 {
8760 return false;
8761 }
8762 self.open_reader(self.run_refs[0].run_id)
8763 .map(|r| r.is_clean())
8764 .unwrap_or(false)
8765 }
8766
8767 fn resolve_footprint(
8774 &self,
8775 conditions: &[crate::query::Condition],
8776 snapshot: Snapshot,
8777 ) -> roaring::RoaringBitmap {
8778 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8779 return roaring::RoaringBitmap::new();
8780 }
8781 if self.run_refs.is_empty() {
8782 return roaring::RoaringBitmap::new();
8783 }
8784 if self.run_refs.len() == 1 {
8786 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8787 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8788 return rids.to_roaring_lossy();
8789 }
8790 }
8791 }
8792 roaring::RoaringBitmap::new()
8793 }
8794
8795 pub fn query_columns_native_cached(
8806 &mut self,
8807 conditions: &[crate::query::Condition],
8808 projection: Option<&[u16]>,
8809 snapshot: Snapshot,
8810 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8811 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8812 }
8813
8814 pub fn query_columns_native_cached_with_control(
8815 &mut self,
8816 conditions: &[crate::query::Condition],
8817 projection: Option<&[u16]>,
8818 snapshot: Snapshot,
8819 control: &crate::ExecutionControl,
8820 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8821 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8822 }
8823
8824 fn query_columns_native_cached_inner(
8825 &mut self,
8826 conditions: &[crate::query::Condition],
8827 projection: Option<&[u16]>,
8828 snapshot: Snapshot,
8829 control: Option<&crate::ExecutionControl>,
8830 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8831 execution_checkpoint(control, 0)?;
8832 if self.ttl.is_some() {
8835 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8836 }
8837 if conditions.is_empty() {
8838 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8839 }
8840 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8844 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8845 crate::trace::QueryTrace::record(|t| {
8846 t.result_cache_hit = true;
8847 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8848 });
8849 return Ok(Some((*hit).clone()));
8850 }
8851 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8852 execution_checkpoint(control, 0)?;
8853 if let Some(cols) = &res {
8854 let footprint = self.resolve_footprint(conditions, snapshot);
8855 let condition_cols = crate::query::condition_columns(conditions);
8856 execution_checkpoint(control, 0)?;
8857 self.result_cache.lock().insert(
8858 key,
8859 CachedEntry {
8860 data: CachedData::Columns(Arc::new(cols.clone())),
8861 footprint,
8862 condition_cols,
8863 },
8864 );
8865 }
8866 Ok(res)
8867 }
8868
8869 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8874 if self.ttl.is_some() {
8875 return self.query(q);
8876 }
8877 if q.conditions.is_empty() {
8878 return self.query(q);
8879 }
8880 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8881 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8882 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8883 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8884 crate::trace::QueryTrace::record(|t| {
8885 t.result_cache_hit = true;
8886 t.scan_mode = crate::trace::ScanMode::Materialized;
8887 });
8888 return Ok((*hit).clone());
8889 }
8890 let rows = self.query(q)?;
8891 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8892 let condition_cols = crate::query::condition_columns(&q.conditions);
8893 self.result_cache.lock().insert(
8894 key,
8895 CachedEntry {
8896 data: CachedData::Rows(Arc::new(rows.clone())),
8897 footprint,
8898 condition_cols,
8899 },
8900 );
8901 Ok(rows)
8902 }
8903
8904 #[allow(clippy::type_complexity)]
8919 pub fn query_columns_native_traced(
8920 &mut self,
8921 conditions: &[crate::query::Condition],
8922 projection: Option<&[u16]>,
8923 snapshot: Snapshot,
8924 ) -> Result<(
8925 Option<Vec<(u16, columnar::NativeColumn)>>,
8926 crate::trace::QueryTrace,
8927 )> {
8928 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8929 self.query_columns_native(conditions, projection, snapshot)
8930 });
8931 Ok((result?, trace))
8932 }
8933
8934 #[allow(clippy::type_complexity)]
8937 pub fn query_columns_native_cached_traced(
8938 &mut self,
8939 conditions: &[crate::query::Condition],
8940 projection: Option<&[u16]>,
8941 snapshot: Snapshot,
8942 ) -> Result<(
8943 Option<Vec<(u16, columnar::NativeColumn)>>,
8944 crate::trace::QueryTrace,
8945 )> {
8946 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8947 self.query_columns_native_cached(conditions, projection, snapshot)
8948 });
8949 Ok((result?, trace))
8950 }
8951
8952 pub fn native_page_cursor_traced(
8954 &self,
8955 snapshot: Snapshot,
8956 projection: Vec<(u16, TypeId)>,
8957 conditions: &[crate::query::Condition],
8958 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8959 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8960 self.native_page_cursor(snapshot, projection, conditions)
8961 });
8962 Ok((result?, trace))
8963 }
8964
8965 pub fn native_multi_run_cursor_traced(
8967 &self,
8968 snapshot: Snapshot,
8969 projection: Vec<(u16, TypeId)>,
8970 conditions: &[crate::query::Condition],
8971 ) -> Result<(
8972 Option<crate::cursor::MultiRunCursor>,
8973 crate::trace::QueryTrace,
8974 )> {
8975 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8976 self.native_multi_run_cursor(snapshot, projection, conditions)
8977 });
8978 Ok((result?, trace))
8979 }
8980
8981 pub fn count_conditions_traced(
8983 &mut self,
8984 conditions: &[crate::query::Condition],
8985 snapshot: Snapshot,
8986 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8987 let (result, trace) =
8988 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8989 Ok((result?, trace))
8990 }
8991
8992 pub fn query_traced(
8994 &mut self,
8995 q: &crate::query::Query,
8996 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8997 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8998 Ok((result?, trace))
8999 }
9000
9001 pub fn query_columns_native(
9006 &mut self,
9007 conditions: &[crate::query::Condition],
9008 projection: Option<&[u16]>,
9009 snapshot: Snapshot,
9010 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9011 self.query_columns_native_inner(conditions, projection, snapshot, None)
9012 }
9013
9014 pub fn query_columns_native_with_control(
9015 &mut self,
9016 conditions: &[crate::query::Condition],
9017 projection: Option<&[u16]>,
9018 snapshot: Snapshot,
9019 control: &crate::ExecutionControl,
9020 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9021 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
9022 }
9023
9024 fn query_columns_native_inner(
9025 &mut self,
9026 conditions: &[crate::query::Condition],
9027 projection: Option<&[u16]>,
9028 snapshot: Snapshot,
9029 control: Option<&crate::ExecutionControl>,
9030 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9031 use crate::query::Condition;
9032 execution_checkpoint(control, 0)?;
9033 if self.ttl.is_some() {
9036 return Ok(None);
9037 }
9038 if conditions.is_empty() {
9039 return Ok(None);
9040 }
9041 self.ensure_indexes_complete()?;
9042
9043 let served = |c: &Condition| {
9048 matches!(
9049 c,
9050 Condition::Pk(_)
9051 | Condition::BitmapEq { .. }
9052 | Condition::BitmapIn { .. }
9053 | Condition::BytesPrefix { .. }
9054 | Condition::FmContains { .. }
9055 | Condition::FmContainsAll { .. }
9056 | Condition::Ann { .. }
9057 | Condition::Range { .. }
9058 | Condition::RangeF64 { .. }
9059 | Condition::SparseMatch { .. }
9060 | Condition::MinHashSimilar { .. }
9061 | Condition::IsNull { .. }
9062 | Condition::IsNotNull { .. }
9063 )
9064 };
9065 if !conditions.iter().all(served) {
9066 return Ok(None);
9067 }
9068 let fast_path =
9069 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9070 crate::trace::QueryTrace::record(|t| {
9071 t.run_count = self.run_refs.len();
9072 t.memtable_rows = self.memtable.len();
9073 t.mutable_run_rows = self.mutable_run.len();
9074 t.conditions_pushed = conditions.len();
9075 t.learned_range_used = conditions.iter().any(|c| match c {
9076 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
9077 self.learned_range.contains_key(column_id)
9078 }
9079 _ => false,
9080 });
9081 });
9082 let col_ids: Vec<u16> = projection
9084 .map(|p| p.to_vec())
9085 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
9086 let proj_pairs: Vec<(u16, TypeId)> = col_ids
9087 .iter()
9088 .map(|&cid| {
9089 let ty = self
9090 .schema
9091 .columns
9092 .iter()
9093 .find(|c| c.id == cid)
9094 .map(|c| c.ty.clone())
9095 .unwrap_or(TypeId::Bytes);
9096 (cid, ty)
9097 })
9098 .collect();
9099
9100 if fast_path {
9106 let needs_column = conditions.iter().any(|c| match c {
9109 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
9110 Condition::RangeF64 { column_id, .. } => {
9111 !self.learned_range.contains_key(column_id)
9112 }
9113 _ => false,
9114 });
9115 let mut reader_opt: Option<RunReader> = if needs_column {
9116 Some(self.open_reader(self.run_refs[0].run_id)?)
9117 } else {
9118 None
9119 };
9120 let mut sets: Vec<RowIdSet> = Vec::new();
9121 for (index, c) in conditions.iter().enumerate() {
9122 execution_checkpoint(control, index)?;
9123 let s = match c {
9124 Condition::Range { column_id, lo, hi }
9125 if !self.learned_range.contains_key(column_id) =>
9126 {
9127 if reader_opt.is_none() {
9128 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9129 }
9130 reader_opt
9131 .as_mut()
9132 .expect("reader opened for range")
9133 .range_row_id_set_i64(*column_id, *lo, *hi)?
9134 }
9135 Condition::RangeF64 {
9136 column_id,
9137 lo,
9138 lo_inclusive,
9139 hi,
9140 hi_inclusive,
9141 } if !self.learned_range.contains_key(column_id) => {
9142 if reader_opt.is_none() {
9143 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
9144 }
9145 reader_opt
9146 .as_mut()
9147 .expect("reader opened for range")
9148 .range_row_id_set_f64(
9149 *column_id,
9150 *lo,
9151 *lo_inclusive,
9152 *hi,
9153 *hi_inclusive,
9154 )?
9155 }
9156 _ => self.resolve_condition(c, snapshot)?,
9157 };
9158 sets.push(s);
9159 }
9160 let candidates = RowIdSet::intersect_many(sets);
9161 crate::trace::QueryTrace::record(|t| {
9162 t.survivor_count = Some(candidates.len());
9163 });
9164 if candidates.is_empty() {
9165 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
9166 .iter()
9167 .map(|&id| {
9168 (
9169 id,
9170 columnar::null_native(
9171 proj_pairs
9172 .iter()
9173 .find(|(c, _)| c == &id)
9174 .map(|(_, t)| t.clone())
9175 .unwrap_or(TypeId::Bytes),
9176 0,
9177 ),
9178 )
9179 })
9180 .collect();
9181 return Ok(Some(cols));
9182 }
9183 let mut reader = match reader_opt.take() {
9184 Some(r) => r,
9185 None => self.open_reader(self.run_refs[0].run_id)?,
9186 };
9187 let candidate_ids = candidates.into_sorted_vec();
9188 let (positions, fast_rid) = if let Some(positions) =
9189 reader.positions_for_row_ids_fast(&candidate_ids)
9190 {
9191 (positions, true)
9192 } else {
9193 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
9194 match col {
9195 columnar::NativeColumn::Int64 { data, .. } => {
9196 let mut p = Vec::with_capacity(candidate_ids.len());
9197 for (index, rid) in candidate_ids.iter().enumerate() {
9198 execution_checkpoint(control, index)?;
9199 if let Ok(position) = data.binary_search(&(*rid as i64)) {
9200 p.push(position);
9201 }
9202 }
9203 p.sort_unstable();
9204 (p, false)
9205 }
9206 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
9207 }
9208 };
9209 crate::trace::QueryTrace::record(|t| {
9210 t.scan_mode = crate::trace::ScanMode::NativePushdown;
9211 t.fast_row_id_map = fast_rid;
9212 });
9213 let mut cols = Vec::with_capacity(col_ids.len());
9214 for (index, cid) in col_ids.iter().enumerate() {
9215 execution_checkpoint(control, index)?;
9216 let col = reader.column_native(*cid)?;
9217 cols.push((*cid, col.gather(&positions)));
9218 }
9219 return Ok(Some(cols));
9220 }
9221
9222 if !self.run_refs.is_empty() {
9235 use crate::cursor::{
9236 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
9237 };
9238 let remaining: usize;
9239 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9240 let c = self
9241 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9242 .expect("single-run cursor should build when run_refs.len() == 1");
9243 remaining = c.remaining_rows();
9244 Box::new(c)
9245 } else {
9246 let c = self
9247 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9248 .expect("multi-run cursor should build when run_refs.len() >= 1");
9249 remaining = c.remaining_rows();
9250 Box::new(c)
9251 };
9252 crate::trace::QueryTrace::record(|t| {
9253 if t.survivor_count.is_none() {
9254 t.survivor_count = Some(remaining);
9255 }
9256 });
9257 let cols = match control {
9258 Some(control) => {
9259 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9260 }
9261 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9262 };
9263 return Ok(Some(cols));
9264 }
9265
9266 crate::trace::QueryTrace::record(|t| {
9271 t.scan_mode = crate::trace::ScanMode::Materialized;
9272 t.row_materialized = true;
9273 });
9274 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9275 for (index, c) in conditions.iter().enumerate() {
9276 execution_checkpoint(control, index)?;
9277 sets.push(self.resolve_condition(c, snapshot)?);
9278 }
9279 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9280 let rows = self.rows_for_rids(&rids, snapshot)?;
9281 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9282 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9283 execution_checkpoint(control, index)?;
9284 let vals: Vec<Value> = rows
9285 .iter()
9286 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9287 .collect();
9288 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9289 }
9290 Ok(Some(cols))
9291 }
9292
9293 pub fn native_page_cursor(
9308 &self,
9309 snapshot: Snapshot,
9310 projection: Vec<(u16, TypeId)>,
9311 conditions: &[crate::query::Condition],
9312 ) -> Result<Option<NativePageCursor>> {
9313 use crate::cursor::build_page_plans;
9314 if self.ttl.is_some() {
9315 return Ok(None);
9316 }
9317 if !conditions.is_empty() && !self.indexes_complete {
9320 return Ok(None);
9321 }
9322 if self.run_refs.len() != 1 {
9323 return Ok(None);
9324 }
9325 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9326 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9327
9328 let overlay_rids: HashSet<u64> = {
9331 let mut s = HashSet::new();
9332 for row in self.memtable.visible_versions(snapshot.epoch) {
9333 s.insert(row.row_id.0);
9334 }
9335 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9336 s.insert(row.row_id.0);
9337 }
9338 s
9339 };
9340
9341 let survivors = if conditions.is_empty() {
9345 None
9346 } else {
9347 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9348 };
9349
9350 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9357 survivors.clone()
9358 } else if let Some(s) = &survivors {
9359 let mut run_set = s.clone();
9360 run_set.remove_many(overlay_rids.iter().copied());
9361 Some(run_set)
9362 } else {
9363 Some(RowIdSet::from_unsorted(
9364 rids.iter()
9365 .map(|&r| r as u64)
9366 .filter(|r| !overlay_rids.contains(r))
9367 .collect(),
9368 ))
9369 };
9370
9371 let overlay_rows = if overlay_rids.is_empty() {
9372 Vec::new()
9373 } else {
9374 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9375 self.overlay_visible_rows(snapshot, bound)
9376 };
9377
9378 let plans = if positions.is_empty() {
9380 Vec::new()
9381 } else {
9382 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9383 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9384 };
9385
9386 let overlay = if overlay_rows.is_empty() {
9388 None
9389 } else {
9390 let filtered =
9391 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9392 if filtered.is_empty() {
9393 None
9394 } else {
9395 Some(self.materialize_overlay(&filtered, &projection))
9396 }
9397 };
9398
9399 let overlay_row_count = overlay
9400 .as_ref()
9401 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9402 .unwrap_or(0);
9403 crate::trace::QueryTrace::record(|t| {
9404 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9405 t.run_count = self.run_refs.len();
9406 t.memtable_rows = self.memtable.len();
9407 t.mutable_run_rows = self.mutable_run.len();
9408 t.overlay_rows = overlay_row_count;
9409 t.conditions_pushed = conditions.len();
9410 t.pages_decoded = plans
9411 .iter()
9412 .map(|p| p.positions.len())
9413 .sum::<usize>()
9414 .min(1);
9415 });
9416
9417 Ok(Some(NativePageCursor::new_with_overlay(
9418 reader, projection, plans, overlay,
9419 )))
9420 }
9421 #[allow(clippy::type_complexity)]
9431 pub fn native_multi_run_cursor(
9432 &self,
9433 snapshot: Snapshot,
9434 projection: Vec<(u16, TypeId)>,
9435 conditions: &[crate::query::Condition],
9436 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9437 use crate::cursor::{MultiRunCursor, RunStream};
9438 use crate::sorted_run::SYS_ROW_ID;
9439 use std::collections::{BinaryHeap, HashMap, HashSet};
9440 if self.ttl.is_some() {
9441 return Ok(None);
9442 }
9443 if !conditions.is_empty() && !self.indexes_complete {
9446 return Ok(None);
9447 }
9448 if self.run_refs.is_empty() {
9449 return Ok(None);
9450 }
9451
9452 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9454 Vec::with_capacity(self.run_refs.len());
9455 for rr in &self.run_refs {
9456 let mut reader = self.open_reader(rr.run_id)?;
9457 let (rids, eps, del) = reader.system_columns_native()?;
9458 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9459 run_meta.push((reader, rids, eps, del, page_rows));
9460 }
9461
9462 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9466 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9467 for i in 0..rids.len() {
9468 let rid = rids[i] as u64;
9469 let e = eps[i] as u64;
9470 if e > snapshot.epoch.0 {
9471 continue;
9472 }
9473 let is_del = del[i] != 0;
9474 best.entry(rid)
9475 .and_modify(|cur| {
9476 if e > cur.0 {
9477 *cur = (e, run_idx, i, is_del);
9478 }
9479 })
9480 .or_insert((e, run_idx, i, is_del));
9481 }
9482 }
9483
9484 let overlay_rids: HashSet<u64> = {
9486 let mut s = HashSet::new();
9487 for row in self.memtable.visible_versions(snapshot.epoch) {
9488 s.insert(row.row_id.0);
9489 }
9490 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9491 s.insert(row.row_id.0);
9492 }
9493 s
9494 };
9495
9496 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9498 None
9499 } else {
9500 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9501 for c in conditions {
9502 sets.push(self.resolve_condition(c, snapshot)?);
9503 }
9504 Some(RowIdSet::intersect_many(sets))
9505 };
9506
9507 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9511 for (rid, (_, run_idx, pos, deleted)) in &best {
9512 if *deleted {
9513 continue;
9514 }
9515 if overlay_rids.contains(rid) {
9516 continue;
9517 }
9518 if let Some(s) = &survivors {
9519 if !s.contains(*rid) {
9520 continue;
9521 }
9522 }
9523 per_run[*run_idx].push((*rid, *pos));
9524 }
9525 for v in per_run.iter_mut() {
9526 v.sort_unstable_by_key(|&(rid, _)| rid);
9527 }
9528
9529 let mut streams = Vec::with_capacity(run_meta.len());
9531 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9532 let mut total = 0usize;
9533 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9534 let mut starts = Vec::with_capacity(page_rows.len());
9535 let mut acc = 0usize;
9536 for &r in &page_rows {
9537 starts.push(acc);
9538 acc += r;
9539 }
9540 let mut survivors_vec: Vec<(u64, usize, usize)> =
9541 Vec::with_capacity(per_run[run_idx].len());
9542 for &(rid, pos) in &per_run[run_idx] {
9543 let page_seq = match starts.partition_point(|&s| s <= pos) {
9544 0 => continue,
9545 p => p - 1,
9546 };
9547 let within = pos - starts[page_seq];
9548 survivors_vec.push((rid, page_seq, within));
9549 }
9550 total += survivors_vec.len();
9551 if let Some(&(rid, _, _)) = survivors_vec.first() {
9552 heap.push(std::cmp::Reverse((rid, run_idx)));
9553 }
9554 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9555 }
9556
9557 let overlay_rows = if overlay_rids.is_empty() {
9559 Vec::new()
9560 } else {
9561 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9562 self.overlay_visible_rows(snapshot, bound)
9563 };
9564 let overlay = if overlay_rows.is_empty() {
9565 None
9566 } else {
9567 let filtered =
9568 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9569 if filtered.is_empty() {
9570 None
9571 } else {
9572 Some(self.materialize_overlay(&filtered, &projection))
9573 }
9574 };
9575
9576 let overlay_row_count = overlay
9577 .as_ref()
9578 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9579 .unwrap_or(0);
9580 crate::trace::QueryTrace::record(|t| {
9581 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9582 t.run_count = self.run_refs.len();
9583 t.memtable_rows = self.memtable.len();
9584 t.mutable_run_rows = self.mutable_run.len();
9585 t.overlay_rows = overlay_row_count;
9586 t.conditions_pushed = conditions.len();
9587 t.survivor_count = Some(total);
9588 });
9589
9590 Ok(Some(MultiRunCursor::new(
9591 streams, projection, heap, total, overlay,
9592 )))
9593 }
9594
9595 fn overlay_materialization_bound<'a>(
9607 conditions: &[crate::query::Condition],
9608 survivors: &'a Option<RowIdSet>,
9609 ) -> Option<&'a RowIdSet> {
9610 use crate::query::Condition;
9611 let has_range = conditions
9612 .iter()
9613 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9614 if has_range {
9615 None
9616 } else {
9617 survivors.as_ref()
9618 }
9619 }
9620
9621 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9633 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9634 let mut fold = |row: Row| {
9635 if let Some(b) = bound {
9636 if !b.contains(row.row_id.0) {
9637 return;
9638 }
9639 }
9640 best.entry(row.row_id.0)
9641 .and_modify(|(be, br)| {
9642 if row.committed_epoch > *be {
9643 *be = row.committed_epoch;
9644 *br = row.clone();
9645 }
9646 })
9647 .or_insert_with(|| (row.committed_epoch, row));
9648 };
9649 for row in self.memtable.visible_versions(snapshot.epoch) {
9650 fold(row);
9651 }
9652 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9653 fold(row);
9654 }
9655 let mut out: Vec<Row> = best
9656 .into_values()
9657 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9658 .collect();
9659 out.sort_by_key(|r| r.row_id);
9660 out
9661 }
9662
9663 fn filter_overlay_rows(
9671 &self,
9672 rows: Vec<Row>,
9673 conditions: &[crate::query::Condition],
9674 survivors: Option<&RowIdSet>,
9675 snapshot: Snapshot,
9676 ) -> Result<Vec<Row>> {
9677 if conditions.is_empty() {
9678 return Ok(rows);
9679 }
9680 use crate::query::Condition;
9681 let all_index_served = !conditions
9685 .iter()
9686 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9687 if all_index_served {
9688 return Ok(rows
9689 .into_iter()
9690 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9691 .collect());
9692 }
9693 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9696 for c in conditions {
9697 let s = match c {
9698 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9699 _ => self.resolve_condition(c, snapshot)?,
9700 };
9701 per_cond_sets.push(s);
9702 }
9703 Ok(rows
9704 .into_iter()
9705 .filter(|row| {
9706 conditions.iter().enumerate().all(|(i, c)| match c {
9707 Condition::Range { column_id, lo, hi } => {
9708 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9709 }
9710 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9711 match row.columns.get(column_id) {
9712 Some(Value::Float64(v)) => {
9713 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9714 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9715 lo_ok && hi_ok
9716 }
9717 _ => false,
9718 }
9719 }
9720 _ => per_cond_sets[i].contains(row.row_id.0),
9721 })
9722 })
9723 .collect())
9724 }
9725
9726 fn materialize_overlay(
9729 &self,
9730 rows: &[Row],
9731 projection: &[(u16, TypeId)],
9732 ) -> Vec<columnar::NativeColumn> {
9733 if projection.is_empty() {
9734 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9735 }
9736 let mut cols = Vec::with_capacity(projection.len());
9737 for (cid, ty) in projection {
9738 let vals: Vec<Value> = rows
9739 .iter()
9740 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9741 .collect();
9742 cols.push(columnar::values_to_native(ty.clone(), &vals));
9743 }
9744 cols
9745 }
9746
9747 fn resolve_survivor_rids(
9752 &self,
9753 conditions: &[crate::query::Condition],
9754 reader: &mut RunReader,
9755 snapshot: Snapshot,
9756 ) -> Result<RowIdSet> {
9757 use crate::query::Condition;
9758 let mut sets: Vec<RowIdSet> = Vec::new();
9759 for c in conditions {
9760 self.validate_condition(c)?;
9761 let s: RowIdSet = match c {
9762 Condition::Pk(key) => {
9763 let lookup = self
9764 .schema
9765 .primary_key()
9766 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9767 .unwrap_or_else(|| key.clone());
9768 self.hot
9769 .get(&lookup)
9770 .map(|r| RowIdSet::one(r.0))
9771 .unwrap_or_else(RowIdSet::empty)
9772 }
9773 Condition::BitmapEq { column_id, value } => {
9774 let lookup = self.index_lookup_key_bytes(*column_id, value);
9775 self.bitmap
9776 .get(column_id)
9777 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9778 .unwrap_or_else(RowIdSet::empty)
9779 }
9780 Condition::BitmapIn { column_id, values } => {
9781 let bm = self.bitmap.get(column_id);
9782 let mut acc = roaring::RoaringBitmap::new();
9783 if let Some(b) = bm {
9784 for v in values {
9785 let lookup = self.index_lookup_key_bytes(*column_id, v);
9786 acc |= b.get(&lookup);
9787 }
9788 }
9789 RowIdSet::from_roaring(acc)
9790 }
9791 Condition::BytesPrefix { column_id, prefix } => {
9792 if let Some(b) = self.bitmap.get(column_id) {
9793 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9794 let mut acc = roaring::RoaringBitmap::new();
9795 for key in b.keys() {
9796 if key.starts_with(&lookup_prefix) {
9797 acc |= b.get(&key);
9798 }
9799 }
9800 RowIdSet::from_roaring(acc)
9801 } else {
9802 RowIdSet::empty()
9803 }
9804 }
9805 Condition::FmContains { column_id, pattern } => self
9806 .fm
9807 .get(column_id)
9808 .map(|f| {
9809 RowIdSet::from_unsorted(
9810 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9811 )
9812 })
9813 .unwrap_or_else(RowIdSet::empty),
9814 Condition::FmContainsAll {
9815 column_id,
9816 patterns,
9817 } => {
9818 if let Some(f) = self.fm.get(column_id) {
9819 let sets: Vec<RowIdSet> = patterns
9820 .iter()
9821 .map(|pat| {
9822 RowIdSet::from_unsorted(
9823 f.locate(pat).into_iter().map(|r| r.0).collect(),
9824 )
9825 })
9826 .collect();
9827 RowIdSet::intersect_many(sets)
9828 } else {
9829 RowIdSet::empty()
9830 }
9831 }
9832 Condition::Ann {
9833 column_id,
9834 query,
9835 k,
9836 } => RowIdSet::from_unsorted(
9837 self.retrieve_filtered(
9838 &crate::query::Retriever::Ann {
9839 column_id: *column_id,
9840 query: query.clone(),
9841 k: *k,
9842 },
9843 snapshot,
9844 None,
9845 None,
9846 None,
9847 None,
9848 )?
9849 .into_iter()
9850 .map(|hit| hit.row_id.0)
9851 .collect(),
9852 ),
9853 Condition::SparseMatch {
9854 column_id,
9855 query,
9856 k,
9857 } => RowIdSet::from_unsorted(
9858 self.retrieve_filtered(
9859 &crate::query::Retriever::Sparse {
9860 column_id: *column_id,
9861 query: query.clone(),
9862 k: *k,
9863 },
9864 snapshot,
9865 None,
9866 None,
9867 None,
9868 None,
9869 )?
9870 .into_iter()
9871 .map(|hit| hit.row_id.0)
9872 .collect(),
9873 ),
9874 Condition::MinHashSimilar {
9875 column_id,
9876 query,
9877 k,
9878 } => match self.minhash.get(column_id) {
9879 Some(index) => {
9880 let candidates = index.candidate_row_ids(query);
9881 let eligible =
9882 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9883 RowIdSet::from_unsorted(
9884 index
9885 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9886 .into_iter()
9887 .map(|(row_id, _)| row_id.0)
9888 .collect(),
9889 )
9890 }
9891 None => RowIdSet::empty(),
9892 },
9893 Condition::Range { column_id, lo, hi } => {
9894 if let Some(li) = self.learned_range.get(column_id) {
9895 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9896 } else {
9897 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9898 }
9899 }
9900 Condition::RangeF64 {
9901 column_id,
9902 lo,
9903 lo_inclusive,
9904 hi,
9905 hi_inclusive,
9906 } => {
9907 if let Some(li) = self.learned_range.get(column_id) {
9908 RowIdSet::from_unsorted(
9909 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9910 .into_iter()
9911 .collect(),
9912 )
9913 } else {
9914 reader.range_row_id_set_f64(
9915 *column_id,
9916 *lo,
9917 *lo_inclusive,
9918 *hi,
9919 *hi_inclusive,
9920 )?
9921 }
9922 }
9923 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9924 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9925 };
9926 sets.push(s);
9927 }
9928 Ok(RowIdSet::intersect_many(sets))
9929 }
9930
9931 pub fn scan_cursor(
9952 &self,
9953 snapshot: Snapshot,
9954 projection: Vec<(u16, TypeId)>,
9955 conditions: &[crate::query::Condition],
9956 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9957 if self.ttl.is_some() {
9958 return Ok(None);
9959 }
9960 if !conditions.is_empty() && !self.indexes_complete {
9966 return Ok(None);
9967 }
9968 if self.run_refs.len() == 1 {
9969 Ok(self
9970 .native_page_cursor(snapshot, projection, conditions)?
9971 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9972 } else {
9973 Ok(self
9974 .native_multi_run_cursor(snapshot, projection, conditions)?
9975 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9976 }
9977 }
9978
9979 pub fn aggregate_native(
9993 &self,
9994 snapshot: Snapshot,
9995 column: Option<u16>,
9996 conditions: &[crate::query::Condition],
9997 agg: NativeAgg,
9998 ) -> Result<Option<NativeAggResult>> {
9999 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
10000 }
10001
10002 pub fn aggregate_native_with_control(
10003 &self,
10004 snapshot: Snapshot,
10005 column: Option<u16>,
10006 conditions: &[crate::query::Condition],
10007 agg: NativeAgg,
10008 control: &crate::ExecutionControl,
10009 ) -> Result<Option<NativeAggResult>> {
10010 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
10011 }
10012
10013 fn aggregate_native_inner(
10014 &self,
10015 snapshot: Snapshot,
10016 column: Option<u16>,
10017 conditions: &[crate::query::Condition],
10018 agg: NativeAgg,
10019 control: Option<&crate::ExecutionControl>,
10020 ) -> Result<Option<NativeAggResult>> {
10021 execution_checkpoint(control, 0)?;
10022 if self.ttl.is_some() {
10023 return Ok(None);
10024 }
10025 if self.run_refs.len() == 1 && conditions.is_empty() {
10027 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
10028 return Ok(Some(res));
10029 }
10030 }
10031 if matches!(agg, NativeAgg::Count) && column.is_none() {
10035 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
10036 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
10037 }
10038 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10039 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
10040 }
10041 let cid = match column {
10044 Some(c) => c,
10045 None => return Ok(None),
10046 };
10047 let ty = self.column_type(cid);
10048 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
10049 execution_checkpoint(control, 0)?;
10050 return match ty {
10051 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10052 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
10053 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10054 }
10055 TypeId::Float64 => {
10056 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
10057 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10058 }
10059 _ => Ok(None),
10060 };
10061 }
10062 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
10064 execution_checkpoint(control, 0)?;
10065 match ty {
10066 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10067 let mut count = 0u64;
10068 let mut sum = 0i128;
10069 let mut mn = i64::MAX;
10070 let mut mx = i64::MIN;
10071 for row in &rows {
10072 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
10073 count += 1;
10074 sum += i128::from(*v);
10075 mn = mn.min(*v);
10076 mx = mx.max(*v);
10077 }
10078 }
10079 Ok(Some(pack_int(agg, count, sum, mn, mx)))
10080 }
10081 TypeId::Float64 => {
10082 let mut count = 0u64;
10083 let mut sum = 0.0f64;
10084 let mut mn = f64::INFINITY;
10085 let mut mx = f64::NEG_INFINITY;
10086 for row in &rows {
10087 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
10088 count += 1;
10089 sum += *v;
10090 mn = mn.min(*v);
10091 mx = mx.max(*v);
10092 }
10093 }
10094 Ok(Some(pack_float(agg, count, sum, mn, mx)))
10095 }
10096 _ => Ok(None),
10097 }
10098 }
10099
10100 fn visible_rows_filtered(
10102 &self,
10103 snapshot: Snapshot,
10104 conditions: &[crate::query::Condition],
10105 control: Option<&crate::ExecutionControl>,
10106 ) -> Result<Vec<Row>> {
10107 let rows = if let Some(control) = control {
10108 self.visible_rows_controlled(snapshot, control)?
10109 } else {
10110 self.visible_rows(snapshot)?
10111 };
10112 if conditions.is_empty() {
10113 return Ok(rows);
10114 }
10115 Ok(rows
10116 .into_iter()
10117 .filter(|row| {
10118 conditions
10119 .iter()
10120 .all(|cond| condition_matches_row(cond, row, &self.schema))
10121 })
10122 .collect())
10123 }
10124
10125 fn aggregate_from_stats(
10133 &self,
10134 snapshot: Snapshot,
10135 column: Option<u16>,
10136 agg: NativeAgg,
10137 ) -> Result<Option<NativeAggResult>> {
10138 let cid = match (agg, column) {
10139 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
10140 _ => return Ok(None), };
10142 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
10143 return Ok(None);
10144 };
10145 let Some(cs) = stats.get(&cid) else {
10146 return Ok(None);
10147 };
10148 match agg {
10149 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
10151 self.live_count.saturating_sub(cs.null_count),
10152 ))),
10153 NativeAgg::Min | NativeAgg::Max => {
10154 let bound = if agg == NativeAgg::Min {
10155 &cs.min
10156 } else {
10157 &cs.max
10158 };
10159 match bound {
10160 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
10161 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
10162 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
10167 None => Ok(None),
10168 }
10169 }
10170 _ => Ok(None),
10171 }
10172 }
10173
10174 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
10183 if self.ttl.is_some() {
10184 return Ok(None);
10185 }
10186 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
10187 return Ok(None);
10188 }
10189 self.ensure_indexes_complete()?;
10192 let reader = self.open_reader(self.run_refs[0].run_id)?;
10193 if self.live_count != reader.row_count() as u64 {
10194 return Ok(None);
10195 }
10196 let Some(bm) = self.bitmap.get(&column_id) else {
10197 return Ok(None); };
10199 let mut distinct = bm.value_count() as u64;
10200 if !bm.get(&Value::Null.encode_key()).is_empty() {
10203 distinct = distinct.saturating_sub(1);
10204 }
10205 Ok(Some(distinct))
10206 }
10207
10208 pub fn aggregate_incremental(
10220 &mut self,
10221 cache_key: u64,
10222 conditions: &[crate::query::Condition],
10223 column: Option<u16>,
10224 agg: NativeAgg,
10225 ) -> Result<IncrementalAggResult> {
10226 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
10227 }
10228
10229 pub fn aggregate_incremental_with_control(
10230 &mut self,
10231 cache_key: u64,
10232 conditions: &[crate::query::Condition],
10233 column: Option<u16>,
10234 agg: NativeAgg,
10235 control: &crate::ExecutionControl,
10236 ) -> Result<IncrementalAggResult> {
10237 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
10238 }
10239
10240 fn aggregate_incremental_inner(
10241 &mut self,
10242 cache_key: u64,
10243 conditions: &[crate::query::Condition],
10244 column: Option<u16>,
10245 agg: NativeAgg,
10246 control: Option<&crate::ExecutionControl>,
10247 ) -> Result<IncrementalAggResult> {
10248 execution_checkpoint(control, 0)?;
10249 let snap = self.snapshot();
10250 let cur_wm = self.allocator.current().0;
10251 let cur_epoch = snap.epoch.0;
10252 let incremental_ok = self.ttl.is_none()
10259 && !self.had_deletes
10260 && self.memtable.is_empty()
10261 && self.mutable_run.is_empty();
10262
10263 if incremental_ok {
10266 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10267 if cached.epoch == cur_epoch {
10268 return Ok(IncrementalAggResult {
10269 state: cached.state,
10270 incremental: true,
10271 delta_rows: 0,
10272 });
10273 }
10274 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10275 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10276 let mut delta_rids = Vec::with_capacity(delta_len);
10277 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10278 execution_checkpoint(control, index)?;
10279 delta_rids.push(row_id);
10280 }
10281 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10282 execution_checkpoint(control, 0)?;
10283 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10284 let delta_state = agg_state_from_rows(
10285 &delta_rows,
10286 conditions,
10287 &index_sets,
10288 column,
10289 agg,
10290 &self.schema,
10291 control,
10292 )?;
10293 let merged = cached.state.merge(delta_state);
10294 let delta_n = delta_rids.len() as u64;
10295 Arc::make_mut(&mut self.agg_cache).insert(
10296 cache_key,
10297 CachedAgg {
10298 state: merged.clone(),
10299 watermark: cur_wm,
10300 epoch: cur_epoch,
10301 },
10302 );
10303 return Ok(IncrementalAggResult {
10304 state: merged,
10305 incremental: true,
10306 delta_rows: delta_n,
10307 });
10308 }
10309 }
10310 }
10311
10312 let cursor_ok =
10317 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10318 let state = if cursor_ok && agg != NativeAgg::Avg {
10319 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10320 Some(result) => {
10321 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10322 }
10323 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10324 }
10325 } else {
10326 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10327 };
10328 if incremental_ok {
10330 Arc::make_mut(&mut self.agg_cache).insert(
10331 cache_key,
10332 CachedAgg {
10333 state: state.clone(),
10334 watermark: cur_wm,
10335 epoch: cur_epoch,
10336 },
10337 );
10338 }
10339 Ok(IncrementalAggResult {
10340 state,
10341 incremental: false,
10342 delta_rows: 0,
10343 })
10344 }
10345
10346 fn agg_state_full_scan(
10349 &self,
10350 conditions: &[crate::query::Condition],
10351 column: Option<u16>,
10352 agg: NativeAgg,
10353 snap: Snapshot,
10354 control: Option<&crate::ExecutionControl>,
10355 ) -> Result<AggState> {
10356 execution_checkpoint(control, 0)?;
10357 let rows = self.visible_rows(snap)?;
10358 execution_checkpoint(control, 0)?;
10359 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10360 agg_state_from_rows(
10361 &rows,
10362 conditions,
10363 &index_sets,
10364 column,
10365 agg,
10366 &self.schema,
10367 control,
10368 )
10369 }
10370
10371 fn resolve_index_conditions(
10374 &self,
10375 conditions: &[crate::query::Condition],
10376 snapshot: Snapshot,
10377 ) -> Result<Vec<RowIdSet>> {
10378 use crate::query::Condition;
10379 let mut sets = Vec::new();
10380 for c in conditions {
10381 if matches!(
10382 c,
10383 Condition::Ann { .. }
10384 | Condition::SparseMatch { .. }
10385 | Condition::MinHashSimilar { .. }
10386 ) {
10387 sets.push(self.resolve_condition(c, snapshot)?);
10388 }
10389 }
10390 Ok(sets)
10391 }
10392
10393 fn column_type(&self, cid: u16) -> TypeId {
10394 self.schema
10395 .columns
10396 .iter()
10397 .find(|c| c.id == cid)
10398 .map(|c| c.ty.clone())
10399 .unwrap_or(TypeId::Bytes)
10400 }
10401
10402 pub fn approx_aggregate(
10411 &mut self,
10412 conditions: &[crate::query::Condition],
10413 column: Option<u16>,
10414 agg: ApproxAgg,
10415 z: f64,
10416 ) -> Result<Option<ApproxResult>> {
10417 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10418 }
10419
10420 pub fn approx_aggregate_with_candidate_authorization(
10423 &mut self,
10424 conditions: &[crate::query::Condition],
10425 column: Option<u16>,
10426 agg: ApproxAgg,
10427 z: f64,
10428 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10429 ) -> Result<Option<ApproxResult>> {
10430 use crate::query::Condition;
10431 self.ensure_reservoir_complete()?;
10432 let mut snapshot = self.snapshot();
10437 let n_pop = self.count();
10438 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10439 if sample_rids.is_empty() {
10440 return Ok(None);
10441 }
10442 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10444 if live_sample.is_empty() {
10445 snapshot = Snapshot::unbounded();
10446 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10447 }
10448 let s = live_sample.len();
10449 if s == 0 {
10450 return Ok(None);
10451 }
10452 let authorized = authorization
10453 .map(|authorization| {
10454 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10455 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10456 })
10457 .transpose()?;
10458
10459 let mut index_sets: Vec<RowIdSet> = Vec::new();
10462 for c in conditions {
10463 if matches!(
10464 c,
10465 Condition::Ann { .. }
10466 | Condition::SparseMatch { .. }
10467 | Condition::MinHashSimilar { .. }
10468 ) {
10469 index_sets.push(self.resolve_condition(c, snapshot)?);
10470 }
10471 }
10472
10473 let cid = match (agg, column) {
10475 (ApproxAgg::Count, _) => None,
10476 (_, Some(c)) => Some(c),
10477 _ => return Ok(None),
10478 };
10479 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10480 for r in &live_sample {
10481 if authorized
10482 .as_ref()
10483 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10484 {
10485 continue;
10486 }
10487 if !conditions
10489 .iter()
10490 .all(|c| condition_matches_row(c, r, &self.schema))
10491 {
10492 continue;
10493 }
10494 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10496 continue;
10497 }
10498 if let Some(cid) = cid {
10499 let mut cells = r
10500 .columns
10501 .get(&cid)
10502 .cloned()
10503 .map(|value| vec![(cid, value)])
10504 .unwrap_or_default();
10505 if let Some(authorization) = authorization {
10506 authorization.security.apply_masks_to_cells(
10507 authorization.table,
10508 &mut cells,
10509 authorization.principal,
10510 );
10511 }
10512 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10513 passing_vals.push(v);
10514 } } else {
10516 passing_vals.push(0.0); }
10518 }
10519 let m = passing_vals.len();
10520
10521 let (point, half) = match agg {
10522 ApproxAgg::Count => {
10523 let p = m as f64 / s as f64;
10525 let point = n_pop as f64 * p;
10526 let var = if s > 1 {
10527 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10528 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10529 } else {
10530 0.0
10531 };
10532 (point, z * var.sqrt())
10533 }
10534 ApproxAgg::Sum => {
10535 let y: Vec<f64> = live_sample
10537 .iter()
10538 .map(|r| {
10539 let passes_row = authorized
10540 .as_ref()
10541 .is_none_or(|authorized| authorized.contains(&r.row_id))
10542 && conditions
10543 .iter()
10544 .all(|c| condition_matches_row(c, r, &self.schema))
10545 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10546 if passes_row {
10547 cid.and_then(|cid| {
10548 let mut cells = r
10549 .columns
10550 .get(&cid)
10551 .cloned()
10552 .map(|value| vec![(cid, value)])
10553 .unwrap_or_default();
10554 if let Some(authorization) = authorization {
10555 authorization.security.apply_masks_to_cells(
10556 authorization.table,
10557 &mut cells,
10558 authorization.principal,
10559 );
10560 }
10561 as_f64(cells.first().map(|(_, value)| value))
10562 })
10563 .unwrap_or(0.0)
10564 } else {
10565 0.0
10566 }
10567 })
10568 .collect();
10569 let mean_y = y.iter().sum::<f64>() / s as f64;
10570 let point = n_pop as f64 * mean_y;
10571 let var = if s > 1 {
10572 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10573 let var_y = ss / (s - 1) as f64;
10574 n_pop as f64 * n_pop as f64 * var_y / s as f64
10575 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10576 } else {
10577 0.0
10578 };
10579 (point, z * var.sqrt())
10580 }
10581 ApproxAgg::Avg => {
10582 if m == 0 {
10583 return Ok(Some(ApproxResult {
10584 point: 0.0,
10585 ci_low: 0.0,
10586 ci_high: 0.0,
10587 n_population: n_pop,
10588 n_sample_live: s,
10589 n_passing: 0,
10590 }));
10591 }
10592 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10593 let half = if m > 1 {
10594 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10595 let sd = (ss / (m - 1) as f64).sqrt();
10596 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10597 z * sd / (m as f64).sqrt() * fpc.sqrt()
10598 } else {
10599 0.0
10600 };
10601 (mean, half)
10602 }
10603 };
10604
10605 Ok(Some(ApproxResult {
10606 point,
10607 ci_low: point - half,
10608 ci_high: point + half,
10609 n_population: n_pop,
10610 n_sample_live: s,
10611 n_passing: m,
10612 }))
10613 }
10614
10615 pub fn exact_column_stats(
10623 &self,
10624 _snapshot: Snapshot,
10625 projection: &[u16],
10626 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10627 if self.ttl.is_some()
10628 || !(self.memtable.is_empty()
10629 && self.mutable_run.is_empty()
10630 && self.run_refs.len() == 1)
10631 {
10632 return Ok(None);
10633 }
10634 let reader = self.open_reader(self.run_refs[0].run_id)?;
10635 if self.live_count != reader.row_count() as u64 {
10636 return Ok(None);
10637 }
10638 let mut out = HashMap::new();
10639 for &cid in projection {
10640 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10641 Some(c) => c,
10642 None => continue,
10643 };
10644 let Some(stats) = reader.column_page_stats(cid) else {
10646 out.insert(
10647 cid,
10648 ColumnStat {
10649 min: None,
10650 max: None,
10651 null_count: self.live_count,
10652 },
10653 );
10654 continue;
10655 };
10656 let stat = match cdef.ty {
10657 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10658 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10659 min: mn.map(Value::Int64),
10660 max: mx.map(Value::Int64),
10661 null_count: n,
10662 })
10663 }
10664 TypeId::Float64 => {
10665 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10666 min: mn.map(Value::Float64),
10667 max: mx.map(Value::Float64),
10668 null_count: n,
10669 })
10670 }
10671 _ => None,
10672 };
10673 if let Some(s) = stat {
10674 out.insert(cid, s);
10675 }
10676 }
10677 Ok(Some(out))
10678 }
10679
10680 pub fn dir(&self) -> &Path {
10681 &self.dir
10682 }
10683
10684 pub fn schema(&self) -> &Schema {
10685 &self.schema
10686 }
10687
10688 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
10689 self.ann.get(&column_id)
10690 }
10691
10692 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
10693 self.ann.get_mut(&column_id)
10694 }
10695
10696 pub(crate) fn set_catalog_name(&mut self, name: String) {
10697 self.name = name;
10698 }
10699
10700 pub(crate) fn prepare_alter_column(
10701 &mut self,
10702 column_name: &str,
10703 change: &AlterColumn,
10704 ) -> Result<(ColumnDef, Option<Schema>)> {
10705 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10706 return Err(MongrelError::InvalidArgument(
10707 "ALTER COLUMN requires committing staged writes first".into(),
10708 ));
10709 }
10710 let old = self
10711 .schema
10712 .columns
10713 .iter()
10714 .find(|c| c.name == column_name)
10715 .cloned()
10716 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10717 let mut next = old.clone();
10718
10719 if let Some(name) = &change.name {
10720 let trimmed = name.trim();
10721 if trimmed.is_empty() {
10722 return Err(MongrelError::InvalidArgument(
10723 "ALTER COLUMN name must not be empty".into(),
10724 ));
10725 }
10726 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10727 return Err(MongrelError::Schema(format!(
10728 "column {trimmed} already exists"
10729 )));
10730 }
10731 next.name = trimmed.to_string();
10732 }
10733
10734 if let Some(ty) = &change.ty {
10735 next.ty = ty.clone();
10736 }
10737 if let Some(flags) = change.flags {
10738 validate_alter_column_flags(old.flags, flags)?;
10739 next.flags = flags;
10740 }
10741
10742 if let Some(default_change) = &change.default_value {
10743 next.default_value = default_change.clone();
10744 }
10745 if let Some(source_change) = &change.embedding_source {
10746 next.embedding_source = source_change.clone();
10747 }
10748
10749 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10750 if old.flags.contains(ColumnFlags::NULLABLE)
10751 && !next.flags.contains(ColumnFlags::NULLABLE)
10752 && self.column_has_nulls(old.id)?
10753 {
10754 return Err(MongrelError::InvalidArgument(format!(
10755 "column '{}' contains NULL values",
10756 old.name
10757 )));
10758 }
10759 if next == old {
10760 return Ok((next, None));
10761 }
10762 let mut schema = self.schema.clone();
10763 let index = schema
10764 .columns
10765 .iter()
10766 .position(|column| column.id == next.id)
10767 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10768 schema.columns[index] = next.clone();
10769 schema.schema_id = schema
10770 .schema_id
10771 .checked_add(1)
10772 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10773 schema.validate_auto_increment()?;
10774 schema.validate_defaults()?;
10775 Ok((next, Some(schema)))
10776 }
10777
10778 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10779 self.schema = schema;
10780 self.auto_inc = resolve_auto_inc(&self.schema);
10781 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10782 self.clear_result_cache();
10783 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10784 }
10785
10786 pub(crate) fn publish_index_schema_change(
10790 &mut self,
10791 schema: Schema,
10792 artifact: SecondaryIndexArtifact,
10793 ) -> Result<()> {
10794 self.apply_altered_schema_prepared(schema);
10795 self.prune_index_maps_to_schema();
10796 match artifact {
10797 SecondaryIndexArtifact::Bitmap(column_id, index) => {
10798 self.bitmap.insert(column_id, index);
10799 }
10800 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
10801 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
10802 }
10803 SecondaryIndexArtifact::Fm(column_id, index) => {
10804 self.fm.insert(column_id, *index);
10805 }
10806 SecondaryIndexArtifact::Ann(column_id, index) => {
10807 self.ann.insert(column_id, *index);
10808 }
10809 SecondaryIndexArtifact::Sparse(column_id, index) => {
10810 self.sparse.insert(column_id, index);
10811 }
10812 SecondaryIndexArtifact::MinHash(column_id, index) => {
10813 self.minhash.insert(column_id, index);
10814 }
10815 }
10816 self.indexes_complete = true;
10817 self.seal_generations();
10818 let view = Arc::new(self.capture_read_generation());
10819 self.published.store(Arc::clone(&view));
10820 checkpoint_current_schema(self)?;
10821 self.invalidate_index_checkpoint();
10824 Ok(())
10825 }
10826
10827 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
10833 self.apply_altered_schema_prepared(schema);
10834 self.prune_index_maps_to_schema();
10835 self.indexes_complete = true;
10836 self.seal_generations();
10837 let view = Arc::new(self.capture_read_generation());
10838 self.published.store(Arc::clone(&view));
10839 checkpoint_current_schema(self)?;
10840 self.invalidate_index_checkpoint();
10842 Ok(())
10843 }
10844
10845 fn prune_index_maps_to_schema(&mut self) {
10846 let keep_bitmap: std::collections::HashSet<u16> = self
10847 .schema
10848 .indexes
10849 .iter()
10850 .filter(|index| index.kind == IndexKind::Bitmap)
10851 .map(|index| index.column_id)
10852 .collect();
10853 let keep_ann: std::collections::HashSet<u16> = self
10854 .schema
10855 .indexes
10856 .iter()
10857 .filter(|index| index.kind == IndexKind::Ann)
10858 .map(|index| index.column_id)
10859 .collect();
10860 let keep_fm: std::collections::HashSet<u16> = self
10861 .schema
10862 .indexes
10863 .iter()
10864 .filter(|index| index.kind == IndexKind::FmIndex)
10865 .map(|index| index.column_id)
10866 .collect();
10867 let keep_sparse: std::collections::HashSet<u16> = self
10868 .schema
10869 .indexes
10870 .iter()
10871 .filter(|index| index.kind == IndexKind::Sparse)
10872 .map(|index| index.column_id)
10873 .collect();
10874 let keep_minhash: std::collections::HashSet<u16> = self
10875 .schema
10876 .indexes
10877 .iter()
10878 .filter(|index| index.kind == IndexKind::MinHash)
10879 .map(|index| index.column_id)
10880 .collect();
10881 let keep_learned: std::collections::HashSet<u16> = self
10882 .schema
10883 .indexes
10884 .iter()
10885 .filter(|index| index.kind == IndexKind::LearnedRange)
10886 .map(|index| index.column_id)
10887 .collect();
10888 self.bitmap
10889 .retain(|column_id, _| keep_bitmap.contains(column_id));
10890 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
10891 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
10892 self.sparse
10893 .retain(|column_id, _| keep_sparse.contains(column_id));
10894 self.minhash
10895 .retain(|column_id, _| keep_minhash.contains(column_id));
10896 {
10897 let learned = Arc::make_mut(&mut self.learned_range);
10898 learned.retain(|column_id, _| keep_learned.contains(column_id));
10899 }
10900 }
10901
10902 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10903 checkpoint_current_schema(self)
10904 }
10905
10906 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10907 self.ensure_writable()?;
10908 let previous_schema = self.schema.clone();
10909 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10910 if let Some(schema) = schema {
10911 self.apply_altered_schema_prepared(schema);
10912 self.checkpoint_standalone_schema_change(previous_schema)?;
10913 }
10914 Ok(column)
10915 }
10916
10917 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10918 if self.live_count == 0 {
10919 return Ok(false);
10920 }
10921 let snap = self.snapshot();
10922 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10923 Ok(columns
10924 .first()
10925 .map(|(_, col)| col.null_count(col.len()) != 0)
10926 .unwrap_or(true))
10927 }
10928
10929 fn has_stored_versions(&self) -> bool {
10930 !self.memtable.is_empty()
10931 || !self.mutable_run.is_empty()
10932 || self.run_refs.iter().any(|r| r.row_count > 0)
10933 || !self.retiring.is_empty()
10934 }
10935
10936 pub fn add_column(
10941 &mut self,
10942 name: &str,
10943 ty: TypeId,
10944 flags: ColumnFlags,
10945 default_value: Option<crate::schema::DefaultExpr>,
10946 ) -> Result<u16> {
10947 self.add_column_with_id(name, ty, flags, default_value, None)
10948 }
10949
10950 pub fn add_column_with_id(
10951 &mut self,
10952 name: &str,
10953 ty: TypeId,
10954 flags: ColumnFlags,
10955 default_value: Option<crate::schema::DefaultExpr>,
10956 requested_id: Option<u16>,
10957 ) -> Result<u16> {
10958 self.ensure_writable()?;
10959 let previous_schema = self.schema.clone();
10960 let (column, schema) =
10961 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
10962 self.apply_altered_schema_prepared(schema);
10963 self.checkpoint_standalone_schema_change(previous_schema)?;
10964 Ok(column.id)
10965 }
10966
10967 pub(crate) fn prepare_add_column(
10968 &mut self,
10969 name: &str,
10970 ty: TypeId,
10971 flags: ColumnFlags,
10972 default_value: Option<crate::schema::DefaultExpr>,
10973 requested_id: Option<u16>,
10974 ) -> Result<(ColumnDef, Schema)> {
10975 self.ensure_writable()?;
10976 if self.schema.columns.iter().any(|c| c.name == name) {
10977 return Err(MongrelError::Schema(format!(
10978 "column {name} already exists"
10979 )));
10980 }
10981 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10982 if self.schema.columns.iter().any(|c| c.id == id) {
10983 return Err(MongrelError::Schema(format!(
10984 "column id {id} already exists"
10985 )));
10986 }
10987 id
10988 } else {
10989 self.schema
10990 .columns
10991 .iter()
10992 .map(|c| c.id)
10993 .max()
10994 .unwrap_or(0)
10995 .checked_add(1)
10996 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10997 };
10998 let column = ColumnDef {
10999 id,
11000 name: name.to_string(),
11001 ty,
11002 flags,
11003 default_value,
11004 embedding_source: None,
11005 };
11006 let mut next_schema = self.schema.clone();
11007 next_schema.columns.push(column.clone());
11008 next_schema.schema_id = next_schema
11009 .schema_id
11010 .checked_add(1)
11011 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11012 next_schema.validate_auto_increment()?;
11013 next_schema.validate_defaults()?;
11014 Ok((column, next_schema))
11015 }
11016
11017 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
11026 self.ensure_writable()?;
11027 let cid = self
11028 .schema
11029 .columns
11030 .iter()
11031 .find(|c| c.name == column_name)
11032 .map(|c| c.id)
11033 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11034 let ty = self
11035 .schema
11036 .columns
11037 .iter()
11038 .find(|c| c.id == cid)
11039 .map(|c| c.ty.clone())
11040 .unwrap_or(TypeId::Int64);
11041 if !matches!(
11042 ty,
11043 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
11044 ) {
11045 return Err(MongrelError::Schema(format!(
11046 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
11047 )));
11048 }
11049 if self
11050 .schema
11051 .indexes
11052 .iter()
11053 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
11054 {
11055 return Ok(()); }
11057 let previous_schema = self.schema.clone();
11058 let previous_learned_range = Arc::clone(&self.learned_range);
11059 let mut next_schema = previous_schema.clone();
11060 next_schema.indexes.push(IndexDef {
11061 name: format!("{}_learned_range", column_name),
11062 column_id: cid,
11063 kind: IndexKind::LearnedRange,
11064 predicate: None,
11065 options: Default::default(),
11066 });
11067 next_schema.schema_id = next_schema
11068 .schema_id
11069 .checked_add(1)
11070 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11071 self.apply_altered_schema_prepared(next_schema);
11072 if let Err(error) = self.build_learned_ranges() {
11073 self.apply_altered_schema_prepared(previous_schema);
11074 self.learned_range = previous_learned_range;
11075 return Err(error);
11076 }
11077 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
11078 if !matches!(
11079 &error,
11080 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
11081 ) {
11082 self.learned_range = previous_learned_range;
11083 }
11084 return Err(error);
11085 }
11086 Ok(())
11087 }
11088
11089 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
11090 let mut schema_published = false;
11091 let schema_result = match self._root_guard.as_deref() {
11092 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
11093 schema_published = true;
11094 }),
11095 None => write_schema_with_after(&self.dir, &self.schema, || {
11096 schema_published = true;
11097 }),
11098 };
11099 if schema_result.is_err() && !schema_published {
11100 self.apply_altered_schema_prepared(previous_schema);
11101 return schema_result;
11102 }
11103
11104 let manifest_result = self.persist_manifest(self.current_epoch());
11105 match (schema_result, manifest_result) {
11106 (_, Ok(())) => Ok(()),
11107 (Ok(()), Err(error)) => {
11108 self.poison_after_maintenance_publish_failure();
11109 Err(MongrelError::DurableCommit {
11110 epoch: self.current_epoch().0,
11111 message: format!(
11112 "schema is durable but matching manifest publication failed: {error}"
11113 ),
11114 })
11115 }
11116 (Err(schema_error), Err(manifest_error)) => {
11117 self.poison_after_maintenance_publish_failure();
11118 Err(MongrelError::CommitOutcomeUnknown {
11119 epoch: self.current_epoch().0,
11120 message: format!(
11121 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11122 ),
11123 })
11124 }
11125 }
11126 }
11127
11128 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
11131 self.sync_byte_threshold = threshold;
11132 if let WalSink::Private(w) = &mut self.wal {
11133 w.set_sync_byte_threshold(threshold);
11134 }
11135 }
11136
11137 pub fn page_cache_flush(&self) {
11141 self.page_cache.flush_to_disk();
11142 }
11143
11144 pub fn page_cache_len(&self) -> usize {
11146 self.page_cache.len()
11147 }
11148
11149 pub fn decoded_cache_len(&self) -> usize {
11152 self.decoded_cache.len()
11153 }
11154
11155 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
11158 self.memtable.drain_sorted()
11159 }
11160
11161 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
11162 self.runs_dir().join(format!("r-{run_id}.sr"))
11163 }
11164
11165 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
11166 match self.runs_root.as_deref() {
11167 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
11168 None => Ok(None),
11169 }
11170 }
11171
11172 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
11173 match self.runs_root.as_deref() {
11174 Some(root) => Ok(Some(root.create_regular_new(name)?)),
11175 None => Ok(None),
11176 }
11177 }
11178
11179 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
11180 match self.runs_root.as_deref() {
11181 Some(root) => match root.remove_file(name) {
11182 Ok(()) => Ok(()),
11183 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11184 Err(error) => Err(error.into()),
11185 },
11186 None => match std::fs::remove_file(self.runs_dir().join(name)) {
11187 Ok(()) => Ok(()),
11188 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
11189 Err(error) => Err(error.into()),
11190 },
11191 }
11192 }
11193
11194 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
11195 match self.runs_root.as_deref() {
11196 Some(root) => root
11197 .rename_file_new(source, destination)
11198 .map_err(Into::into),
11199 None => crate::durable_file::rename(
11200 &self.runs_dir().join(source),
11201 &self.runs_dir().join(destination),
11202 )
11203 .map_err(Into::into),
11204 }
11205 }
11206
11207 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11208 self.run_refs.iter().map(|run| run.run_id)
11209 }
11210
11211 pub(crate) fn table_dir(&self) -> &Path {
11212 &self.dir
11213 }
11214
11215 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
11216 &self.schema
11217 }
11218
11219 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
11220 let id = self.next_run_id;
11221 self.next_run_id = self
11222 .next_run_id
11223 .checked_add(1)
11224 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
11225 Ok(id)
11226 }
11227
11228 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
11229 self.run_refs.push(run_ref);
11230 }
11231
11232 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
11242 self.retiring.push(crate::manifest::RetiredRun {
11243 run_id,
11244 retire_epoch,
11245 });
11246 }
11247
11248 pub(crate) fn reap_retiring(
11252 &mut self,
11253 min_active: Epoch,
11254 backup_pinned: &std::collections::HashSet<u128>,
11255 ) -> Result<usize> {
11256 if self.retiring.is_empty() {
11257 return Ok(0);
11258 }
11259 let mut reaped = 0;
11260 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11261 for r in std::mem::take(&mut self.retiring) {
11267 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11268 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11269 reaped += 1;
11270 } else {
11271 kept.push(r);
11272 }
11273 }
11274 self.retiring = kept;
11275 if reaped > 0 {
11276 self.persist_manifest(self.current_epoch())?;
11277 }
11278 Ok(reaped)
11279 }
11280
11281 pub(crate) fn has_reapable_retiring(
11282 &self,
11283 min_active: Epoch,
11284 backup_pinned: &std::collections::HashSet<u128>,
11285 ) -> bool {
11286 self.retiring
11287 .iter()
11288 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11289 }
11290
11291 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11292 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11293 return false;
11294 }
11295 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11296 self.run_refs.push(run_ref);
11297 self.indexes_complete = false;
11298 true
11299 }
11300
11301 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11302 self.kek.as_ref()
11303 }
11304
11305 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11306 let mut reader = match self.runs_root.as_deref() {
11307 Some(root) => RunReader::open_file_with_cache(
11308 root.open_regular(format!("r-{run_id}.sr"))?,
11309 self.schema.clone(),
11310 self.kek.clone(),
11311 Some(self.page_cache.clone()),
11312 Some(self.decoded_cache.clone()),
11313 self.table_id,
11314 Some(&self.verified_runs),
11315 None,
11316 )?,
11317 None => RunReader::open_with_cache(
11318 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11319 self.schema.clone(),
11320 self.kek.clone(),
11321 Some(self.page_cache.clone()),
11322 Some(self.decoded_cache.clone()),
11323 self.table_id,
11324 Some(&self.verified_runs),
11325 )?,
11326 };
11327 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11331 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11332 }
11333 Ok(reader)
11334 }
11335
11336 pub(crate) fn run_refs(&self) -> &[RunRef] {
11337 &self.run_refs
11338 }
11339
11340 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11341 self.retiring.iter().map(|run| run.run_id)
11342 }
11343
11344 pub(crate) fn runs_dir(&self) -> PathBuf {
11345 self.runs_root
11346 .as_deref()
11347 .and_then(|root| root.io_path().ok())
11348 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11349 }
11350
11351 pub(crate) fn wal_dir(&self) -> PathBuf {
11352 self.dir.join(WAL_DIR)
11353 }
11354
11355 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11356 self.run_refs = refs;
11357 }
11358
11359 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11360 self.compaction_zstd_level
11361 }
11362
11363 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11364 self.kek.clone()
11365 }
11366
11367 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11371 self.kek.as_ref().map(|k| k.derive_idx_key())
11372 }
11373
11374 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11378 self.kek.as_ref().map(|k| *k.derive_meta_key())
11379 }
11380
11381 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11384 self.column_keys
11385 .iter()
11386 .map(|(&id, &(_, scheme))| (id, scheme))
11387 .collect()
11388 }
11389
11390 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11395 self.tokenize_value_enc(column_id, v)
11396 }
11397
11398 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11399 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11400 let (key, scheme) = self.column_keys.get(&column_id)?;
11401 let token: Vec<u8> = match (*scheme, v) {
11402 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11403 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11404 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11405 _ => hmac_token(key, &v.encode_key()).to_vec(),
11406 };
11407 Some(Value::Bytes(token))
11408 }
11409
11410 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11412 self.index_lookup_key_bytes(column_id, &v.encode_key())
11413 }
11414
11415 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11418 {
11419 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11420 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11421 if *scheme == SCHEME_HMAC_EQ {
11422 return hmac_token(key, encoded).to_vec();
11423 }
11424 }
11425 }
11426 let _ = column_id;
11427 encoded.to_vec()
11428 }
11429}
11430
11431fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11432 let columnar::NativeColumn::Int64 { data, validity } = col else {
11433 return false;
11434 };
11435 if data.len() < n || !columnar::all_non_null(validity, n) {
11436 return false;
11437 }
11438 data.iter()
11439 .take(n)
11440 .zip(data.iter().skip(1))
11441 .all(|(a, b)| a < b)
11442}
11443
11444#[derive(Debug, Clone)]
11448pub struct ColumnStat {
11449 pub min: Option<Value>,
11450 pub max: Option<Value>,
11451 pub null_count: u64,
11452}
11453
11454#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11456pub enum NativeAgg {
11457 Count,
11458 Sum,
11459 Min,
11460 Max,
11461 Avg,
11462}
11463
11464#[derive(Debug, Clone, PartialEq)]
11466pub enum NativeAggResult {
11467 Count(u64),
11468 Int(i64),
11469 Float(f64),
11470 Null,
11472}
11473
11474#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11476pub enum ApproxAgg {
11477 Count,
11478 Sum,
11479 Avg,
11480}
11481
11482#[derive(Debug, Clone)]
11486pub struct ApproxResult {
11487 pub point: f64,
11489 pub ci_low: f64,
11491 pub ci_high: f64,
11493 pub n_population: u64,
11495 pub n_sample_live: usize,
11497 pub n_passing: usize,
11499}
11500
11501#[derive(Debug, Clone, PartialEq)]
11506pub enum AggState {
11507 Count(u64),
11509 SumI {
11511 sum: i128,
11512 count: u64,
11513 },
11514 SumF {
11516 sum: f64,
11517 count: u64,
11518 },
11519 AvgI {
11521 sum: i128,
11522 count: u64,
11523 },
11524 AvgF {
11526 sum: f64,
11527 count: u64,
11528 },
11529 MinI(i64),
11531 MaxI(i64),
11532 MinF(f64),
11534 MaxF(f64),
11535 Empty,
11537}
11538
11539impl AggState {
11540 pub fn merge(self, other: AggState) -> AggState {
11542 use AggState::*;
11543 match (self, other) {
11544 (Empty, x) | (x, Empty) => x,
11545 (Count(a), Count(b)) => Count(a + b),
11546 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11547 sum: sa + sb,
11548 count: ca + cb,
11549 },
11550 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11551 sum: sa + sb,
11552 count: ca + cb,
11553 },
11554 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11555 sum: sa + sb,
11556 count: ca + cb,
11557 },
11558 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11559 sum: sa + sb,
11560 count: ca + cb,
11561 },
11562 (MinI(a), MinI(b)) => MinI(a.min(b)),
11563 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11564 (MinF(a), MinF(b)) => MinF(a.min(b)),
11565 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11566 _ => Empty, }
11568 }
11569
11570 pub fn point(&self) -> Option<f64> {
11572 match self {
11573 AggState::Count(n) => Some(*n as f64),
11574 AggState::SumI { sum, .. } => Some(*sum as f64),
11575 AggState::SumF { sum, .. } => Some(*sum),
11576 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11577 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11578 AggState::MinI(n) => Some(*n as f64),
11579 AggState::MaxI(n) => Some(*n as f64),
11580 AggState::MinF(n) => Some(*n),
11581 AggState::MaxF(n) => Some(*n),
11582 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11583 }
11584 }
11585
11586 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11590 let is_float = matches!(ty, Some(TypeId::Float64));
11591 match (agg, result) {
11592 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11593 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11594 sum: x as i128,
11595 count: 1, },
11597 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11598 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11599 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11600 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11601 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11602 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11603 (NativeAgg::Count, _) => AggState::Empty,
11604 (_, NativeAggResult::Null) => AggState::Empty,
11605 _ => {
11606 let _ = is_float;
11607 AggState::Empty
11608 }
11609 }
11610 }
11611}
11612
11613#[derive(Debug, Clone)]
11616pub struct CachedAgg {
11617 pub state: AggState,
11618 pub watermark: u64,
11619 pub epoch: u64,
11620}
11621
11622#[derive(Debug, Clone)]
11624pub struct IncrementalAggResult {
11625 pub state: AggState,
11627 pub incremental: bool,
11630 pub delta_rows: u64,
11632}
11633
11634fn agg_state_from_rows(
11638 rows: &[Row],
11639 conditions: &[crate::query::Condition],
11640 index_sets: &[RowIdSet],
11641 column: Option<u16>,
11642 agg: NativeAgg,
11643 schema: &Schema,
11644 control: Option<&crate::ExecutionControl>,
11645) -> Result<AggState> {
11646 let mut count: u64 = 0;
11647 let mut sum_i: i128 = 0;
11648 let mut sum_f: f64 = 0.0;
11649 let mut mn_i: i64 = i64::MAX;
11650 let mut mx_i: i64 = i64::MIN;
11651 let mut mn_f: f64 = f64::INFINITY;
11652 let mut mx_f: f64 = f64::NEG_INFINITY;
11653 let mut saw_int = false;
11654 let mut saw_float = false;
11655 for (index, r) in rows.iter().enumerate() {
11656 execution_checkpoint(control, index)?;
11657 if !conditions
11658 .iter()
11659 .all(|c| condition_matches_row(c, r, schema))
11660 {
11661 continue;
11662 }
11663 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11664 continue;
11665 }
11666 match agg {
11667 NativeAgg::Count => match column {
11668 None => count += 1,
11670 Some(cid) => match r.columns.get(&cid) {
11673 None | Some(Value::Null) => {}
11674 Some(_) => count += 1,
11675 },
11676 },
11677 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11678 Some(Value::Int64(n)) => {
11679 count += 1;
11680 sum_i += *n as i128;
11681 mn_i = mn_i.min(*n);
11682 mx_i = mx_i.max(*n);
11683 saw_int = true;
11684 }
11685 Some(Value::Float64(f)) => {
11686 count += 1;
11687 sum_f += f;
11688 mn_f = mn_f.min(*f);
11689 mx_f = mx_f.max(*f);
11690 saw_float = true;
11691 }
11692 _ => {}
11693 },
11694 }
11695 }
11696 Ok(match agg {
11697 NativeAgg::Count => {
11698 if count == 0 {
11699 AggState::Empty
11700 } else {
11701 AggState::Count(count)
11702 }
11703 }
11704 NativeAgg::Sum => {
11705 if count == 0 {
11706 AggState::Empty
11707 } else if saw_int {
11708 AggState::SumI { sum: sum_i, count }
11709 } else {
11710 AggState::SumF { sum: sum_f, count }
11711 }
11712 }
11713 NativeAgg::Avg => {
11714 if count == 0 {
11715 AggState::Empty
11716 } else if saw_int {
11717 AggState::AvgI { sum: sum_i, count }
11718 } else {
11719 AggState::AvgF { sum: sum_f, count }
11720 }
11721 }
11722 NativeAgg::Min => {
11723 if !saw_int && !saw_float {
11724 AggState::Empty
11725 } else if saw_int {
11726 AggState::MinI(mn_i)
11727 } else {
11728 AggState::MinF(mn_f)
11729 }
11730 }
11731 NativeAgg::Max => {
11732 if !saw_int && !saw_float {
11733 AggState::Empty
11734 } else if saw_int {
11735 AggState::MaxI(mx_i)
11736 } else {
11737 AggState::MaxF(mx_f)
11738 }
11739 }
11740 })
11741}
11742
11743fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11747 use crate::query::Condition;
11748 match c {
11749 Condition::Pk(key) => match schema.primary_key() {
11750 Some(pk) => row
11751 .columns
11752 .get(&pk.id)
11753 .map(|v| v.encode_key() == *key)
11754 .unwrap_or(false),
11755 None => false,
11756 },
11757 Condition::BitmapEq { column_id, value } => row
11758 .columns
11759 .get(column_id)
11760 .map(|v| v.encode_key() == *value)
11761 .unwrap_or(false),
11762 Condition::BitmapIn { column_id, values } => {
11763 let key = row.columns.get(column_id).map(|v| v.encode_key());
11764 match key {
11765 Some(k) => values.contains(&k),
11766 None => false,
11767 }
11768 }
11769 Condition::BytesPrefix { column_id, prefix } => row
11770 .columns
11771 .get(column_id)
11772 .map(|v| v.encode_key().starts_with(prefix))
11773 .unwrap_or(false),
11774 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11775 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11776 _ => false,
11777 },
11778 Condition::RangeF64 {
11779 column_id,
11780 lo,
11781 lo_inclusive,
11782 hi,
11783 hi_inclusive,
11784 } => match row.columns.get(column_id) {
11785 Some(Value::Float64(n)) => {
11786 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11787 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11788 lo_ok && hi_ok
11789 }
11790 _ => false,
11791 },
11792 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11793 Some(Value::Bytes(b)) => {
11794 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11795 }
11796 _ => false,
11797 },
11798 Condition::FmContainsAll {
11799 column_id,
11800 patterns,
11801 } => match row.columns.get(column_id) {
11802 Some(Value::Bytes(b)) => patterns
11803 .iter()
11804 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11805 _ => false,
11806 },
11807 Condition::Ann { .. }
11808 | Condition::SparseMatch { .. }
11809 | Condition::MinHashSimilar { .. } => true,
11810 Condition::IsNull { column_id } => {
11811 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11812 }
11813 Condition::IsNotNull { column_id } => {
11814 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11815 }
11816 }
11817}
11818
11819fn as_f64(v: Option<&Value>) -> Option<f64> {
11821 match v {
11822 Some(Value::Int64(n)) => Some(*n as f64),
11823 Some(Value::Float64(f)) => Some(*f),
11824 _ => None,
11825 }
11826}
11827
11828fn accumulate_int(
11832 cursor: &mut dyn crate::cursor::Cursor,
11833 control: Option<&crate::ExecutionControl>,
11834) -> Result<(u64, i128, i64, i64)> {
11835 let mut count: u64 = 0;
11836 let mut sum: i128 = 0;
11837 let mut mn: i64 = i64::MAX;
11838 let mut mx: i64 = i64::MIN;
11839 while let Some(cols) = cursor.next_batch()? {
11840 execution_checkpoint(control, 0)?;
11841 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11842 if crate::columnar::all_non_null(validity, data.len()) {
11843 count += data.len() as u64;
11845 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11846 execution_checkpoint(control, chunk_index * 1024)?;
11847 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11848 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11849 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11850 }
11851 } else {
11852 for (i, &v) in data.iter().enumerate() {
11853 execution_checkpoint(control, i)?;
11854 if crate::columnar::validity_bit(validity, i) {
11855 count += 1;
11856 sum += v as i128;
11857 mn = mn.min(v);
11858 mx = mx.max(v);
11859 }
11860 }
11861 }
11862 }
11863 }
11864 Ok((count, sum, mn, mx))
11865}
11866
11867fn accumulate_float(
11869 cursor: &mut dyn crate::cursor::Cursor,
11870 control: Option<&crate::ExecutionControl>,
11871) -> Result<(u64, f64, f64, f64)> {
11872 let mut count: u64 = 0;
11873 let mut sum: f64 = 0.0;
11874 let mut mn: f64 = f64::INFINITY;
11875 let mut mx: f64 = f64::NEG_INFINITY;
11876 while let Some(cols) = cursor.next_batch()? {
11877 execution_checkpoint(control, 0)?;
11878 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11879 if crate::columnar::all_non_null(validity, data.len()) {
11880 count += data.len() as u64;
11881 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11882 execution_checkpoint(control, chunk_index * 1024)?;
11883 sum += chunk.iter().sum::<f64>();
11884 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11885 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11886 }
11887 } else {
11888 for (i, &v) in data.iter().enumerate() {
11889 execution_checkpoint(control, i)?;
11890 if crate::columnar::validity_bit(validity, i) {
11891 count += 1;
11892 sum += v;
11893 mn = mn.min(v);
11894 mx = mx.max(v);
11895 }
11896 }
11897 }
11898 }
11899 }
11900 Ok((count, sum, mn, mx))
11901}
11902
11903#[inline]
11904fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11905 if index.is_multiple_of(256) {
11906 control
11907 .map(crate::ExecutionControl::checkpoint)
11908 .transpose()?;
11909 }
11910 Ok(())
11911}
11912
11913fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11914 if count == 0 && !matches!(agg, NativeAgg::Count) {
11915 return NativeAggResult::Null;
11916 }
11917 match agg {
11918 NativeAgg::Count => NativeAggResult::Count(count),
11919 NativeAgg::Sum => match sum.try_into() {
11922 Ok(v) => NativeAggResult::Int(v),
11923 Err(_) => NativeAggResult::Null,
11924 },
11925 NativeAgg::Min => NativeAggResult::Int(mn),
11926 NativeAgg::Max => NativeAggResult::Int(mx),
11927 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11928 }
11929}
11930
11931fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11932 if count == 0 && !matches!(agg, NativeAgg::Count) {
11933 return NativeAggResult::Null;
11934 }
11935 match agg {
11936 NativeAgg::Count => NativeAggResult::Count(count),
11937 NativeAgg::Sum => NativeAggResult::Float(sum),
11938 NativeAgg::Min => NativeAggResult::Float(mn),
11939 NativeAgg::Max => NativeAggResult::Float(mx),
11940 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11941 }
11942}
11943
11944fn agg_int(
11947 stats: &[crate::page::PageStat],
11948 decode: fn(Option<&[u8]>) -> Option<i64>,
11949) -> Option<(Option<i64>, Option<i64>, u64)> {
11950 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11951 let mut any = false;
11952 for s in stats {
11953 if let Some(v) = decode(s.min.as_deref()) {
11954 mn = mn.min(v);
11955 any = true;
11956 }
11957 if let Some(v) = decode(s.max.as_deref()) {
11958 mx = mx.max(v);
11959 any = true;
11960 }
11961 nulls += s.null_count;
11962 }
11963 any.then_some((Some(mn), Some(mx), nulls))
11964}
11965
11966fn agg_float(
11968 stats: &[crate::page::PageStat],
11969 decode: fn(Option<&[u8]>) -> Option<f64>,
11970) -> Option<(Option<f64>, Option<f64>, u64)> {
11971 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11972 let mut any = false;
11973 for s in stats {
11974 if let Some(v) = decode(s.min.as_deref()) {
11975 mn = mn.min(v);
11976 any = true;
11977 }
11978 if let Some(v) = decode(s.max.as_deref()) {
11979 mx = mx.max(v);
11980 any = true;
11981 }
11982 nulls += s.null_count;
11983 }
11984 any.then_some((Some(mn), Some(mx), nulls))
11985}
11986
11987type SecondaryIndexes = (
11989 HashMap<u16, BitmapIndex>,
11990 HashMap<u16, AnnIndex>,
11991 HashMap<u16, FmIndex>,
11992 HashMap<u16, SparseIndex>,
11993 HashMap<u16, MinHashIndex>,
11994);
11995
11996fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11997 let mut bitmap = HashMap::new();
11998 let mut ann = HashMap::new();
11999 let mut fm = HashMap::new();
12000 let mut sparse = HashMap::new();
12001 let mut minhash = HashMap::new();
12002 for idef in &schema.indexes {
12003 match idef.kind {
12004 IndexKind::Bitmap => {
12005 bitmap.insert(idef.column_id, BitmapIndex::new());
12006 }
12007 IndexKind::Ann => {
12008 let dim = schema
12009 .columns
12010 .iter()
12011 .find(|c| c.id == idef.column_id)
12012 .and_then(|c| match c.ty {
12013 TypeId::Embedding { dim } => Some(dim as usize),
12014 _ => None,
12015 })
12016 .unwrap_or(0);
12017 let options = idef.options.ann.clone().unwrap_or_default();
12018 ann.insert(
12019 idef.column_id,
12020 AnnIndex::with_full_options(
12021 dim,
12022 options.m,
12023 options.ef_construction,
12024 options.ef_search,
12025 &options,
12026 ),
12027 );
12028 }
12029 IndexKind::FmIndex => {
12030 fm.insert(idef.column_id, FmIndex::new());
12031 }
12032 IndexKind::Sparse => {
12033 sparse.insert(idef.column_id, SparseIndex::new());
12034 }
12035 IndexKind::MinHash => {
12036 let options = idef.options.minhash.clone().unwrap_or_default();
12037 minhash.insert(
12038 idef.column_id,
12039 MinHashIndex::with_options(options.permutations, options.bands),
12040 );
12041 }
12042 _ => {}
12043 }
12044 }
12045 (bitmap, ann, fm, sparse, minhash)
12046}
12047
12048const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
12049 | ColumnFlags::AUTO_INCREMENT
12050 | ColumnFlags::ENCRYPTED
12051 | ColumnFlags::ENCRYPTED_INDEXABLE
12052 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
12053
12054fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
12055 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
12056 return Err(MongrelError::Schema(
12057 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
12058 ));
12059 }
12060 Ok(())
12061}
12062
12063fn validate_alter_column_type(
12064 schema: &Schema,
12065 old: &ColumnDef,
12066 next: &ColumnDef,
12067 has_stored_versions: bool,
12068) -> Result<()> {
12069 if old.ty == next.ty {
12070 return Ok(());
12071 }
12072 if schema.indexes.iter().any(|i| i.column_id == old.id) {
12073 return Err(MongrelError::Schema(format!(
12074 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
12075 old.name
12076 )));
12077 }
12078 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
12079 return Ok(());
12080 }
12081 Err(MongrelError::Schema(format!(
12082 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
12083 old.ty, next.ty
12084 )))
12085}
12086
12087fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
12088 matches!(
12089 (old, new),
12090 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
12091 )
12092}
12093
12094fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
12100 let mut prev: Option<i64> = None;
12101 for r in rows {
12102 match r.columns.get(&pk_id) {
12103 Some(Value::Int64(v)) => {
12104 if prev.is_some_and(|p| p >= *v) {
12105 return false;
12106 }
12107 prev = Some(*v);
12108 }
12109 _ => return false,
12110 }
12111 }
12112 true
12113}
12114
12115#[allow(clippy::too_many_arguments)]
12116fn index_into(
12117 schema: &Schema,
12118 row: &Row,
12119 hot: &mut HotIndex,
12120 bitmap: &mut HashMap<u16, BitmapIndex>,
12121 ann: &mut HashMap<u16, AnnIndex>,
12122 fm: &mut HashMap<u16, FmIndex>,
12123 sparse: &mut HashMap<u16, SparseIndex>,
12124 minhash: &mut HashMap<u16, MinHashIndex>,
12125) {
12126 for idef in &schema.indexes {
12127 let Some(val) = row.columns.get(&idef.column_id) else {
12128 continue;
12129 };
12130 match idef.kind {
12131 IndexKind::Bitmap => {
12132 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12133 b.insert(val.encode_key(), row.row_id);
12134 }
12135 }
12136 IndexKind::Ann => {
12137 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12138 if let Some(meta) = val.generated_embedding_metadata() {
12139 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12141 continue;
12142 }
12143 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12144 .is_err()
12145 {
12146 continue;
12147 }
12148 }
12149 a.insert_validated(v, row.row_id);
12150 }
12151 }
12152 IndexKind::FmIndex => {
12153 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12154 f.insert(b.clone(), row.row_id);
12155 }
12156 }
12157 IndexKind::Sparse => {
12158 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12159 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12162 s.insert(&terms, row.row_id);
12163 }
12164 }
12165 }
12166 IndexKind::MinHash => {
12167 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12168 let tokens = crate::index::token_hashes_from_bytes(b);
12171 mh.insert(&tokens, row.row_id);
12172 }
12173 }
12174 _ => {}
12175 }
12176 }
12177 if let Some(pk_col) = schema.primary_key() {
12178 if let Some(pk_val) = row.columns.get(&pk_col.id) {
12179 hot.insert(pk_val.encode_key(), row.row_id);
12180 }
12181 }
12182}
12183
12184#[allow(clippy::too_many_arguments)]
12187fn index_into_single(
12188 idef: &IndexDef,
12189 _schema: &Schema,
12190 row: &Row,
12191 _hot: &mut HotIndex,
12192 bitmap: &mut HashMap<u16, BitmapIndex>,
12193 ann: &mut HashMap<u16, AnnIndex>,
12194 fm: &mut HashMap<u16, FmIndex>,
12195 sparse: &mut HashMap<u16, SparseIndex>,
12196 minhash: &mut HashMap<u16, MinHashIndex>,
12197) {
12198 let Some(val) = row.columns.get(&idef.column_id) else {
12199 return;
12200 };
12201 match idef.kind {
12202 IndexKind::Bitmap => {
12203 if let Some(b) = bitmap.get_mut(&idef.column_id) {
12204 b.insert(val.encode_key(), row.row_id);
12205 }
12206 }
12207 IndexKind::Ann => {
12208 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
12209 if let Some(meta) = val.generated_embedding_metadata() {
12210 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
12212 return;
12213 }
12214 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
12215 .is_err()
12216 {
12217 return;
12218 }
12219 }
12220 a.insert_validated(v, row.row_id);
12221 }
12222 }
12223 IndexKind::FmIndex => {
12224 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
12225 f.insert(b.clone(), row.row_id);
12226 }
12227 }
12228 IndexKind::Sparse => {
12229 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
12230 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
12231 s.insert(&terms, row.row_id);
12232 }
12233 }
12234 }
12235 IndexKind::MinHash => {
12236 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
12237 let tokens = crate::index::token_hashes_from_bytes(b);
12238 mh.insert(&tokens, row.row_id);
12239 }
12240 }
12241 _ => {}
12242 }
12243}
12244
12245fn eval_partial_predicate(
12251 pred: &str,
12252 columns_map: &HashMap<u16, &Value>,
12253 name_to_id: &HashMap<&str, u16>,
12254) -> bool {
12255 let lower = pred.trim().to_ascii_lowercase();
12256 if let Some(rest) = lower.strip_suffix(" is not null") {
12258 let col_name = rest.trim();
12259 if let Some(col_id) = name_to_id.get(col_name) {
12260 return columns_map
12261 .get(col_id)
12262 .is_some_and(|v| !matches!(v, Value::Null));
12263 }
12264 }
12265 if let Some(rest) = lower.strip_suffix(" is null") {
12267 let col_name = rest.trim();
12268 if let Some(col_id) = name_to_id.get(col_name) {
12269 return columns_map
12270 .get(col_id)
12271 .is_none_or(|v| matches!(v, Value::Null));
12272 }
12273 }
12274 true
12277}
12278
12279#[allow(dead_code)]
12285fn bulk_index_key(
12286 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12287 column_id: u16,
12288 ty: TypeId,
12289 col: &columnar::NativeColumn,
12290 i: usize,
12291) -> Option<Vec<u8>> {
12292 let encoded = columnar::encode_key_native(ty, col, i)?;
12293 {
12294 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12295 if let Some((key, scheme)) = column_keys.get(&column_id) {
12296 return Some(match (*scheme, col) {
12297 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12298 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12299 ope_token_i64(key, data[i]).to_vec()
12300 }
12301 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12302 ope_token_f64(key, data[i]).to_vec()
12303 }
12304 _ => hmac_token(key, &encoded).to_vec(),
12305 });
12306 }
12307 }
12308 Some(encoded)
12309}
12310
12311pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12312 write_schema_with_after(dir, schema, || {})
12313}
12314
12315pub(crate) fn write_schema_durable(
12316 root: &crate::durable_file::DurableRoot,
12317 schema: &Schema,
12318) -> Result<()> {
12319 write_schema_durable_with_after(root, schema, || {})
12320}
12321
12322fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12323where
12324 F: FnOnce(),
12325{
12326 let json = serde_json::to_string_pretty(schema)
12327 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12328 crate::durable_file::write_atomic_with_after(
12329 &dir.join(SCHEMA_FILENAME),
12330 json.as_bytes(),
12331 after_publish,
12332 )?;
12333 Ok(())
12334}
12335
12336fn write_schema_durable_with_after<F>(
12337 root: &crate::durable_file::DurableRoot,
12338 schema: &Schema,
12339 after_publish: F,
12340) -> Result<()>
12341where
12342 F: FnOnce(),
12343{
12344 let json = serde_json::to_string_pretty(schema)
12345 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12346 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12347 Ok(())
12348}
12349
12350fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12351 let mut schema_published = false;
12352 let schema_result = match table._root_guard.as_deref() {
12353 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12354 schema_published = true;
12355 }),
12356 None => write_schema_with_after(&table.dir, &table.schema, || {
12357 schema_published = true;
12358 }),
12359 };
12360 if schema_result.is_err() && !schema_published {
12361 return schema_result;
12362 }
12363 match table.persist_manifest(table.current_epoch()) {
12364 Ok(()) => Ok(()),
12365 Err(manifest_error) => Err(match schema_result {
12366 Ok(()) => manifest_error,
12367 Err(schema_error) => MongrelError::Other(format!(
12368 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12369 )),
12370 }),
12371 }
12372}
12373
12374fn read_schema(dir: &Path) -> Result<Schema> {
12375 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12376 read_schema_file(file)
12377}
12378
12379fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12380 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12381 use std::io::Read;
12382
12383 let length = file.metadata()?.len();
12384 if length > MAX_SCHEMA_BYTES {
12385 return Err(MongrelError::ResourceLimitExceeded {
12386 resource: "schema bytes",
12387 requested: usize::try_from(length).unwrap_or(usize::MAX),
12388 limit: MAX_SCHEMA_BYTES as usize,
12389 });
12390 }
12391 let mut bytes = Vec::with_capacity(length as usize);
12392 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12393 if bytes.len() as u64 != length {
12394 return Err(MongrelError::Schema(
12395 "schema length changed while reading".into(),
12396 ));
12397 }
12398 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12399}
12400
12401fn preflight_standalone_open(
12402 dir: &Path,
12403 runs_root: Option<&crate::durable_file::DurableRoot>,
12404 idx_root: Option<&crate::durable_file::DurableRoot>,
12405 manifest: &Manifest,
12406 schema: &Schema,
12407 records: &[crate::wal::Record],
12408 kek: Option<Arc<Kek>>,
12409) -> Result<()> {
12410 crate::wal::validate_shared_transaction_framing(records)?;
12411 if manifest.schema_id > schema.schema_id
12412 || manifest.flushed_epoch > manifest.current_epoch
12413 || manifest.global_idx_epoch > manifest.current_epoch
12414 || manifest.next_row_id == u64::MAX
12415 || manifest.auto_inc_next < 0
12416 || manifest.auto_inc_next == i64::MAX
12417 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12418 {
12419 return Err(MongrelError::InvalidArgument(
12420 "manifest counters or schema identity are invalid".into(),
12421 ));
12422 }
12423 let mut run_ids = HashSet::new();
12424 let mut maximum_row_id = None::<u64>;
12425 for run in &manifest.runs {
12426 if run.run_id >= u64::MAX as u128
12427 || !run_ids.insert(run.run_id)
12428 || run.epoch_created > manifest.current_epoch
12429 {
12430 return Err(MongrelError::InvalidArgument(
12431 "manifest contains an invalid or duplicate active run".into(),
12432 ));
12433 }
12434 let mut reader = match runs_root {
12435 Some(root) => RunReader::open_file(
12436 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12437 schema.clone(),
12438 kek.clone(),
12439 )?,
12440 None => RunReader::open(
12441 dir.join(RUNS_DIR)
12442 .join(format!("r-{}.sr", run.run_id as u64)),
12443 schema.clone(),
12444 kek.clone(),
12445 )?,
12446 };
12447 let header = reader.header();
12448 if header.run_id != run.run_id
12449 || header.level != run.level
12450 || header.row_count != run.row_count
12451 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12452 || header.is_uniform_epoch() && header.epoch_created != 0
12453 || header.schema_id > schema.schema_id
12454 {
12455 return Err(MongrelError::InvalidArgument(format!(
12456 "run {} differs from its manifest",
12457 run.run_id
12458 )));
12459 }
12460 if header.row_count != 0 {
12461 maximum_row_id = Some(
12462 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12463 );
12464 }
12465 reader.validate_all_pages()?;
12466 }
12467 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12468 return Err(MongrelError::InvalidArgument(
12469 "manifest next_row_id does not advance beyond persisted rows".into(),
12470 ));
12471 }
12472 for run in &manifest.retiring {
12473 if run.run_id >= u64::MAX as u128
12474 || run.retire_epoch > manifest.current_epoch
12475 || !run_ids.insert(run.run_id)
12476 {
12477 return Err(MongrelError::InvalidArgument(
12478 "manifest contains an invalid or duplicate retired run".into(),
12479 ));
12480 }
12481 }
12482 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12483 match idx_root {
12484 Some(root) => {
12485 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12486 }
12487 None => {
12488 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12489 }
12490 }
12491
12492 let committed = records
12493 .iter()
12494 .filter_map(|record| match record.op {
12495 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12496 _ => None,
12497 })
12498 .collect::<HashMap<_, _>>();
12499 for record in records {
12500 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12501 continue;
12502 };
12503 match &record.op {
12504 Op::Put { table_id, rows } => {
12505 if *table_id != manifest.table_id {
12506 return Err(MongrelError::CorruptWal {
12507 offset: record.seq.0,
12508 reason: format!(
12509 "private WAL record references table {table_id}, expected {}",
12510 manifest.table_id
12511 ),
12512 });
12513 }
12514 let rows: Vec<Row> =
12515 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12516 offset: record.seq.0,
12517 reason: format!("committed Put payload could not be decoded: {error}"),
12518 })?;
12519 for row in rows {
12520 if row.deleted || row.row_id.0 == u64::MAX {
12521 return Err(MongrelError::CorruptWal {
12522 offset: record.seq.0,
12523 reason: "committed Put contains an invalid row identity".into(),
12524 });
12525 }
12526 let cells = row.columns.into_iter().collect::<Vec<_>>();
12527 schema
12528 .validate_values(&cells)
12529 .map_err(|error| MongrelError::CorruptWal {
12530 offset: record.seq.0,
12531 reason: format!("committed Put violates table schema: {error}"),
12532 })?;
12533 if schema.auto_increment_column().is_some_and(|column| {
12534 matches!(
12535 cells.iter().find(|(id, _)| *id == column.id),
12536 Some((_, Value::Int64(value))) if *value == i64::MAX
12537 )
12538 }) {
12539 return Err(MongrelError::CorruptWal {
12540 offset: record.seq.0,
12541 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12542 });
12543 }
12544 }
12545 }
12546 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12547 if *table_id != manifest.table_id =>
12548 {
12549 return Err(MongrelError::CorruptWal {
12550 offset: record.seq.0,
12551 reason: format!(
12552 "private WAL record references table {table_id}, expected {}",
12553 manifest.table_id
12554 ),
12555 });
12556 }
12557 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12558 return Err(MongrelError::CorruptWal {
12559 offset: record.seq.0,
12560 reason: "private WAL contains shared spilled-run metadata".into(),
12561 });
12562 }
12563 _ => {}
12564 }
12565 }
12566 Ok(())
12567}
12568
12569fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12570 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12571}
12572
12573fn wal_segment_number(path: &Path) -> Option<u64> {
12574 path.file_stem()
12575 .and_then(|stem| stem.to_str())
12576 .and_then(|stem| stem.strip_prefix("seg-"))
12577 .and_then(|number| number.parse().ok())
12578}
12579
12580fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12581 let n = list_wal_numbers(wal_dir)?;
12582 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12583}
12584
12585fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12586 list_wal_numbers(wal_dir)?
12587 .map(|maximum| {
12588 maximum
12589 .checked_add(1)
12590 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12591 })
12592 .unwrap_or(Ok(0))
12593}
12594
12595fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12596 let mut max_n = None;
12597 let entries = match std::fs::read_dir(wal_dir) {
12598 Ok(entries) => entries,
12599 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12600 Err(error) => return Err(error.into()),
12601 };
12602 for entry in entries {
12603 let entry = entry?;
12604 let fname = entry.file_name();
12605 let Some(s) = fname.to_str() else {
12606 continue;
12607 };
12608 let Some(stripped) = s.strip_prefix("seg-") else {
12609 continue;
12610 };
12611 let Some(number) = stripped.strip_suffix(".wal") else {
12612 return Err(MongrelError::CorruptWal {
12613 offset: 0,
12614 reason: format!("malformed WAL segment name {s:?}"),
12615 });
12616 };
12617 let n = number
12618 .parse::<u32>()
12619 .map_err(|_| MongrelError::CorruptWal {
12620 offset: 0,
12621 reason: format!("malformed WAL segment name {s:?}"),
12622 })?;
12623 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12624 return Err(MongrelError::CorruptWal {
12625 offset: n as u64,
12626 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12627 });
12628 }
12629 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12630 }
12631 Ok(max_n)
12632}