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
384const 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;
391
392#[derive(Clone, Copy, Debug)]
407struct AutoIncState {
408 column_id: u16,
409 next: i64,
410 seeded: bool,
411}
412
413pub(crate) struct RecoveryMetadataPlan {
414 live_count: u64,
415 auto_inc: Option<AutoIncState>,
416 changed: bool,
417}
418
419type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
420
421fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
424 schema.auto_increment_column().map(|c| AutoIncState {
425 column_id: c.id,
426 next: 0,
427 seeded: false,
428 })
429}
430
431#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
444pub enum IndexBuildPolicy {
445 #[default]
447 Deferred,
448 Eager,
450}
451
452#[derive(Clone)]
453struct ReversePkSegment {
454 values: HashMap<RowId, Vec<u8>>,
455 removed: HashSet<RowId>,
456}
457
458#[derive(Clone)]
459struct ReversePkMap {
460 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
461 active: ReversePkSegment,
462}
463
464impl ReversePkMap {
465 fn new() -> Self {
466 Self {
467 frozen: Arc::new(Vec::new()),
468 active: ReversePkSegment {
469 values: HashMap::new(),
470 removed: HashSet::new(),
471 },
472 }
473 }
474
475 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
476 let mut map = Self::new();
477 map.active.values.extend(entries);
478 map
479 }
480
481 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
482 self.active.removed.remove(&row_id);
483 self.active.values.insert(row_id, key);
484 }
485
486 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
487 if let Some(key) = self.active.values.get(row_id) {
488 return Some(key);
489 }
490 if self.active.removed.contains(row_id) {
491 return None;
492 }
493 for segment in self.frozen.iter().rev() {
494 if let Some(key) = segment.values.get(row_id) {
495 return Some(key);
496 }
497 if segment.removed.contains(row_id) {
498 return None;
499 }
500 }
501 None
502 }
503
504 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
505 let previous = self.get(row_id).cloned();
506 self.active.values.remove(row_id);
507 self.active.removed.insert(*row_id);
508 previous
509 }
510
511 fn clear(&mut self) {
512 *self = Self::new();
513 }
514
515 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
516 let mut entries = HashMap::new();
517 for segment in self
518 .frozen
519 .iter()
520 .map(Arc::as_ref)
521 .chain(std::iter::once(&self.active))
522 {
523 for row_id in &segment.removed {
524 entries.remove(row_id);
525 }
526 entries.extend(
527 segment
528 .values
529 .iter()
530 .map(|(row_id, key)| (*row_id, key.clone())),
531 );
532 }
533 entries
534 }
535
536 fn seal(&mut self) {
537 if self.active.values.is_empty() && self.active.removed.is_empty() {
538 return;
539 }
540 let active = std::mem::replace(
541 &mut self.active,
542 ReversePkSegment {
543 values: HashMap::new(),
544 removed: HashSet::new(),
545 },
546 );
547 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
548 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
549 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
550 values: self.entries(),
551 removed: HashSet::new(),
552 })]);
553 }
554 }
555}
556
557#[derive(Clone)]
570pub struct ReadGeneration {
571 schema: Arc<Schema>,
572 base_runs: Arc<Vec<RunRef>>,
573 deltas: TableDeltas,
574 indexes: Arc<IndexGeneration>,
575 visible_through: Epoch,
576}
577
578#[derive(Clone)]
584pub struct TableDeltas {
585 memtable: Memtable,
586 mutable_run: MutableRun,
587 hot: HotIndex,
588 pk_by_row: ReversePkMap,
589}
590
591impl ReadGeneration {
592 fn empty(schema: &Schema) -> Self {
595 Self {
596 schema: Arc::new(schema.clone()),
597 base_runs: Arc::new(Vec::new()),
598 deltas: TableDeltas {
599 memtable: Memtable::new(),
600 mutable_run: MutableRun::new(),
601 hot: HotIndex::new(),
602 pk_by_row: ReversePkMap::new(),
603 },
604 indexes: Arc::new(IndexGeneration::default()),
605 visible_through: Epoch(0),
606 }
607 }
608
609 pub fn schema(&self) -> &Arc<Schema> {
611 &self.schema
612 }
613
614 pub fn base_runs(&self) -> &[RunRef] {
616 &self.base_runs
617 }
618
619 pub fn indexes(&self) -> &Arc<IndexGeneration> {
621 &self.indexes
622 }
623
624 pub fn visible_through(&self) -> Epoch {
626 self.visible_through
627 }
628
629 pub fn deltas(&self) -> &TableDeltas {
633 &self.deltas
634 }
635}
636
637impl TableDeltas {
638 pub fn approx_bytes(&self) -> u64 {
641 self.memtable
642 .approx_bytes()
643 .saturating_add(self.mutable_run.approx_bytes())
644 }
645
646 pub fn memtable_len(&self) -> usize {
648 self.memtable.len()
649 }
650
651 pub fn mutable_run_len(&self) -> usize {
653 self.mutable_run.len()
654 }
655
656 pub fn hot_len(&self) -> usize {
658 self.hot.len()
659 }
660
661 pub fn reverse_pk_len(&self) -> usize {
663 self.pk_by_row.entries().len()
664 }
665}
666
667#[derive(Clone)]
669pub struct Table {
670 dir: PathBuf,
671 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
672 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
673 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
674 table_id: u64,
675 name: String,
679 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
684 read_only: bool,
687 durable_commit_failed: bool,
691 wal: WalSink,
692 memtable: Memtable,
693 mutable_run: MutableRun,
698 mutable_run_spill_bytes: u64,
700 compaction_zstd_level: i32,
703 allocator: RowIdAllocator,
704 epoch: Arc<EpochAuthority>,
705 data_generation: u64,
708 schema: Schema,
709 hot: HotIndex,
710 kek: Option<Arc<Kek>>,
713 column_keys: HashMap<u16, ([u8; 32], u8)>,
717 run_refs: Vec<RunRef>,
718 retiring: Vec<crate::manifest::RetiredRun>,
721 next_run_id: u64,
722 sync_byte_threshold: u64,
723 current_txn_id: u64,
728 pending_private_mutations: bool,
732 bitmap: HashMap<u16, BitmapIndex>,
733 ann: HashMap<u16, AnnIndex>,
734 fm: HashMap<u16, FmIndex>,
735 sparse: HashMap<u16, SparseIndex>,
736 minhash: HashMap<u16, MinHashIndex>,
737 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
740 pk_by_row: ReversePkMap,
742 pinned: BTreeMap<Epoch, usize>,
745 pub(crate) live_count: u64,
748 reservoir: crate::reservoir::Reservoir,
751 reservoir_complete: bool,
759 had_deletes: bool,
763 agg_cache: Arc<HashMap<u64, CachedAgg>>,
767 global_idx_epoch: u64,
771 indexes_complete: bool,
776 index_build_policy: IndexBuildPolicy,
778 pk_by_row_complete: bool,
785 flushed_epoch: u64,
788 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
791 snapshots: Arc<crate::retention::SnapshotRegistry>,
794 commit_lock: Arc<parking_lot::Mutex<()>>,
796 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
799 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
809 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
818 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
820 pending_delete_rids: roaring::RoaringBitmap,
823 pending_put_cols: std::collections::HashSet<u16>,
826 pending_rows: Vec<Row>,
832 pending_rows_auto_inc: Vec<bool>,
833 pending_dels: Vec<RowId>,
836 pending_truncate: Option<Epoch>,
840 auto_inc: Option<AutoIncState>,
843 ttl: Option<TtlPolicy>,
846 pins: Arc<crate::retention::PinRegistry>,
853 published: Arc<ArcSwap<ReadGeneration>>,
858 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
863}
864
865const _: () = {
872 const fn assert_sync<T: ?Sized + Sync>() {}
873 assert_sync::<Table>();
874};
875
876enum CachedData {
882 Rows(Arc<Vec<Row>>),
883 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
884}
885
886impl CachedData {
887 fn approx_bytes(&self) -> u64 {
888 match self {
889 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
890 CachedData::Columns(c) => c
891 .iter()
892 .map(|(_, c)| c.approx_bytes())
893 .sum::<u64>()
894 .saturating_add(c.len() as u64 * 16),
895 }
896 }
897}
898
899struct CachedEntry {
903 data: CachedData,
904 footprint: roaring::RoaringBitmap,
905 condition_cols: Vec<u16>,
906}
907
908struct ResultCache {
919 entries: std::collections::HashMap<u64, CachedEntry>,
920 order: std::collections::VecDeque<u64>,
921 bytes: u64,
922 max_bytes: u64,
923 dir: Option<std::path::PathBuf>,
924 #[allow(dead_code)]
925 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
926}
927
928#[derive(serde::Serialize, serde::Deserialize)]
930struct SerializedEntry {
931 condition_cols: Vec<u16>,
932 footprint_bits: Vec<u32>,
933 data: SerializedData,
934}
935
936#[derive(serde::Serialize, serde::Deserialize)]
937enum SerializedData {
938 Rows(Vec<Row>),
939 Columns(Vec<(u16, columnar::NativeColumn)>),
940}
941
942impl SerializedEntry {
943 fn from_entry(entry: &CachedEntry) -> Self {
944 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
945 let data = match &entry.data {
946 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
947 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
948 };
949 Self {
950 condition_cols: entry.condition_cols.clone(),
951 footprint_bits,
952 data,
953 }
954 }
955
956 fn into_entry(self) -> Option<CachedEntry> {
957 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
958 let data = match self.data {
959 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
960 SerializedData::Columns(c) => {
961 if !c.iter().all(|(_, col)| col.validate()) {
964 return None;
965 }
966 CachedData::Columns(Arc::new(c))
967 }
968 };
969 Some(CachedEntry {
970 data,
971 footprint,
972 condition_cols: self.condition_cols,
973 })
974 }
975}
976
977impl ResultCache {
978 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
979
980 fn new() -> Self {
981 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
982 }
983
984 fn with_max_bytes(max_bytes: u64) -> Self {
985 Self {
986 entries: std::collections::HashMap::new(),
987 order: std::collections::VecDeque::new(),
988 bytes: 0,
989 max_bytes,
990 dir: None,
991 cache_dek: None,
992 }
993 }
994
995 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
996 let _ = std::fs::create_dir_all(&dir);
997 self.dir = Some(dir);
998 self
999 }
1000
1001 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1002 self.cache_dek = dek;
1003 self
1004 }
1005
1006 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1007 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1008 }
1009
1010 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1014 let Some(path) = self.disk_path(key) else {
1015 return;
1016 };
1017 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
1018 Ok(s) => s,
1019 Err(_) => return,
1020 };
1021 let on_disk = if let Some(dek) = &self.cache_dek {
1023 match self.encrypt_cache(&serialized, dek) {
1024 Some(b) => b,
1025 None => return,
1026 }
1027 } else {
1028 serialized
1029 };
1030 let tmp = path.with_extension("tmp");
1031 use std::io::Write;
1032 let write = || -> std::io::Result<()> {
1033 let mut f = std::fs::File::create(&tmp)?;
1034 f.write_all(&on_disk)?;
1035 f.flush()?;
1036 Ok(())
1037 };
1038 if write().is_err() {
1039 let _ = std::fs::remove_file(&tmp);
1040 return;
1041 }
1042 let _ = std::fs::rename(&tmp, &path);
1043 }
1044
1045 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1047 let path = self.disk_path(key)?;
1048 let bytes = std::fs::read(&path).ok()?;
1049 let plaintext = if let Some(dek) = &self.cache_dek {
1050 self.decrypt_cache(&bytes, dek)?
1051 } else {
1052 bytes
1053 };
1054 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1055 serialized.into_entry()
1056 }
1057
1058 fn remove_from_disk(&self, key: u64) {
1060 if let Some(path) = self.disk_path(key) {
1061 let _ = std::fs::remove_file(&path);
1062 }
1063 }
1064
1065 #[cfg(feature = "encryption")]
1067 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1068 use crate::encryption::Cipher;
1069 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1070 let mut nonce = [0u8; 12];
1071 crate::encryption::fill_random(&mut nonce).ok()?;
1072 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1073 let mut out = Vec::with_capacity(12 + ct.len());
1074 out.extend_from_slice(&nonce);
1075 out.extend_from_slice(&ct);
1076 Some(out)
1077 }
1078
1079 #[cfg(not(feature = "encryption"))]
1080 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1081 None
1082 }
1083
1084 #[cfg(feature = "encryption")]
1086 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1087 use crate::encryption::Cipher;
1088 if bytes.len() < 28 {
1089 return None;
1090 }
1091 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1092 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1093 let ct = &bytes[12..];
1094 cipher.decrypt_page(&nonce, ct).ok()
1095 }
1096
1097 #[cfg(not(feature = "encryption"))]
1098 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1099 None
1100 }
1101
1102 fn load_persistent(&mut self) {
1105 let Some(dir) = self.dir.as_ref().cloned() else {
1106 return;
1107 };
1108 let entries = match std::fs::read_dir(&dir) {
1109 Ok(e) => e,
1110 Err(_) => return,
1111 };
1112 for entry in entries.flatten() {
1113 let path = entry.path();
1114 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1116 let _ = std::fs::remove_file(&path);
1117 continue;
1118 }
1119 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1120 continue;
1121 }
1122 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1123 Some(s) => s,
1124 None => continue,
1125 };
1126 let key = match u64::from_str_radix(stem, 16) {
1127 Ok(k) => k,
1128 Err(_) => continue,
1129 };
1130 let bytes = match std::fs::read(&path) {
1131 Ok(b) => b,
1132 Err(_) => continue,
1133 };
1134 let plaintext = if let Some(dek) = &self.cache_dek {
1136 match self.decrypt_cache(&bytes, dek) {
1137 Some(p) => p,
1138 None => {
1139 let _ = std::fs::remove_file(&path);
1140 continue;
1141 }
1142 }
1143 } else {
1144 bytes
1145 };
1146 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1147 Ok(serialized) => {
1148 if let Some(entry) = serialized.into_entry() {
1149 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1150 self.entries.insert(key, entry);
1151 self.order.push_back(key);
1152 } else {
1153 let _ = std::fs::remove_file(&path);
1154 }
1155 }
1156 Err(_) => {
1157 let _ = std::fs::remove_file(&path);
1158 }
1159 }
1160 }
1161 self.evict();
1162 }
1163
1164 fn set_max_bytes(&mut self, max_bytes: u64) {
1165 self.max_bytes = max_bytes;
1166 self.evict();
1167 }
1168
1169 fn touch(&mut self, key: u64) {
1171 self.order.retain(|k| *k != key);
1172 self.order.push_back(key);
1173 }
1174
1175 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1176 let res = self.entries.get(&key).and_then(|e| match &e.data {
1177 CachedData::Rows(r) => Some(r.clone()),
1178 CachedData::Columns(_) => None,
1179 });
1180 if res.is_some() {
1181 self.touch(key);
1182 return res;
1183 }
1184 if let Some(entry) = self.load_from_disk(key) {
1186 let res = match &entry.data {
1187 CachedData::Rows(r) => Some(r.clone()),
1188 CachedData::Columns(_) => None,
1189 };
1190 if res.is_some() {
1191 let approx = entry.data.approx_bytes();
1192 self.bytes = self.bytes.saturating_add(approx);
1193 self.entries.insert(key, entry);
1194 self.order.push_back(key);
1195 self.evict();
1196 return res;
1197 }
1198 }
1199 None
1200 }
1201
1202 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1203 let res = self.entries.get(&key).and_then(|e| match &e.data {
1204 CachedData::Columns(c) => Some(c.clone()),
1205 CachedData::Rows(_) => None,
1206 });
1207 if res.is_some() {
1208 self.touch(key);
1209 return res;
1210 }
1211 if let Some(entry) = self.load_from_disk(key) {
1213 let res = match &entry.data {
1214 CachedData::Columns(c) => Some(c.clone()),
1215 CachedData::Rows(_) => None,
1216 };
1217 if res.is_some() {
1218 let approx = entry.data.approx_bytes();
1219 self.bytes = self.bytes.saturating_add(approx);
1220 self.entries.insert(key, entry);
1221 self.order.push_back(key);
1222 self.evict();
1223 return res;
1224 }
1225 }
1226 None
1227 }
1228
1229 fn insert(&mut self, key: u64, entry: CachedEntry) {
1230 let approx = entry.data.approx_bytes();
1231 if self.entries.remove(&key).is_some() {
1232 self.order.retain(|k| *k != key);
1233 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1234 }
1235 self.store_to_disk(key, &entry);
1237 self.bytes = self.bytes.saturating_add(approx);
1238 self.entries.insert(key, entry);
1239 self.order.push_back(key);
1240 self.evict();
1241 }
1242
1243 fn invalidate(
1252 &mut self,
1253 delete_rids: &roaring::RoaringBitmap,
1254 put_cols: &std::collections::HashSet<u16>,
1255 ) {
1256 if self.entries.is_empty() {
1257 return;
1258 }
1259 let has_deletes = !delete_rids.is_empty();
1260 let to_remove: std::collections::HashSet<u64> = self
1261 .entries
1262 .iter()
1263 .filter(|(_, e)| {
1264 let delete_hit = if e.footprint.is_empty() {
1265 has_deletes
1266 } else {
1267 e.footprint.intersection_len(delete_rids) > 0
1268 };
1269 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1270 delete_hit || col_hit
1271 })
1272 .map(|(&k, _)| k)
1273 .collect();
1274 for key in &to_remove {
1275 if let Some(e) = self.entries.remove(key) {
1276 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1277 }
1278 self.remove_from_disk(*key);
1279 }
1280 if !to_remove.is_empty() {
1281 self.order.retain(|k| !to_remove.contains(k));
1282 }
1283 }
1284
1285 fn clear(&mut self) {
1286 if let Some(dir) = &self.dir {
1288 if let Ok(entries) = std::fs::read_dir(dir) {
1289 for entry in entries.flatten() {
1290 let path = entry.path();
1291 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1292 let _ = std::fs::remove_file(&path);
1293 }
1294 }
1295 }
1296 }
1297 self.entries.clear();
1298 self.order.clear();
1299 self.bytes = 0;
1300 }
1301
1302 fn evict(&mut self) {
1303 while self.bytes > self.max_bytes {
1304 let Some(k) = self.order.pop_front() else {
1305 break;
1306 };
1307 if let Some(e) = self.entries.remove(&k) {
1308 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1309 self.remove_from_disk(k);
1313 }
1314 }
1315 }
1316}
1317
1318type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1325
1326fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1327 let _ = kek;
1328 #[cfg(feature = "encryption")]
1329 {
1330 if let Some(k) = kek {
1331 return (
1332 Some(k.derive_table_wal_key(_table_id)),
1333 Some(k.derive_cache_key()),
1334 );
1335 }
1336 }
1337 (None, None)
1338}
1339
1340#[cfg(feature = "encryption")]
1341fn read_table_encryption_salt_root(
1342 root: &crate::durable_file::DurableRoot,
1343) -> Result<[u8; crate::encryption::SALT_LEN]> {
1344 use std::io::Read;
1345
1346 let mut file = root
1347 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1348 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1349 let length = file.metadata()?.len();
1350 if length != crate::encryption::SALT_LEN as u64 {
1351 return Err(MongrelError::InvalidArgument(format!(
1352 "salt file is {length} bytes, expected {}",
1353 crate::encryption::SALT_LEN
1354 )));
1355 }
1356 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1357 file.read_exact(&mut salt)?;
1358 Ok(salt)
1359}
1360
1361#[cfg(feature = "encryption")]
1363fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1364 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1365}
1366
1367#[cfg(not(feature = "encryption"))]
1368fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1369 Box::new(crate::encryption::PlaintextCipher)
1370}
1371
1372fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1373 let Some(kek) = kek else {
1374 return HashMap::new();
1375 };
1376 #[cfg(feature = "encryption")]
1377 {
1378 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1379 schema
1380 .columns
1381 .iter()
1382 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1383 .map(|c| {
1384 let scheme = if schema
1385 .indexes
1386 .iter()
1387 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1388 {
1389 SCHEME_OPE_RANGE
1390 } else {
1391 SCHEME_HMAC_EQ
1392 };
1393 let key: [u8; 32] = *kek.derive_column_key(c.id);
1394 (c.id, (key, scheme))
1395 })
1396 .collect()
1397 }
1398 #[cfg(not(feature = "encryption"))]
1399 {
1400 let _ = (kek, schema);
1401 HashMap::new()
1402 }
1403}
1404
1405pub(crate) struct SharedCtx {
1410 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1411 pub epoch: Arc<EpochAuthority>,
1412 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1413 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1414 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1415 pub kek: Option<Arc<Kek>>,
1416 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1422 pub shared: Option<SharedWalCtx>,
1426 pub table_name: Option<String>,
1429 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1432 pub read_only: bool,
1434}
1435
1436#[derive(Clone)]
1442pub(crate) struct SharedWalCtx {
1443 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1444 pub group: Arc<GroupCommit>,
1445 pub poisoned: Arc<AtomicBool>,
1446 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1447 pub change_wake: tokio::sync::broadcast::Sender<()>,
1448 pub lifecycle: Arc<crate::core::LifecycleController>,
1451}
1452
1453enum WalSink {
1456 Private(Wal),
1457 Shared(SharedWalCtx),
1458 ReadOnly,
1459}
1460
1461impl Clone for WalSink {
1462 fn clone(&self) -> Self {
1463 match self {
1464 Self::Shared(shared) => Self::Shared(shared.clone()),
1465 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1466 }
1467 }
1468}
1469
1470impl SharedCtx {
1471 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1475 let n_shards = if cache_dir.is_some() {
1479 1
1480 } else {
1481 crate::cache::CACHE_SHARDS
1482 };
1483 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1484 let page_cache = if let Some(d) = cache_dir {
1485 Arc::new(crate::cache::Sharded::new(1, || {
1486 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1487 }))
1488 } else {
1489 Arc::new(crate::cache::Sharded::new(n_shards, || {
1490 crate::cache::PageCache::new(per_shard)
1491 }))
1492 };
1493 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1494 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1495 crate::cache::CACHE_SHARDS,
1496 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1497 ));
1498 Self {
1499 root_guard: None,
1500 epoch: Arc::new(EpochAuthority::new(0)),
1501 page_cache,
1502 decoded_cache,
1503 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1504 kek,
1505 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1506 shared: None,
1507 table_name: None,
1508 auth: None,
1509 read_only: false,
1510 }
1511 }
1512}
1513
1514fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1518 use crate::query::Condition;
1519 match c {
1520 Condition::Pk(_)
1522 | Condition::BitmapEq { .. }
1523 | Condition::BitmapIn { .. }
1524 | Condition::BytesPrefix { .. }
1525 | Condition::IsNull { .. }
1526 | Condition::IsNotNull { .. } => 0,
1527 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1529 1
1530 }
1531 Condition::FmContains { .. }
1533 | Condition::FmContainsAll { .. }
1534 | Condition::Ann { .. }
1535 | Condition::SparseMatch { .. } => 2,
1536 }
1537}
1538
1539impl Table {
1540 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1541 let dir = dir.as_ref().to_path_buf();
1542 crate::durable_file::create_directory_all(&dir)?;
1543 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1544 let pinned = root.io_path()?;
1545 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1546 ctx.root_guard = Some(root);
1547 Self::create_in(&pinned, schema, table_id, ctx)
1548 }
1549
1550 #[cfg(feature = "encryption")]
1561 pub fn create_encrypted(
1562 dir: impl AsRef<Path>,
1563 schema: Schema,
1564 table_id: u64,
1565 passphrase: &str,
1566 ) -> Result<Self> {
1567 let dir = dir.as_ref().to_path_buf();
1568 crate::durable_file::create_directory_all(&dir)?;
1569 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1570 root.create_directory_all(META_DIR)?;
1571 let salt = crate::encryption::random_salt()?;
1572 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1573 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1574 let pinned = root.io_path()?;
1575 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1576 ctx.root_guard = Some(root);
1577 Self::create_in(&pinned, schema, table_id, ctx)
1578 }
1579
1580 #[cfg(feature = "encryption")]
1585 pub fn create_with_key(
1586 dir: impl AsRef<Path>,
1587 schema: Schema,
1588 table_id: u64,
1589 key: &[u8],
1590 ) -> Result<Self> {
1591 let dir = dir.as_ref().to_path_buf();
1592 crate::durable_file::create_directory_all(&dir)?;
1593 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1594 root.create_directory_all(META_DIR)?;
1595 let salt = crate::encryption::random_salt()?;
1596 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1597 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1598 let pinned = root.io_path()?;
1599 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1600 ctx.root_guard = Some(root);
1601 Self::create_in(&pinned, schema, table_id, ctx)
1602 }
1603
1604 #[cfg(feature = "encryption")]
1606 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1607 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1608 let salt = read_table_encryption_salt_root(&root)?;
1609 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1610 let pinned = root.io_path()?;
1611 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1612 ctx.root_guard = Some(root);
1613 Self::open_in(&pinned, ctx)
1614 }
1615
1616 pub(crate) fn create_in(
1617 dir: impl AsRef<Path>,
1618 schema: Schema,
1619 table_id: u64,
1620 ctx: SharedCtx,
1621 ) -> Result<Self> {
1622 schema.validate_auto_increment()?;
1623 schema.validate_defaults()?;
1624 schema.validate_ai()?;
1625 for index in &schema.indexes {
1626 index.validate_options()?;
1627 }
1628 let dir = dir.as_ref().to_path_buf();
1629 let runs_root = match ctx.root_guard.as_ref() {
1630 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1631 None => {
1632 crate::durable_file::create_directory_all(&dir)?;
1633 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1634 None
1635 }
1636 };
1637 match ctx.root_guard.as_deref() {
1638 Some(root) => write_schema_durable(root, &schema)?,
1639 None => write_schema(&dir, &schema)?,
1640 }
1641 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1642 let (wal, current_txn_id) = match ctx.shared.clone() {
1645 Some(s) => (WalSink::Shared(s), 0),
1646 None => {
1647 let pinned_wal_root = match ctx.root_guard.as_deref() {
1648 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1649 None => None,
1650 };
1651 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1652 root.io_path()?
1653 } else {
1654 let wal_dir = dir.join(WAL_DIR);
1655 crate::durable_file::create_directory_all(&wal_dir)?;
1656 wal_dir
1657 };
1658 let mut w = if let Some(ref dk) = wal_dek {
1659 Wal::create_with_cipher(
1660 wal_dir.join("seg-000000.wal"),
1661 Epoch(0),
1662 Some(make_cipher(dk)),
1663 0,
1664 )?
1665 } else {
1666 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1667 };
1668 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1669 (WalSink::Private(w), 1)
1670 }
1671 };
1672 let mut manifest = Manifest::new(table_id, schema.schema_id);
1673 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1678 match ctx.root_guard.as_deref() {
1679 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1680 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1681 }
1682 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1683 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1684 let auto_inc = resolve_auto_inc(&schema);
1685 let rcache_dir = dir.join(RCACHE_DIR);
1686 let initial_view = ReadGeneration::empty(&schema);
1687 Ok(Self {
1688 dir,
1689 _root_guard: ctx.root_guard,
1690 runs_root,
1691 idx_root: None,
1692 table_id,
1693 name: ctx.table_name.unwrap_or_default(),
1694 auth: ctx.auth,
1695 read_only: ctx.read_only,
1696 durable_commit_failed: false,
1697 wal,
1698 memtable: Memtable::new(),
1699 mutable_run: MutableRun::new(),
1700 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1701 compaction_zstd_level: 3,
1702 allocator: RowIdAllocator::new(0),
1703 epoch: ctx.epoch,
1704 data_generation: 0,
1705 schema,
1706 hot: HotIndex::new(),
1707 kek: ctx.kek,
1708 column_keys,
1709 run_refs: Vec::new(),
1710 retiring: Vec::new(),
1711 next_run_id: 1,
1712 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1713 current_txn_id,
1714 pending_private_mutations: false,
1715 bitmap,
1716 ann,
1717 fm,
1718 sparse,
1719 minhash,
1720 learned_range: Arc::new(HashMap::new()),
1721 pk_by_row: ReversePkMap::new(),
1722 pinned: BTreeMap::new(),
1723 live_count: 0,
1724 reservoir: crate::reservoir::Reservoir::default(),
1725 reservoir_complete: true,
1726 had_deletes: false,
1727 agg_cache: Arc::new(HashMap::new()),
1728 global_idx_epoch: 0,
1729 indexes_complete: true,
1730 index_build_policy: IndexBuildPolicy::default(),
1731 pk_by_row_complete: false,
1732 flushed_epoch: 0,
1733 page_cache: ctx.page_cache,
1734 decoded_cache: ctx.decoded_cache,
1735 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1736 snapshots: ctx.snapshots,
1737 commit_lock: ctx.commit_lock,
1738 result_cache: Arc::new(parking_lot::Mutex::new(
1739 ResultCache::new()
1740 .with_dir(rcache_dir)
1741 .with_cache_dek(cache_dek.clone()),
1742 )),
1743 pending_delete_rids: roaring::RoaringBitmap::new(),
1744 pending_put_cols: std::collections::HashSet::new(),
1745 pending_rows: Vec::new(),
1746 pending_rows_auto_inc: Vec::new(),
1747 pending_dels: Vec::new(),
1748 pending_truncate: None,
1749 wal_dek,
1750 auto_inc,
1751 ttl: None,
1752 pins: Arc::new(crate::retention::PinRegistry::new()),
1753 published: Arc::new(ArcSwap::from_pointee(initial_view)),
1754 read_generation_pin: None,
1755 })
1756 }
1757
1758 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1762 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1763 let pinned = root.io_path()?;
1764 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1765 ctx.root_guard = Some(root);
1766 Self::open_in(&pinned, ctx)
1767 }
1768
1769 #[cfg(feature = "encryption")]
1772 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1773 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1774 let salt = read_table_encryption_salt_root(&root)?;
1775 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1776 let pinned = root.io_path()?;
1777 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1778 ctx.root_guard = Some(root);
1779 let t = Self::open_in(&pinned, ctx)?;
1780 Ok(t)
1781 }
1782
1783 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1784 let dir = dir.as_ref().to_path_buf();
1785 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1786 let mut manifest = match ctx.root_guard.as_ref() {
1787 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1788 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1789 };
1790 let schema: Schema = match ctx.root_guard.as_ref() {
1791 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1792 None => read_schema(&dir)?,
1793 };
1794 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1800 let runs_root = match ctx.root_guard.as_ref() {
1801 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1802 None => None,
1803 };
1804 let idx_root = match ctx.root_guard.as_ref() {
1805 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1806 Ok(root) => Some(Arc::new(root)),
1807 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1808 Err(error) => return Err(error.into()),
1809 },
1810 None => None,
1811 };
1812 schema.validate_auto_increment()?;
1813 schema.validate_defaults()?;
1814 schema.validate_ai()?;
1815 for index in &schema.indexes {
1816 index.validate_options()?;
1817 }
1818 let replay_epoch = Epoch(manifest.current_epoch);
1819 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1820 let private_replayed = if ctx.shared.is_none() {
1821 match latest_wal_segment(&dir.join(WAL_DIR))? {
1822 Some(path) => {
1823 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1824 crate::wal::replay_with_cipher(path, cipher)?
1825 }
1826 None => Vec::new(),
1827 }
1828 } else {
1829 Vec::new()
1830 };
1831 if ctx.shared.is_none() {
1832 preflight_standalone_open(
1833 &dir,
1834 runs_root.as_deref(),
1835 idx_root.as_deref(),
1836 &manifest,
1837 &schema,
1838 &private_replayed,
1839 ctx.kek.clone(),
1840 )?;
1841 }
1842 let next_run_id = derive_next_run_id(
1843 &dir,
1844 runs_root.as_deref(),
1845 &manifest.runs,
1846 &manifest.retiring,
1847 )?;
1848 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1852 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1853 None => {
1854 let replayed = private_replayed;
1855 let wal_dir = dir.join(WAL_DIR);
1861 crate::durable_file::create_directory_all(&wal_dir)?;
1862 let segment = next_wal_segment(&wal_dir)?;
1863 let segment_no = wal_segment_number(&segment).unwrap_or(0);
1864 let temporary = wal_dir.join(format!(
1865 ".recovery-{}-{}-{segment_no:06}.tmp",
1866 std::process::id(),
1867 std::time::SystemTime::now()
1868 .duration_since(std::time::UNIX_EPOCH)
1869 .unwrap_or_default()
1870 .as_nanos()
1871 ));
1872 let mut w = Wal::create_with_cipher(
1873 &temporary,
1874 replay_epoch,
1875 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1876 segment_no,
1877 )?;
1878 for record in &replayed {
1879 w.append_txn(record.txn_id, record.op.clone())?;
1880 }
1881 let mut w = w.publish_as(segment)?;
1882 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1883 let next_txn_id = replayed
1884 .iter()
1885 .map(|record| record.txn_id)
1886 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
1887 .max()
1888 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
1889 .unwrap_or(1);
1890 (WalSink::Private(w), replayed, next_txn_id)
1891 }
1892 };
1893
1894 let mut memtable = Memtable::new();
1895 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1896 let persisted_epoch = manifest.current_epoch;
1897 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1904 s.next = manifest.auto_inc_next;
1905 s.seeded = manifest.auto_inc_next > 0;
1906 s
1907 });
1908
1909 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1916 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1917 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
1918 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1919 std::collections::BTreeMap::new();
1920 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1921 let mut recovered_epoch = manifest.current_epoch;
1922 let mut recovered_manifest_dirty = schema_manifest_repair;
1923 let mut saw_delete = false;
1924 for record in replayed {
1925 let txn_id = record.txn_id;
1926 match record.op {
1927 Op::Put { rows, .. } => {
1928 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1929 for row in &rows {
1930 allocator.advance_to(row.row_id)?;
1931 if let Some(ai) = auto_inc.as_mut() {
1932 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1933 let next = n.checked_add(1).ok_or_else(|| {
1934 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
1935 })?;
1936 if next > ai.next {
1937 ai.next = next;
1938 }
1939 }
1940 }
1941 }
1942 staged_puts.entry(txn_id).or_default().extend(rows);
1943 }
1944 Op::Delete { row_ids, .. } => {
1945 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1946 }
1947 Op::TxnCommit { epoch, .. } => {
1948 let commit_epoch = Epoch(epoch);
1949 recovered_epoch = recovered_epoch.max(epoch);
1950 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
1951 memtable = Memtable::new();
1952 replayed_puts.clear();
1953 replayed_deletes.clear();
1954 manifest.runs.clear();
1955 manifest.retiring.clear();
1956 manifest.live_count = 0;
1957 manifest.global_idx_epoch = 0;
1958 manifest.current_epoch = manifest.current_epoch.max(epoch);
1959 recovered_manifest_dirty = true;
1960 saw_delete = true;
1961 }
1962 if let Some(puts) = staged_puts.remove(&txn_id) {
1963 if commit_epoch.0 > manifest.flushed_epoch {
1964 for row in &puts {
1965 memtable.upsert(row.clone());
1966 }
1967 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1968 }
1969 }
1970 if let Some(dels) = staged_deletes.remove(&txn_id) {
1971 saw_delete = true;
1972 if commit_epoch.0 > manifest.flushed_epoch {
1973 for rid in dels {
1974 memtable.tombstone(rid, commit_epoch);
1975 replayed_deletes.push((rid, commit_epoch));
1976 }
1977 }
1978 }
1979 }
1980 Op::TxnAbort => {
1981 staged_puts.remove(&txn_id);
1982 staged_deletes.remove(&txn_id);
1983 staged_truncates.remove(&txn_id);
1984 }
1985 Op::TruncateTable { .. } => {
1986 staged_puts.remove(&txn_id);
1987 staged_deletes.remove(&txn_id);
1988 staged_truncates.insert(txn_id);
1989 }
1990 Op::ExternalTableState { .. }
1991 | Op::Flush { .. }
1992 | Op::Ddl(_)
1993 | Op::BeforeImage { .. }
1994 | Op::CommitTimestamp { .. }
1995 | Op::SpilledRows { .. } => {}
1996 }
1997 }
1998
1999 let rcache_dir = dir.join(RCACHE_DIR);
2000 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2001 let initial_view = ReadGeneration::empty(&schema);
2002 let mut db = Self {
2003 dir,
2004 _root_guard: ctx.root_guard,
2005 runs_root,
2006 idx_root,
2007 table_id: manifest.table_id,
2008 name: ctx.table_name.unwrap_or_default(),
2009 auth: ctx.auth,
2010 read_only: ctx.read_only,
2011 durable_commit_failed: false,
2012 wal,
2013 memtable,
2014 mutable_run: MutableRun::new(),
2015 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2016 compaction_zstd_level: 3,
2017 allocator,
2018 epoch: ctx.epoch,
2019 data_generation: persisted_epoch,
2020 schema,
2021 hot: HotIndex::new(),
2022 kek: ctx.kek,
2023 column_keys,
2024 run_refs: manifest.runs.clone(),
2025 retiring: manifest.retiring.clone(),
2026 next_run_id,
2027 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2028 current_txn_id,
2029 pending_private_mutations: false,
2030 bitmap: HashMap::new(),
2031 ann: HashMap::new(),
2032 fm: HashMap::new(),
2033 sparse: HashMap::new(),
2034 minhash: HashMap::new(),
2035 learned_range: Arc::new(HashMap::new()),
2036 pk_by_row: ReversePkMap::new(),
2037 pinned: BTreeMap::new(),
2038 live_count: manifest.live_count,
2039 reservoir: crate::reservoir::Reservoir::default(),
2040 reservoir_complete: false,
2041 had_deletes: saw_delete
2042 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2043 != manifest.live_count,
2044 agg_cache: Arc::new(HashMap::new()),
2045 global_idx_epoch: manifest.global_idx_epoch,
2046 indexes_complete: true,
2047 index_build_policy: IndexBuildPolicy::default(),
2048 pk_by_row_complete: false,
2049 flushed_epoch: manifest.flushed_epoch,
2050 page_cache: ctx.page_cache,
2051 decoded_cache: ctx.decoded_cache,
2052 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2053 snapshots: ctx.snapshots,
2054 commit_lock: ctx.commit_lock,
2055 result_cache: Arc::new(parking_lot::Mutex::new(
2056 ResultCache::new()
2057 .with_dir(rcache_dir)
2058 .with_cache_dek(cache_dek.clone()),
2059 )),
2060 pending_delete_rids: roaring::RoaringBitmap::new(),
2061 pending_put_cols: std::collections::HashSet::new(),
2062 pending_rows: Vec::new(),
2063 pending_rows_auto_inc: Vec::new(),
2064 pending_dels: Vec::new(),
2065 pending_truncate: None,
2066 wal_dek,
2067 auto_inc,
2068 ttl: manifest.ttl,
2069 pins: Arc::new(crate::retention::PinRegistry::new()),
2070 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2071 read_generation_pin: None,
2072 };
2073
2074 db.epoch.advance_recovered(Epoch(recovered_epoch));
2077
2078 let checkpoint = match db.idx_root.as_deref() {
2083 Some(root) => {
2084 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2085 }
2086 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2087 };
2088 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2089 c.epoch_built == manifest.global_idx_epoch
2090 && manifest.global_idx_epoch > 0
2091 && manifest
2092 .runs
2093 .iter()
2094 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2095 });
2096 if let Some(loaded) = checkpoint {
2097 if checkpoint_valid {
2098 db.hot = loaded.hot;
2099 db.bitmap = loaded.bitmap;
2100 db.ann = loaded.ann;
2101 db.fm = loaded.fm;
2102 db.sparse = loaded.sparse;
2103 db.minhash = loaded.minhash;
2104 db.learned_range = Arc::new(loaded.learned_range);
2105 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2110 for (cid, idx) in bitmap0 {
2111 db.bitmap.entry(cid).or_insert(idx);
2112 }
2113 for (cid, idx) in ann0 {
2114 db.ann.entry(cid).or_insert(idx);
2115 }
2116 for (cid, idx) in fm0 {
2117 db.fm.entry(cid).or_insert(idx);
2118 }
2119 for (cid, idx) in sparse0 {
2120 db.sparse.entry(cid).or_insert(idx);
2121 }
2122 for (cid, idx) in minhash0 {
2123 db.minhash.entry(cid).or_insert(idx);
2124 }
2125 }
2128 }
2129 if !checkpoint_valid {
2130 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2131 db.bitmap = bitmap;
2132 db.ann = ann;
2133 db.fm = fm;
2134 db.sparse = sparse;
2135 db.minhash = minhash;
2136 db.rebuild_indexes_from_runs()?;
2137 db.build_learned_ranges()?;
2138 }
2139
2140 for (epoch, group) in replayed_puts {
2145 let (losers, winner_pks) = db.partition_pk_winners(&group);
2146 for (key, &row_id) in &winner_pks {
2147 if let Some(old_rid) = db.hot.get(key) {
2148 if old_rid != row_id {
2149 db.tombstone_row(old_rid, epoch, false);
2150 }
2151 }
2152 }
2153 for &loser_rid in &losers {
2154 db.tombstone_row(loser_rid, epoch, false);
2155 }
2156 for (key, row_id) in winner_pks {
2157 db.insert_hot_pk(key, row_id);
2158 }
2159 if db.schema.primary_key().is_none() {
2160 for r in &group {
2161 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2162 }
2163 }
2164 for r in &group {
2165 if !losers.contains(&r.row_id) {
2166 db.index_row(r);
2167 }
2168 }
2169 }
2170 for (rid, epoch) in &replayed_deletes {
2174 db.remove_hot_for_row(*rid, *epoch);
2175 }
2176
2177 if recovered_manifest_dirty {
2178 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2179 db.live_count = rows.len() as u64;
2180 db.persist_manifest(Epoch(recovered_epoch))?;
2181 }
2182
2183 db.result_cache.lock().load_persistent();
2190 Ok(db)
2191 }
2192
2193 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2199 if self.reservoir_complete {
2200 return Ok(());
2201 }
2202 self.rebuild_reservoir()?;
2203 self.reservoir_complete = true;
2204 Ok(())
2205 }
2206
2207 fn rebuild_reservoir(&mut self) -> Result<()> {
2210 let snap = self.snapshot();
2211 let rows = self.visible_rows(snap)?;
2212 self.reservoir.reset();
2213 for r in rows {
2214 self.reservoir.offer(r.row_id.0);
2215 }
2216 Ok(())
2217 }
2218
2219 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2220 self.rebuild_indexes_from_runs_inner(None)
2221 }
2222
2223 fn rebuild_indexes_from_runs_inner(
2224 &mut self,
2225 control: Option<&crate::ExecutionControl>,
2226 ) -> Result<()> {
2227 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2230 crate::retention::PinSource::OnlineIndexBuild,
2231 self.current_epoch(),
2232 );
2233 self.hot = HotIndex::new();
2234 self.pk_by_row.clear();
2235 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2236 self.bitmap = bitmap;
2237 self.ann = ann;
2238 self.fm = fm;
2239 self.sparse = sparse;
2240 self.minhash = minhash;
2241 let snapshot = Epoch(u64::MAX);
2242 let ttl_now = unix_nanos_now();
2243 let mut scanned = 0_usize;
2244 for rr in self.run_refs.clone() {
2245 if let Some(control) = control {
2246 control.checkpoint()?;
2247 }
2248 let mut reader = self.open_reader(rr.run_id)?;
2249 for row in reader.visible_rows(snapshot)? {
2250 if scanned.is_multiple_of(256) {
2251 if let Some(control) = control {
2252 control.checkpoint()?;
2253 }
2254 }
2255 scanned += 1;
2256 if self.row_expired_at(&row, ttl_now) {
2257 continue;
2258 }
2259 let tok_row = self.tokenized_for_indexes(&row);
2260 index_into(
2261 &self.schema,
2262 &tok_row,
2263 &mut self.hot,
2264 &mut self.bitmap,
2265 &mut self.ann,
2266 &mut self.fm,
2267 &mut self.sparse,
2268 &mut self.minhash,
2269 );
2270 }
2271 }
2272 for row in self.mutable_run.visible_versions(snapshot) {
2273 if scanned.is_multiple_of(256) {
2274 if let Some(control) = control {
2275 control.checkpoint()?;
2276 }
2277 }
2278 scanned += 1;
2279 if row.deleted {
2280 self.remove_hot_for_row(row.row_id, snapshot);
2281 } else if !self.row_expired_at(&row, ttl_now) {
2282 self.index_row(&row);
2283 }
2284 }
2285 for row in self.memtable.visible_versions(snapshot) {
2286 if scanned.is_multiple_of(256) {
2287 if let Some(control) = control {
2288 control.checkpoint()?;
2289 }
2290 }
2291 scanned += 1;
2292 if row.deleted {
2293 self.remove_hot_for_row(row.row_id, snapshot);
2294 } else if !self.row_expired_at(&row, ttl_now) {
2295 self.index_row(&row);
2296 }
2297 }
2298 self.refresh_pk_by_row_from_hot();
2299 Ok(())
2300 }
2301
2302 fn refresh_pk_by_row_from_hot(&mut self) {
2303 self.pk_by_row_complete = true;
2304 if self.schema.primary_key().is_none() {
2305 self.pk_by_row.clear();
2306 return;
2307 }
2308 self.pk_by_row = ReversePkMap::from_entries(
2314 self.hot
2315 .entries()
2316 .into_iter()
2317 .map(|(key, row_id)| (row_id, key)),
2318 );
2319 }
2320
2321 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2322 if self.schema.primary_key().is_some() {
2323 self.pk_by_row.insert(row_id, key.clone());
2324 }
2325 self.hot.insert(key, row_id);
2326 }
2327
2328 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2332 self.build_learned_ranges_inner(None)
2333 }
2334
2335 fn build_learned_ranges_inner(
2336 &mut self,
2337 control: Option<&crate::ExecutionControl>,
2338 ) -> Result<()> {
2339 self.learned_range = Arc::new(HashMap::new());
2340 if self.run_refs.len() != 1 {
2341 return Ok(());
2342 }
2343 let cols: Vec<(u16, usize)> = self
2344 .schema
2345 .indexes
2346 .iter()
2347 .filter(|i| i.kind == IndexKind::LearnedRange)
2348 .map(|i| {
2349 (
2350 i.column_id,
2351 i.options
2352 .learned_range
2353 .as_ref()
2354 .map(|options| options.epsilon)
2355 .unwrap_or(16),
2356 )
2357 })
2358 .collect();
2359 if cols.is_empty() {
2360 return Ok(());
2361 }
2362 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2363 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2364 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2365 _ => return Ok(()),
2366 };
2367 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2368 if column_index % 256 == 0 {
2369 if let Some(control) = control {
2370 control.checkpoint()?;
2371 }
2372 }
2373 let ty = self
2374 .schema
2375 .columns
2376 .iter()
2377 .find(|c| c.id == cid)
2378 .map(|c| c.ty.clone())
2379 .unwrap_or(TypeId::Int64);
2380 match ty {
2381 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2382 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2383 let pairs: Vec<(i64, u64)> = data
2384 .iter()
2385 .zip(row_ids.iter())
2386 .map(|(v, r)| (*v, *r))
2387 .collect();
2388 Arc::make_mut(&mut self.learned_range).insert(
2389 cid,
2390 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2391 );
2392 }
2393 }
2394 TypeId::Float64 => {
2395 if let columnar::NativeColumn::Float64 { data, .. } =
2396 reader.column_native(cid)?
2397 {
2398 let pairs: Vec<(f64, u64)> = data
2399 .iter()
2400 .zip(row_ids.iter())
2401 .map(|(v, r)| (*v, *r))
2402 .collect();
2403 Arc::make_mut(&mut self.learned_range).insert(
2404 cid,
2405 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2406 );
2407 }
2408 }
2409 _ => {}
2410 }
2411 }
2412 Ok(())
2413 }
2414
2415 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2422 if self.indexes_complete {
2423 crate::trace::QueryTrace::record(|t| {
2424 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2425 });
2426 return Ok(());
2427 }
2428 crate::trace::QueryTrace::record(|t| {
2429 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2430 });
2431 self.rebuild_indexes_from_runs()?;
2432 self.build_learned_ranges()?;
2433 self.indexes_complete = true;
2434 let epoch = self.current_epoch();
2435 self.checkpoint_indexes(epoch);
2436 Ok(())
2437 }
2438
2439 #[doc(hidden)]
2442 pub fn ensure_indexes_complete_controlled<F>(
2443 &mut self,
2444 control: &crate::ExecutionControl,
2445 before_publish: F,
2446 ) -> Result<bool>
2447 where
2448 F: FnOnce() -> bool,
2449 {
2450 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2451 .map(|(changed, _)| changed)
2452 }
2453
2454 #[doc(hidden)]
2457 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2458 &mut self,
2459 control: &crate::ExecutionControl,
2460 before_publish: F,
2461 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2462 where
2463 F: FnOnce() -> bool,
2464 {
2465 if self.indexes_complete {
2466 crate::trace::QueryTrace::record(|trace| {
2467 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2468 });
2469 return Ok((false, None));
2470 }
2471 crate::trace::QueryTrace::record(|trace| {
2472 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2473 });
2474 control.checkpoint()?;
2475 let maintenance_epoch = self.current_epoch();
2476 self.rebuild_indexes_from_runs_inner(Some(control))?;
2477 self.build_learned_ranges_inner(Some(control))?;
2478 control.checkpoint()?;
2479 if !before_publish() {
2480 return Err(MongrelError::Cancelled);
2481 }
2482 self.indexes_complete = true;
2483 self.checkpoint_indexes(maintenance_epoch);
2484 Ok((
2485 true,
2486 Some(MaintenanceReceipt {
2487 epoch: maintenance_epoch,
2488 }),
2489 ))
2490 }
2491
2492 fn pending_epoch(&self) -> Epoch {
2493 Epoch(self.epoch.visible().0 + 1)
2494 }
2495
2496 fn is_shared(&self) -> bool {
2499 matches!(self.wal, WalSink::Shared(_))
2500 }
2501
2502 fn ensure_txn_id(&mut self) -> Result<u64> {
2506 if self.current_txn_id == 0 {
2507 let id = match &self.wal {
2508 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2509 WalSink::Private(_) => {
2510 return Err(MongrelError::Full(
2511 "standalone transaction id namespace exhausted".into(),
2512 ))
2513 }
2514 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2515 };
2516 self.current_txn_id = id;
2517 }
2518 Ok(self.current_txn_id)
2519 }
2520
2521 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2524 self.ensure_writable()?;
2525 let txn_id = self.ensure_txn_id()?;
2526 let table_id = self.table_id;
2527 match &mut self.wal {
2528 WalSink::Private(w) => {
2529 w.append_txn(txn_id, op)?;
2530 self.pending_private_mutations = true;
2531 }
2532 WalSink::Shared(s) => {
2533 s.wal.lock().append(txn_id, table_id, op)?;
2534 }
2535 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2536 }
2537 Ok(())
2538 }
2539
2540 fn ensure_writable(&self) -> Result<()> {
2541 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2542 return Err(MongrelError::ReadOnlyReplica);
2543 }
2544 if self.durable_commit_failed {
2545 return Err(MongrelError::Other(
2546 "table poisoned by post-commit failure; reopen required".into(),
2547 ));
2548 }
2549 Ok(())
2550 }
2551
2552 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2563 match &self.auth {
2564 Some(checker) => checker.check(&self.name, perm),
2565 None => Ok(()),
2566 }
2567 }
2568 pub fn require_select(&self) -> Result<()> {
2573 self.require(crate::auth_state::RequiredPermission::Select)
2574 }
2575 fn require_insert(&self) -> Result<()> {
2576 self.require(crate::auth_state::RequiredPermission::Insert)
2577 }
2578 #[allow(dead_code)]
2582 fn require_update(&self) -> Result<()> {
2583 self.require(crate::auth_state::RequiredPermission::Update)
2584 }
2585 fn require_delete(&self) -> Result<()> {
2586 self.require(crate::auth_state::RequiredPermission::Delete)
2587 }
2588
2589 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2592 self.require_insert()?;
2593 Ok(self.put_returning(columns)?.0)
2594 }
2595
2596 pub fn put_returning(
2601 &mut self,
2602 mut columns: Vec<(u16, Value)>,
2603 ) -> Result<(RowId, Option<i64>)> {
2604 self.require_insert()?;
2605 let assigned = self.fill_auto_inc(&mut columns)?;
2606 self.apply_defaults(&mut columns)?;
2607 self.schema.validate_values(&columns)?;
2608 let row_id = if self.schema.clustered {
2613 self.derive_clustered_row_id(&columns)?
2614 } else {
2615 self.allocator.alloc()?
2616 };
2617 let epoch = self.pending_epoch();
2618 let mut row = Row::new(row_id, epoch);
2619 for (col_id, val) in columns {
2620 row.columns.insert(col_id, val);
2621 }
2622 self.commit_rows(vec![row], assigned.is_some())?;
2623 Ok((row_id, assigned))
2624 }
2625
2626 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2629 self.require_insert()?;
2630 Ok(self
2631 .put_batch_returning(batch)?
2632 .into_iter()
2633 .map(|(r, _)| r)
2634 .collect())
2635 }
2636
2637 pub fn put_batch_returning(
2640 &mut self,
2641 batch: Vec<Vec<(u16, Value)>>,
2642 ) -> Result<Vec<(RowId, Option<i64>)>> {
2643 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2644 for mut cols in batch {
2645 let assigned = self.fill_auto_inc(&mut cols)?;
2646 self.apply_defaults(&mut cols)?;
2647 filled.push((cols, assigned));
2648 }
2649 for (cols, _) in &filled {
2650 self.schema.validate_values(cols)?;
2651 }
2652 let epoch = self.pending_epoch();
2653 let mut rows = Vec::with_capacity(filled.len());
2654 let mut ids = Vec::with_capacity(filled.len());
2655 let first_row_id = if self.schema.clustered {
2656 None
2657 } else {
2658 let count = u64::try_from(filled.len())
2659 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2660 Some(self.allocator.alloc_range(count)?.0)
2661 };
2662 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2663 let row_id = match first_row_id {
2664 Some(first) => RowId(first + row_index as u64),
2665 None => self.derive_clustered_row_id(&cols)?,
2666 };
2667 let mut row = Row::new(row_id, epoch);
2668 for (c, v) in cols {
2669 row.columns.insert(c, v);
2670 }
2671 ids.push((row_id, assigned));
2672 rows.push(row);
2673 }
2674 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2675 self.commit_rows(rows, all_auto_generated)?;
2676 Ok(ids)
2677 }
2678
2679 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2685 self.ensure_writable()?;
2686 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2687 return Ok(None);
2688 };
2689 let pos = columns.iter().position(|(c, _)| *c == cid);
2690 let assigned = match pos {
2691 Some(i) => match &columns[i].1 {
2692 Value::Null => {
2693 let next = self.alloc_auto_inc_value()?;
2694 columns[i].1 = Value::Int64(next);
2695 Some(next)
2696 }
2697 Value::Int64(n) => {
2698 self.advance_auto_inc_past(*n)?;
2699 None
2700 }
2701 other => {
2702 return Err(MongrelError::InvalidArgument(format!(
2703 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2704 other
2705 )))
2706 }
2707 },
2708 None => {
2709 let next = self.alloc_auto_inc_value()?;
2710 columns.push((cid, Value::Int64(next)));
2711 Some(next)
2712 }
2713 };
2714 Ok(assigned)
2715 }
2716
2717 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2723 for col in &self.schema.columns {
2724 let Some(expr) = &col.default_value else {
2725 continue;
2726 };
2727 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2729 continue;
2730 }
2731 let pos = columns.iter().position(|(c, _)| *c == col.id);
2732 let needs_default = match pos {
2733 None => true,
2734 Some(i) => matches!(columns[i].1, Value::Null),
2735 };
2736 if !needs_default {
2737 continue;
2738 }
2739 let v = match expr {
2740 crate::schema::DefaultExpr::Static(v) => v.clone(),
2741 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2742 crate::schema::DefaultExpr::Uuid => {
2743 let mut buf = [0u8; 16];
2744 getrandom::getrandom(&mut buf)
2745 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2746 Value::Uuid(buf)
2747 }
2748 };
2749 match pos {
2750 None => columns.push((col.id, v)),
2751 Some(i) => columns[i].1 = v,
2752 }
2753 }
2754 Ok(())
2755 }
2756
2757 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2759 self.ensure_auto_inc_seeded()?;
2760 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2762 let v = ai.next;
2763 ai.next = ai
2764 .next
2765 .checked_add(1)
2766 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2767 Ok(v)
2768 }
2769
2770 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2773 self.ensure_auto_inc_seeded()?;
2774 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2775 let floor = used
2776 .checked_add(1)
2777 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2778 .max(1);
2779 if ai.next < floor {
2780 ai.next = floor;
2781 }
2782 Ok(())
2783 }
2784
2785 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2790 let needs_seed = match self.auto_inc {
2791 Some(ai) => !ai.seeded,
2792 None => return Ok(()),
2793 };
2794 if !needs_seed {
2795 return Ok(());
2796 }
2797 if self.seed_empty_auto_inc() {
2798 return Ok(());
2799 }
2800 let cid = self
2801 .auto_inc
2802 .as_ref()
2803 .expect("auto-inc column present")
2804 .column_id;
2805 let max = self.scan_max_int64(cid)?;
2806 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2807 let floor = max
2808 .checked_add(1)
2809 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2810 .max(1);
2811 if ai.next < floor {
2812 ai.next = floor;
2813 }
2814 ai.seeded = true;
2815 Ok(())
2816 }
2817
2818 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2819 if n == 0 || self.auto_inc.is_none() {
2820 return Ok(None);
2821 }
2822 self.ensure_auto_inc_seeded()?;
2823 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2824 let start = ai.next;
2825 let count = i64::try_from(n)
2826 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2827 ai.next = ai
2828 .next
2829 .checked_add(count)
2830 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2831 Ok(Some(start))
2832 }
2833
2834 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2838 let mut max: i64 = 0;
2839 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2840 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2841 if *n > max {
2842 max = *n;
2843 }
2844 }
2845 }
2846 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2847 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2848 if *n > max {
2849 max = *n;
2850 }
2851 }
2852 }
2853 for rr in self.run_refs.clone() {
2854 let reader = self.open_reader(rr.run_id)?;
2855 if let Some(stats) = reader.column_page_stats(column_id) {
2856 for s in stats {
2857 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2858 if n > max {
2859 max = n;
2860 }
2861 }
2862 }
2863 } else if reader.has_column(column_id) {
2864 if let columnar::NativeColumn::Int64 { data, validity } =
2865 reader.column_native_shared(column_id)?
2866 {
2867 for (i, n) in data.iter().enumerate() {
2868 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2869 {
2870 max = *n;
2871 }
2872 }
2873 }
2874 }
2875 }
2876 Ok(max)
2877 }
2878
2879 fn seed_empty_auto_inc(&mut self) -> bool {
2880 let Some(ai) = self.auto_inc.as_mut() else {
2881 return false;
2882 };
2883 if ai.seeded || self.live_count != 0 {
2884 return false;
2885 }
2886 if ai.next < 1 {
2887 ai.next = 1;
2888 }
2889 ai.seeded = true;
2890 true
2891 }
2892
2893 fn advance_auto_inc_from_native_columns(
2894 &mut self,
2895 columns: &[(u16, columnar::NativeColumn)],
2896 n: usize,
2897 live_before: u64,
2898 ) -> Result<()> {
2899 let Some(ai) = self.auto_inc.as_mut() else {
2900 return Ok(());
2901 };
2902 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2903 return Ok(());
2904 };
2905 let columnar::NativeColumn::Int64 { data, validity } = col else {
2906 return Err(MongrelError::InvalidArgument(format!(
2907 "AUTO_INCREMENT column {} must be Int64",
2908 ai.column_id
2909 )));
2910 };
2911 let max = if native_int64_strictly_increasing(col, n) {
2912 data.get(n.saturating_sub(1)).copied()
2913 } else {
2914 data.iter()
2915 .take(n)
2916 .enumerate()
2917 .filter_map(|(i, v)| {
2918 if validity.is_empty() || columnar::validity_bit(validity, i) {
2919 Some(*v)
2920 } else {
2921 None
2922 }
2923 })
2924 .max()
2925 };
2926 if let Some(max) = max {
2927 let floor = max
2928 .checked_add(1)
2929 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2930 .max(1);
2931 if ai.next < floor {
2932 ai.next = floor;
2933 }
2934 if ai.seeded || live_before == 0 {
2935 ai.seeded = true;
2936 }
2937 }
2938 Ok(())
2939 }
2940
2941 fn fill_auto_inc_native_columns(
2942 &mut self,
2943 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2944 n: usize,
2945 ) -> Result<()> {
2946 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2947 return Ok(());
2948 };
2949 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2950 if let Some(start) = self.alloc_auto_inc_range(n)? {
2951 columns.push((
2952 cid,
2953 columnar::NativeColumn::Int64 {
2954 data: (start..start.saturating_add(n as i64)).collect(),
2955 validity: vec![0xFF; n.div_ceil(8)],
2956 },
2957 ));
2958 }
2959 return Ok(());
2960 };
2961
2962 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2963 return Err(MongrelError::InvalidArgument(format!(
2964 "AUTO_INCREMENT column {cid} must be Int64"
2965 )));
2966 };
2967 if data.len() < n {
2968 return Err(MongrelError::InvalidArgument(format!(
2969 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2970 data.len()
2971 )));
2972 }
2973 if columnar::all_non_null(validity, n) {
2974 return Ok(());
2975 }
2976 if validity.iter().all(|b| *b == 0) {
2977 if let Some(start) = self.alloc_auto_inc_range(n)? {
2978 for (i, slot) in data.iter_mut().take(n).enumerate() {
2979 *slot = start.saturating_add(i as i64);
2980 }
2981 *validity = vec![0xFF; n.div_ceil(8)];
2982 }
2983 return Ok(());
2984 }
2985
2986 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2987 for (i, slot) in data.iter_mut().enumerate().take(n) {
2988 if columnar::validity_bit(validity, i) {
2989 self.advance_auto_inc_past(*slot)?;
2990 } else {
2991 *slot = self.alloc_auto_inc_value()?;
2992 }
2993 }
2994 *validity = new_validity;
2995 Ok(())
2996 }
2997
2998 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3012 self.ensure_writable()?;
3013 if self.auto_inc.is_none() {
3014 return Ok(None);
3015 }
3016 Ok(Some(self.alloc_auto_inc_value()?))
3017 }
3018
3019 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3025 let payload = bincode::serialize(&rows)?;
3026 self.wal_append_data(Op::Put {
3027 table_id: self.table_id,
3028 rows: payload,
3029 })?;
3030 if self.is_shared() {
3031 self.pending_rows_auto_inc
3032 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3033 self.pending_rows.extend(rows);
3034 } else {
3035 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3036 }
3037 Ok(())
3038 }
3039
3040 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3043 self.ensure_indexes_complete()
3044 }
3045
3046 pub(crate) fn prepare_durable_publish_controlled(
3047 &mut self,
3048 control: &crate::ExecutionControl,
3049 ) -> Result<()> {
3050 self.ensure_indexes_complete_controlled(control, || true)?;
3051 Ok(())
3052 }
3053
3054 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3055 self.apply_put_rows_inner_prepared(rows, true);
3056 }
3057
3058 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3059 if check_existing_pk {
3060 self.ensure_indexes_complete()?;
3061 }
3062 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3063 Ok(())
3064 }
3065
3066 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3070 if rows.len() == 1 {
3074 let row = rows.into_iter().next().expect("len checked");
3075 self.apply_put_row_single(row, check_existing_pk);
3076 return;
3077 }
3078 let pk_id = self.schema.primary_key().map(|c| c.id);
3095 let probe = match pk_id {
3096 Some(pid) => {
3097 check_existing_pk
3098 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3099 }
3100 None => false,
3101 };
3102 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3105 for r in rows {
3106 for &cid in r.columns.keys() {
3107 self.pending_put_cols.insert(cid);
3108 }
3109 match pk_id {
3110 Some(pid) if probe || maintain_pk_by_row => {
3111 if let Some(pk_val) = r.columns.get(&pid) {
3112 let key = self.index_lookup_key(pid, pk_val);
3113 if probe {
3114 if let Some(old_rid) = self.hot.get(&key) {
3115 if old_rid != r.row_id {
3116 self.tombstone_row(old_rid, r.committed_epoch, true);
3117 }
3118 }
3119 }
3120 if maintain_pk_by_row {
3121 self.pk_by_row.insert(r.row_id, key);
3122 }
3123 }
3124 }
3125 Some(_) => {}
3126 None => {
3127 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3128 }
3129 }
3130 self.index_row(&r);
3131 self.reservoir.offer(r.row_id.0);
3132 self.memtable.upsert(r);
3133 self.live_count = self.live_count.saturating_add(1);
3136 }
3137 self.data_generation = self.data_generation.wrapping_add(1);
3138 }
3139
3140 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3144 for &cid in row.columns.keys() {
3145 self.pending_put_cols.insert(cid);
3146 }
3147 let epoch = row.committed_epoch;
3148 if let Some(pk_col) = self.schema.primary_key() {
3149 let pk_id = pk_col.id;
3150 if let Some(pk_val) = row.columns.get(&pk_id) {
3151 let maintain_pk_by_row = self.pk_by_row_complete;
3155 if check_existing_pk || maintain_pk_by_row {
3156 let key = self.index_lookup_key(pk_id, pk_val);
3157 if check_existing_pk {
3158 if let Some(old_rid) = self.hot.get(&key) {
3159 if old_rid != row.row_id {
3160 self.tombstone_row(old_rid, epoch, true);
3161 }
3162 }
3163 }
3164 if maintain_pk_by_row {
3165 self.pk_by_row.insert(row.row_id, key);
3166 }
3167 }
3168 }
3169 } else {
3170 self.hot
3171 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3172 }
3173 self.index_row(&row);
3174 self.reservoir.offer(row.row_id.0);
3175 self.memtable.upsert(row);
3176 self.live_count = self.live_count.saturating_add(1);
3177 self.data_generation = self.data_generation.wrapping_add(1);
3178 }
3179
3180 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3183 self.allocator.alloc()
3184 }
3185
3186 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3192 let pk = self.schema.primary_key().ok_or_else(|| {
3193 MongrelError::Schema("clustered table requires a single-column primary key".into())
3194 })?;
3195 let pk_val = columns
3196 .iter()
3197 .find(|(id, _)| *id == pk.id)
3198 .map(|(_, v)| v)
3199 .ok_or_else(|| {
3200 MongrelError::Schema(format!(
3201 "clustered table missing primary key column {} ({})",
3202 pk.id, pk.name
3203 ))
3204 })?;
3205 let key_bytes = pk_val.encode_key();
3206 let mut hash: u64 = 0xcbf29ce484222325;
3208 for &b in &key_bytes {
3209 hash ^= b as u64;
3210 hash = hash.wrapping_mul(0x100000001b3);
3211 }
3212 Ok(RowId(hash.max(1)))
3215 }
3216
3217 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3225 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3226 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3227 }
3228 let n = rows.len();
3229 for r in rows {
3230 for &cid in r.columns.keys() {
3231 self.pending_put_cols.insert(cid);
3232 }
3233 }
3234 let (losers, winner_pks) = self.partition_pk_winners(rows);
3235 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3236 for (key, &row_id) in &winner_pks {
3238 if let Some(old_rid) = self.hot.get(key) {
3239 if old_rid != row_id {
3240 self.tombstone_row(old_rid, epoch, true);
3241 }
3242 }
3243 }
3244 for &loser_rid in &losers {
3247 self.tombstone_row(loser_rid, epoch, false);
3248 }
3249 for (key, row_id) in winner_pks {
3251 self.insert_hot_pk(key, row_id);
3252 }
3253 if self.schema.primary_key().is_none() {
3254 for r in rows {
3255 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3256 }
3257 }
3258 for r in rows {
3259 self.allocator.advance_to(r.row_id)?;
3260 if !losers.contains(&r.row_id) {
3261 self.index_row(r);
3262 }
3263 }
3264 for r in rows {
3265 if !losers.contains(&r.row_id) {
3266 self.reservoir.offer(r.row_id.0);
3267 }
3268 }
3269 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3270 self.data_generation = self.data_generation.wrapping_add(1);
3271 Ok(())
3272 }
3273
3274 pub(crate) fn recover_apply(
3279 &mut self,
3280 rows: Vec<Row>,
3281 deletes: Vec<(RowId, Epoch)>,
3282 ) -> Result<()> {
3283 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3287 std::collections::BTreeMap::new();
3288 for row in rows {
3289 if row.row_id.0 >= u64::MAX - 1 {
3290 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3291 }
3292 self.allocator.advance_to(row.row_id)?;
3293 if let Some(ai) = self.auto_inc.as_mut() {
3298 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3299 let next = n.checked_add(1).ok_or_else(|| {
3300 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3301 })?;
3302 if next > ai.next {
3303 ai.next = next;
3304 }
3305 }
3306 }
3307 by_epoch.entry(row.committed_epoch).or_default().push(row);
3308 }
3309 for (epoch, group) in by_epoch {
3310 let (losers, winner_pks) = self.partition_pk_winners(&group);
3311 for (key, &row_id) in &winner_pks {
3313 if let Some(old_rid) = self.hot.get(key) {
3314 if old_rid != row_id {
3315 self.tombstone_row(old_rid, epoch, false);
3316 }
3317 }
3318 }
3319 for (key, row_id) in winner_pks {
3320 self.insert_hot_pk(key, row_id);
3321 }
3322 if self.schema.primary_key().is_none() {
3323 for r in &group {
3324 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3325 }
3326 }
3327 for r in &group {
3328 if !losers.contains(&r.row_id) {
3329 self.memtable.upsert(r.clone());
3330 self.index_row(r);
3331 }
3332 }
3333 }
3334 for (rid, epoch) in deletes {
3335 self.memtable.tombstone(rid, epoch);
3336 self.remove_hot_for_row(rid, epoch);
3337 }
3338 self.reservoir_complete = false;
3341 Ok(())
3342 }
3343
3344 pub(crate) fn flushed_epoch(&self) -> u64 {
3346 self.flushed_epoch
3347 }
3348
3349 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3350 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3351 }
3352
3353 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3355 self.schema.validate_values(cells)
3356 }
3357
3358 fn validate_columns_not_null(
3362 &self,
3363 columns: &[(u16, columnar::NativeColumn)],
3364 n: usize,
3365 ) -> Result<()> {
3366 let by_id: HashMap<u16, &columnar::NativeColumn> =
3367 columns.iter().map(|(id, c)| (*id, c)).collect();
3368 for col in &self.schema.columns {
3369 if !col.flags.contains(ColumnFlags::NULLABLE) {
3370 match by_id.get(&col.id) {
3371 None => {
3372 return Err(MongrelError::InvalidArgument(format!(
3373 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3374 col.name, col.id
3375 )));
3376 }
3377 Some(c) => {
3378 if c.null_count(n) != 0 {
3379 return Err(MongrelError::InvalidArgument(format!(
3380 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3381 col.name, col.id
3382 )));
3383 }
3384 }
3385 }
3386 }
3387 if let TypeId::Enum { variants } = &col.ty {
3388 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3389 if by_id.contains_key(&col.id) {
3390 return Err(MongrelError::InvalidArgument(format!(
3391 "column '{}' ({}) enum requires a bytes column",
3392 col.name, col.id
3393 )));
3394 }
3395 continue;
3396 };
3397 for index in 0..n {
3398 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3399 continue;
3400 };
3401 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3402 return Err(MongrelError::InvalidArgument(format!(
3403 "column '{}' ({}) enum value {:?} is not one of {:?}",
3404 col.name,
3405 col.id,
3406 String::from_utf8_lossy(value),
3407 variants
3408 )));
3409 }
3410 }
3411 }
3412 }
3413 Ok(())
3414 }
3415
3416 fn bulk_pk_winner_indices(
3421 &self,
3422 columns: &[(u16, columnar::NativeColumn)],
3423 n: usize,
3424 ) -> Option<Vec<usize>> {
3425 let pk_col = self.schema.primary_key()?;
3426 let pk_id = pk_col.id;
3427 let pk_ty = pk_col.ty.clone();
3428 let by_id: HashMap<u16, &columnar::NativeColumn> =
3429 columns.iter().map(|(id, c)| (*id, c)).collect();
3430 let pk_native = by_id.get(&pk_id)?;
3431 if native_int64_strictly_increasing(pk_native, n) {
3432 return None;
3433 }
3434 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3436 let mut null_pk_rows: Vec<usize> = Vec::new();
3437 for i in 0..n {
3438 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3439 Some(key) => {
3440 last.insert(key, i);
3441 }
3442 None => null_pk_rows.push(i),
3443 }
3444 }
3445 let mut winners: HashSet<usize> = last.values().copied().collect();
3446 for i in null_pk_rows {
3447 winners.insert(i);
3448 }
3449 Some((0..n).filter(|i| winners.contains(i)).collect())
3450 }
3451
3452 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3454 self.require_delete()?;
3455 let epoch = self.pending_epoch();
3456 self.wal_append_data(Op::Delete {
3457 table_id: self.table_id,
3458 row_ids: vec![row_id],
3459 })?;
3460 if self.is_shared() {
3461 self.pending_dels.push(row_id);
3462 } else {
3463 self.apply_delete(row_id, epoch);
3464 }
3465 Ok(())
3466 }
3467
3468 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3469 let pre = self.get(row_id, self.snapshot());
3470 self.delete(row_id)?;
3471 Ok(pre.map(|row| {
3472 let mut columns: Vec<_> = row.columns.into_iter().collect();
3473 columns.sort_by_key(|(id, _)| *id);
3474 OwnedRow { columns }
3475 }))
3476 }
3477
3478 pub fn truncate(&mut self) -> Result<()> {
3480 self.require_delete()?;
3481 let epoch = self.pending_epoch();
3482 self.wal_append_data(Op::TruncateTable {
3483 table_id: self.table_id,
3484 })?;
3485 self.pending_rows.clear();
3486 self.pending_rows_auto_inc.clear();
3487 self.pending_dels.clear();
3488 self.pending_truncate = Some(epoch);
3489 Ok(())
3490 }
3491
3492 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3494 self.run_refs.clear();
3499 self.retiring.clear();
3500 self.memtable = Memtable::new();
3501 self.mutable_run = MutableRun::new();
3502 self.hot = HotIndex::new();
3503 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3504 self.bitmap = bitmap;
3505 self.ann = ann;
3506 self.fm = fm;
3507 self.sparse = sparse;
3508 self.minhash = minhash;
3509 self.learned_range = Arc::new(HashMap::new());
3510 self.pk_by_row.clear();
3511 self.pk_by_row_complete = false;
3512 self.live_count = 0;
3513 self.reservoir = crate::reservoir::Reservoir::default();
3514 self.reservoir_complete = true;
3515 self.had_deletes = true;
3516 self.agg_cache = Arc::new(HashMap::new());
3517 self.global_idx_epoch = 0;
3518 self.indexes_complete = true;
3519 self.pending_delete_rids.clear();
3520 self.pending_put_cols.clear();
3521 self.pending_rows.clear();
3522 self.pending_rows_auto_inc.clear();
3523 self.pending_dels.clear();
3524 self.clear_result_cache();
3525 self.invalidate_index_checkpoint();
3526 self.data_generation = self.data_generation.wrapping_add(1);
3527 }
3528
3529 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3532 self.remove_hot_for_row(row_id, epoch);
3533 self.tombstone_row(row_id, epoch, true);
3534 self.data_generation = self.data_generation.wrapping_add(1);
3535 }
3536
3537 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3541 let tombstone = Row {
3542 row_id,
3543 committed_epoch: epoch,
3544 columns: std::collections::HashMap::new(),
3545 deleted: true,
3546 };
3547 self.memtable.upsert(tombstone);
3548 self.pk_by_row.remove(&row_id);
3549 if adjust_live_count {
3550 self.live_count = self.live_count.saturating_sub(1);
3551 }
3552 self.pending_delete_rids.insert(row_id.0 as u32);
3554 self.had_deletes = true;
3557 self.agg_cache = Arc::new(HashMap::new());
3558 }
3559
3560 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3564 let Some(pk_col) = self.schema.primary_key() else {
3565 return;
3566 };
3567 if self.pk_by_row_complete {
3570 if let Some(key) = self.pk_by_row.remove(&row_id) {
3571 if self.hot.get(&key) == Some(row_id) {
3572 self.hot.remove(&key);
3573 }
3574 }
3575 return;
3576 }
3577 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3596 if self.indexes_complete {
3597 let pk_val = self
3598 .memtable
3599 .get_version(row_id, lookup_epoch)
3600 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3601 .or_else(|| {
3602 self.mutable_run
3603 .get_version(row_id, lookup_epoch)
3604 .filter(|(_, r)| !r.deleted)
3605 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3606 })
3607 .or_else(|| {
3608 self.run_refs.iter().find_map(|rr| {
3609 let mut reader = self.open_reader(rr.run_id).ok()?;
3610 let (_, deleted, val) = reader
3611 .get_version_column(row_id, lookup_epoch, pk_col.id)
3612 .ok()??;
3613 if deleted {
3614 return None;
3615 }
3616 val
3617 })
3618 });
3619 if let Some(pk_val) = pk_val {
3620 let key = self.index_lookup_key(pk_col.id, &pk_val);
3621 if self.hot.get(&key) == Some(row_id) {
3622 self.hot.remove(&key);
3623 }
3624 return;
3625 }
3626 }
3627 self.refresh_pk_by_row_from_hot();
3632 if let Some(key) = self.pk_by_row.remove(&row_id) {
3633 if self.hot.get(&key) == Some(row_id) {
3634 self.hot.remove(&key);
3635 }
3636 }
3637 }
3638
3639 fn partition_pk_winners(
3644 &self,
3645 rows: &[Row],
3646 ) -> (
3647 std::collections::HashSet<RowId>,
3648 std::collections::HashMap<Vec<u8>, RowId>,
3649 ) {
3650 let mut losers = std::collections::HashSet::new();
3651 let Some(pk_col) = self.schema.primary_key() else {
3652 return (losers, std::collections::HashMap::new());
3653 };
3654 let pk_id = pk_col.id;
3655 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3656 std::collections::HashMap::new();
3657 for r in rows {
3658 let Some(pk_val) = r.columns.get(&pk_id) else {
3659 continue;
3660 };
3661 let key = self.index_lookup_key(pk_id, pk_val);
3662 if let Some(&old_rid) = winners.get(&key) {
3663 losers.insert(old_rid);
3664 }
3665 winners.insert(key, r.row_id);
3666 }
3667 (losers, winners)
3668 }
3669
3670 fn index_row(&mut self, row: &Row) {
3671 if row.deleted {
3672 return;
3673 }
3674 let any_predicate = self
3682 .schema
3683 .indexes
3684 .iter()
3685 .any(|idx| idx.predicate.is_some());
3686 if any_predicate {
3687 let columns_map: HashMap<u16, &Value> =
3688 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3689 let name_to_id: HashMap<&str, u16> = self
3690 .schema
3691 .columns
3692 .iter()
3693 .map(|c| (c.name.as_str(), c.id))
3694 .collect();
3695 for idx in &self.schema.indexes {
3696 if let Some(pred) = &idx.predicate {
3697 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3698 continue; }
3700 }
3701 index_into_single(
3703 idx,
3704 &self.schema,
3705 row,
3706 &mut self.hot,
3707 &mut self.bitmap,
3708 &mut self.ann,
3709 &mut self.fm,
3710 &mut self.sparse,
3711 &mut self.minhash,
3712 );
3713 }
3714 return;
3715 }
3716 if self.column_keys.is_empty() {
3720 index_into(
3721 &self.schema,
3722 row,
3723 &mut self.hot,
3724 &mut self.bitmap,
3725 &mut self.ann,
3726 &mut self.fm,
3727 &mut self.sparse,
3728 &mut self.minhash,
3729 );
3730 return;
3731 }
3732 let effective_row = self.tokenized_for_indexes(row);
3733 index_into(
3734 &self.schema,
3735 &effective_row,
3736 &mut self.hot,
3737 &mut self.bitmap,
3738 &mut self.ann,
3739 &mut self.fm,
3740 &mut self.sparse,
3741 &mut self.minhash,
3742 );
3743 }
3744
3745 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3751 if self.column_keys.is_empty() {
3752 return row.clone();
3753 }
3754 #[cfg(feature = "encryption")]
3755 {
3756 use crate::encryption::SCHEME_HMAC_EQ;
3757 let mut tok = row.clone();
3758 for (&cid, &(_, scheme)) in &self.column_keys {
3759 if scheme != SCHEME_HMAC_EQ {
3760 continue;
3761 }
3762 if let Some(v) = tok.columns.get(&cid).cloned() {
3763 if let Some(t) = self.tokenize_value(cid, &v) {
3764 tok.columns.insert(cid, t);
3765 }
3766 }
3767 }
3768 tok
3769 }
3770 #[cfg(not(feature = "encryption"))]
3771 {
3772 row.clone()
3773 }
3774 }
3775
3776 pub fn commit(&mut self) -> Result<Epoch> {
3781 self.commit_inner(None)
3782 }
3783
3784 #[doc(hidden)]
3787 pub fn commit_controlled<F>(
3788 &mut self,
3789 control: &crate::ExecutionControl,
3790 mut before_commit: F,
3791 ) -> Result<Epoch>
3792 where
3793 F: FnMut() -> Result<()>,
3794 {
3795 self.commit_inner(Some((control, &mut before_commit)))
3796 }
3797
3798 fn commit_inner(
3799 &mut self,
3800 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3801 ) -> Result<Epoch> {
3802 self.ensure_writable()?;
3803 if !self.has_pending_mutations() {
3804 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3805 return Err(MongrelError::Full(
3806 "standalone transaction id namespace exhausted".into(),
3807 ));
3808 }
3809 return Ok(self.epoch.visible());
3810 }
3811 self.commit_new_epoch_inner(controlled)
3812 }
3813
3814 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3818 self.commit_new_epoch_inner(None)
3819 }
3820
3821 fn commit_new_epoch_inner(
3822 &mut self,
3823 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3824 ) -> Result<Epoch> {
3825 self.ensure_writable()?;
3826 if self.is_shared() {
3827 self.commit_shared(controlled)
3828 } else {
3829 self.commit_private(controlled)
3830 }
3831 }
3832
3833 fn commit_private(
3835 &mut self,
3836 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3837 ) -> Result<Epoch> {
3838 let commit_lock = Arc::clone(&self.commit_lock);
3842 let _g = commit_lock.lock();
3843 let txn_id = self.ensure_txn_id()?;
3846 if let Some((control, before_commit)) = controlled {
3847 control.checkpoint()?;
3848 before_commit()?;
3849 }
3850 let new_epoch = self.epoch.bump_assigned();
3851 let epoch_authority = Arc::clone(&self.epoch);
3852 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3853 let wal_result = match &mut self.wal {
3857 WalSink::Private(w) => w
3858 .append_txn(
3859 txn_id,
3860 Op::TxnCommit {
3861 epoch: new_epoch.0,
3862 added_runs: Vec::new(),
3863 },
3864 )
3865 .and_then(|_| w.sync()),
3866 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3867 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
3868 };
3869 if let Err(error) = wal_result {
3870 self.durable_commit_failed = true;
3871 return Err(MongrelError::CommitOutcomeUnknown {
3872 epoch: new_epoch.0,
3873 message: error.to_string(),
3874 });
3875 }
3876 if let Some(epoch) = self.pending_truncate.take() {
3879 self.apply_truncate(epoch);
3880 }
3881 self.invalidate_pending_cache();
3882 let publish_result = self.persist_manifest(new_epoch);
3883 self.epoch.publish_in_order(new_epoch);
3887 epoch_guard.disarm();
3888 if let Err(error) = publish_result {
3889 self.durable_commit_failed = true;
3890 return Err(MongrelError::DurableCommit {
3891 epoch: new_epoch.0,
3892 message: error.to_string(),
3893 });
3894 }
3895 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
3896 self.pending_private_mutations = false;
3897 self.data_generation = self.data_generation.wrapping_add(1);
3898 Ok(new_epoch)
3899 }
3900
3901 fn commit_shared(
3907 &mut self,
3908 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3909 ) -> Result<Epoch> {
3910 use std::sync::atomic::Ordering;
3911 let s = match &self.wal {
3912 WalSink::Shared(s) => s.clone(),
3913 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3914 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3915 };
3916 if s.poisoned.load(Ordering::Relaxed) {
3917 return Err(MongrelError::Other(
3918 "database poisoned by fsync error".into(),
3919 ));
3920 }
3921 let commit_lock = Arc::clone(&self.commit_lock);
3928 let _g = commit_lock.lock();
3929 if !self.pending_rows.is_empty() {
3930 match controlled.as_ref() {
3931 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
3932 None => self.prepare_durable_publish()?,
3933 }
3934 }
3935 let txn_id = self.ensure_txn_id()?;
3938 let mut wal = s.wal.lock();
3939 if let Some((control, before_commit)) = controlled {
3940 control.checkpoint()?;
3941 before_commit()?;
3942 }
3943 let new_epoch = self.epoch.bump_assigned();
3944 let epoch_authority = Arc::clone(&self.epoch);
3945 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3946 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
3947 Ok(commit_seq) => commit_seq,
3948 Err(error) => {
3949 s.poisoned.store(true, Ordering::Relaxed);
3950 s.lifecycle.poison();
3951 return Err(MongrelError::CommitOutcomeUnknown {
3952 epoch: new_epoch.0,
3953 message: error.to_string(),
3954 });
3955 }
3956 };
3957 drop(wal);
3958 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
3959 s.poisoned.store(true, Ordering::Relaxed);
3960 s.lifecycle.poison();
3961 return Err(MongrelError::CommitOutcomeUnknown {
3962 epoch: new_epoch.0,
3963 message: error.to_string(),
3964 });
3965 }
3966
3967 if self.pending_truncate.take().is_some() {
3970 self.apply_truncate(new_epoch);
3971 }
3972 let mut rows = std::mem::take(&mut self.pending_rows);
3973 if !rows.is_empty() {
3974 for r in &mut rows {
3975 r.committed_epoch = new_epoch;
3976 }
3977 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3978 let all_auto_generated =
3979 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3980 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
3981 } else {
3982 self.pending_rows_auto_inc.clear();
3983 }
3984 let dels = std::mem::take(&mut self.pending_dels);
3985 for rid in dels {
3986 self.apply_delete(rid, new_epoch);
3987 }
3988
3989 self.invalidate_pending_cache();
3990 let publish_result = self.persist_manifest(new_epoch);
3991 self.epoch.publish_in_order(new_epoch);
3992 epoch_guard.disarm();
3993 let _ = s.change_wake.send(());
3994 if let Err(error) = publish_result {
3995 self.durable_commit_failed = true;
3996 s.poisoned.store(true, Ordering::Relaxed);
3997 s.lifecycle.poison();
3998 return Err(MongrelError::DurableCommit {
3999 epoch: new_epoch.0,
4000 message: error.to_string(),
4001 });
4002 }
4003 self.current_txn_id = 0;
4005 self.data_generation = self.data_generation.wrapping_add(1);
4006 Ok(new_epoch)
4007 }
4008
4009 pub fn flush(&mut self) -> Result<Epoch> {
4017 self.flush_with_outcome().map(|(epoch, _)| epoch)
4018 }
4019
4020 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4022 self.flush_with_outcome_inner(None)
4023 }
4024
4025 #[doc(hidden)]
4028 pub fn flush_with_outcome_controlled<F>(
4029 &mut self,
4030 control: &crate::ExecutionControl,
4031 mut before_commit: F,
4032 ) -> Result<(Epoch, bool)>
4033 where
4034 F: FnMut() -> Result<()>,
4035 {
4036 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4037 }
4038
4039 fn flush_with_outcome_inner(
4040 &mut self,
4041 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4042 ) -> Result<(Epoch, bool)> {
4043 match controlled.as_ref() {
4044 Some((control, _)) => {
4045 self.ensure_indexes_complete_controlled(control, || true)?;
4046 }
4047 None => self.ensure_indexes_complete()?,
4048 }
4049 let committed = self.has_pending_mutations();
4050 let epoch = self.commit_inner(controlled)?;
4051 let finish: Result<(Epoch, bool)> = (|| {
4052 let rows = self.memtable.drain_sorted();
4053 if !rows.is_empty() {
4054 self.mutable_run.insert_many(rows);
4055 }
4056 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4057 self.spill_mutable_run(epoch)?;
4058 self.mark_flushed(epoch)?;
4062 self.persist_manifest(epoch)?;
4063 self.build_learned_ranges()?;
4064 self.checkpoint_indexes(epoch);
4067 }
4068 Ok((epoch, committed))
4071 })();
4072 let outcome = match finish {
4073 Err(error) if committed => Err(MongrelError::DurableCommit {
4074 epoch: epoch.0,
4075 message: error.to_string(),
4076 }),
4077 result => result,
4078 };
4079 if outcome.is_ok() {
4080 let _ = self.publish_read_generation();
4086 }
4087 outcome
4088 }
4089
4090 fn has_pending_mutations(&self) -> bool {
4091 self.pending_private_mutations
4092 || !self.pending_rows.is_empty()
4093 || !self.pending_dels.is_empty()
4094 || self.pending_truncate.is_some()
4095 }
4096
4097 pub fn has_pending_writes(&self) -> bool {
4098 self.has_pending_mutations()
4099 }
4100
4101 pub fn force_flush(&mut self) -> Result<Epoch> {
4110 let saved = self.mutable_run_spill_bytes;
4111 self.mutable_run_spill_bytes = 1;
4112 let result = self.flush();
4113 self.mutable_run_spill_bytes = saved;
4114 result
4115 }
4116
4117 pub fn close(&mut self) -> Result<()> {
4124 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4125 self.force_flush()?;
4126 }
4127 Ok(())
4128 }
4129
4130 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4137 let op = Op::Flush {
4138 table_id: self.table_id,
4139 flushed_epoch: epoch.0,
4140 };
4141 match &mut self.wal {
4142 WalSink::Private(w) => {
4143 w.append_system(op)?;
4144 w.sync()?;
4145 }
4146 WalSink::Shared(s) => {
4147 s.wal.lock().append_system(op)?;
4152 }
4153 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4154 }
4155 self.flushed_epoch = epoch.0;
4156 if matches!(self.wal, WalSink::Private(_)) {
4157 self.rotate_wal(epoch)?;
4158 }
4159 Ok(())
4160 }
4161
4162 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4166 if self.mutable_run.is_empty() {
4167 return Ok(());
4168 }
4169 let run_id = self.alloc_run_id()?;
4170 let rows = self.mutable_run.drain_sorted();
4171 if rows.is_empty() {
4172 return Ok(());
4173 }
4174 let path = self.run_path(run_id);
4175 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4176 if let Some(kek) = &self.kek {
4177 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4178 }
4179 let header = match self.create_run_file(run_id)? {
4180 Some(file) => writer.write_file(file, &rows)?,
4181 None => writer.write(&path, &rows)?,
4182 };
4183 self.run_refs.push(RunRef {
4184 run_id: run_id as u128,
4185 level: 0,
4186 epoch_created: epoch.0,
4187 row_count: header.row_count,
4188 });
4189 Ok(())
4190 }
4191
4192 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4196 self.mutable_run_spill_bytes = bytes.max(1);
4197 }
4198
4199 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4203 self.compaction_zstd_level = level;
4204 }
4205
4206 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4210 self.result_cache.lock().set_max_bytes(max_bytes);
4211 }
4212
4213 pub(crate) fn clear_result_cache(&mut self) {
4217 self.result_cache.lock().clear();
4218 }
4219
4220 pub fn mutable_run_len(&self) -> usize {
4222 self.mutable_run.len()
4223 }
4224
4225 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4228 self.mutable_run.drain_sorted()
4229 }
4230
4231 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4233 let mut snapshot = self.mutable_run.clone();
4234 snapshot.drain_sorted()
4235 }
4236
4237 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4242 self.ensure_writable()?;
4243 let n = batch.len();
4244 if n == 0 {
4245 return Ok(self.current_epoch());
4246 }
4247 for row in &batch {
4248 self.schema.validate_values(row)?;
4249 }
4250 let epoch = self.commit_new_epoch()?;
4251 let live_before = self.live_count;
4252 self.spill_mutable_run(epoch)?;
4256 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4257 && self.indexes_complete
4258 && self.run_refs.is_empty()
4259 && self.memtable.is_empty()
4260 && self.mutable_run.is_empty();
4261 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4267 use rayon::prelude::*;
4268 self.schema
4269 .columns
4270 .par_iter()
4271 .map(|cdef| {
4272 (
4273 cdef.id,
4274 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4275 )
4276 })
4277 .collect::<Vec<_>>()
4278 };
4279 drop(batch);
4280 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4285 self.validate_columns_not_null(&user_columns, n)?;
4286 let winner_idx = self
4287 .bulk_pk_winner_indices(&user_columns, n)
4288 .filter(|idx| idx.len() != n);
4289 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4290 match winner_idx.as_deref() {
4291 Some(idx) => {
4292 let compacted = user_columns
4293 .iter()
4294 .map(|(id, c)| (*id, c.gather(idx)))
4295 .collect();
4296 (compacted, idx.len())
4297 }
4298 None => (std::mem::take(&mut user_columns), n),
4299 };
4300 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4301 let first = self.allocator.alloc_range(write_n as u64)?.0;
4302 for rid in first..first + write_n as u64 {
4303 self.reservoir.offer(rid);
4304 }
4305 let run_id = self.alloc_run_id()?;
4306 let path = self.run_path(run_id);
4307 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4308 .clean(true)
4309 .with_lz4()
4310 .with_native_endian();
4311 if let Some(kek) = &self.kek {
4312 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4313 }
4314 let header = match self.create_run_file(run_id)? {
4315 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4316 None => writer.write_native(&path, &write_columns, write_n, first)?,
4317 };
4318 self.run_refs.push(RunRef {
4319 run_id: run_id as u128,
4320 level: 0,
4321 epoch_created: epoch.0,
4322 row_count: header.row_count,
4323 });
4324 self.live_count = self.live_count.saturating_add(write_n as u64);
4325 if eager_index_build {
4326 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4327 self.index_columns_bulk(&write_columns, &row_ids);
4328 self.indexes_complete = true;
4329 self.build_learned_ranges()?;
4330 } else {
4331 self.indexes_complete = false;
4332 }
4333 self.mark_flushed(epoch)?;
4334 self.persist_manifest(epoch)?;
4335 if eager_index_build {
4336 self.checkpoint_indexes(epoch);
4337 }
4338 self.clear_result_cache();
4339 Ok(epoch)
4340 }
4341
4342 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4345 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4346 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4347 let segment_no = segment
4350 .file_stem()
4351 .and_then(|s| s.to_str())
4352 .and_then(|s| s.strip_prefix("seg-"))
4353 .and_then(|s| s.parse::<u64>().ok())
4354 .unwrap_or(0);
4355 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4356 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4357 wal.sync()?;
4358 self.wal = WalSink::Private(wal);
4359 Ok(())
4360 }
4361
4362 pub(crate) fn invalidate_pending_cache(&mut self) {
4367 self.result_cache
4368 .lock()
4369 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4370 self.pending_delete_rids.clear();
4371 self.pending_put_cols.clear();
4372 }
4373
4374 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4375 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4376 m.current_epoch = epoch.0;
4377 m.next_row_id = self.allocator.current().0;
4378 m.runs = self.run_refs.clone();
4379 m.live_count = self.live_count;
4380 m.global_idx_epoch = self.global_idx_epoch;
4381 m.flushed_epoch = self.flushed_epoch;
4382 m.retiring = self.retiring.clone();
4383 m.auto_inc_next = match self.auto_inc {
4387 Some(ai) if ai.seeded => ai.next,
4388 _ => 0,
4389 };
4390 m.ttl = self.ttl;
4391 let meta_dek = self.manifest_meta_dek();
4392 match self._root_guard.as_deref() {
4393 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4394 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4395 }
4396 Ok(())
4397 }
4398
4399 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4400 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4404 let live_count = u64::try_from(rows.len())
4405 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4406 let auto_inc = match self.auto_inc {
4407 Some(mut state) => {
4408 let maximum = self.scan_max_int64(state.column_id)?;
4409 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4410 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4411 })?;
4412 state.next = state.next.max(after_maximum).max(1);
4413 state.seeded = true;
4414 Some(state)
4415 }
4416 None => None,
4417 };
4418 Ok(RecoveryMetadataPlan {
4419 live_count,
4420 auto_inc,
4421 changed: live_count != self.live_count
4422 || auto_inc.is_some_and(|planned| {
4423 self.auto_inc.is_none_or(|current| {
4424 current.next != planned.next || current.seeded != planned.seeded
4425 })
4426 }),
4427 })
4428 }
4429
4430 pub(crate) fn apply_recovered_metadata(
4431 &mut self,
4432 plan: RecoveryMetadataPlan,
4433 epoch: Epoch,
4434 ) -> Result<()> {
4435 if !plan.changed {
4436 return Ok(());
4437 }
4438 self.live_count = plan.live_count;
4439 self.auto_inc = plan.auto_inc;
4440 self.persist_manifest(epoch)
4441 }
4442
4443 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4449 if !self.indexes_complete {
4452 return;
4453 }
4454 if crate::catalog::inject_hook("index.publish.before").is_err() {
4457 return;
4458 }
4459 if self.idx_root.is_none() {
4460 if let Some(root) = self._root_guard.as_ref() {
4461 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4462 return;
4463 };
4464 self.idx_root = Some(Arc::new(idx_root));
4465 }
4466 }
4467 let snap = global_idx::IndexSnapshot {
4468 hot: &self.hot,
4469 bitmap: &self.bitmap,
4470 ann: &self.ann,
4471 fm: &self.fm,
4472 sparse: &self.sparse,
4473 minhash: &self.minhash,
4474 learned_range: &self.learned_range,
4475 };
4476 let idx_dek = self.idx_dek();
4478 let written = match self.idx_root.as_deref() {
4479 Some(root) => global_idx::write_atomic_root(
4480 root,
4481 self.table_id,
4482 epoch.0,
4483 snap,
4484 idx_dek.as_deref(),
4485 ),
4486 None => global_idx::write_atomic(
4487 &self.dir,
4488 self.table_id,
4489 epoch.0,
4490 snap,
4491 idx_dek.as_deref(),
4492 ),
4493 };
4494 if written.is_ok() {
4495 self.global_idx_epoch = epoch.0;
4496 let _ = self.persist_manifest(epoch);
4497 let _ = crate::catalog::inject_hook("index.publish.after");
4499 }
4500 }
4501
4502 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4505 self.global_idx_epoch = 0;
4506 if let Some(root) = self.idx_root.as_deref() {
4507 let _ = root.remove_file(global_idx::IDX_FILENAME);
4508 } else {
4509 global_idx::remove(&self.dir);
4510 }
4511 let _ = self.persist_manifest(self.epoch.visible());
4512 }
4513
4514 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4519 self.indexes_complete = false;
4520 self.global_idx_epoch = 0;
4521 if let Some(root) = self.idx_root.as_deref() {
4522 let _ = root.remove_file(global_idx::IDX_FILENAME);
4523 } else {
4524 global_idx::remove(&self.dir);
4525 }
4526 }
4527
4528 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4529 self.indexes_complete = true;
4530 }
4531
4532 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4538 self.durable_commit_failed = true;
4539 if let WalSink::Shared(shared) = &self.wal {
4540 shared
4541 .poisoned
4542 .store(true, std::sync::atomic::Ordering::Relaxed);
4543 }
4544 }
4545
4546 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4550 self.durable_commit_failed = true;
4551 }
4552
4553 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4556 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4557 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4558 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4559 best = Some((epoch, row));
4560 }
4561 }
4562 for rr in &self.run_refs {
4563 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4564 continue;
4565 };
4566 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4567 continue;
4568 };
4569 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4570 best = Some((epoch, row));
4571 }
4572 }
4573 let now_nanos = unix_nanos_now();
4574 match best {
4575 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4576 Some((_, r)) => Some(r),
4577 None => None,
4578 }
4579 }
4580
4581 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4585 self.visible_rows_at_time(snapshot, unix_nanos_now())
4586 }
4587
4588 #[doc(hidden)]
4591 pub fn visible_rows_controlled(
4592 &self,
4593 snapshot: Snapshot,
4594 control: &crate::ExecutionControl,
4595 ) -> Result<Vec<Row>> {
4596 let mut out = Vec::new();
4597 self.for_each_visible_row_controlled(snapshot, control, |row| {
4598 out.push(row);
4599 Ok(())
4600 })?;
4601 Ok(out)
4602 }
4603
4604 #[doc(hidden)]
4607 pub fn for_each_visible_row_controlled<F>(
4608 &self,
4609 snapshot: Snapshot,
4610 control: &crate::ExecutionControl,
4611 visit: F,
4612 ) -> Result<()>
4613 where
4614 F: FnMut(Row) -> Result<()>,
4615 {
4616 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4617 control.checkpoint()?;
4618 let memtable = self.memtable.visible_versions(snapshot.epoch);
4619 if !memtable.is_empty() {
4620 sources.push(ControlledVisibleSource::memory(memtable));
4621 }
4622 control.checkpoint()?;
4623 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4624 if !mutable.is_empty() {
4625 sources.push(ControlledVisibleSource::memory(mutable));
4626 }
4627 for run in &self.run_refs {
4628 control.checkpoint()?;
4629 let reader = self.open_reader(run.run_id)?;
4630 sources.push(ControlledVisibleSource::run(
4631 reader.into_visible_version_cursor(snapshot.epoch)?,
4632 ));
4633 }
4634 let now_nanos = unix_nanos_now();
4635 merge_controlled_visible_sources(
4636 &mut sources,
4637 control,
4638 |row| self.row_expired_at(row, now_nanos),
4639 visit,
4640 )
4641 }
4642
4643 #[doc(hidden)]
4644 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4645 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4646 let mut fold = |row: Row| {
4647 best.entry(row.row_id.0)
4648 .and_modify(|e| {
4649 if row.committed_epoch > e.0 {
4650 *e = (row.committed_epoch, row.clone());
4651 }
4652 })
4653 .or_insert_with(|| (row.committed_epoch, row));
4654 };
4655 for row in self.memtable.visible_versions(snapshot.epoch) {
4656 fold(row);
4657 }
4658 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4659 fold(row);
4660 }
4661 for rr in &self.run_refs {
4662 let mut reader = self.open_reader(rr.run_id)?;
4663 for row in reader.visible_versions(snapshot.epoch)? {
4664 fold(row);
4665 }
4666 }
4667 let mut out: Vec<Row> = best
4668 .into_values()
4669 .filter_map(|(_, r)| {
4670 if r.deleted || self.row_expired_at(&r, now_nanos) {
4671 None
4672 } else {
4673 Some(r)
4674 }
4675 })
4676 .collect();
4677 out.sort_by_key(|r| r.row_id);
4678 Ok(out)
4679 }
4680
4681 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4688 if self.ttl.is_none()
4689 && self.memtable.is_empty()
4690 && self.mutable_run.is_empty()
4691 && self.run_refs.len() == 1
4692 {
4693 let rr = self.run_refs[0].clone();
4694 let mut reader = self.open_reader(rr.run_id)?;
4695 let idxs = reader.visible_indices(snapshot.epoch)?;
4696 let mut cols = Vec::with_capacity(self.schema.columns.len());
4697 for cdef in &self.schema.columns {
4698 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4699 }
4700 return Ok(cols);
4701 }
4702 let rows = self.visible_rows(snapshot)?;
4704 let mut cols: Vec<(u16, Vec<Value>)> = self
4705 .schema
4706 .columns
4707 .iter()
4708 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4709 .collect();
4710 for r in &rows {
4711 for (cid, vec) in cols.iter_mut() {
4712 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4713 }
4714 }
4715 Ok(cols)
4716 }
4717
4718 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4720 let row_id = self.hot.get(key)?;
4721 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4722 Some(row_id)
4723 } else {
4724 None
4725 }
4726 }
4727
4728 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4733 self.query_at_with_allowed(q, self.snapshot(), None)
4734 }
4735
4736 pub fn query_controlled(
4739 &mut self,
4740 q: &crate::query::Query,
4741 control: &crate::ExecutionControl,
4742 ) -> Result<Vec<Row>> {
4743 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4744 }
4745
4746 pub fn query_at_with_allowed(
4749 &mut self,
4750 q: &crate::query::Query,
4751 snapshot: Snapshot,
4752 allowed: Option<&std::collections::HashSet<RowId>>,
4753 ) -> Result<Vec<Row>> {
4754 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4755 }
4756
4757 #[doc(hidden)]
4758 pub fn query_at_with_allowed_controlled(
4759 &mut self,
4760 q: &crate::query::Query,
4761 snapshot: Snapshot,
4762 allowed: Option<&std::collections::HashSet<RowId>>,
4763 control: &crate::ExecutionControl,
4764 ) -> Result<Vec<Row>> {
4765 self.require_select()?;
4766 self.ensure_indexes_complete_controlled(control, || true)?;
4767 self.validate_native_query(q)?;
4768 self.query_conditions_at(
4769 &q.conditions,
4770 snapshot,
4771 allowed,
4772 q.limit,
4773 q.offset,
4774 None,
4775 unix_nanos_now(),
4776 Some(control),
4777 )
4778 }
4779
4780 #[doc(hidden)]
4781 pub fn query_at_with_allowed_after(
4782 &mut self,
4783 q: &crate::query::Query,
4784 snapshot: Snapshot,
4785 allowed: Option<&std::collections::HashSet<RowId>>,
4786 after_row_id: Option<RowId>,
4787 ) -> Result<Vec<Row>> {
4788 self.query_at_with_allowed_after_at_time(
4789 q,
4790 snapshot,
4791 allowed,
4792 after_row_id,
4793 unix_nanos_now(),
4794 )
4795 }
4796
4797 #[doc(hidden)]
4798 pub fn query_at_with_allowed_after_at_time(
4799 &mut self,
4800 q: &crate::query::Query,
4801 snapshot: Snapshot,
4802 allowed: Option<&std::collections::HashSet<RowId>>,
4803 after_row_id: Option<RowId>,
4804 query_time_nanos: i64,
4805 ) -> Result<Vec<Row>> {
4806 self.require_select()?;
4807 self.ensure_indexes_complete()?;
4808 self.validate_native_query(q)?;
4809 self.query_conditions_at(
4810 &q.conditions,
4811 snapshot,
4812 allowed,
4813 q.limit,
4814 q.offset,
4815 after_row_id,
4816 query_time_nanos,
4817 None,
4818 )
4819 }
4820
4821 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4822 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4823 return Err(MongrelError::InvalidArgument(format!(
4824 "query exceeds {} conditions",
4825 crate::query::MAX_HARD_CONDITIONS
4826 )));
4827 }
4828 if let Some(limit) = q.limit {
4829 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4830 return Err(MongrelError::InvalidArgument(format!(
4831 "query limit must be between 1 and {}",
4832 crate::query::MAX_FINAL_LIMIT
4833 )));
4834 }
4835 }
4836 if q.offset > crate::query::MAX_QUERY_OFFSET {
4837 return Err(MongrelError::InvalidArgument(format!(
4838 "query offset exceeds {}",
4839 crate::query::MAX_QUERY_OFFSET
4840 )));
4841 }
4842 Ok(())
4843 }
4844
4845 #[doc(hidden)]
4848 pub fn query_all_at(
4849 &mut self,
4850 conditions: &[crate::query::Condition],
4851 snapshot: Snapshot,
4852 ) -> Result<Vec<Row>> {
4853 self.require_select()?;
4854 self.ensure_indexes_complete()?;
4855 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4856 return Err(MongrelError::InvalidArgument(format!(
4857 "query exceeds {} conditions",
4858 crate::query::MAX_HARD_CONDITIONS
4859 )));
4860 }
4861 self.query_conditions_at(
4862 conditions,
4863 snapshot,
4864 None,
4865 None,
4866 0,
4867 None,
4868 unix_nanos_now(),
4869 None,
4870 )
4871 }
4872
4873 #[allow(clippy::too_many_arguments)]
4874 fn query_conditions_at(
4875 &self,
4876 conditions: &[crate::query::Condition],
4877 snapshot: Snapshot,
4878 allowed: Option<&std::collections::HashSet<RowId>>,
4879 limit: Option<usize>,
4880 offset: usize,
4881 after_row_id: Option<RowId>,
4882 query_time_nanos: i64,
4883 control: Option<&crate::ExecutionControl>,
4884 ) -> Result<Vec<Row>> {
4885 control
4886 .map(crate::ExecutionControl::checkpoint)
4887 .transpose()?;
4888 crate::trace::QueryTrace::record(|t| {
4889 t.run_count = self.run_refs.len();
4890 t.memtable_rows = self.memtable.len();
4891 t.mutable_run_rows = self.mutable_run.len();
4892 });
4893 if conditions.is_empty() {
4897 crate::trace::QueryTrace::record(|t| {
4898 t.scan_mode = crate::trace::ScanMode::Materialized;
4899 t.row_materialized = true;
4900 });
4901 let mut rows = match control {
4902 Some(control) => self.visible_rows_controlled(snapshot, control)?,
4903 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
4904 };
4905 if let Some(allowed) = allowed {
4906 let mut filtered = Vec::with_capacity(rows.len());
4907 for (index, row) in rows.into_iter().enumerate() {
4908 if index & 255 == 0 {
4909 control
4910 .map(crate::ExecutionControl::checkpoint)
4911 .transpose()?;
4912 }
4913 if allowed.contains(&row.row_id) {
4914 filtered.push(row);
4915 }
4916 }
4917 rows = filtered;
4918 }
4919 if let Some(after_row_id) = after_row_id {
4920 rows.retain(|row| row.row_id > after_row_id);
4921 }
4922 rows.drain(..offset.min(rows.len()));
4923 if let Some(limit) = limit {
4924 rows.truncate(limit);
4925 }
4926 return Ok(rows);
4927 }
4928 crate::trace::QueryTrace::record(|t| {
4929 t.conditions_pushed = conditions.len();
4930 t.scan_mode = crate::trace::ScanMode::Materialized;
4931 t.row_materialized = true;
4932 });
4933 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
4940 ordered.sort_by_key(|c| condition_cost_rank(c));
4941 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
4942 for c in &ordered {
4943 control
4944 .map(crate::ExecutionControl::checkpoint)
4945 .transpose()?;
4946 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
4947 let empty = s.is_empty();
4948 sets.push(s);
4949 if empty {
4950 break;
4951 }
4952 }
4953 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4954 if let Some(allowed) = allowed {
4955 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
4956 }
4957 if let Some(after_row_id) = after_row_id {
4958 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
4959 rids.drain(..first);
4960 }
4961 rids.drain(..offset.min(rids.len()));
4962 if let Some(limit) = limit {
4963 rids.truncate(limit);
4964 }
4965 control
4966 .map(crate::ExecutionControl::checkpoint)
4967 .transpose()?;
4968 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
4969 }
4970
4971 pub fn retrieve(
4973 &mut self,
4974 retriever: &crate::query::Retriever,
4975 ) -> Result<Vec<crate::query::RetrieverHit>> {
4976 self.retrieve_with_allowed(retriever, None)
4977 }
4978
4979 pub fn retrieve_at(
4980 &mut self,
4981 retriever: &crate::query::Retriever,
4982 snapshot: Snapshot,
4983 allowed: Option<&std::collections::HashSet<RowId>>,
4984 ) -> Result<Vec<crate::query::RetrieverHit>> {
4985 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
4986 }
4987
4988 pub fn retrieve_with_allowed(
4991 &mut self,
4992 retriever: &crate::query::Retriever,
4993 allowed: Option<&std::collections::HashSet<RowId>>,
4994 ) -> Result<Vec<crate::query::RetrieverHit>> {
4995 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
4996 }
4997
4998 pub fn retrieve_at_with_allowed(
4999 &mut self,
5000 retriever: &crate::query::Retriever,
5001 snapshot: Snapshot,
5002 allowed: Option<&std::collections::HashSet<RowId>>,
5003 ) -> Result<Vec<crate::query::RetrieverHit>> {
5004 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5005 }
5006
5007 pub fn retrieve_at_with_allowed_and_context(
5008 &mut self,
5009 retriever: &crate::query::Retriever,
5010 snapshot: Snapshot,
5011 allowed: Option<&std::collections::HashSet<RowId>>,
5012 context: Option<&crate::query::AiExecutionContext>,
5013 ) -> Result<Vec<crate::query::RetrieverHit>> {
5014 self.require_select()?;
5015 self.ensure_indexes_complete()?;
5016 self.validate_retriever(retriever)?;
5017 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5018 }
5019
5020 pub fn retrieve_at_with_candidate_authorization_and_context(
5021 &mut self,
5022 retriever: &crate::query::Retriever,
5023 snapshot: Snapshot,
5024 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5025 context: Option<&crate::query::AiExecutionContext>,
5026 ) -> Result<Vec<crate::query::RetrieverHit>> {
5027 self.require_select()?;
5028 self.ensure_indexes_complete()?;
5029 self.retrieve_at_with_candidate_authorization_on_generation(
5030 retriever,
5031 snapshot,
5032 authorization,
5033 context,
5034 )
5035 }
5036
5037 #[doc(hidden)]
5038 pub fn retrieve_at_with_candidate_authorization_on_generation(
5039 &self,
5040 retriever: &crate::query::Retriever,
5041 snapshot: Snapshot,
5042 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5043 context: Option<&crate::query::AiExecutionContext>,
5044 ) -> Result<Vec<crate::query::RetrieverHit>> {
5045 self.require_select()?;
5046 self.validate_retriever(retriever)?;
5047 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5048 }
5049
5050 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5051 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5052 let (column_id, k) = match retriever {
5053 Retriever::Ann {
5054 column_id,
5055 query,
5056 k,
5057 } => {
5058 let index = self.ann.get(column_id).ok_or_else(|| {
5059 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5060 })?;
5061 if query.len() != index.dim() {
5062 return Err(MongrelError::InvalidArgument(format!(
5063 "ANN query dimension must be {}, got {}",
5064 index.dim(),
5065 query.len()
5066 )));
5067 }
5068 if query.iter().any(|value| !value.is_finite()) {
5069 return Err(MongrelError::InvalidArgument(
5070 "ANN query values must be finite".into(),
5071 ));
5072 }
5073 (*column_id, *k)
5074 }
5075 Retriever::Sparse {
5076 column_id,
5077 query,
5078 k,
5079 } => {
5080 if !self.sparse.contains_key(column_id) {
5081 return Err(MongrelError::InvalidArgument(format!(
5082 "column {column_id} has no Sparse index"
5083 )));
5084 }
5085 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5086 return Err(MongrelError::InvalidArgument(
5087 "Sparse query must be non-empty with finite weights".into(),
5088 ));
5089 }
5090 if query.len() > MAX_SPARSE_TERMS {
5091 return Err(MongrelError::InvalidArgument(format!(
5092 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5093 )));
5094 }
5095 (*column_id, *k)
5096 }
5097 Retriever::MinHash {
5098 column_id,
5099 members,
5100 k,
5101 } => {
5102 if !self.minhash.contains_key(column_id) {
5103 return Err(MongrelError::InvalidArgument(format!(
5104 "column {column_id} has no MinHash index"
5105 )));
5106 }
5107 if members.is_empty() {
5108 return Err(MongrelError::InvalidArgument(
5109 "MinHash members must not be empty".into(),
5110 ));
5111 }
5112 if members.len() > MAX_SET_MEMBERS {
5113 return Err(MongrelError::InvalidArgument(format!(
5114 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5115 )));
5116 }
5117 let mut total_bytes = 0usize;
5118 for member in members {
5119 let bytes = member.encoded_len();
5120 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5121 return Err(MongrelError::InvalidArgument(format!(
5122 "MinHash member exceeds {} bytes",
5123 crate::query::MAX_SET_MEMBER_BYTES
5124 )));
5125 }
5126 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5127 MongrelError::InvalidArgument("MinHash input size overflow".into())
5128 })?;
5129 }
5130 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5131 return Err(MongrelError::InvalidArgument(format!(
5132 "MinHash input exceeds {} bytes",
5133 crate::query::MAX_SET_INPUT_BYTES
5134 )));
5135 }
5136 (*column_id, *k)
5137 }
5138 };
5139 if k == 0 {
5140 return Err(MongrelError::InvalidArgument(
5141 "retriever k must be > 0".into(),
5142 ));
5143 }
5144 if k > MAX_RETRIEVER_K {
5145 return Err(MongrelError::InvalidArgument(format!(
5146 "retriever k exceeds {MAX_RETRIEVER_K}"
5147 )));
5148 }
5149 debug_assert!(self
5150 .schema
5151 .columns
5152 .iter()
5153 .any(|column| column.id == column_id));
5154 Ok(())
5155 }
5156
5157 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5158 use crate::query::Condition;
5159 match condition {
5160 Condition::Ann {
5161 column_id,
5162 query,
5163 k,
5164 } => self.validate_retriever(&crate::query::Retriever::Ann {
5165 column_id: *column_id,
5166 query: query.clone(),
5167 k: *k,
5168 }),
5169 Condition::SparseMatch {
5170 column_id,
5171 query,
5172 k,
5173 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5174 column_id: *column_id,
5175 query: query.clone(),
5176 k: *k,
5177 }),
5178 Condition::MinHashSimilar {
5179 column_id,
5180 query,
5181 k,
5182 } => {
5183 if !self.minhash.contains_key(column_id) {
5184 return Err(MongrelError::InvalidArgument(format!(
5185 "column {column_id} has no MinHash index"
5186 )));
5187 }
5188 if query.is_empty() || *k == 0 {
5189 return Err(MongrelError::InvalidArgument(
5190 "MinHash query must be non-empty and k must be > 0".into(),
5191 ));
5192 }
5193 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5194 {
5195 return Err(MongrelError::InvalidArgument(format!(
5196 "MinHash query must have <= {} members and k <= {}",
5197 crate::query::MAX_SET_MEMBERS,
5198 crate::query::MAX_RETRIEVER_K
5199 )));
5200 }
5201 Ok(())
5202 }
5203 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5204 Err(MongrelError::InvalidArgument(format!(
5205 "bitmap IN exceeds {} values",
5206 crate::query::MAX_SET_MEMBERS
5207 )))
5208 }
5209 Condition::FmContainsAll { patterns, .. }
5210 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5211 {
5212 Err(MongrelError::InvalidArgument(format!(
5213 "FM query exceeds {} patterns",
5214 crate::query::MAX_HARD_CONDITIONS
5215 )))
5216 }
5217 _ => Ok(()),
5218 }
5219 }
5220
5221 fn retrieve_filtered(
5222 &self,
5223 retriever: &crate::query::Retriever,
5224 snapshot: Snapshot,
5225 hard_filter: Option<&RowIdSet>,
5226 allowed: Option<&std::collections::HashSet<RowId>>,
5227 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5228 context: Option<&crate::query::AiExecutionContext>,
5229 ) -> Result<Vec<crate::query::RetrieverHit>> {
5230 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5231 let started = std::time::Instant::now();
5232 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5233 Retriever::Ann {
5234 column_id,
5235 query,
5236 k,
5237 } => {
5238 let Some(index) = self.ann.get(column_id) else {
5239 return Ok(Vec::new());
5240 };
5241 let cap = ann_candidate_cap(index.len(), context);
5242 if cap == 0 {
5243 return Ok(Vec::new());
5244 }
5245 let mut breadth = (*k).max(1).min(cap);
5246 let mut eligibility = std::collections::HashMap::new();
5247 let mut filtered = loop {
5248 let mut seen = std::collections::HashSet::new();
5249 if let Some(context) = context {
5250 context.checkpoint()?;
5251 }
5252 let raw = index.search_with_context(query, breadth, context)?;
5253 let unchecked: Vec<_> = raw
5254 .iter()
5255 .map(|(row_id, _)| *row_id)
5256 .filter(|row_id| !eligibility.contains_key(row_id))
5257 .filter(|row_id| {
5258 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5259 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5260 })
5261 .collect();
5262 let eligible = self.eligible_and_authorized_candidate_ids(
5263 &unchecked,
5264 *column_id,
5265 snapshot,
5266 candidate_authorization,
5267 context,
5268 )?;
5269 for row_id in unchecked {
5270 eligibility.insert(row_id, eligible.contains(&row_id));
5271 }
5272 let filtered: Vec<_> = raw
5273 .into_iter()
5274 .filter(|(row_id, _)| {
5275 seen.insert(*row_id)
5276 && eligibility.get(row_id).copied().unwrap_or(false)
5277 })
5278 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
5279 .collect();
5280 if filtered.len() >= *k || breadth >= cap {
5281 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5282 crate::trace::QueryTrace::record(|trace| {
5283 trace.ann_candidate_cap_hit = true;
5284 });
5285 }
5286 break filtered;
5287 }
5288 breadth = breadth.saturating_mul(2).min(cap);
5289 };
5290 filtered.truncate(*k);
5291 filtered
5292 }
5293 Retriever::Sparse {
5294 column_id,
5295 query,
5296 k,
5297 } => self
5298 .sparse
5299 .get(column_id)
5300 .map(|index| -> Result<Vec<_>> {
5301 let mut breadth = (*k).max(1);
5302 let mut eligibility = std::collections::HashMap::new();
5303 loop {
5304 if let Some(context) = context {
5305 context.checkpoint()?;
5306 }
5307 let raw = index.search_with_context(query, breadth, context)?;
5308 let unchecked: Vec<_> = raw
5309 .iter()
5310 .map(|(row_id, _)| *row_id)
5311 .filter(|row_id| !eligibility.contains_key(row_id))
5312 .filter(|row_id| {
5313 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5314 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5315 })
5316 .collect();
5317 let eligible = self.eligible_and_authorized_candidate_ids(
5318 &unchecked,
5319 *column_id,
5320 snapshot,
5321 candidate_authorization,
5322 context,
5323 )?;
5324 for row_id in unchecked {
5325 eligibility.insert(row_id, eligible.contains(&row_id));
5326 }
5327 let filtered: Vec<_> = raw
5328 .iter()
5329 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5330 .take(*k)
5331 .map(|(row_id, score)| {
5332 (*row_id, RetrieverScore::SparseDotProduct(*score))
5333 })
5334 .collect();
5335 if filtered.len() >= *k || raw.len() < breadth {
5336 break Ok(filtered);
5337 }
5338 let next = breadth.saturating_mul(2);
5339 if next == breadth {
5340 break Ok(filtered);
5341 }
5342 breadth = next;
5343 }
5344 })
5345 .transpose()?
5346 .unwrap_or_default(),
5347 Retriever::MinHash {
5348 column_id,
5349 members,
5350 k,
5351 } => self
5352 .minhash
5353 .get(column_id)
5354 .map(|index| -> Result<Vec<_>> {
5355 let mut hashes = Vec::with_capacity(members.len());
5356 for member in members {
5357 if let Some(context) = context {
5358 context.consume(crate::query::work_units(
5359 member.encoded_len(),
5360 crate::query::PARSE_WORK_QUANTUM,
5361 ))?;
5362 }
5363 hashes.push(member.hash_v1());
5364 }
5365 let mut breadth = (*k).max(1);
5366 let mut eligibility = std::collections::HashMap::new();
5367 loop {
5368 if let Some(context) = context {
5369 context.checkpoint()?;
5370 }
5371 let raw = index.search_with_context(&hashes, breadth, context)?;
5372 let unchecked: Vec<_> = raw
5373 .iter()
5374 .map(|(row_id, _)| *row_id)
5375 .filter(|row_id| !eligibility.contains_key(row_id))
5376 .filter(|row_id| {
5377 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5378 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5379 })
5380 .collect();
5381 let eligible = self.eligible_and_authorized_candidate_ids(
5382 &unchecked,
5383 *column_id,
5384 snapshot,
5385 candidate_authorization,
5386 context,
5387 )?;
5388 for row_id in unchecked {
5389 eligibility.insert(row_id, eligible.contains(&row_id));
5390 }
5391 let filtered: Vec<_> = raw
5392 .iter()
5393 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5394 .take(*k)
5395 .map(|(row_id, score)| {
5396 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5397 })
5398 .collect();
5399 if filtered.len() >= *k || raw.len() < breadth {
5400 break Ok(filtered);
5401 }
5402 let next = breadth.saturating_mul(2);
5403 if next == breadth {
5404 break Ok(filtered);
5405 }
5406 breadth = next;
5407 }
5408 })
5409 .transpose()?
5410 .unwrap_or_default(),
5411 };
5412 let elapsed = started.elapsed().as_nanos() as u64;
5413 crate::trace::QueryTrace::record(|trace| {
5414 match retriever {
5415 Retriever::Ann { .. } => {
5416 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5417 }
5418 Retriever::Sparse { .. } => {
5419 trace.sparse_candidate_nanos =
5420 trace.sparse_candidate_nanos.saturating_add(elapsed)
5421 }
5422 Retriever::MinHash { .. } => {
5423 trace.minhash_candidate_nanos =
5424 trace.minhash_candidate_nanos.saturating_add(elapsed)
5425 }
5426 }
5427 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5428 });
5429 Ok(scored
5430 .into_iter()
5431 .enumerate()
5432 .map(|(rank, (row_id, score))| RetrieverHit {
5433 row_id,
5434 rank: rank + 1,
5435 score,
5436 })
5437 .collect())
5438 }
5439
5440 fn eligible_candidate_ids(
5441 &self,
5442 candidates: &[RowId],
5443 _column_id: u16,
5444 snapshot: Snapshot,
5445 context: Option<&crate::query::AiExecutionContext>,
5446 ) -> Result<std::collections::HashSet<RowId>> {
5447 if !self.had_deletes
5448 && self.ttl.is_none()
5449 && self.pending_put_cols.is_empty()
5450 && snapshot.epoch == self.snapshot().epoch
5451 {
5452 return Ok(candidates.iter().copied().collect());
5453 }
5454 let mut readers: Vec<_> = self
5455 .run_refs
5456 .iter()
5457 .map(|run| self.open_reader(run.run_id))
5458 .collect::<Result<_>>()?;
5459 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5460 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5461 for &row_id in candidates {
5462 if let Some(context) = context {
5463 context.consume(1)?;
5464 }
5465 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5466 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5467 let overlay = match (mem, mutable) {
5468 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5469 (Some(value), None) | (None, Some(value)) => Some(value),
5470 (None, None) => None,
5471 };
5472 if let Some((_, row)) = overlay {
5473 if !row.deleted && !self.row_expired_at(&row, now) {
5474 eligible.insert(row_id);
5475 }
5476 continue;
5477 }
5478 let mut best: Option<(Epoch, bool, usize)> = None;
5479 for (index, reader) in readers.iter_mut().enumerate() {
5480 if let Some((epoch, deleted)) =
5481 reader.get_version_visibility(row_id, snapshot.epoch)?
5482 {
5483 if best
5484 .as_ref()
5485 .map(|(best_epoch, ..)| epoch > *best_epoch)
5486 .unwrap_or(true)
5487 {
5488 best = Some((epoch, deleted, index));
5489 }
5490 }
5491 }
5492 let Some((_, false, reader_index)) = best else {
5493 continue;
5494 };
5495 if let Some(ttl) = self.ttl {
5496 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5497 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5498 {
5499 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5500 continue;
5501 }
5502 }
5503 }
5504 eligible.insert(row_id);
5505 }
5506 Ok(eligible)
5507 }
5508
5509 fn eligible_and_authorized_candidate_ids(
5510 &self,
5511 candidates: &[RowId],
5512 column_id: u16,
5513 snapshot: Snapshot,
5514 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5515 context: Option<&crate::query::AiExecutionContext>,
5516 ) -> Result<std::collections::HashSet<RowId>> {
5517 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5518 let Some(authorization) = authorization else {
5519 return Ok(eligible);
5520 };
5521 let candidates: Vec<_> = eligible.into_iter().collect();
5522 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5523 }
5524
5525 fn policy_allowed_candidate_ids(
5526 &self,
5527 candidates: &[RowId],
5528 snapshot: Snapshot,
5529 authorization: &crate::security::CandidateAuthorization<'_>,
5530 context: Option<&crate::query::AiExecutionContext>,
5531 ) -> Result<std::collections::HashSet<RowId>> {
5532 let started = std::time::Instant::now();
5533 if candidates.is_empty()
5534 || authorization.principal.is_admin
5535 || !authorization.security.rls_enabled(authorization.table)
5536 {
5537 return Ok(candidates.iter().copied().collect());
5538 }
5539 if let Some(context) = context {
5540 context.checkpoint()?;
5541 }
5542 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5543 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5544 .iter()
5545 .map(|row_id| {
5546 (
5547 *row_id,
5548 Row {
5549 row_id: *row_id,
5550 committed_epoch: snapshot.epoch,
5551 columns: std::collections::HashMap::new(),
5552 deleted: false,
5553 },
5554 )
5555 })
5556 .collect();
5557 let columns = authorization
5558 .security
5559 .select_policy_columns(authorization.table, authorization.principal);
5560 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5561 let mut decoded = 0usize;
5562 for column_id in &columns {
5563 if let Some(context) = context {
5564 context.checkpoint()?;
5565 }
5566 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5567 &row_ids, *column_id, snapshot, query_now, context,
5568 )? {
5569 if let Some(row) = rows.get_mut(&row_id) {
5570 row.columns.insert(*column_id, value);
5571 decoded = decoded.saturating_add(1);
5572 }
5573 }
5574 }
5575 if let Some(context) = context {
5576 context.consume(candidates.len().saturating_add(decoded))?;
5577 }
5578 let allowed = rows
5579 .into_values()
5580 .filter_map(|row| {
5581 authorization
5582 .security
5583 .row_allowed(
5584 authorization.table,
5585 crate::security::PolicyCommand::Select,
5586 &row,
5587 authorization.principal,
5588 false,
5589 )
5590 .then_some(row.row_id)
5591 })
5592 .collect();
5593 crate::trace::QueryTrace::record(|trace| {
5594 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5595 trace.rls_policy_columns_decoded =
5596 trace.rls_policy_columns_decoded.saturating_add(decoded);
5597 trace.authorization_nanos = trace
5598 .authorization_nanos
5599 .saturating_add(started.elapsed().as_nanos() as u64);
5600 });
5601 Ok(allowed)
5602 }
5603
5604 pub fn search(
5606 &mut self,
5607 request: &crate::query::SearchRequest,
5608 ) -> Result<Vec<crate::query::SearchHit>> {
5609 self.search_with_allowed(request, None)
5610 }
5611
5612 pub fn search_at(
5613 &mut self,
5614 request: &crate::query::SearchRequest,
5615 snapshot: Snapshot,
5616 authorized: Option<&std::collections::HashSet<RowId>>,
5617 ) -> Result<Vec<crate::query::SearchHit>> {
5618 self.search_at_with_allowed(request, snapshot, authorized)
5619 }
5620
5621 pub fn search_with_allowed(
5622 &mut self,
5623 request: &crate::query::SearchRequest,
5624 authorized: Option<&std::collections::HashSet<RowId>>,
5625 ) -> Result<Vec<crate::query::SearchHit>> {
5626 self.search_at_with_allowed(request, self.snapshot(), authorized)
5627 }
5628
5629 pub fn search_at_with_allowed(
5630 &mut self,
5631 request: &crate::query::SearchRequest,
5632 snapshot: Snapshot,
5633 authorized: Option<&std::collections::HashSet<RowId>>,
5634 ) -> Result<Vec<crate::query::SearchHit>> {
5635 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5636 }
5637
5638 pub fn search_at_with_allowed_and_context(
5639 &mut self,
5640 request: &crate::query::SearchRequest,
5641 snapshot: Snapshot,
5642 authorized: Option<&std::collections::HashSet<RowId>>,
5643 context: Option<&crate::query::AiExecutionContext>,
5644 ) -> Result<Vec<crate::query::SearchHit>> {
5645 self.ensure_indexes_complete()?;
5646 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5647 }
5648
5649 pub fn search_at_with_candidate_authorization_and_context(
5650 &mut self,
5651 request: &crate::query::SearchRequest,
5652 snapshot: Snapshot,
5653 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5654 context: Option<&crate::query::AiExecutionContext>,
5655 ) -> Result<Vec<crate::query::SearchHit>> {
5656 self.ensure_indexes_complete()?;
5657 self.search_at_with_filters_and_context(
5658 request,
5659 snapshot,
5660 None,
5661 authorization,
5662 context,
5663 None,
5664 )
5665 }
5666
5667 #[doc(hidden)]
5668 pub fn search_at_with_candidate_authorization_on_generation(
5669 &self,
5670 request: &crate::query::SearchRequest,
5671 snapshot: Snapshot,
5672 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5673 context: Option<&crate::query::AiExecutionContext>,
5674 ) -> Result<Vec<crate::query::SearchHit>> {
5675 self.search_at_with_filters_and_context(
5676 request,
5677 snapshot,
5678 None,
5679 authorization,
5680 context,
5681 None,
5682 )
5683 }
5684
5685 #[doc(hidden)]
5686 pub fn search_at_with_candidate_authorization_on_generation_after(
5687 &self,
5688 request: &crate::query::SearchRequest,
5689 snapshot: Snapshot,
5690 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5691 context: Option<&crate::query::AiExecutionContext>,
5692 after: Option<crate::query::SearchAfter>,
5693 ) -> Result<Vec<crate::query::SearchHit>> {
5694 self.search_at_with_filters_and_context(
5695 request,
5696 snapshot,
5697 None,
5698 authorization,
5699 context,
5700 after,
5701 )
5702 }
5703
5704 fn search_at_with_filters_and_context(
5705 &self,
5706 request: &crate::query::SearchRequest,
5707 snapshot: Snapshot,
5708 authorized: Option<&std::collections::HashSet<RowId>>,
5709 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5710 context: Option<&crate::query::AiExecutionContext>,
5711 after: Option<crate::query::SearchAfter>,
5712 ) -> Result<Vec<crate::query::SearchHit>> {
5713 use crate::query::{
5714 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5715 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5716 };
5717 let total_started = std::time::Instant::now();
5718 let rank_offset = after.map_or(0, |after| after.returned_count);
5719 self.require_select()?;
5720 if request.limit == 0 {
5721 return Err(MongrelError::InvalidArgument(
5722 "search limit must be > 0".into(),
5723 ));
5724 }
5725 if request.limit > MAX_FINAL_LIMIT {
5726 return Err(MongrelError::InvalidArgument(format!(
5727 "search limit exceeds {MAX_FINAL_LIMIT}"
5728 )));
5729 }
5730 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5731 return Err(MongrelError::InvalidArgument(
5732 "search-after score must be finite".into(),
5733 ));
5734 }
5735 if request.retrievers.is_empty() {
5736 return Err(MongrelError::InvalidArgument(
5737 "search requires at least one retriever".into(),
5738 ));
5739 }
5740 if request.retrievers.len() > MAX_RETRIEVERS {
5741 return Err(MongrelError::InvalidArgument(format!(
5742 "search exceeds {MAX_RETRIEVERS} retrievers"
5743 )));
5744 }
5745 if request.must.len() > MAX_HARD_CONDITIONS {
5746 return Err(MongrelError::InvalidArgument(format!(
5747 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5748 )));
5749 }
5750 for condition in &request.must {
5751 self.validate_condition(condition)?;
5752 }
5753 if request.must.iter().any(|condition| {
5754 matches!(
5755 condition,
5756 Condition::Ann { .. }
5757 | Condition::SparseMatch { .. }
5758 | Condition::MinHashSimilar { .. }
5759 )
5760 }) {
5761 return Err(MongrelError::InvalidArgument(
5762 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5763 .into(),
5764 ));
5765 }
5766 let mut names = std::collections::HashSet::new();
5767 for named in &request.retrievers {
5768 if named.name.is_empty()
5769 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5770 || !names.insert(named.name.as_str())
5771 {
5772 return Err(MongrelError::InvalidArgument(format!(
5773 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5774 crate::query::MAX_RETRIEVER_NAME_BYTES
5775 )));
5776 }
5777 if !named.weight.is_finite()
5778 || named.weight < 0.0
5779 || named.weight > MAX_RETRIEVER_WEIGHT
5780 {
5781 return Err(MongrelError::InvalidArgument(format!(
5782 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5783 )));
5784 }
5785 self.validate_retriever(&named.retriever)?;
5786 }
5787 let projection = request
5788 .projection
5789 .clone()
5790 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5791 if projection.len() > MAX_PROJECTION_COLUMNS {
5792 return Err(MongrelError::InvalidArgument(format!(
5793 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5794 )));
5795 }
5796 for column_id in &projection {
5797 if !self
5798 .schema
5799 .columns
5800 .iter()
5801 .any(|column| column.id == *column_id)
5802 {
5803 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5804 }
5805 }
5806 if let Some(crate::query::Rerank::ExactVector {
5807 embedding_column,
5808 query,
5809 candidate_limit,
5810 weight,
5811 ..
5812 }) = &request.rerank
5813 {
5814 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5815 {
5816 return Err(MongrelError::InvalidArgument(format!(
5817 "rerank candidate_limit must be between search limit and {}",
5818 crate::query::MAX_RETRIEVER_K
5819 )));
5820 }
5821 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5822 return Err(MongrelError::InvalidArgument(format!(
5823 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5824 )));
5825 }
5826 let column = self
5827 .schema
5828 .columns
5829 .iter()
5830 .find(|column| column.id == *embedding_column)
5831 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5832 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5833 return Err(MongrelError::InvalidArgument(format!(
5834 "rerank column {embedding_column} is not an embedding"
5835 )));
5836 };
5837 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5838 return Err(MongrelError::InvalidArgument(format!(
5839 "rerank query must contain {dim} finite values"
5840 )));
5841 }
5842 }
5843
5844 let hard_filter_started = std::time::Instant::now();
5845 let hard_filter = if request.must.is_empty() {
5846 None
5847 } else {
5848 let mut sets = Vec::with_capacity(request.must.len());
5849 for condition in &request.must {
5850 if let Some(context) = context {
5851 context.checkpoint()?;
5852 }
5853 sets.push(self.resolve_condition(condition, snapshot)?);
5854 }
5855 Some(RowIdSet::intersect_many(sets))
5856 };
5857 crate::trace::QueryTrace::record(|trace| {
5858 trace.hard_filter_nanos = trace
5859 .hard_filter_nanos
5860 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
5861 });
5862 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
5863 return Ok(Vec::new());
5864 }
5865
5866 let constant = match request.fusion {
5867 Fusion::ReciprocalRank { constant } => constant,
5868 };
5869 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
5870 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
5871 let mut fusion_nanos = 0u64;
5872 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
5873 std::collections::HashMap::new();
5874 for named in retrievers {
5875 if named.weight == 0.0 {
5876 continue;
5877 }
5878 if let Some(context) = context {
5879 context.checkpoint()?;
5880 }
5881 let hits = self.retrieve_filtered(
5882 &named.retriever,
5883 snapshot,
5884 hard_filter.as_ref(),
5885 authorized,
5886 candidate_authorization,
5887 context,
5888 )?;
5889 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
5890 let fusion_started = std::time::Instant::now();
5891 for hit in hits {
5892 if let Some(context) = context {
5893 context.consume(1)?;
5894 }
5895 let contribution = named.weight / (constant as f64 + hit.rank as f64);
5896 if !contribution.is_finite() {
5897 return Err(MongrelError::InvalidArgument(
5898 "retriever contribution must be finite".into(),
5899 ));
5900 }
5901 let max_fused_candidates = context.map_or(
5902 crate::query::MAX_FUSED_CANDIDATES,
5903 crate::query::AiExecutionContext::max_fused_candidates,
5904 );
5905 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
5906 return Err(MongrelError::WorkBudgetExceeded);
5907 }
5908 let entry = fused.entry(hit.row_id).or_default();
5909 entry.0 += contribution;
5910 if !entry.0.is_finite() {
5911 return Err(MongrelError::InvalidArgument(
5912 "fused score must be finite".into(),
5913 ));
5914 }
5915 entry.1.push(ComponentScore {
5916 retriever_name: retriever_name.clone(),
5917 rank: hit.rank,
5918 raw_score: hit.score,
5919 contribution,
5920 });
5921 }
5922 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
5923 }
5924 let union_size = fused.len();
5925 let mut ranked: Vec<_> = fused
5926 .into_iter()
5927 .map(|(row_id, (fused_score, components))| {
5928 (row_id, fused_score, components, None, fused_score)
5929 })
5930 .collect();
5931 let order = |(a_row, _, _, _, a_score): &(
5932 RowId,
5933 f64,
5934 Vec<ComponentScore>,
5935 Option<f32>,
5936 f64,
5937 ),
5938 (b_row, _, _, _, b_score): &(
5939 RowId,
5940 f64,
5941 Vec<ComponentScore>,
5942 Option<f32>,
5943 f64,
5944 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
5945 if let Some(crate::query::Rerank::ExactVector {
5946 embedding_column,
5947 query,
5948 metric,
5949 candidate_limit,
5950 weight,
5951 }) = &request.rerank
5952 {
5953 let fused_order = |(a_row, a_score, ..): &(
5954 RowId,
5955 f64,
5956 Vec<ComponentScore>,
5957 Option<f32>,
5958 f64,
5959 ),
5960 (b_row, b_score, ..): &(
5961 RowId,
5962 f64,
5963 Vec<ComponentScore>,
5964 Option<f32>,
5965 f64,
5966 )| {
5967 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
5968 };
5969 let selection_started = std::time::Instant::now();
5970 if let Some(context) = context {
5971 context.consume(ranked.len())?;
5972 }
5973 if ranked.len() > *candidate_limit {
5974 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
5975 ranked.truncate(*candidate_limit);
5976 }
5977 ranked.sort_by(fused_order);
5978 fusion_nanos =
5979 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5980 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5981 if let Some(context) = context {
5982 context.consume(row_ids.len())?;
5983 }
5984 let query_now =
5985 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5986 let gather_started = std::time::Instant::now();
5987 let vectors = self.values_for_rids_batch_at_with_context(
5988 &row_ids,
5989 *embedding_column,
5990 snapshot,
5991 query_now,
5992 context,
5993 )?;
5994 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5995 let vector_work =
5996 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
5997 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
5998 if let Some(context) = context {
5999 context.consume(vector_work)?;
6000 }
6001 query
6002 .iter()
6003 .map(|value| f64::from(*value).powi(2))
6004 .sum::<f64>()
6005 .sqrt()
6006 } else {
6007 0.0
6008 };
6009 let score_started = std::time::Instant::now();
6010 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6011 for (row_id, value) in vectors {
6012 let Value::Embedding(vector) = value else {
6013 continue;
6014 };
6015 let score = match metric {
6016 crate::query::VectorMetric::DotProduct => {
6017 if let Some(context) = context {
6018 context.consume(vector_work)?;
6019 }
6020 query
6021 .iter()
6022 .zip(&vector)
6023 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6024 .sum::<f64>()
6025 }
6026 crate::query::VectorMetric::Cosine => {
6027 if let Some(context) = context {
6028 context.consume(vector_work.saturating_mul(2))?;
6029 }
6030 let dot = query
6031 .iter()
6032 .zip(&vector)
6033 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6034 .sum::<f64>();
6035 let norm = vector
6036 .iter()
6037 .map(|value| f64::from(*value).powi(2))
6038 .sum::<f64>()
6039 .sqrt();
6040 if query_norm == 0.0 || norm == 0.0 {
6041 0.0
6042 } else {
6043 dot / (query_norm * norm)
6044 }
6045 }
6046 crate::query::VectorMetric::Euclidean => {
6047 if let Some(context) = context {
6048 context.consume(vector_work)?;
6049 }
6050 query
6051 .iter()
6052 .zip(&vector)
6053 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6054 .sum::<f64>()
6055 .sqrt()
6056 }
6057 };
6058 if !score.is_finite() {
6059 return Err(MongrelError::InvalidArgument(
6060 "exact rerank score must be finite".into(),
6061 ));
6062 }
6063 scores.insert(row_id, score as f32);
6064 }
6065 let mut reranked = Vec::with_capacity(ranked.len());
6066 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6067 let Some(score) = scores.get(&row_id).copied() else {
6068 continue;
6069 };
6070 let ordering_score = match metric {
6071 crate::query::VectorMetric::Euclidean => -f64::from(score),
6072 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6073 f64::from(score)
6074 }
6075 };
6076 let final_score = fused_score + *weight * ordering_score;
6077 if !final_score.is_finite() {
6078 return Err(MongrelError::InvalidArgument(
6079 "final rerank score must be finite".into(),
6080 ));
6081 }
6082 reranked.push((row_id, fused_score, components, Some(score), final_score));
6083 }
6084 ranked = reranked;
6085 ranked.sort_by(order);
6086 crate::trace::QueryTrace::record(|trace| {
6087 trace.exact_vector_gather_nanos =
6088 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6089 trace.exact_vector_score_nanos = trace
6090 .exact_vector_score_nanos
6091 .saturating_add(score_started.elapsed().as_nanos() as u64);
6092 });
6093 }
6094 if let Some(after) = after {
6095 ranked.retain(|(row_id, _, _, _, final_score)| {
6096 final_score.total_cmp(&after.final_score).is_lt()
6097 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6098 });
6099 }
6100 let projection_started = std::time::Instant::now();
6101 let sentinel = projection
6102 .first()
6103 .copied()
6104 .or_else(|| self.schema.columns.first().map(|column| column.id));
6105 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6106 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6107 let mut projection_rows = 0usize;
6108 let mut projection_cells = 0usize;
6109 while out.len() < request.limit && !ranked.is_empty() {
6110 if let Some(context) = context {
6111 context.checkpoint()?;
6112 context.consume(ranked.len())?;
6113 }
6114 let needed = request.limit - out.len();
6115 let window_size = ranked
6116 .len()
6117 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6118 let selection_started = std::time::Instant::now();
6119 let mut remainder = if ranked.len() > window_size {
6120 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6121 ranked.split_off(window_size)
6122 } else {
6123 Vec::new()
6124 };
6125 ranked.sort_by(order);
6126 fusion_nanos =
6127 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6128 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6129 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6130 if let Some(context) = context {
6131 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6132 }
6133 projection_rows = projection_rows.saturating_add(row_ids.len());
6134 projection_cells =
6135 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6136 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6137 std::collections::HashMap::new();
6138 if let Some(column_id) = sentinel {
6139 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6140 &row_ids, column_id, snapshot, query_now, context,
6141 )? {
6142 cells.entry(row_id).or_default().insert(column_id, value);
6143 }
6144 }
6145 for &column_id in &projection {
6146 if Some(column_id) == sentinel {
6147 continue;
6148 }
6149 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6150 &row_ids, column_id, snapshot, query_now, context,
6151 )? {
6152 cells.entry(row_id).or_default().insert(column_id, value);
6153 }
6154 }
6155 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6156 ranked.drain(..)
6157 {
6158 let Some(row_cells) = cells.remove(&row_id) else {
6159 continue;
6160 };
6161 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6162 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6163 out.push(SearchHit {
6164 row_id,
6165 cells: projection
6166 .iter()
6167 .filter_map(|column_id| {
6168 row_cells
6169 .get(column_id)
6170 .cloned()
6171 .map(|value| (*column_id, value))
6172 })
6173 .collect(),
6174 components,
6175 fused_score,
6176 exact_rerank_score,
6177 final_score,
6178 final_rank,
6179 });
6180 if out.len() == request.limit {
6181 break;
6182 }
6183 }
6184 ranked.append(&mut remainder);
6185 }
6186 crate::trace::QueryTrace::record(|trace| {
6187 trace.union_size = union_size;
6188 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6189 trace.projection_nanos = trace
6190 .projection_nanos
6191 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6192 trace.total_nanos = trace
6193 .total_nanos
6194 .saturating_add(total_started.elapsed().as_nanos() as u64);
6195 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6196 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6197 if let Some(context) = context {
6198 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6199 }
6200 });
6201 Ok(out)
6202 }
6203
6204 pub fn set_similarity(
6207 &mut self,
6208 request: &crate::query::SetSimilarityRequest,
6209 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6210 self.set_similarity_with_allowed(request, None)
6211 }
6212
6213 pub fn set_similarity_at(
6214 &mut self,
6215 request: &crate::query::SetSimilarityRequest,
6216 snapshot: Snapshot,
6217 allowed: Option<&std::collections::HashSet<RowId>>,
6218 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6219 self.set_similarity_explained_at(request, snapshot, allowed)
6220 .map(|(hits, _)| hits)
6221 }
6222
6223 pub fn ann_rerank(
6225 &mut self,
6226 request: &crate::query::AnnRerankRequest,
6227 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6228 self.ann_rerank_with_allowed(request, None)
6229 }
6230
6231 pub fn ann_rerank_with_allowed(
6232 &mut self,
6233 request: &crate::query::AnnRerankRequest,
6234 allowed: Option<&std::collections::HashSet<RowId>>,
6235 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6236 self.ann_rerank_at(request, self.snapshot(), allowed)
6237 }
6238
6239 pub fn ann_rerank_at(
6240 &mut self,
6241 request: &crate::query::AnnRerankRequest,
6242 snapshot: Snapshot,
6243 allowed: Option<&std::collections::HashSet<RowId>>,
6244 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6245 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6246 }
6247
6248 pub fn ann_rerank_at_with_context(
6249 &mut self,
6250 request: &crate::query::AnnRerankRequest,
6251 snapshot: Snapshot,
6252 allowed: Option<&std::collections::HashSet<RowId>>,
6253 context: Option<&crate::query::AiExecutionContext>,
6254 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6255 self.ensure_indexes_complete()?;
6256 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6257 }
6258
6259 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6260 &mut self,
6261 request: &crate::query::AnnRerankRequest,
6262 snapshot: Snapshot,
6263 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6264 context: Option<&crate::query::AiExecutionContext>,
6265 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6266 self.ensure_indexes_complete()?;
6267 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6268 }
6269
6270 #[doc(hidden)]
6271 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6272 &self,
6273 request: &crate::query::AnnRerankRequest,
6274 snapshot: Snapshot,
6275 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6276 context: Option<&crate::query::AiExecutionContext>,
6277 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6278 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6279 }
6280
6281 fn ann_rerank_at_with_filters_and_context(
6282 &self,
6283 request: &crate::query::AnnRerankRequest,
6284 snapshot: Snapshot,
6285 allowed: Option<&std::collections::HashSet<RowId>>,
6286 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6287 context: Option<&crate::query::AiExecutionContext>,
6288 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6289 use crate::query::{
6290 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6291 };
6292 if request.candidate_k == 0 || request.limit == 0 {
6293 return Err(MongrelError::InvalidArgument(
6294 "candidate_k and limit must be > 0".into(),
6295 ));
6296 }
6297 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6298 return Err(MongrelError::InvalidArgument(format!(
6299 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6300 )));
6301 }
6302 let retriever = Retriever::Ann {
6303 column_id: request.column_id,
6304 query: request.query.clone(),
6305 k: request.candidate_k,
6306 };
6307 self.require_select()?;
6308 self.validate_retriever(&retriever)?;
6309 let hits = self.retrieve_filtered(
6310 &retriever,
6311 snapshot,
6312 None,
6313 allowed,
6314 candidate_authorization,
6315 context,
6316 )?;
6317 let distances: std::collections::HashMap<_, _> = hits
6318 .iter()
6319 .filter_map(|hit| match hit.score {
6320 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
6321 _ => None,
6322 })
6323 .collect();
6324 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6325 if let Some(context) = context {
6326 context.consume(row_ids.len())?;
6327 }
6328 let gather_started = std::time::Instant::now();
6329 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6330 let values = self.values_for_rids_batch_at_with_context(
6331 &row_ids,
6332 request.column_id,
6333 snapshot,
6334 query_now,
6335 context,
6336 )?;
6337 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6338 let score_started = std::time::Instant::now();
6339 let vector_work =
6340 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6341 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6342 if let Some(context) = context {
6343 context.consume(vector_work)?;
6344 }
6345 request
6346 .query
6347 .iter()
6348 .map(|value| f64::from(*value).powi(2))
6349 .sum::<f64>()
6350 .sqrt()
6351 } else {
6352 0.0
6353 };
6354 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6355 for (row_id, value) in values {
6356 let Value::Embedding(vector) = value else {
6357 continue;
6358 };
6359 let exact_score = match request.metric {
6360 VectorMetric::DotProduct => {
6361 if let Some(context) = context {
6362 context.consume(vector_work)?;
6363 }
6364 request
6365 .query
6366 .iter()
6367 .zip(&vector)
6368 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6369 .sum::<f64>()
6370 }
6371 VectorMetric::Cosine => {
6372 if let Some(context) = context {
6373 context.consume(vector_work.saturating_mul(2))?;
6374 }
6375 let dot = request
6376 .query
6377 .iter()
6378 .zip(&vector)
6379 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6380 .sum::<f64>();
6381 let norm = vector
6382 .iter()
6383 .map(|value| f64::from(*value).powi(2))
6384 .sum::<f64>()
6385 .sqrt();
6386 if query_norm == 0.0 || norm == 0.0 {
6387 0.0
6388 } else {
6389 dot / (query_norm * norm)
6390 }
6391 }
6392 VectorMetric::Euclidean => {
6393 if let Some(context) = context {
6394 context.consume(vector_work)?;
6395 }
6396 request
6397 .query
6398 .iter()
6399 .zip(&vector)
6400 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6401 .sum::<f64>()
6402 .sqrt()
6403 }
6404 };
6405 let exact_score = exact_score as f32;
6406 if !exact_score.is_finite() {
6407 return Err(MongrelError::InvalidArgument(
6408 "exact ANN score must be finite".into(),
6409 ));
6410 }
6411 reranked.push(AnnRerankHit {
6412 row_id,
6413 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
6414 exact_score,
6415 });
6416 }
6417 reranked.sort_by(|left, right| {
6418 let score = match request.metric {
6419 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6420 VectorMetric::Cosine | VectorMetric::DotProduct => {
6421 right.exact_score.total_cmp(&left.exact_score)
6422 }
6423 };
6424 score.then_with(|| left.row_id.cmp(&right.row_id))
6425 });
6426 reranked.truncate(request.limit);
6427 crate::trace::QueryTrace::record(|trace| {
6428 trace.exact_vector_gather_nanos =
6429 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6430 trace.exact_vector_score_nanos = trace
6431 .exact_vector_score_nanos
6432 .saturating_add(score_started.elapsed().as_nanos() as u64);
6433 });
6434 Ok(reranked)
6435 }
6436
6437 pub fn set_similarity_with_allowed(
6438 &mut self,
6439 request: &crate::query::SetSimilarityRequest,
6440 allowed: Option<&std::collections::HashSet<RowId>>,
6441 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6442 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6443 .map(|(hits, _)| hits)
6444 }
6445
6446 pub fn set_similarity_explained(
6447 &mut self,
6448 request: &crate::query::SetSimilarityRequest,
6449 ) -> Result<(
6450 Vec<crate::query::SetSimilarityHit>,
6451 crate::query::SetSimilarityTrace,
6452 )> {
6453 self.set_similarity_explained_at(request, self.snapshot(), None)
6454 }
6455
6456 fn set_similarity_explained_at(
6457 &mut self,
6458 request: &crate::query::SetSimilarityRequest,
6459 snapshot: Snapshot,
6460 allowed: Option<&std::collections::HashSet<RowId>>,
6461 ) -> Result<(
6462 Vec<crate::query::SetSimilarityHit>,
6463 crate::query::SetSimilarityTrace,
6464 )> {
6465 self.ensure_indexes_complete()?;
6466 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6467 }
6468
6469 pub fn set_similarity_at_with_context(
6470 &mut self,
6471 request: &crate::query::SetSimilarityRequest,
6472 snapshot: Snapshot,
6473 allowed: Option<&std::collections::HashSet<RowId>>,
6474 context: Option<&crate::query::AiExecutionContext>,
6475 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6476 self.ensure_indexes_complete()?;
6477 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6478 .map(|(hits, _)| hits)
6479 }
6480
6481 pub fn set_similarity_at_with_candidate_authorization_and_context(
6482 &mut self,
6483 request: &crate::query::SetSimilarityRequest,
6484 snapshot: Snapshot,
6485 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6486 context: Option<&crate::query::AiExecutionContext>,
6487 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6488 self.ensure_indexes_complete()?;
6489 self.set_similarity_explained_at_with_context(
6490 request,
6491 snapshot,
6492 None,
6493 authorization,
6494 context,
6495 )
6496 .map(|(hits, _)| hits)
6497 }
6498
6499 #[doc(hidden)]
6500 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6501 &self,
6502 request: &crate::query::SetSimilarityRequest,
6503 snapshot: Snapshot,
6504 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6505 context: Option<&crate::query::AiExecutionContext>,
6506 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6507 self.set_similarity_explained_at_with_context(
6508 request,
6509 snapshot,
6510 None,
6511 authorization,
6512 context,
6513 )
6514 .map(|(hits, _)| hits)
6515 }
6516
6517 fn set_similarity_explained_at_with_context(
6518 &self,
6519 request: &crate::query::SetSimilarityRequest,
6520 snapshot: Snapshot,
6521 allowed: Option<&std::collections::HashSet<RowId>>,
6522 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6523 context: Option<&crate::query::AiExecutionContext>,
6524 ) -> Result<(
6525 Vec<crate::query::SetSimilarityHit>,
6526 crate::query::SetSimilarityTrace,
6527 )> {
6528 use crate::query::{
6529 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6530 MAX_SET_MEMBERS,
6531 };
6532 let mut trace = crate::query::SetSimilarityTrace::default();
6533 if request.members.is_empty() {
6534 return Ok((Vec::new(), trace));
6535 }
6536 if request.candidate_k == 0 || request.limit == 0 {
6537 return Err(MongrelError::InvalidArgument(
6538 "candidate_k and limit must be > 0".into(),
6539 ));
6540 }
6541 if request.candidate_k > MAX_RETRIEVER_K
6542 || request.limit > MAX_FINAL_LIMIT
6543 || request.members.len() > MAX_SET_MEMBERS
6544 {
6545 return Err(MongrelError::InvalidArgument(format!(
6546 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6547 )));
6548 }
6549 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6550 return Err(MongrelError::InvalidArgument(
6551 "min_jaccard must be finite and between 0 and 1".into(),
6552 ));
6553 }
6554 let started = std::time::Instant::now();
6555 let retriever = Retriever::MinHash {
6556 column_id: request.column_id,
6557 members: request.members.clone(),
6558 k: request.candidate_k,
6559 };
6560 self.require_select()?;
6561 self.validate_retriever(&retriever)?;
6562 let hits = self.retrieve_filtered(
6563 &retriever,
6564 snapshot,
6565 None,
6566 allowed,
6567 candidate_authorization,
6568 context,
6569 )?;
6570 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6571 trace.candidate_count = hits.len();
6572 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6573 if let Some(context) = context {
6574 context.consume(row_ids.len())?;
6575 }
6576 let started = std::time::Instant::now();
6577 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6578 let values = self.values_for_rids_batch_at_with_context(
6579 &row_ids,
6580 request.column_id,
6581 snapshot,
6582 query_now,
6583 context,
6584 )?;
6585 trace.gather_us = started.elapsed().as_micros() as u64;
6586 if let Some(context) = context {
6587 context.consume(request.members.len())?;
6588 }
6589 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6590 let estimates: std::collections::HashMap<_, _> = hits
6591 .into_iter()
6592 .filter_map(|hit| match hit.score {
6593 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6594 _ => None,
6595 })
6596 .collect();
6597 let started = std::time::Instant::now();
6598 let mut parsed = Vec::with_capacity(values.len());
6599 for (row_id, value) in values {
6600 let Value::Bytes(bytes) = value else {
6601 continue;
6602 };
6603 if let Some(context) = context {
6604 context.consume(crate::query::work_units(
6605 bytes.len(),
6606 crate::query::PARSE_WORK_QUANTUM,
6607 ))?;
6608 }
6609 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6610 continue;
6611 };
6612 if let Some(context) = context {
6613 context.consume(members.len())?;
6614 }
6615 let stored = members
6616 .into_iter()
6617 .filter_map(|member| match member {
6618 serde_json::Value::String(value) => {
6619 Some(crate::query::SetMember::String(value))
6620 }
6621 serde_json::Value::Number(value) => {
6622 Some(crate::query::SetMember::Number(value))
6623 }
6624 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6625 _ => None,
6626 })
6627 .collect::<std::collections::HashSet<_>>();
6628 parsed.push((row_id, stored));
6629 }
6630 trace.parse_us = started.elapsed().as_micros() as u64;
6631 trace.verified_count = parsed.len();
6632 let started = std::time::Instant::now();
6633 let mut exact = Vec::new();
6634 for (row_id, stored) in parsed {
6635 if let Some(context) = context {
6636 context.consume(query.len().saturating_add(stored.len()))?;
6637 }
6638 let union = query.union(&stored).count();
6639 let score = if union == 0 {
6640 1.0
6641 } else {
6642 query.intersection(&stored).count() as f32 / union as f32
6643 };
6644 if score >= request.min_jaccard {
6645 exact.push(SetSimilarityHit {
6646 row_id,
6647 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6648 exact_jaccard: score,
6649 });
6650 }
6651 }
6652 exact.sort_by(|a, b| {
6653 b.exact_jaccard
6654 .total_cmp(&a.exact_jaccard)
6655 .then_with(|| a.row_id.cmp(&b.row_id))
6656 });
6657 exact.truncate(request.limit);
6658 trace.score_us = started.elapsed().as_micros() as u64;
6659 crate::trace::QueryTrace::record(|query_trace| {
6660 query_trace.exact_set_gather_nanos = query_trace
6661 .exact_set_gather_nanos
6662 .saturating_add(trace.gather_us.saturating_mul(1_000));
6663 query_trace.exact_set_parse_nanos = query_trace
6664 .exact_set_parse_nanos
6665 .saturating_add(trace.parse_us.saturating_mul(1_000));
6666 query_trace.exact_set_score_nanos = query_trace
6667 .exact_set_score_nanos
6668 .saturating_add(trace.score_us.saturating_mul(1_000));
6669 });
6670 Ok((exact, trace))
6671 }
6672
6673 fn values_for_rids_batch_at(
6675 &self,
6676 row_ids: &[u64],
6677 column_id: u16,
6678 snapshot: Snapshot,
6679 now: i64,
6680 ) -> Result<Vec<(RowId, Value)>> {
6681 if self.ttl.is_none()
6682 && self.memtable.is_empty()
6683 && self.mutable_run.is_empty()
6684 && self.run_refs.len() == 1
6685 {
6686 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6687 if row_ids.len().saturating_mul(24) < reader.row_count() {
6692 let mut values = Vec::with_capacity(row_ids.len());
6693 for &raw_row_id in row_ids {
6694 let row_id = RowId(raw_row_id);
6695 if let Some((_, false, Some(value))) =
6696 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6697 {
6698 values.push((row_id, value));
6699 }
6700 }
6701 return Ok(values);
6702 }
6703 let (positions, visible_row_ids) =
6704 reader.visible_positions_with_rids(snapshot.epoch)?;
6705 let requested: Vec<(RowId, usize)> = row_ids
6706 .iter()
6707 .filter_map(|raw| {
6708 visible_row_ids
6709 .binary_search(&(*raw as i64))
6710 .ok()
6711 .map(|index| (RowId(*raw), positions[index]))
6712 })
6713 .collect();
6714 let values = reader.gather_column(
6715 column_id,
6716 &requested
6717 .iter()
6718 .map(|(_, position)| *position)
6719 .collect::<Vec<_>>(),
6720 )?;
6721 return Ok(requested
6722 .into_iter()
6723 .zip(values)
6724 .map(|((row_id, _), value)| (row_id, value))
6725 .collect());
6726 }
6727 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6728 }
6729
6730 fn values_for_rids_batch_at_with_context(
6731 &self,
6732 row_ids: &[u64],
6733 column_id: u16,
6734 snapshot: Snapshot,
6735 now: i64,
6736 context: Option<&crate::query::AiExecutionContext>,
6737 ) -> Result<Vec<(RowId, Value)>> {
6738 let Some(context) = context else {
6739 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6740 };
6741 let mut values = Vec::with_capacity(row_ids.len());
6742 for chunk in row_ids.chunks(256) {
6743 context.checkpoint()?;
6744 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6745 }
6746 Ok(values)
6747 }
6748
6749 fn values_for_rids_at(
6751 &self,
6752 row_ids: &[u64],
6753 column_id: u16,
6754 snapshot: Snapshot,
6755 now: i64,
6756 ) -> Result<Vec<(RowId, Value)>> {
6757 let mut readers: Vec<_> = self
6758 .run_refs
6759 .iter()
6760 .map(|run| self.open_reader(run.run_id))
6761 .collect::<Result<_>>()?;
6762 let mut out = Vec::with_capacity(row_ids.len());
6763 for &raw_row_id in row_ids {
6764 let row_id = RowId(raw_row_id);
6765 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6766 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6767 let overlay = match (mem, mutable) {
6768 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6769 (a_epoch, a)
6770 } else {
6771 (b_epoch, b)
6772 }),
6773 (Some(value), None) | (None, Some(value)) => Some(value),
6774 (None, None) => None,
6775 };
6776 if let Some((_, row)) = overlay {
6777 if !row.deleted && !self.row_expired_at(&row, now) {
6778 if let Some(value) = row.columns.get(&column_id) {
6779 out.push((row_id, value.clone()));
6780 }
6781 }
6782 continue;
6783 }
6784
6785 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6786 for (index, reader) in readers.iter_mut().enumerate() {
6787 if let Some((epoch, deleted, value)) =
6788 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6789 {
6790 if best
6791 .as_ref()
6792 .map(|(best_epoch, ..)| epoch > *best_epoch)
6793 .unwrap_or(true)
6794 {
6795 best = Some((epoch, deleted, value, index));
6796 }
6797 }
6798 }
6799 let Some((_, false, Some(value), reader_index)) = best else {
6800 continue;
6801 };
6802 if let Some(ttl) = self.ttl {
6803 if ttl.column_id != column_id {
6804 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6805 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6806 {
6807 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6808 continue;
6809 }
6810 }
6811 } else if let Value::Int64(timestamp) = value {
6812 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6813 continue;
6814 }
6815 }
6816 }
6817 out.push((row_id, value));
6818 }
6819 Ok(out)
6820 }
6821
6822 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6827 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6828 }
6829
6830 pub fn rows_for_rids_with_context(
6831 &self,
6832 rids: &[u64],
6833 snapshot: Snapshot,
6834 context: &crate::query::AiExecutionContext,
6835 ) -> Result<Vec<Row>> {
6836 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6837 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
6838 }
6839
6840 fn rows_for_rids_at_time(
6841 &self,
6842 rids: &[u64],
6843 snapshot: Snapshot,
6844 ttl_now: i64,
6845 control: Option<&crate::ExecutionControl>,
6846 ) -> Result<Vec<Row>> {
6847 use std::collections::HashMap;
6848 let mut rows = Vec::with_capacity(rids.len());
6849 let tier_size = self.memtable.len() + self.mutable_run.len();
6866 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
6867 if rids.len().saturating_mul(24) < tier_size {
6868 for &rid in rids {
6869 if overlay.len() & 255 == 0 {
6870 control
6871 .map(crate::ExecutionControl::checkpoint)
6872 .transpose()?;
6873 }
6874 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
6875 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
6876 let newest = match (mem, mrun) {
6877 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
6878 (Some((_, mr)), None) => Some(mr),
6879 (None, Some((_, rr))) => Some(rr),
6880 (None, None) => None,
6881 };
6882 if let Some(row) = newest {
6883 overlay.insert(rid, row);
6884 }
6885 }
6886 } else {
6887 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
6888 overlay
6889 .entry(row.row_id.0)
6890 .and_modify(|e| {
6891 if row.committed_epoch > e.committed_epoch {
6892 *e = row.clone();
6893 }
6894 })
6895 .or_insert(row);
6896 };
6897 for (index, row) in self
6898 .memtable
6899 .visible_versions(snapshot.epoch)
6900 .into_iter()
6901 .enumerate()
6902 {
6903 if index & 255 == 0 {
6904 control
6905 .map(crate::ExecutionControl::checkpoint)
6906 .transpose()?;
6907 }
6908 fold_newest(row, &mut overlay);
6909 }
6910 for (index, row) in self
6911 .mutable_run
6912 .visible_versions(snapshot.epoch)
6913 .into_iter()
6914 .enumerate()
6915 {
6916 if index & 255 == 0 {
6917 control
6918 .map(crate::ExecutionControl::checkpoint)
6919 .transpose()?;
6920 }
6921 fold_newest(row, &mut overlay);
6922 }
6923 }
6924 if self.run_refs.len() == 1 {
6925 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6926 if rids.len().saturating_mul(24) < reader.row_count() {
6934 for (index, &rid) in rids.iter().enumerate() {
6935 if index & 255 == 0 {
6936 control
6937 .map(crate::ExecutionControl::checkpoint)
6938 .transpose()?;
6939 }
6940 if let Some(r) = overlay.get(&rid) {
6941 if !r.deleted {
6942 rows.push(r.clone());
6943 }
6944 continue;
6945 }
6946 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
6947 if !row.deleted {
6948 rows.push(row);
6949 }
6950 }
6951 }
6952 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6953 return Ok(rows);
6954 }
6955 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6964 enum Src {
6967 Overlay,
6968 Run,
6969 }
6970 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
6971 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
6972 for (index, rid) in rids.iter().enumerate() {
6973 if index & 255 == 0 {
6974 control
6975 .map(crate::ExecutionControl::checkpoint)
6976 .transpose()?;
6977 }
6978 if overlay.contains_key(rid) {
6979 plan.push(Src::Overlay);
6980 continue;
6981 }
6982 match vis_rids.binary_search(&(*rid as i64)) {
6983 Ok(i) => {
6984 plan.push(Src::Run);
6985 fetch.push(positions[i]);
6986 }
6987 Err(_) => { }
6988 }
6989 }
6990 let fetched = reader.materialize_batch(&fetch)?;
6991 let mut fetched_iter = fetched.into_iter();
6992 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
6993 if index & 255 == 0 {
6994 control
6995 .map(crate::ExecutionControl::checkpoint)
6996 .transpose()?;
6997 }
6998 match src {
6999 Src::Overlay => {
7000 if let Some(r) = overlay.get(rid) {
7001 if !r.deleted {
7002 rows.push(r.clone());
7003 }
7004 }
7005 }
7006 Src::Run => {
7007 if let Some(row) = fetched_iter.next() {
7008 if !row.deleted {
7009 rows.push(row);
7010 }
7011 }
7012 }
7013 }
7014 }
7015 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7016 return Ok(rows);
7017 }
7018 let mut readers: Vec<_> = self
7022 .run_refs
7023 .iter()
7024 .map(|rr| self.open_reader(rr.run_id))
7025 .collect::<Result<Vec<_>>>()?;
7026 for (index, rid) in rids.iter().enumerate() {
7027 if index & 255 == 0 {
7028 control
7029 .map(crate::ExecutionControl::checkpoint)
7030 .transpose()?;
7031 }
7032 if let Some(r) = overlay.get(rid) {
7033 if !r.deleted {
7034 rows.push(r.clone());
7035 }
7036 continue;
7037 }
7038 let mut best: Option<(Epoch, Row)> = None;
7039 for reader in readers.iter_mut() {
7040 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7041 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7042 best = Some((epoch, row));
7043 }
7044 }
7045 }
7046 if let Some((_, r)) = best {
7047 if !r.deleted {
7048 rows.push(r);
7049 }
7050 }
7051 }
7052 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7053 Ok(rows)
7054 }
7055
7056 pub fn indexes_complete(&self) -> bool {
7066 self.indexes_complete
7067 }
7068
7069 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7071 self.index_build_policy
7072 }
7073
7074 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7078 self.index_build_policy = policy;
7079 }
7080
7081 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7086 if !self.indexes_complete {
7090 return None;
7091 }
7092 let b = self.bitmap.get(&column_id)?;
7093 let result: Vec<Vec<u8>> = b
7094 .keys()
7095 .into_iter()
7096 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7097 .collect();
7098 Some(result)
7099 }
7100
7101 pub fn fk_join_row_ids(
7102 &self,
7103 fk_column_id: u16,
7104 pk_values: &[Vec<u8>],
7105 fk_conditions: &[crate::query::Condition],
7106 snapshot: Snapshot,
7107 ) -> Result<Vec<u64>> {
7108 let Some(b) = self.bitmap.get(&fk_column_id) else {
7109 return Ok(Vec::new());
7110 };
7111 let mut join_set = {
7112 let mut acc = roaring::RoaringBitmap::new();
7113 for v in pk_values {
7114 acc |= b.get(v);
7115 }
7116 RowIdSet::from_roaring(acc)
7117 };
7118 if !fk_conditions.is_empty() {
7119 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7120 sets.push(join_set);
7121 for c in fk_conditions {
7122 sets.push(self.resolve_condition(c, snapshot)?);
7123 }
7124 join_set = RowIdSet::intersect_many(sets);
7125 }
7126 Ok(join_set.into_sorted_vec())
7127 }
7128
7129 pub fn fk_join_count(
7135 &self,
7136 fk_column_id: u16,
7137 pk_values: &[Vec<u8>],
7138 fk_conditions: &[crate::query::Condition],
7139 snapshot: Snapshot,
7140 ) -> Result<u64> {
7141 let Some(b) = self.bitmap.get(&fk_column_id) else {
7142 return Ok(0);
7143 };
7144 let mut acc = roaring::RoaringBitmap::new();
7145 for v in pk_values {
7146 acc |= b.get(v);
7147 }
7148 if fk_conditions.is_empty() {
7149 return Ok(acc.len());
7150 }
7151 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7152 sets.push(RowIdSet::from_roaring(acc));
7153 for c in fk_conditions {
7154 sets.push(self.resolve_condition(c, snapshot)?);
7155 }
7156 Ok(RowIdSet::intersect_many(sets).len() as u64)
7157 }
7158
7159 fn resolve_condition(
7164 &self,
7165 c: &crate::query::Condition,
7166 snapshot: Snapshot,
7167 ) -> Result<RowIdSet> {
7168 self.resolve_condition_with_allowed(c, snapshot, None)
7169 }
7170
7171 fn resolve_condition_with_allowed(
7172 &self,
7173 c: &crate::query::Condition,
7174 snapshot: Snapshot,
7175 allowed: Option<&std::collections::HashSet<RowId>>,
7176 ) -> Result<RowIdSet> {
7177 use crate::query::Condition;
7178 self.validate_condition(c)?;
7179 Ok(match c {
7180 Condition::Pk(key) => {
7181 let lookup = self
7182 .schema
7183 .primary_key()
7184 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7185 .unwrap_or_else(|| key.clone());
7186 self.hot
7187 .get(&lookup)
7188 .map(|r| RowIdSet::one(r.0))
7189 .unwrap_or_else(RowIdSet::empty)
7190 }
7191 Condition::BitmapEq { column_id, value } => {
7192 let lookup = self.index_lookup_key_bytes(*column_id, value);
7193 self.bitmap
7194 .get(column_id)
7195 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7196 .unwrap_or_else(RowIdSet::empty)
7197 }
7198 Condition::BitmapIn { column_id, values } => {
7199 let bm = self.bitmap.get(column_id);
7200 let mut acc = roaring::RoaringBitmap::new();
7201 if let Some(b) = bm {
7202 for v in values {
7203 let lookup = self.index_lookup_key_bytes(*column_id, v);
7204 acc |= b.get(&lookup);
7205 }
7206 }
7207 RowIdSet::from_roaring(acc)
7208 }
7209 Condition::BytesPrefix { column_id, prefix } => {
7210 if let Some(b) = self.bitmap.get(column_id) {
7215 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7216 let mut acc = roaring::RoaringBitmap::new();
7217 for key in b.keys() {
7218 if key.starts_with(&lookup_prefix) {
7219 acc |= b.get(&key);
7220 }
7221 }
7222 RowIdSet::from_roaring(acc)
7223 } else {
7224 RowIdSet::empty()
7225 }
7226 }
7227 Condition::FmContains { column_id, pattern } => self
7228 .fm
7229 .get(column_id)
7230 .map(|f| {
7231 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7232 })
7233 .unwrap_or_else(RowIdSet::empty),
7234 Condition::FmContainsAll {
7235 column_id,
7236 patterns,
7237 } => {
7238 if let Some(f) = self.fm.get(column_id) {
7241 let sets: Vec<RowIdSet> = patterns
7242 .iter()
7243 .map(|pat| {
7244 RowIdSet::from_unsorted(
7245 f.locate(pat).into_iter().map(|r| r.0).collect(),
7246 )
7247 })
7248 .collect();
7249 RowIdSet::intersect_many(sets)
7250 } else {
7251 RowIdSet::empty()
7252 }
7253 }
7254 Condition::Ann {
7255 column_id,
7256 query,
7257 k,
7258 } => RowIdSet::from_unsorted(
7259 self.retrieve_filtered(
7260 &crate::query::Retriever::Ann {
7261 column_id: *column_id,
7262 query: query.clone(),
7263 k: *k,
7264 },
7265 snapshot,
7266 None,
7267 allowed,
7268 None,
7269 None,
7270 )?
7271 .into_iter()
7272 .map(|hit| hit.row_id.0)
7273 .collect(),
7274 ),
7275 Condition::SparseMatch {
7276 column_id,
7277 query,
7278 k,
7279 } => RowIdSet::from_unsorted(
7280 self.retrieve_filtered(
7281 &crate::query::Retriever::Sparse {
7282 column_id: *column_id,
7283 query: query.clone(),
7284 k: *k,
7285 },
7286 snapshot,
7287 None,
7288 allowed,
7289 None,
7290 None,
7291 )?
7292 .into_iter()
7293 .map(|hit| hit.row_id.0)
7294 .collect(),
7295 ),
7296 Condition::MinHashSimilar {
7297 column_id,
7298 query,
7299 k,
7300 } => match self.minhash.get(column_id) {
7301 Some(index) => {
7302 let candidates = index.candidate_row_ids(query);
7303 let eligible =
7304 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7305 RowIdSet::from_unsorted(
7306 index
7307 .search_filtered(query, *k, |row_id| {
7308 eligible.contains(&row_id)
7309 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7310 })
7311 .into_iter()
7312 .map(|(row_id, _)| row_id.0)
7313 .collect(),
7314 )
7315 }
7316 None => RowIdSet::empty(),
7317 },
7318 Condition::Range { column_id, lo, hi } => {
7319 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7328 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7329 } else if self.run_refs.len() == 1 {
7330 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7331 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7332 } else {
7333 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7334 };
7335 set.remove_many(self.overlay_rid_set(snapshot));
7336 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7337 set
7338 }
7339 Condition::RangeF64 {
7340 column_id,
7341 lo,
7342 lo_inclusive,
7343 hi,
7344 hi_inclusive,
7345 } => {
7346 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7349 RowIdSet::from_unsorted(
7350 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7351 .into_iter()
7352 .collect(),
7353 )
7354 } else if self.run_refs.len() == 1 {
7355 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7356 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7357 } else {
7358 return self.range_scan_f64(
7359 *column_id,
7360 *lo,
7361 *lo_inclusive,
7362 *hi,
7363 *hi_inclusive,
7364 snapshot,
7365 );
7366 };
7367 set.remove_many(self.overlay_rid_set(snapshot));
7368 self.range_scan_overlay_f64(
7369 &mut set,
7370 *column_id,
7371 *lo,
7372 *lo_inclusive,
7373 *hi,
7374 *hi_inclusive,
7375 snapshot,
7376 );
7377 set
7378 }
7379 Condition::IsNull { column_id } => {
7380 let mut set = if self.run_refs.len() == 1 {
7381 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7382 r.null_row_id_set(*column_id, true)?
7383 } else {
7384 return self.null_scan(*column_id, true, snapshot);
7385 };
7386 set.remove_many(self.overlay_rid_set(snapshot));
7387 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7388 set
7389 }
7390 Condition::IsNotNull { column_id } => {
7391 let mut set = if self.run_refs.len() == 1 {
7392 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7393 r.null_row_id_set(*column_id, false)?
7394 } else {
7395 return self.null_scan(*column_id, false, snapshot);
7396 };
7397 set.remove_many(self.overlay_rid_set(snapshot));
7398 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7399 set
7400 }
7401 })
7402 }
7403
7404 fn range_scan_i64(
7412 &self,
7413 column_id: u16,
7414 lo: i64,
7415 hi: i64,
7416 snapshot: Snapshot,
7417 ) -> Result<RowIdSet> {
7418 let mut row_ids = Vec::new();
7419 let overlay_rids = self.overlay_rid_set(snapshot);
7420 for rr in &self.run_refs {
7421 let mut reader = self.open_reader(rr.run_id)?;
7422 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7423 for rid in matched {
7424 if !overlay_rids.contains(&rid) {
7425 row_ids.push(rid);
7426 }
7427 }
7428 }
7429 let mut s = RowIdSet::from_unsorted(row_ids);
7430 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7431 Ok(s)
7432 }
7433
7434 fn range_scan_f64(
7437 &self,
7438 column_id: u16,
7439 lo: f64,
7440 lo_inclusive: bool,
7441 hi: f64,
7442 hi_inclusive: bool,
7443 snapshot: Snapshot,
7444 ) -> Result<RowIdSet> {
7445 let mut row_ids = Vec::new();
7446 let overlay_rids = self.overlay_rid_set(snapshot);
7447 for rr in &self.run_refs {
7448 let mut reader = self.open_reader(rr.run_id)?;
7449 let matched = reader.range_row_ids_visible_f64(
7450 column_id,
7451 lo,
7452 lo_inclusive,
7453 hi,
7454 hi_inclusive,
7455 snapshot.epoch,
7456 )?;
7457 for rid in matched {
7458 if !overlay_rids.contains(&rid) {
7459 row_ids.push(rid);
7460 }
7461 }
7462 }
7463 let mut s = RowIdSet::from_unsorted(row_ids);
7464 self.range_scan_overlay_f64(
7465 &mut s,
7466 column_id,
7467 lo,
7468 lo_inclusive,
7469 hi,
7470 hi_inclusive,
7471 snapshot,
7472 );
7473 Ok(s)
7474 }
7475
7476 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7478 let mut s = HashSet::new();
7479 for row in self.memtable.visible_versions(snapshot.epoch) {
7480 s.insert(row.row_id.0);
7481 }
7482 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7483 s.insert(row.row_id.0);
7484 }
7485 s
7486 }
7487
7488 fn range_scan_overlay_i64(
7489 &self,
7490 s: &mut RowIdSet,
7491 column_id: u16,
7492 lo: i64,
7493 hi: i64,
7494 snapshot: Snapshot,
7495 ) {
7496 let mut newest: HashMap<u64, &Row> = HashMap::new();
7501 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7502 let memtable = self.memtable.visible_versions(snapshot.epoch);
7503 for r in &mutable {
7504 newest.entry(r.row_id.0).or_insert(r);
7505 }
7506 for r in &memtable {
7507 newest.insert(r.row_id.0, r);
7508 }
7509 for row in newest.values() {
7510 if !row.deleted {
7511 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7512 if *v >= lo && *v <= hi {
7513 s.insert(row.row_id.0);
7514 }
7515 }
7516 }
7517 }
7518 }
7519
7520 #[allow(clippy::too_many_arguments)]
7521 fn range_scan_overlay_f64(
7522 &self,
7523 s: &mut RowIdSet,
7524 column_id: u16,
7525 lo: f64,
7526 lo_inclusive: bool,
7527 hi: f64,
7528 hi_inclusive: bool,
7529 snapshot: Snapshot,
7530 ) {
7531 let mut newest: HashMap<u64, &Row> = HashMap::new();
7534 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7535 let memtable = self.memtable.visible_versions(snapshot.epoch);
7536 for r in &mutable {
7537 newest.entry(r.row_id.0).or_insert(r);
7538 }
7539 for r in &memtable {
7540 newest.insert(r.row_id.0, r);
7541 }
7542 for row in newest.values() {
7543 if !row.deleted {
7544 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7545 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7546 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7547 if ok_lo && ok_hi {
7548 s.insert(row.row_id.0);
7549 }
7550 }
7551 }
7552 }
7553 }
7554
7555 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7558 let mut row_ids = Vec::new();
7559 let overlay_rids = self.overlay_rid_set(snapshot);
7560 for rr in &self.run_refs {
7561 let mut reader = self.open_reader(rr.run_id)?;
7562 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7563 for rid in matched {
7564 if !overlay_rids.contains(&rid) {
7565 row_ids.push(rid);
7566 }
7567 }
7568 }
7569 let mut s = RowIdSet::from_unsorted(row_ids);
7570 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7571 Ok(s)
7572 }
7573
7574 fn null_scan_overlay(
7578 &self,
7579 s: &mut RowIdSet,
7580 column_id: u16,
7581 want_nulls: bool,
7582 snapshot: Snapshot,
7583 ) {
7584 let mut newest: HashMap<u64, &Row> = HashMap::new();
7585 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7586 let memtable = self.memtable.visible_versions(snapshot.epoch);
7587 for r in &mutable {
7588 newest.entry(r.row_id.0).or_insert(r);
7589 }
7590 for r in &memtable {
7591 newest.insert(r.row_id.0, r);
7592 }
7593 for row in newest.values() {
7594 if row.deleted {
7595 continue;
7596 }
7597 let is_null = !row.columns.contains_key(&column_id)
7598 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7599 if is_null == want_nulls {
7600 s.insert(row.row_id.0);
7601 }
7602 }
7603 }
7604
7605 pub fn snapshot(&self) -> Snapshot {
7606 Snapshot::at(self.epoch.visible())
7607 }
7608
7609 pub fn data_generation(&self) -> u64 {
7611 self.data_generation
7612 }
7613
7614 pub(crate) fn bump_data_generation(&mut self) {
7615 self.data_generation = self.data_generation.wrapping_add(1);
7616 }
7617
7618 pub fn table_id(&self) -> u64 {
7620 self.table_id
7621 }
7622
7623 fn seal_generations(&mut self) {
7629 self.memtable.seal();
7630 self.mutable_run.seal();
7631 self.hot.seal();
7632 for index in self.bitmap.values_mut() {
7633 index.seal();
7634 }
7635 for index in self.ann.values_mut() {
7636 index.seal();
7637 }
7638 for index in self.fm.values_mut() {
7639 index.seal();
7640 }
7641 for index in self.sparse.values_mut() {
7642 index.seal();
7643 }
7644 for index in self.minhash.values_mut() {
7645 index.seal();
7646 }
7647 self.pk_by_row.seal();
7648 }
7649
7650 fn capture_read_generation(&self) -> ReadGeneration {
7655 let visible_through = self.current_epoch();
7656 ReadGeneration {
7657 schema: Arc::new(self.schema.clone()),
7658 base_runs: Arc::new(self.run_refs.clone()),
7659 deltas: TableDeltas {
7660 memtable: self.memtable.clone(),
7661 mutable_run: self.mutable_run.clone(),
7662 hot: self.hot.clone(),
7663 pk_by_row: self.pk_by_row.clone(),
7664 },
7665 indexes: Arc::new(IndexGeneration::capture(
7666 &self.bitmap,
7667 &self.learned_range,
7668 &self.fm,
7669 &self.ann,
7670 &self.sparse,
7671 &self.minhash,
7672 visible_through,
7673 )),
7674 visible_through,
7675 }
7676 }
7677
7678 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
7683 self.ensure_indexes_complete()?;
7684 self.seal_generations();
7685 let view = Arc::new(self.capture_read_generation());
7686 self.published.store(Arc::clone(&view));
7687 Ok(view)
7688 }
7689
7690 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
7696 self.published.load_full()
7697 }
7698
7699 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
7701 &self.pins
7702 }
7703
7704 pub fn version_gc_floor(&self) -> Epoch {
7709 self.min_active_snapshot()
7710 .unwrap_or_else(|| self.current_epoch())
7711 }
7712
7713 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
7720 let mut report = self.pins.report();
7721 let transaction_floor = [
7722 self.pinned.keys().next().copied(),
7723 self.snapshots.min_pinned(),
7724 ]
7725 .into_iter()
7726 .flatten()
7727 .min();
7728 if let Some(epoch) = transaction_floor {
7729 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
7730 }
7731 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
7732 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
7733 }
7734 report
7735 }
7736
7737 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7738 self.publish_read_generation()?;
7739 let mut generation = self.clone();
7740 generation.read_only = true;
7741 generation.wal = WalSink::ReadOnly;
7742 generation.pending_delete_rids.clear();
7743 generation.pending_put_cols.clear();
7744 generation.pending_rows.clear();
7745 generation.pending_rows_auto_inc.clear();
7746 generation.pending_dels.clear();
7747 generation.pending_truncate = None;
7748 generation.agg_cache = Arc::new(HashMap::new());
7749 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
7752 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
7755 crate::retention::PinSource::ReadGeneration,
7756 self.current_epoch(),
7757 )));
7758 Ok(generation)
7759 }
7760
7761 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7762 (std::mem::size_of::<Self>() as u64)
7763 .saturating_add(self.memtable.approx_bytes())
7764 .saturating_add(self.mutable_run.approx_bytes())
7765 .saturating_add(self.live_count.saturating_mul(64))
7766 }
7767
7768 pub fn pin_snapshot(&mut self) -> Snapshot {
7771 let e = self.epoch.visible();
7772 *self.pinned.entry(e).or_insert(0) += 1;
7773 Snapshot::at(e)
7774 }
7775
7776 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7778 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7779 *count -= 1;
7780 if *count == 0 {
7781 self.pinned.remove(&snap.epoch);
7782 }
7783 }
7784 }
7785
7786 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7798 let local = self.pinned.keys().next().copied();
7799 let global = self.snapshots.min_pinned();
7800 let history = self.snapshots.history_floor(self.current_epoch());
7801 let pinned = self.pins.oldest_pinned();
7802 [local, global, history, pinned].into_iter().flatten().min()
7803 }
7804
7805 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7809 self.ensure_writable()?;
7810 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7811 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7812 }
7813
7814 pub fn clear_ttl(&mut self) -> Result<()> {
7815 self.ensure_writable()?;
7816 self.apply_ttl_policy_at(None, self.current_epoch())
7817 }
7818
7819 pub fn ttl(&self) -> Option<TtlPolicy> {
7820 self.ttl
7821 }
7822
7823 pub(crate) fn prepare_ttl_policy(
7824 &self,
7825 column_name: &str,
7826 duration_nanos: u64,
7827 ) -> Result<TtlPolicy> {
7828 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7829 return Err(MongrelError::InvalidArgument(
7830 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7831 ));
7832 }
7833 let column = self
7834 .schema
7835 .columns
7836 .iter()
7837 .find(|column| column.name == column_name)
7838 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7839 if column.ty != TypeId::TimestampNanos {
7840 return Err(MongrelError::Schema(format!(
7841 "TTL column {column_name} must be TimestampNanos, is {:?}",
7842 column.ty
7843 )));
7844 }
7845 Ok(TtlPolicy {
7846 column_id: column.id,
7847 duration_nanos,
7848 })
7849 }
7850
7851 pub(crate) fn apply_ttl_policy_at(
7852 &mut self,
7853 policy: Option<TtlPolicy>,
7854 epoch: Epoch,
7855 ) -> Result<()> {
7856 if let Some(policy) = policy {
7857 let column = self
7858 .schema
7859 .columns
7860 .iter()
7861 .find(|column| column.id == policy.column_id)
7862 .ok_or_else(|| {
7863 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
7864 })?;
7865 if column.ty != TypeId::TimestampNanos
7866 || policy.duration_nanos == 0
7867 || policy.duration_nanos > i64::MAX as u64
7868 {
7869 return Err(MongrelError::Schema("invalid TTL policy".into()));
7870 }
7871 }
7872 self.ttl = policy;
7873 self.agg_cache = Arc::new(HashMap::new());
7874 self.clear_result_cache();
7875 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7876 self.persist_manifest(epoch)
7877 }
7878
7879 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
7880 let Some(policy) = self.ttl else {
7881 return false;
7882 };
7883 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
7884 return false;
7885 };
7886 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
7887 }
7888
7889 pub fn current_epoch(&self) -> Epoch {
7890 self.epoch.visible()
7891 }
7892
7893 pub fn memtable_len(&self) -> usize {
7894 self.memtable.len()
7895 }
7896
7897 pub fn count(&self) -> u64 {
7900 if self.ttl.is_none()
7901 && self.pending_put_cols.is_empty()
7902 && self.pending_delete_rids.is_empty()
7903 && self.pending_rows.is_empty()
7904 && self.pending_dels.is_empty()
7905 && self.pending_truncate.is_none()
7906 {
7907 self.live_count
7908 } else {
7909 self.visible_rows(self.snapshot())
7910 .map(|rows| rows.len() as u64)
7911 .unwrap_or(self.live_count)
7912 }
7913 }
7914
7915 pub fn count_conditions(
7919 &mut self,
7920 conditions: &[crate::query::Condition],
7921 snapshot: Snapshot,
7922 ) -> Result<Option<u64>> {
7923 use crate::query::Condition;
7924 if self.ttl.is_some() {
7925 if conditions.is_empty() {
7926 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
7927 }
7928 let mut sets = Vec::with_capacity(conditions.len());
7929 for condition in conditions {
7930 sets.push(self.resolve_condition(condition, snapshot)?);
7931 }
7932 let survivors = RowIdSet::intersect_many(sets);
7933 let rows = self.visible_rows(snapshot)?;
7934 return Ok(Some(
7935 rows.into_iter()
7936 .filter(|row| survivors.contains(row.row_id.0))
7937 .count() as u64,
7938 ));
7939 }
7940 if conditions.is_empty() {
7941 return Ok(Some(self.count()));
7942 }
7943 let served = |c: &Condition| {
7944 matches!(
7945 c,
7946 Condition::Pk(_)
7947 | Condition::BitmapEq { .. }
7948 | Condition::BitmapIn { .. }
7949 | Condition::BytesPrefix { .. }
7950 | Condition::FmContains { .. }
7951 | Condition::FmContainsAll { .. }
7952 | Condition::Ann { .. }
7953 | Condition::Range { .. }
7954 | Condition::RangeF64 { .. }
7955 | Condition::SparseMatch { .. }
7956 | Condition::MinHashSimilar { .. }
7957 | Condition::IsNull { .. }
7958 | Condition::IsNotNull { .. }
7959 )
7960 };
7961 if !conditions.iter().all(served) {
7962 return Ok(None);
7963 }
7964 self.ensure_indexes_complete()?;
7965 if !self.pending_put_cols.is_empty()
7966 || !self.pending_delete_rids.is_empty()
7967 || !self.pending_rows.is_empty()
7968 || !self.pending_dels.is_empty()
7969 || self.pending_truncate.is_some()
7970 {
7971 let mut sets = Vec::with_capacity(conditions.len());
7972 for condition in conditions {
7973 sets.push(self.resolve_condition(condition, snapshot)?);
7974 }
7975 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7976 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
7977 }
7978 let mut sets = Vec::with_capacity(conditions.len());
7979 for condition in conditions {
7980 sets.push(self.resolve_condition(condition, snapshot)?);
7981 }
7982 let mut rids = RowIdSet::intersect_many(sets);
7983 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7993 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
7994 }
7995 let count = rids.len() as u64;
7996 crate::trace::QueryTrace::record(|t| {
7997 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
7998 t.survivor_count = Some(count as usize);
7999 t.conditions_pushed = conditions.len();
8000 });
8001 Ok(Some(count))
8002 }
8003
8004 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8009 let mut out = Vec::new();
8010 for row in self.memtable.visible_versions(snapshot.epoch) {
8011 if row.deleted {
8012 out.push(row.row_id.0);
8013 }
8014 }
8015 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8016 if row.deleted {
8017 out.push(row.row_id.0);
8018 }
8019 }
8020 out
8021 }
8022
8023 pub fn bulk_load_columns(
8032 &mut self,
8033 user_columns: Vec<(u16, columnar::NativeColumn)>,
8034 ) -> Result<Epoch> {
8035 self.bulk_load_columns_with(user_columns, 3, false, true)
8036 }
8037
8038 pub fn bulk_load_fast(
8045 &mut self,
8046 user_columns: Vec<(u16, columnar::NativeColumn)>,
8047 ) -> Result<Epoch> {
8048 self.bulk_load_columns_with(user_columns, -1, true, false)
8049 }
8050
8051 fn bulk_load_columns_with(
8052 &mut self,
8053 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8054 zstd_level: i32,
8055 force_plain: bool,
8056 lz4: bool,
8057 ) -> Result<Epoch> {
8058 self.ensure_writable()?;
8059 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8060 if n == 0 {
8061 return Ok(self.current_epoch());
8062 }
8063 let epoch = self.commit_new_epoch()?;
8064 let live_before = self.live_count;
8065 self.spill_mutable_run(epoch)?;
8067 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8068 && self.indexes_complete
8069 && self.run_refs.is_empty()
8070 && self.memtable.is_empty()
8071 && self.mutable_run.is_empty();
8072 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8075 self.validate_columns_not_null(&user_columns, n)?;
8076 let winner_idx = self
8077 .bulk_pk_winner_indices(&user_columns, n)
8078 .filter(|idx| idx.len() != n);
8079 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8080 match winner_idx.as_deref() {
8081 Some(idx) => {
8082 let compacted = user_columns
8083 .iter()
8084 .map(|(id, c)| (*id, c.gather(idx)))
8085 .collect();
8086 (compacted, idx.len())
8087 }
8088 None => (user_columns, n),
8089 };
8090 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8091 let first = self.allocator.alloc_range(write_n as u64)?.0;
8092 for rid in first..first + write_n as u64 {
8093 self.reservoir.offer(rid);
8094 }
8095 let run_id = self.alloc_run_id()?;
8096 let path = self.run_path(run_id);
8097 let mut writer =
8098 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8099 if force_plain {
8100 writer = writer.with_plain();
8101 } else if lz4 {
8102 writer = writer.with_lz4();
8105 } else {
8106 writer = writer.with_zstd_level(zstd_level);
8107 }
8108 if let Some(kek) = &self.kek {
8109 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8110 }
8111 let header = match self.create_run_file(run_id)? {
8112 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8113 None => writer.write_native(&path, &write_columns, write_n, first)?,
8114 };
8115 self.run_refs.push(RunRef {
8116 run_id: run_id as u128,
8117 level: 0,
8118 epoch_created: epoch.0,
8119 row_count: header.row_count,
8120 });
8121 self.live_count = self.live_count.saturating_add(write_n as u64);
8122 if eager_index_build {
8123 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8124 self.index_columns_bulk(&write_columns, &row_ids);
8125 self.indexes_complete = true;
8126 self.build_learned_ranges()?;
8127 } else {
8128 self.indexes_complete = false;
8132 }
8133 self.mark_flushed(epoch)?;
8134 self.persist_manifest(epoch)?;
8135 if eager_index_build {
8136 self.checkpoint_indexes(epoch);
8137 }
8138 self.clear_result_cache();
8139 self.data_generation = self.data_generation.wrapping_add(1);
8140 Ok(epoch)
8141 }
8142
8143 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8161 let n = row_ids.len();
8162 if n == 0 {
8163 return;
8164 }
8165 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8166 columns.iter().map(|(id, c)| (*id, c)).collect();
8167 let ty_of: std::collections::HashMap<u16, TypeId> = self
8168 .schema
8169 .columns
8170 .iter()
8171 .map(|c| (c.id, c.ty.clone()))
8172 .collect();
8173 let pk_id = self.schema.primary_key().map(|c| c.id);
8174
8175 for (i, &rid) in row_ids.iter().enumerate() {
8176 let row_id = RowId(rid);
8177 if let Some(pid) = pk_id {
8178 if let Some(col) = by_id.get(&pid) {
8179 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8180 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8181 self.insert_hot_pk(key, row_id);
8182 }
8183 }
8184 }
8185 for idef in &self.schema.indexes {
8186 let Some(col) = by_id.get(&idef.column_id) else {
8187 continue;
8188 };
8189 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8190 match idef.kind {
8191 IndexKind::Bitmap => {
8192 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8193 if let Some(key) =
8194 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8195 {
8196 b.insert(key, row_id);
8197 }
8198 }
8199 }
8200 IndexKind::FmIndex => {
8201 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8202 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8203 f.insert(bytes.to_vec(), row_id);
8204 }
8205 }
8206 }
8207 IndexKind::Sparse => {
8208 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8209 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8210 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8211 s.insert(&terms, row_id);
8212 }
8213 }
8214 }
8215 }
8216 IndexKind::MinHash => {
8217 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8218 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8219 let tokens = crate::index::token_hashes_from_bytes(bytes);
8220 mh.insert(&tokens, row_id);
8221 }
8222 }
8223 }
8224 _ => {}
8225 }
8226 }
8227 }
8228 }
8229
8230 pub fn visible_columns_native(
8235 &self,
8236 snapshot: Snapshot,
8237 projection: Option<&[u16]>,
8238 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8239 self.visible_columns_native_inner(snapshot, projection, None)
8240 }
8241
8242 pub fn visible_columns_native_with_control(
8243 &self,
8244 snapshot: Snapshot,
8245 projection: Option<&[u16]>,
8246 control: &crate::ExecutionControl,
8247 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8248 self.visible_columns_native_inner(snapshot, projection, Some(control))
8249 }
8250
8251 fn visible_columns_native_inner(
8252 &self,
8253 snapshot: Snapshot,
8254 projection: Option<&[u16]>,
8255 control: Option<&crate::ExecutionControl>,
8256 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8257 execution_checkpoint(control, 0)?;
8258 let wanted: Vec<u16> = match projection {
8259 Some(p) => p.to_vec(),
8260 None => self.schema.columns.iter().map(|c| c.id).collect(),
8261 };
8262 if self.ttl.is_none()
8263 && self.memtable.is_empty()
8264 && self.mutable_run.is_empty()
8265 && self.run_refs.len() == 1
8266 {
8267 let rr = self.run_refs[0].clone();
8268 let mut reader = self.open_reader(rr.run_id)?;
8269 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8270 execution_checkpoint(control, 0)?;
8271 let all_visible = idxs.len() == reader.row_count();
8272 if reader.has_mmap() && control.is_none() {
8278 use rayon::prelude::*;
8279 let valid: Vec<u16> = wanted
8282 .iter()
8283 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8284 .copied()
8285 .collect();
8286 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8288 .par_iter()
8289 .filter_map(|cid| {
8290 reader
8291 .column_native_shared(*cid)
8292 .ok()
8293 .map(|col| (*cid, col))
8294 })
8295 .collect();
8296 let cols = decoded
8297 .into_iter()
8298 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8299 .collect();
8300 return Ok(cols);
8301 }
8302 let mut cols = Vec::with_capacity(wanted.len());
8303 for (index, cid) in wanted.iter().enumerate() {
8304 execution_checkpoint(control, index)?;
8305 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8306 Some(c) => c,
8307 None => continue,
8308 };
8309 let col = reader.column_native(cdef.id)?;
8310 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8311 }
8312 return Ok(cols);
8313 }
8314 let vcols = self.visible_columns(snapshot)?;
8315 execution_checkpoint(control, 0)?;
8316 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8317 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8318 .into_iter()
8319 .filter(|(id, _)| want_set.contains(id))
8320 .map(|(id, vals)| {
8321 let ty = self
8322 .schema
8323 .columns
8324 .iter()
8325 .find(|c| c.id == id)
8326 .map(|c| c.ty.clone())
8327 .unwrap_or(TypeId::Bytes);
8328 (id, columnar::values_to_native(ty, &vals))
8329 })
8330 .collect();
8331 Ok(out)
8332 }
8333
8334 pub fn run_count(&self) -> usize {
8335 self.run_refs.len()
8336 }
8337
8338 pub fn memtable_is_empty(&self) -> bool {
8340 self.memtable.is_empty()
8341 }
8342
8343 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8347 self.page_cache.stats()
8348 }
8349
8350 pub fn reset_page_cache_stats(&self) {
8352 self.page_cache.reset_stats();
8353 }
8354
8355 pub fn run_ids(&self) -> Vec<u128> {
8358 self.run_refs.iter().map(|r| r.run_id).collect()
8359 }
8360
8361 pub fn single_run_is_clean(&self) -> bool {
8365 if self.ttl.is_some() || self.run_refs.len() != 1 {
8366 return false;
8367 }
8368 self.open_reader(self.run_refs[0].run_id)
8369 .map(|r| r.is_clean())
8370 .unwrap_or(false)
8371 }
8372
8373 fn resolve_footprint(
8380 &self,
8381 conditions: &[crate::query::Condition],
8382 snapshot: Snapshot,
8383 ) -> roaring::RoaringBitmap {
8384 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8385 return roaring::RoaringBitmap::new();
8386 }
8387 if self.run_refs.is_empty() {
8388 return roaring::RoaringBitmap::new();
8389 }
8390 if self.run_refs.len() == 1 {
8392 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8393 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8394 return rids.to_roaring_lossy();
8395 }
8396 }
8397 }
8398 roaring::RoaringBitmap::new()
8399 }
8400
8401 pub fn query_columns_native_cached(
8412 &mut self,
8413 conditions: &[crate::query::Condition],
8414 projection: Option<&[u16]>,
8415 snapshot: Snapshot,
8416 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8417 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8418 }
8419
8420 pub fn query_columns_native_cached_with_control(
8421 &mut self,
8422 conditions: &[crate::query::Condition],
8423 projection: Option<&[u16]>,
8424 snapshot: Snapshot,
8425 control: &crate::ExecutionControl,
8426 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8427 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8428 }
8429
8430 fn query_columns_native_cached_inner(
8431 &mut self,
8432 conditions: &[crate::query::Condition],
8433 projection: Option<&[u16]>,
8434 snapshot: Snapshot,
8435 control: Option<&crate::ExecutionControl>,
8436 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8437 execution_checkpoint(control, 0)?;
8438 if self.ttl.is_some() {
8441 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8442 }
8443 if conditions.is_empty() {
8444 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8445 }
8446 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8450 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8451 crate::trace::QueryTrace::record(|t| {
8452 t.result_cache_hit = true;
8453 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8454 });
8455 return Ok(Some((*hit).clone()));
8456 }
8457 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8458 execution_checkpoint(control, 0)?;
8459 if let Some(cols) = &res {
8460 let footprint = self.resolve_footprint(conditions, snapshot);
8461 let condition_cols = crate::query::condition_columns(conditions);
8462 execution_checkpoint(control, 0)?;
8463 self.result_cache.lock().insert(
8464 key,
8465 CachedEntry {
8466 data: CachedData::Columns(Arc::new(cols.clone())),
8467 footprint,
8468 condition_cols,
8469 },
8470 );
8471 }
8472 Ok(res)
8473 }
8474
8475 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8480 if self.ttl.is_some() {
8481 return self.query(q);
8482 }
8483 if q.conditions.is_empty() {
8484 return self.query(q);
8485 }
8486 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8487 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8488 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8489 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8490 crate::trace::QueryTrace::record(|t| {
8491 t.result_cache_hit = true;
8492 t.scan_mode = crate::trace::ScanMode::Materialized;
8493 });
8494 return Ok((*hit).clone());
8495 }
8496 let rows = self.query(q)?;
8497 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8498 let condition_cols = crate::query::condition_columns(&q.conditions);
8499 self.result_cache.lock().insert(
8500 key,
8501 CachedEntry {
8502 data: CachedData::Rows(Arc::new(rows.clone())),
8503 footprint,
8504 condition_cols,
8505 },
8506 );
8507 Ok(rows)
8508 }
8509
8510 #[allow(clippy::type_complexity)]
8525 pub fn query_columns_native_traced(
8526 &mut self,
8527 conditions: &[crate::query::Condition],
8528 projection: Option<&[u16]>,
8529 snapshot: Snapshot,
8530 ) -> Result<(
8531 Option<Vec<(u16, columnar::NativeColumn)>>,
8532 crate::trace::QueryTrace,
8533 )> {
8534 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8535 self.query_columns_native(conditions, projection, snapshot)
8536 });
8537 Ok((result?, trace))
8538 }
8539
8540 #[allow(clippy::type_complexity)]
8543 pub fn query_columns_native_cached_traced(
8544 &mut self,
8545 conditions: &[crate::query::Condition],
8546 projection: Option<&[u16]>,
8547 snapshot: Snapshot,
8548 ) -> Result<(
8549 Option<Vec<(u16, columnar::NativeColumn)>>,
8550 crate::trace::QueryTrace,
8551 )> {
8552 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8553 self.query_columns_native_cached(conditions, projection, snapshot)
8554 });
8555 Ok((result?, trace))
8556 }
8557
8558 pub fn native_page_cursor_traced(
8560 &self,
8561 snapshot: Snapshot,
8562 projection: Vec<(u16, TypeId)>,
8563 conditions: &[crate::query::Condition],
8564 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8565 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8566 self.native_page_cursor(snapshot, projection, conditions)
8567 });
8568 Ok((result?, trace))
8569 }
8570
8571 pub fn native_multi_run_cursor_traced(
8573 &self,
8574 snapshot: Snapshot,
8575 projection: Vec<(u16, TypeId)>,
8576 conditions: &[crate::query::Condition],
8577 ) -> Result<(
8578 Option<crate::cursor::MultiRunCursor>,
8579 crate::trace::QueryTrace,
8580 )> {
8581 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8582 self.native_multi_run_cursor(snapshot, projection, conditions)
8583 });
8584 Ok((result?, trace))
8585 }
8586
8587 pub fn count_conditions_traced(
8589 &mut self,
8590 conditions: &[crate::query::Condition],
8591 snapshot: Snapshot,
8592 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8593 let (result, trace) =
8594 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8595 Ok((result?, trace))
8596 }
8597
8598 pub fn query_traced(
8600 &mut self,
8601 q: &crate::query::Query,
8602 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8603 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8604 Ok((result?, trace))
8605 }
8606
8607 pub fn query_columns_native(
8612 &mut self,
8613 conditions: &[crate::query::Condition],
8614 projection: Option<&[u16]>,
8615 snapshot: Snapshot,
8616 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8617 self.query_columns_native_inner(conditions, projection, snapshot, None)
8618 }
8619
8620 pub fn query_columns_native_with_control(
8621 &mut self,
8622 conditions: &[crate::query::Condition],
8623 projection: Option<&[u16]>,
8624 snapshot: Snapshot,
8625 control: &crate::ExecutionControl,
8626 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8627 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8628 }
8629
8630 fn query_columns_native_inner(
8631 &mut self,
8632 conditions: &[crate::query::Condition],
8633 projection: Option<&[u16]>,
8634 snapshot: Snapshot,
8635 control: Option<&crate::ExecutionControl>,
8636 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8637 use crate::query::Condition;
8638 execution_checkpoint(control, 0)?;
8639 if self.ttl.is_some() {
8642 return Ok(None);
8643 }
8644 if conditions.is_empty() {
8645 return Ok(None);
8646 }
8647 self.ensure_indexes_complete()?;
8648
8649 let served = |c: &Condition| {
8654 matches!(
8655 c,
8656 Condition::Pk(_)
8657 | Condition::BitmapEq { .. }
8658 | Condition::BitmapIn { .. }
8659 | Condition::BytesPrefix { .. }
8660 | Condition::FmContains { .. }
8661 | Condition::FmContainsAll { .. }
8662 | Condition::Ann { .. }
8663 | Condition::Range { .. }
8664 | Condition::RangeF64 { .. }
8665 | Condition::SparseMatch { .. }
8666 | Condition::MinHashSimilar { .. }
8667 | Condition::IsNull { .. }
8668 | Condition::IsNotNull { .. }
8669 )
8670 };
8671 if !conditions.iter().all(served) {
8672 return Ok(None);
8673 }
8674 let fast_path =
8675 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8676 crate::trace::QueryTrace::record(|t| {
8677 t.run_count = self.run_refs.len();
8678 t.memtable_rows = self.memtable.len();
8679 t.mutable_run_rows = self.mutable_run.len();
8680 t.conditions_pushed = conditions.len();
8681 t.learned_range_used = conditions.iter().any(|c| match c {
8682 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8683 self.learned_range.contains_key(column_id)
8684 }
8685 _ => false,
8686 });
8687 });
8688 let col_ids: Vec<u16> = projection
8690 .map(|p| p.to_vec())
8691 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8692 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8693 .iter()
8694 .map(|&cid| {
8695 let ty = self
8696 .schema
8697 .columns
8698 .iter()
8699 .find(|c| c.id == cid)
8700 .map(|c| c.ty.clone())
8701 .unwrap_or(TypeId::Bytes);
8702 (cid, ty)
8703 })
8704 .collect();
8705
8706 if fast_path {
8712 let needs_column = conditions.iter().any(|c| match c {
8715 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8716 Condition::RangeF64 { column_id, .. } => {
8717 !self.learned_range.contains_key(column_id)
8718 }
8719 _ => false,
8720 });
8721 let mut reader_opt: Option<RunReader> = if needs_column {
8722 Some(self.open_reader(self.run_refs[0].run_id)?)
8723 } else {
8724 None
8725 };
8726 let mut sets: Vec<RowIdSet> = Vec::new();
8727 for (index, c) in conditions.iter().enumerate() {
8728 execution_checkpoint(control, index)?;
8729 let s = match c {
8730 Condition::Range { column_id, lo, hi }
8731 if !self.learned_range.contains_key(column_id) =>
8732 {
8733 if reader_opt.is_none() {
8734 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8735 }
8736 reader_opt
8737 .as_mut()
8738 .expect("reader opened for range")
8739 .range_row_id_set_i64(*column_id, *lo, *hi)?
8740 }
8741 Condition::RangeF64 {
8742 column_id,
8743 lo,
8744 lo_inclusive,
8745 hi,
8746 hi_inclusive,
8747 } if !self.learned_range.contains_key(column_id) => {
8748 if reader_opt.is_none() {
8749 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8750 }
8751 reader_opt
8752 .as_mut()
8753 .expect("reader opened for range")
8754 .range_row_id_set_f64(
8755 *column_id,
8756 *lo,
8757 *lo_inclusive,
8758 *hi,
8759 *hi_inclusive,
8760 )?
8761 }
8762 _ => self.resolve_condition(c, snapshot)?,
8763 };
8764 sets.push(s);
8765 }
8766 let candidates = RowIdSet::intersect_many(sets);
8767 crate::trace::QueryTrace::record(|t| {
8768 t.survivor_count = Some(candidates.len());
8769 });
8770 if candidates.is_empty() {
8771 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8772 .iter()
8773 .map(|&id| {
8774 (
8775 id,
8776 columnar::null_native(
8777 proj_pairs
8778 .iter()
8779 .find(|(c, _)| c == &id)
8780 .map(|(_, t)| t.clone())
8781 .unwrap_or(TypeId::Bytes),
8782 0,
8783 ),
8784 )
8785 })
8786 .collect();
8787 return Ok(Some(cols));
8788 }
8789 let mut reader = match reader_opt.take() {
8790 Some(r) => r,
8791 None => self.open_reader(self.run_refs[0].run_id)?,
8792 };
8793 let candidate_ids = candidates.into_sorted_vec();
8794 let (positions, fast_rid) = if let Some(positions) =
8795 reader.positions_for_row_ids_fast(&candidate_ids)
8796 {
8797 (positions, true)
8798 } else {
8799 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8800 match col {
8801 columnar::NativeColumn::Int64 { data, .. } => {
8802 let mut p = Vec::with_capacity(candidate_ids.len());
8803 for (index, rid) in candidate_ids.iter().enumerate() {
8804 execution_checkpoint(control, index)?;
8805 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8806 p.push(position);
8807 }
8808 }
8809 p.sort_unstable();
8810 (p, false)
8811 }
8812 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8813 }
8814 };
8815 crate::trace::QueryTrace::record(|t| {
8816 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8817 t.fast_row_id_map = fast_rid;
8818 });
8819 let mut cols = Vec::with_capacity(col_ids.len());
8820 for (index, cid) in col_ids.iter().enumerate() {
8821 execution_checkpoint(control, index)?;
8822 let col = reader.column_native(*cid)?;
8823 cols.push((*cid, col.gather(&positions)));
8824 }
8825 return Ok(Some(cols));
8826 }
8827
8828 if !self.run_refs.is_empty() {
8841 use crate::cursor::{
8842 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8843 };
8844 let remaining: usize;
8845 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
8846 let c = self
8847 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
8848 .expect("single-run cursor should build when run_refs.len() == 1");
8849 remaining = c.remaining_rows();
8850 Box::new(c)
8851 } else {
8852 let c = self
8853 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
8854 .expect("multi-run cursor should build when run_refs.len() >= 1");
8855 remaining = c.remaining_rows();
8856 Box::new(c)
8857 };
8858 crate::trace::QueryTrace::record(|t| {
8859 if t.survivor_count.is_none() {
8860 t.survivor_count = Some(remaining);
8861 }
8862 });
8863 let cols = match control {
8864 Some(control) => {
8865 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
8866 }
8867 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
8868 };
8869 return Ok(Some(cols));
8870 }
8871
8872 crate::trace::QueryTrace::record(|t| {
8877 t.scan_mode = crate::trace::ScanMode::Materialized;
8878 t.row_materialized = true;
8879 });
8880 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8881 for (index, c) in conditions.iter().enumerate() {
8882 execution_checkpoint(control, index)?;
8883 sets.push(self.resolve_condition(c, snapshot)?);
8884 }
8885 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8886 let rows = self.rows_for_rids(&rids, snapshot)?;
8887 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
8888 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
8889 execution_checkpoint(control, index)?;
8890 let vals: Vec<Value> = rows
8891 .iter()
8892 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8893 .collect();
8894 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
8895 }
8896 Ok(Some(cols))
8897 }
8898
8899 pub fn native_page_cursor(
8914 &self,
8915 snapshot: Snapshot,
8916 projection: Vec<(u16, TypeId)>,
8917 conditions: &[crate::query::Condition],
8918 ) -> Result<Option<NativePageCursor>> {
8919 use crate::cursor::build_page_plans;
8920 if self.ttl.is_some() {
8921 return Ok(None);
8922 }
8923 if !conditions.is_empty() && !self.indexes_complete {
8926 return Ok(None);
8927 }
8928 if self.run_refs.len() != 1 {
8929 return Ok(None);
8930 }
8931 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8932 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8933
8934 let overlay_rids: HashSet<u64> = {
8937 let mut s = HashSet::new();
8938 for row in self.memtable.visible_versions(snapshot.epoch) {
8939 s.insert(row.row_id.0);
8940 }
8941 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8942 s.insert(row.row_id.0);
8943 }
8944 s
8945 };
8946
8947 let survivors = if conditions.is_empty() {
8951 None
8952 } else {
8953 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
8954 };
8955
8956 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
8963 survivors.clone()
8964 } else if let Some(s) = &survivors {
8965 let mut run_set = s.clone();
8966 run_set.remove_many(overlay_rids.iter().copied());
8967 Some(run_set)
8968 } else {
8969 Some(RowIdSet::from_unsorted(
8970 rids.iter()
8971 .map(|&r| r as u64)
8972 .filter(|r| !overlay_rids.contains(r))
8973 .collect(),
8974 ))
8975 };
8976
8977 let overlay_rows = if overlay_rids.is_empty() {
8978 Vec::new()
8979 } else {
8980 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8981 self.overlay_visible_rows(snapshot, bound)
8982 };
8983
8984 let plans = if positions.is_empty() {
8986 Vec::new()
8987 } else {
8988 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
8989 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
8990 };
8991
8992 let overlay = if overlay_rows.is_empty() {
8994 None
8995 } else {
8996 let filtered =
8997 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8998 if filtered.is_empty() {
8999 None
9000 } else {
9001 Some(self.materialize_overlay(&filtered, &projection))
9002 }
9003 };
9004
9005 let overlay_row_count = overlay
9006 .as_ref()
9007 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9008 .unwrap_or(0);
9009 crate::trace::QueryTrace::record(|t| {
9010 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9011 t.run_count = self.run_refs.len();
9012 t.memtable_rows = self.memtable.len();
9013 t.mutable_run_rows = self.mutable_run.len();
9014 t.overlay_rows = overlay_row_count;
9015 t.conditions_pushed = conditions.len();
9016 t.pages_decoded = plans
9017 .iter()
9018 .map(|p| p.positions.len())
9019 .sum::<usize>()
9020 .min(1);
9021 });
9022
9023 Ok(Some(NativePageCursor::new_with_overlay(
9024 reader, projection, plans, overlay,
9025 )))
9026 }
9027 #[allow(clippy::type_complexity)]
9037 pub fn native_multi_run_cursor(
9038 &self,
9039 snapshot: Snapshot,
9040 projection: Vec<(u16, TypeId)>,
9041 conditions: &[crate::query::Condition],
9042 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9043 use crate::cursor::{MultiRunCursor, RunStream};
9044 use crate::sorted_run::SYS_ROW_ID;
9045 use std::collections::{BinaryHeap, HashMap, HashSet};
9046 if self.ttl.is_some() {
9047 return Ok(None);
9048 }
9049 if !conditions.is_empty() && !self.indexes_complete {
9052 return Ok(None);
9053 }
9054 if self.run_refs.is_empty() {
9055 return Ok(None);
9056 }
9057
9058 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9060 Vec::with_capacity(self.run_refs.len());
9061 for rr in &self.run_refs {
9062 let mut reader = self.open_reader(rr.run_id)?;
9063 let (rids, eps, del) = reader.system_columns_native()?;
9064 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9065 run_meta.push((reader, rids, eps, del, page_rows));
9066 }
9067
9068 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9072 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9073 for i in 0..rids.len() {
9074 let rid = rids[i] as u64;
9075 let e = eps[i] as u64;
9076 if e > snapshot.epoch.0 {
9077 continue;
9078 }
9079 let is_del = del[i] != 0;
9080 best.entry(rid)
9081 .and_modify(|cur| {
9082 if e > cur.0 {
9083 *cur = (e, run_idx, i, is_del);
9084 }
9085 })
9086 .or_insert((e, run_idx, i, is_del));
9087 }
9088 }
9089
9090 let overlay_rids: HashSet<u64> = {
9092 let mut s = HashSet::new();
9093 for row in self.memtable.visible_versions(snapshot.epoch) {
9094 s.insert(row.row_id.0);
9095 }
9096 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9097 s.insert(row.row_id.0);
9098 }
9099 s
9100 };
9101
9102 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9104 None
9105 } else {
9106 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9107 for c in conditions {
9108 sets.push(self.resolve_condition(c, snapshot)?);
9109 }
9110 Some(RowIdSet::intersect_many(sets))
9111 };
9112
9113 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9117 for (rid, (_, run_idx, pos, deleted)) in &best {
9118 if *deleted {
9119 continue;
9120 }
9121 if overlay_rids.contains(rid) {
9122 continue;
9123 }
9124 if let Some(s) = &survivors {
9125 if !s.contains(*rid) {
9126 continue;
9127 }
9128 }
9129 per_run[*run_idx].push((*rid, *pos));
9130 }
9131 for v in per_run.iter_mut() {
9132 v.sort_unstable_by_key(|&(rid, _)| rid);
9133 }
9134
9135 let mut streams = Vec::with_capacity(run_meta.len());
9137 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9138 let mut total = 0usize;
9139 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9140 let mut starts = Vec::with_capacity(page_rows.len());
9141 let mut acc = 0usize;
9142 for &r in &page_rows {
9143 starts.push(acc);
9144 acc += r;
9145 }
9146 let mut survivors_vec: Vec<(u64, usize, usize)> =
9147 Vec::with_capacity(per_run[run_idx].len());
9148 for &(rid, pos) in &per_run[run_idx] {
9149 let page_seq = match starts.partition_point(|&s| s <= pos) {
9150 0 => continue,
9151 p => p - 1,
9152 };
9153 let within = pos - starts[page_seq];
9154 survivors_vec.push((rid, page_seq, within));
9155 }
9156 total += survivors_vec.len();
9157 if let Some(&(rid, _, _)) = survivors_vec.first() {
9158 heap.push(std::cmp::Reverse((rid, run_idx)));
9159 }
9160 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9161 }
9162
9163 let overlay_rows = if overlay_rids.is_empty() {
9165 Vec::new()
9166 } else {
9167 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9168 self.overlay_visible_rows(snapshot, bound)
9169 };
9170 let overlay = if overlay_rows.is_empty() {
9171 None
9172 } else {
9173 let filtered =
9174 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9175 if filtered.is_empty() {
9176 None
9177 } else {
9178 Some(self.materialize_overlay(&filtered, &projection))
9179 }
9180 };
9181
9182 let overlay_row_count = overlay
9183 .as_ref()
9184 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9185 .unwrap_or(0);
9186 crate::trace::QueryTrace::record(|t| {
9187 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9188 t.run_count = self.run_refs.len();
9189 t.memtable_rows = self.memtable.len();
9190 t.mutable_run_rows = self.mutable_run.len();
9191 t.overlay_rows = overlay_row_count;
9192 t.conditions_pushed = conditions.len();
9193 t.survivor_count = Some(total);
9194 });
9195
9196 Ok(Some(MultiRunCursor::new(
9197 streams, projection, heap, total, overlay,
9198 )))
9199 }
9200
9201 fn overlay_materialization_bound<'a>(
9213 conditions: &[crate::query::Condition],
9214 survivors: &'a Option<RowIdSet>,
9215 ) -> Option<&'a RowIdSet> {
9216 use crate::query::Condition;
9217 let has_range = conditions
9218 .iter()
9219 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9220 if has_range {
9221 None
9222 } else {
9223 survivors.as_ref()
9224 }
9225 }
9226
9227 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9239 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9240 let mut fold = |row: Row| {
9241 if let Some(b) = bound {
9242 if !b.contains(row.row_id.0) {
9243 return;
9244 }
9245 }
9246 best.entry(row.row_id.0)
9247 .and_modify(|(be, br)| {
9248 if row.committed_epoch > *be {
9249 *be = row.committed_epoch;
9250 *br = row.clone();
9251 }
9252 })
9253 .or_insert_with(|| (row.committed_epoch, row));
9254 };
9255 for row in self.memtable.visible_versions(snapshot.epoch) {
9256 fold(row);
9257 }
9258 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9259 fold(row);
9260 }
9261 let mut out: Vec<Row> = best
9262 .into_values()
9263 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9264 .collect();
9265 out.sort_by_key(|r| r.row_id);
9266 out
9267 }
9268
9269 fn filter_overlay_rows(
9277 &self,
9278 rows: Vec<Row>,
9279 conditions: &[crate::query::Condition],
9280 survivors: Option<&RowIdSet>,
9281 snapshot: Snapshot,
9282 ) -> Result<Vec<Row>> {
9283 if conditions.is_empty() {
9284 return Ok(rows);
9285 }
9286 use crate::query::Condition;
9287 let all_index_served = !conditions
9291 .iter()
9292 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9293 if all_index_served {
9294 return Ok(rows
9295 .into_iter()
9296 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9297 .collect());
9298 }
9299 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9302 for c in conditions {
9303 let s = match c {
9304 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9305 _ => self.resolve_condition(c, snapshot)?,
9306 };
9307 per_cond_sets.push(s);
9308 }
9309 Ok(rows
9310 .into_iter()
9311 .filter(|row| {
9312 conditions.iter().enumerate().all(|(i, c)| match c {
9313 Condition::Range { column_id, lo, hi } => {
9314 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9315 }
9316 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9317 match row.columns.get(column_id) {
9318 Some(Value::Float64(v)) => {
9319 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9320 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9321 lo_ok && hi_ok
9322 }
9323 _ => false,
9324 }
9325 }
9326 _ => per_cond_sets[i].contains(row.row_id.0),
9327 })
9328 })
9329 .collect())
9330 }
9331
9332 fn materialize_overlay(
9335 &self,
9336 rows: &[Row],
9337 projection: &[(u16, TypeId)],
9338 ) -> Vec<columnar::NativeColumn> {
9339 if projection.is_empty() {
9340 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9341 }
9342 let mut cols = Vec::with_capacity(projection.len());
9343 for (cid, ty) in projection {
9344 let vals: Vec<Value> = rows
9345 .iter()
9346 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9347 .collect();
9348 cols.push(columnar::values_to_native(ty.clone(), &vals));
9349 }
9350 cols
9351 }
9352
9353 fn resolve_survivor_rids(
9358 &self,
9359 conditions: &[crate::query::Condition],
9360 reader: &mut RunReader,
9361 snapshot: Snapshot,
9362 ) -> Result<RowIdSet> {
9363 use crate::query::Condition;
9364 let mut sets: Vec<RowIdSet> = Vec::new();
9365 for c in conditions {
9366 self.validate_condition(c)?;
9367 let s: RowIdSet = match c {
9368 Condition::Pk(key) => {
9369 let lookup = self
9370 .schema
9371 .primary_key()
9372 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9373 .unwrap_or_else(|| key.clone());
9374 self.hot
9375 .get(&lookup)
9376 .map(|r| RowIdSet::one(r.0))
9377 .unwrap_or_else(RowIdSet::empty)
9378 }
9379 Condition::BitmapEq { column_id, value } => {
9380 let lookup = self.index_lookup_key_bytes(*column_id, value);
9381 self.bitmap
9382 .get(column_id)
9383 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9384 .unwrap_or_else(RowIdSet::empty)
9385 }
9386 Condition::BitmapIn { column_id, values } => {
9387 let bm = self.bitmap.get(column_id);
9388 let mut acc = roaring::RoaringBitmap::new();
9389 if let Some(b) = bm {
9390 for v in values {
9391 let lookup = self.index_lookup_key_bytes(*column_id, v);
9392 acc |= b.get(&lookup);
9393 }
9394 }
9395 RowIdSet::from_roaring(acc)
9396 }
9397 Condition::BytesPrefix { column_id, prefix } => {
9398 if let Some(b) = self.bitmap.get(column_id) {
9399 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9400 let mut acc = roaring::RoaringBitmap::new();
9401 for key in b.keys() {
9402 if key.starts_with(&lookup_prefix) {
9403 acc |= b.get(&key);
9404 }
9405 }
9406 RowIdSet::from_roaring(acc)
9407 } else {
9408 RowIdSet::empty()
9409 }
9410 }
9411 Condition::FmContains { column_id, pattern } => self
9412 .fm
9413 .get(column_id)
9414 .map(|f| {
9415 RowIdSet::from_unsorted(
9416 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9417 )
9418 })
9419 .unwrap_or_else(RowIdSet::empty),
9420 Condition::FmContainsAll {
9421 column_id,
9422 patterns,
9423 } => {
9424 if let Some(f) = self.fm.get(column_id) {
9425 let sets: Vec<RowIdSet> = patterns
9426 .iter()
9427 .map(|pat| {
9428 RowIdSet::from_unsorted(
9429 f.locate(pat).into_iter().map(|r| r.0).collect(),
9430 )
9431 })
9432 .collect();
9433 RowIdSet::intersect_many(sets)
9434 } else {
9435 RowIdSet::empty()
9436 }
9437 }
9438 Condition::Ann {
9439 column_id,
9440 query,
9441 k,
9442 } => RowIdSet::from_unsorted(
9443 self.retrieve_filtered(
9444 &crate::query::Retriever::Ann {
9445 column_id: *column_id,
9446 query: query.clone(),
9447 k: *k,
9448 },
9449 snapshot,
9450 None,
9451 None,
9452 None,
9453 None,
9454 )?
9455 .into_iter()
9456 .map(|hit| hit.row_id.0)
9457 .collect(),
9458 ),
9459 Condition::SparseMatch {
9460 column_id,
9461 query,
9462 k,
9463 } => RowIdSet::from_unsorted(
9464 self.retrieve_filtered(
9465 &crate::query::Retriever::Sparse {
9466 column_id: *column_id,
9467 query: query.clone(),
9468 k: *k,
9469 },
9470 snapshot,
9471 None,
9472 None,
9473 None,
9474 None,
9475 )?
9476 .into_iter()
9477 .map(|hit| hit.row_id.0)
9478 .collect(),
9479 ),
9480 Condition::MinHashSimilar {
9481 column_id,
9482 query,
9483 k,
9484 } => match self.minhash.get(column_id) {
9485 Some(index) => {
9486 let candidates = index.candidate_row_ids(query);
9487 let eligible =
9488 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9489 RowIdSet::from_unsorted(
9490 index
9491 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9492 .into_iter()
9493 .map(|(row_id, _)| row_id.0)
9494 .collect(),
9495 )
9496 }
9497 None => RowIdSet::empty(),
9498 },
9499 Condition::Range { column_id, lo, hi } => {
9500 if let Some(li) = self.learned_range.get(column_id) {
9501 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9502 } else {
9503 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9504 }
9505 }
9506 Condition::RangeF64 {
9507 column_id,
9508 lo,
9509 lo_inclusive,
9510 hi,
9511 hi_inclusive,
9512 } => {
9513 if let Some(li) = self.learned_range.get(column_id) {
9514 RowIdSet::from_unsorted(
9515 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9516 .into_iter()
9517 .collect(),
9518 )
9519 } else {
9520 reader.range_row_id_set_f64(
9521 *column_id,
9522 *lo,
9523 *lo_inclusive,
9524 *hi,
9525 *hi_inclusive,
9526 )?
9527 }
9528 }
9529 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9530 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9531 };
9532 sets.push(s);
9533 }
9534 Ok(RowIdSet::intersect_many(sets))
9535 }
9536
9537 pub fn scan_cursor(
9558 &self,
9559 snapshot: Snapshot,
9560 projection: Vec<(u16, TypeId)>,
9561 conditions: &[crate::query::Condition],
9562 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9563 if self.ttl.is_some() {
9564 return Ok(None);
9565 }
9566 if !conditions.is_empty() && !self.indexes_complete {
9572 return Ok(None);
9573 }
9574 if self.run_refs.len() == 1 {
9575 Ok(self
9576 .native_page_cursor(snapshot, projection, conditions)?
9577 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9578 } else {
9579 Ok(self
9580 .native_multi_run_cursor(snapshot, projection, conditions)?
9581 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9582 }
9583 }
9584
9585 pub fn aggregate_native(
9599 &self,
9600 snapshot: Snapshot,
9601 column: Option<u16>,
9602 conditions: &[crate::query::Condition],
9603 agg: NativeAgg,
9604 ) -> Result<Option<NativeAggResult>> {
9605 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9606 }
9607
9608 pub fn aggregate_native_with_control(
9609 &self,
9610 snapshot: Snapshot,
9611 column: Option<u16>,
9612 conditions: &[crate::query::Condition],
9613 agg: NativeAgg,
9614 control: &crate::ExecutionControl,
9615 ) -> Result<Option<NativeAggResult>> {
9616 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9617 }
9618
9619 fn aggregate_native_inner(
9620 &self,
9621 snapshot: Snapshot,
9622 column: Option<u16>,
9623 conditions: &[crate::query::Condition],
9624 agg: NativeAgg,
9625 control: Option<&crate::ExecutionControl>,
9626 ) -> Result<Option<NativeAggResult>> {
9627 execution_checkpoint(control, 0)?;
9628 if self.ttl.is_some() {
9629 return Ok(None);
9630 }
9631 if self.run_refs.len() == 1 && conditions.is_empty() {
9633 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9634 return Ok(Some(res));
9635 }
9636 }
9637 if matches!(agg, NativeAgg::Count) && column.is_none() {
9641 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
9642 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
9643 }
9644 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9645 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
9646 }
9647 let cid = match column {
9650 Some(c) => c,
9651 None => return Ok(None),
9652 };
9653 let ty = self.column_type(cid);
9654 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
9655 execution_checkpoint(control, 0)?;
9656 return match ty {
9657 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9658 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9659 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9660 }
9661 TypeId::Float64 => {
9662 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9663 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9664 }
9665 _ => Ok(None),
9666 };
9667 }
9668 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9670 execution_checkpoint(control, 0)?;
9671 match ty {
9672 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9673 let mut count = 0u64;
9674 let mut sum = 0i128;
9675 let mut mn = i64::MAX;
9676 let mut mx = i64::MIN;
9677 for row in &rows {
9678 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
9679 count += 1;
9680 sum += i128::from(*v);
9681 mn = mn.min(*v);
9682 mx = mx.max(*v);
9683 }
9684 }
9685 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9686 }
9687 TypeId::Float64 => {
9688 let mut count = 0u64;
9689 let mut sum = 0.0f64;
9690 let mut mn = f64::INFINITY;
9691 let mut mx = f64::NEG_INFINITY;
9692 for row in &rows {
9693 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
9694 count += 1;
9695 sum += *v;
9696 mn = mn.min(*v);
9697 mx = mx.max(*v);
9698 }
9699 }
9700 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9701 }
9702 _ => Ok(None),
9703 }
9704 }
9705
9706 fn visible_rows_filtered(
9708 &self,
9709 snapshot: Snapshot,
9710 conditions: &[crate::query::Condition],
9711 control: Option<&crate::ExecutionControl>,
9712 ) -> Result<Vec<Row>> {
9713 let rows = if let Some(control) = control {
9714 self.visible_rows_controlled(snapshot, control)?
9715 } else {
9716 self.visible_rows(snapshot)?
9717 };
9718 if conditions.is_empty() {
9719 return Ok(rows);
9720 }
9721 Ok(rows
9722 .into_iter()
9723 .filter(|row| {
9724 conditions
9725 .iter()
9726 .all(|cond| condition_matches_row(cond, row, &self.schema))
9727 })
9728 .collect())
9729 }
9730
9731 fn aggregate_from_stats(
9739 &self,
9740 snapshot: Snapshot,
9741 column: Option<u16>,
9742 agg: NativeAgg,
9743 ) -> Result<Option<NativeAggResult>> {
9744 let cid = match (agg, column) {
9745 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9746 _ => return Ok(None), };
9748 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9749 return Ok(None);
9750 };
9751 let Some(cs) = stats.get(&cid) else {
9752 return Ok(None);
9753 };
9754 match agg {
9755 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9757 self.live_count.saturating_sub(cs.null_count),
9758 ))),
9759 NativeAgg::Min | NativeAgg::Max => {
9760 let bound = if agg == NativeAgg::Min {
9761 &cs.min
9762 } else {
9763 &cs.max
9764 };
9765 match bound {
9766 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9767 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9768 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9773 None => Ok(None),
9774 }
9775 }
9776 _ => Ok(None),
9777 }
9778 }
9779
9780 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9789 if self.ttl.is_some() {
9790 return Ok(None);
9791 }
9792 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9793 return Ok(None);
9794 }
9795 self.ensure_indexes_complete()?;
9798 let reader = self.open_reader(self.run_refs[0].run_id)?;
9799 if self.live_count != reader.row_count() as u64 {
9800 return Ok(None);
9801 }
9802 let Some(bm) = self.bitmap.get(&column_id) else {
9803 return Ok(None); };
9805 let mut distinct = bm.value_count() as u64;
9806 if !bm.get(&Value::Null.encode_key()).is_empty() {
9809 distinct = distinct.saturating_sub(1);
9810 }
9811 Ok(Some(distinct))
9812 }
9813
9814 pub fn aggregate_incremental(
9826 &mut self,
9827 cache_key: u64,
9828 conditions: &[crate::query::Condition],
9829 column: Option<u16>,
9830 agg: NativeAgg,
9831 ) -> Result<IncrementalAggResult> {
9832 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9833 }
9834
9835 pub fn aggregate_incremental_with_control(
9836 &mut self,
9837 cache_key: u64,
9838 conditions: &[crate::query::Condition],
9839 column: Option<u16>,
9840 agg: NativeAgg,
9841 control: &crate::ExecutionControl,
9842 ) -> Result<IncrementalAggResult> {
9843 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9844 }
9845
9846 fn aggregate_incremental_inner(
9847 &mut self,
9848 cache_key: u64,
9849 conditions: &[crate::query::Condition],
9850 column: Option<u16>,
9851 agg: NativeAgg,
9852 control: Option<&crate::ExecutionControl>,
9853 ) -> Result<IncrementalAggResult> {
9854 execution_checkpoint(control, 0)?;
9855 let snap = self.snapshot();
9856 let cur_wm = self.allocator.current().0;
9857 let cur_epoch = snap.epoch.0;
9858 let incremental_ok = self.ttl.is_none()
9865 && !self.had_deletes
9866 && self.memtable.is_empty()
9867 && self.mutable_run.is_empty();
9868
9869 if incremental_ok {
9872 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
9873 if cached.epoch == cur_epoch {
9874 return Ok(IncrementalAggResult {
9875 state: cached.state,
9876 incremental: true,
9877 delta_rows: 0,
9878 });
9879 }
9880 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
9881 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
9882 let mut delta_rids = Vec::with_capacity(delta_len);
9883 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
9884 execution_checkpoint(control, index)?;
9885 delta_rids.push(row_id);
9886 }
9887 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
9888 execution_checkpoint(control, 0)?;
9889 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9890 let delta_state = agg_state_from_rows(
9891 &delta_rows,
9892 conditions,
9893 &index_sets,
9894 column,
9895 agg,
9896 &self.schema,
9897 control,
9898 )?;
9899 let merged = cached.state.merge(delta_state);
9900 let delta_n = delta_rids.len() as u64;
9901 Arc::make_mut(&mut self.agg_cache).insert(
9902 cache_key,
9903 CachedAgg {
9904 state: merged.clone(),
9905 watermark: cur_wm,
9906 epoch: cur_epoch,
9907 },
9908 );
9909 return Ok(IncrementalAggResult {
9910 state: merged,
9911 incremental: true,
9912 delta_rows: delta_n,
9913 });
9914 }
9915 }
9916 }
9917
9918 let cursor_ok =
9923 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9924 let state = if cursor_ok && agg != NativeAgg::Avg {
9925 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
9926 Some(result) => {
9927 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
9928 }
9929 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
9930 }
9931 } else {
9932 self.agg_state_full_scan(conditions, column, agg, snap, control)?
9933 };
9934 if incremental_ok {
9936 Arc::make_mut(&mut self.agg_cache).insert(
9937 cache_key,
9938 CachedAgg {
9939 state: state.clone(),
9940 watermark: cur_wm,
9941 epoch: cur_epoch,
9942 },
9943 );
9944 }
9945 Ok(IncrementalAggResult {
9946 state,
9947 incremental: false,
9948 delta_rows: 0,
9949 })
9950 }
9951
9952 fn agg_state_full_scan(
9955 &self,
9956 conditions: &[crate::query::Condition],
9957 column: Option<u16>,
9958 agg: NativeAgg,
9959 snap: Snapshot,
9960 control: Option<&crate::ExecutionControl>,
9961 ) -> Result<AggState> {
9962 execution_checkpoint(control, 0)?;
9963 let rows = self.visible_rows(snap)?;
9964 execution_checkpoint(control, 0)?;
9965 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9966 agg_state_from_rows(
9967 &rows,
9968 conditions,
9969 &index_sets,
9970 column,
9971 agg,
9972 &self.schema,
9973 control,
9974 )
9975 }
9976
9977 fn resolve_index_conditions(
9980 &self,
9981 conditions: &[crate::query::Condition],
9982 snapshot: Snapshot,
9983 ) -> Result<Vec<RowIdSet>> {
9984 use crate::query::Condition;
9985 let mut sets = Vec::new();
9986 for c in conditions {
9987 if matches!(
9988 c,
9989 Condition::Ann { .. }
9990 | Condition::SparseMatch { .. }
9991 | Condition::MinHashSimilar { .. }
9992 ) {
9993 sets.push(self.resolve_condition(c, snapshot)?);
9994 }
9995 }
9996 Ok(sets)
9997 }
9998
9999 fn column_type(&self, cid: u16) -> TypeId {
10000 self.schema
10001 .columns
10002 .iter()
10003 .find(|c| c.id == cid)
10004 .map(|c| c.ty.clone())
10005 .unwrap_or(TypeId::Bytes)
10006 }
10007
10008 pub fn approx_aggregate(
10017 &mut self,
10018 conditions: &[crate::query::Condition],
10019 column: Option<u16>,
10020 agg: ApproxAgg,
10021 z: f64,
10022 ) -> Result<Option<ApproxResult>> {
10023 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10024 }
10025
10026 pub fn approx_aggregate_with_candidate_authorization(
10029 &mut self,
10030 conditions: &[crate::query::Condition],
10031 column: Option<u16>,
10032 agg: ApproxAgg,
10033 z: f64,
10034 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10035 ) -> Result<Option<ApproxResult>> {
10036 use crate::query::Condition;
10037 self.ensure_reservoir_complete()?;
10038 let snapshot = self.snapshot();
10039 let n_pop = self.count();
10040 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10041 if sample_rids.is_empty() {
10042 return Ok(None);
10043 }
10044 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10046 let s = live_sample.len();
10047 if s == 0 {
10048 return Ok(None);
10049 }
10050 let authorized = authorization
10051 .map(|authorization| {
10052 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10053 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10054 })
10055 .transpose()?;
10056
10057 let mut index_sets: Vec<RowIdSet> = Vec::new();
10060 for c in conditions {
10061 if matches!(
10062 c,
10063 Condition::Ann { .. }
10064 | Condition::SparseMatch { .. }
10065 | Condition::MinHashSimilar { .. }
10066 ) {
10067 index_sets.push(self.resolve_condition(c, snapshot)?);
10068 }
10069 }
10070
10071 let cid = match (agg, column) {
10073 (ApproxAgg::Count, _) => None,
10074 (_, Some(c)) => Some(c),
10075 _ => return Ok(None),
10076 };
10077 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10078 for r in &live_sample {
10079 if authorized
10080 .as_ref()
10081 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10082 {
10083 continue;
10084 }
10085 if !conditions
10087 .iter()
10088 .all(|c| condition_matches_row(c, r, &self.schema))
10089 {
10090 continue;
10091 }
10092 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10094 continue;
10095 }
10096 if let Some(cid) = cid {
10097 let mut cells = r
10098 .columns
10099 .get(&cid)
10100 .cloned()
10101 .map(|value| vec![(cid, value)])
10102 .unwrap_or_default();
10103 if let Some(authorization) = authorization {
10104 authorization.security.apply_masks_to_cells(
10105 authorization.table,
10106 &mut cells,
10107 authorization.principal,
10108 );
10109 }
10110 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10111 passing_vals.push(v);
10112 } } else {
10114 passing_vals.push(0.0); }
10116 }
10117 let m = passing_vals.len();
10118
10119 let (point, half) = match agg {
10120 ApproxAgg::Count => {
10121 let p = m as f64 / s as f64;
10123 let point = n_pop as f64 * p;
10124 let var = if s > 1 {
10125 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10126 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10127 } else {
10128 0.0
10129 };
10130 (point, z * var.sqrt())
10131 }
10132 ApproxAgg::Sum => {
10133 let y: Vec<f64> = live_sample
10135 .iter()
10136 .map(|r| {
10137 let passes_row = authorized
10138 .as_ref()
10139 .is_none_or(|authorized| authorized.contains(&r.row_id))
10140 && conditions
10141 .iter()
10142 .all(|c| condition_matches_row(c, r, &self.schema))
10143 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10144 if passes_row {
10145 cid.and_then(|cid| {
10146 let mut cells = r
10147 .columns
10148 .get(&cid)
10149 .cloned()
10150 .map(|value| vec![(cid, value)])
10151 .unwrap_or_default();
10152 if let Some(authorization) = authorization {
10153 authorization.security.apply_masks_to_cells(
10154 authorization.table,
10155 &mut cells,
10156 authorization.principal,
10157 );
10158 }
10159 as_f64(cells.first().map(|(_, value)| value))
10160 })
10161 .unwrap_or(0.0)
10162 } else {
10163 0.0
10164 }
10165 })
10166 .collect();
10167 let mean_y = y.iter().sum::<f64>() / s as f64;
10168 let point = n_pop as f64 * mean_y;
10169 let var = if s > 1 {
10170 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10171 let var_y = ss / (s - 1) as f64;
10172 n_pop as f64 * n_pop as f64 * var_y / s as f64
10173 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10174 } else {
10175 0.0
10176 };
10177 (point, z * var.sqrt())
10178 }
10179 ApproxAgg::Avg => {
10180 if m == 0 {
10181 return Ok(Some(ApproxResult {
10182 point: 0.0,
10183 ci_low: 0.0,
10184 ci_high: 0.0,
10185 n_population: n_pop,
10186 n_sample_live: s,
10187 n_passing: 0,
10188 }));
10189 }
10190 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10191 let half = if m > 1 {
10192 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10193 let sd = (ss / (m - 1) as f64).sqrt();
10194 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10195 z * sd / (m as f64).sqrt() * fpc.sqrt()
10196 } else {
10197 0.0
10198 };
10199 (mean, half)
10200 }
10201 };
10202
10203 Ok(Some(ApproxResult {
10204 point,
10205 ci_low: point - half,
10206 ci_high: point + half,
10207 n_population: n_pop,
10208 n_sample_live: s,
10209 n_passing: m,
10210 }))
10211 }
10212
10213 pub fn exact_column_stats(
10221 &self,
10222 _snapshot: Snapshot,
10223 projection: &[u16],
10224 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10225 if self.ttl.is_some()
10226 || !(self.memtable.is_empty()
10227 && self.mutable_run.is_empty()
10228 && self.run_refs.len() == 1)
10229 {
10230 return Ok(None);
10231 }
10232 let reader = self.open_reader(self.run_refs[0].run_id)?;
10233 if self.live_count != reader.row_count() as u64 {
10234 return Ok(None);
10235 }
10236 let mut out = HashMap::new();
10237 for &cid in projection {
10238 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10239 Some(c) => c,
10240 None => continue,
10241 };
10242 let Some(stats) = reader.column_page_stats(cid) else {
10244 out.insert(
10245 cid,
10246 ColumnStat {
10247 min: None,
10248 max: None,
10249 null_count: self.live_count,
10250 },
10251 );
10252 continue;
10253 };
10254 let stat = match cdef.ty {
10255 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10256 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10257 min: mn.map(Value::Int64),
10258 max: mx.map(Value::Int64),
10259 null_count: n,
10260 })
10261 }
10262 TypeId::Float64 => {
10263 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10264 min: mn.map(Value::Float64),
10265 max: mx.map(Value::Float64),
10266 null_count: n,
10267 })
10268 }
10269 _ => None,
10270 };
10271 if let Some(s) = stat {
10272 out.insert(cid, s);
10273 }
10274 }
10275 Ok(Some(out))
10276 }
10277
10278 pub fn dir(&self) -> &Path {
10279 &self.dir
10280 }
10281
10282 pub fn schema(&self) -> &Schema {
10283 &self.schema
10284 }
10285
10286 pub(crate) fn set_catalog_name(&mut self, name: String) {
10287 self.name = name;
10288 }
10289
10290 pub(crate) fn prepare_alter_column(
10291 &mut self,
10292 column_name: &str,
10293 change: &AlterColumn,
10294 ) -> Result<(ColumnDef, Option<Schema>)> {
10295 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10296 return Err(MongrelError::InvalidArgument(
10297 "ALTER COLUMN requires committing staged writes first".into(),
10298 ));
10299 }
10300 let old = self
10301 .schema
10302 .columns
10303 .iter()
10304 .find(|c| c.name == column_name)
10305 .cloned()
10306 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10307 let mut next = old.clone();
10308
10309 if let Some(name) = &change.name {
10310 let trimmed = name.trim();
10311 if trimmed.is_empty() {
10312 return Err(MongrelError::InvalidArgument(
10313 "ALTER COLUMN name must not be empty".into(),
10314 ));
10315 }
10316 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10317 return Err(MongrelError::Schema(format!(
10318 "column {trimmed} already exists"
10319 )));
10320 }
10321 next.name = trimmed.to_string();
10322 }
10323
10324 if let Some(ty) = &change.ty {
10325 next.ty = ty.clone();
10326 }
10327 if let Some(flags) = change.flags {
10328 validate_alter_column_flags(old.flags, flags)?;
10329 next.flags = flags;
10330 }
10331
10332 if let Some(default_change) = &change.default_value {
10333 next.default_value = default_change.clone();
10334 }
10335 if let Some(source_change) = &change.embedding_source {
10336 next.embedding_source = source_change.clone();
10337 }
10338
10339 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10340 if old.flags.contains(ColumnFlags::NULLABLE)
10341 && !next.flags.contains(ColumnFlags::NULLABLE)
10342 && self.column_has_nulls(old.id)?
10343 {
10344 return Err(MongrelError::InvalidArgument(format!(
10345 "column '{}' contains NULL values",
10346 old.name
10347 )));
10348 }
10349 if next == old {
10350 return Ok((next, None));
10351 }
10352 let mut schema = self.schema.clone();
10353 let index = schema
10354 .columns
10355 .iter()
10356 .position(|column| column.id == next.id)
10357 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10358 schema.columns[index] = next.clone();
10359 schema.schema_id = schema
10360 .schema_id
10361 .checked_add(1)
10362 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10363 schema.validate_auto_increment()?;
10364 schema.validate_defaults()?;
10365 Ok((next, Some(schema)))
10366 }
10367
10368 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10369 self.schema = schema;
10370 self.auto_inc = resolve_auto_inc(&self.schema);
10371 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10372 self.clear_result_cache();
10373 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10374 }
10375
10376 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10377 checkpoint_current_schema(self)
10378 }
10379
10380 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10381 self.ensure_writable()?;
10382 let previous_schema = self.schema.clone();
10383 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10384 if let Some(schema) = schema {
10385 self.apply_altered_schema_prepared(schema);
10386 self.checkpoint_standalone_schema_change(previous_schema)?;
10387 }
10388 Ok(column)
10389 }
10390
10391 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10392 if self.live_count == 0 {
10393 return Ok(false);
10394 }
10395 let snap = self.snapshot();
10396 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10397 Ok(columns
10398 .first()
10399 .map(|(_, col)| col.null_count(col.len()) != 0)
10400 .unwrap_or(true))
10401 }
10402
10403 fn has_stored_versions(&self) -> bool {
10404 !self.memtable.is_empty()
10405 || !self.mutable_run.is_empty()
10406 || self.run_refs.iter().any(|r| r.row_count > 0)
10407 || !self.retiring.is_empty()
10408 }
10409
10410 pub fn add_column(
10415 &mut self,
10416 name: &str,
10417 ty: TypeId,
10418 flags: ColumnFlags,
10419 default_value: Option<crate::schema::DefaultExpr>,
10420 ) -> Result<u16> {
10421 self.add_column_with_id(name, ty, flags, default_value, None)
10422 }
10423
10424 pub fn add_column_with_id(
10425 &mut self,
10426 name: &str,
10427 ty: TypeId,
10428 flags: ColumnFlags,
10429 default_value: Option<crate::schema::DefaultExpr>,
10430 requested_id: Option<u16>,
10431 ) -> Result<u16> {
10432 self.ensure_writable()?;
10433 if self.schema.columns.iter().any(|c| c.name == name) {
10434 return Err(MongrelError::Schema(format!(
10435 "column {name} already exists"
10436 )));
10437 }
10438 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10439 if self.schema.columns.iter().any(|c| c.id == id) {
10440 return Err(MongrelError::Schema(format!(
10441 "column id {id} already exists"
10442 )));
10443 }
10444 id
10445 } else {
10446 self.schema
10447 .columns
10448 .iter()
10449 .map(|c| c.id)
10450 .max()
10451 .unwrap_or(0)
10452 .checked_add(1)
10453 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10454 };
10455 let previous_schema = self.schema.clone();
10456 let mut next_schema = previous_schema.clone();
10457 next_schema.columns.push(ColumnDef {
10458 id,
10459 name: name.to_string(),
10460 ty,
10461 flags,
10462 default_value,
10463 embedding_source: None,
10464 });
10465 next_schema.schema_id = next_schema
10466 .schema_id
10467 .checked_add(1)
10468 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10469 next_schema.validate_auto_increment()?;
10470 next_schema.validate_defaults()?;
10471 self.apply_altered_schema_prepared(next_schema);
10472 self.checkpoint_standalone_schema_change(previous_schema)?;
10473 Ok(id)
10474 }
10475
10476 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10485 self.ensure_writable()?;
10486 let cid = self
10487 .schema
10488 .columns
10489 .iter()
10490 .find(|c| c.name == column_name)
10491 .map(|c| c.id)
10492 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10493 let ty = self
10494 .schema
10495 .columns
10496 .iter()
10497 .find(|c| c.id == cid)
10498 .map(|c| c.ty.clone())
10499 .unwrap_or(TypeId::Int64);
10500 if !matches!(
10501 ty,
10502 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10503 ) {
10504 return Err(MongrelError::Schema(format!(
10505 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10506 )));
10507 }
10508 if self
10509 .schema
10510 .indexes
10511 .iter()
10512 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10513 {
10514 return Ok(()); }
10516 let previous_schema = self.schema.clone();
10517 let previous_learned_range = Arc::clone(&self.learned_range);
10518 let mut next_schema = previous_schema.clone();
10519 next_schema.indexes.push(IndexDef {
10520 name: format!("{}_learned_range", column_name),
10521 column_id: cid,
10522 kind: IndexKind::LearnedRange,
10523 predicate: None,
10524 options: Default::default(),
10525 });
10526 next_schema.schema_id = next_schema
10527 .schema_id
10528 .checked_add(1)
10529 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10530 self.apply_altered_schema_prepared(next_schema);
10531 if let Err(error) = self.build_learned_ranges() {
10532 self.apply_altered_schema_prepared(previous_schema);
10533 self.learned_range = previous_learned_range;
10534 return Err(error);
10535 }
10536 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10537 if !matches!(
10538 &error,
10539 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10540 ) {
10541 self.learned_range = previous_learned_range;
10542 }
10543 return Err(error);
10544 }
10545 Ok(())
10546 }
10547
10548 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10549 let mut schema_published = false;
10550 let schema_result = match self._root_guard.as_deref() {
10551 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10552 schema_published = true;
10553 }),
10554 None => write_schema_with_after(&self.dir, &self.schema, || {
10555 schema_published = true;
10556 }),
10557 };
10558 if schema_result.is_err() && !schema_published {
10559 self.apply_altered_schema_prepared(previous_schema);
10560 return schema_result;
10561 }
10562
10563 let manifest_result = self.persist_manifest(self.current_epoch());
10564 match (schema_result, manifest_result) {
10565 (_, Ok(())) => Ok(()),
10566 (Ok(()), Err(error)) => {
10567 self.poison_after_maintenance_publish_failure();
10568 Err(MongrelError::DurableCommit {
10569 epoch: self.current_epoch().0,
10570 message: format!(
10571 "schema is durable but matching manifest publication failed: {error}"
10572 ),
10573 })
10574 }
10575 (Err(schema_error), Err(manifest_error)) => {
10576 self.poison_after_maintenance_publish_failure();
10577 Err(MongrelError::CommitOutcomeUnknown {
10578 epoch: self.current_epoch().0,
10579 message: format!(
10580 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10581 ),
10582 })
10583 }
10584 }
10585 }
10586
10587 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10590 self.sync_byte_threshold = threshold;
10591 if let WalSink::Private(w) = &mut self.wal {
10592 w.set_sync_byte_threshold(threshold);
10593 }
10594 }
10595
10596 pub fn page_cache_flush(&self) {
10600 self.page_cache.flush_to_disk();
10601 }
10602
10603 pub fn page_cache_len(&self) -> usize {
10605 self.page_cache.len()
10606 }
10607
10608 pub fn decoded_cache_len(&self) -> usize {
10611 self.decoded_cache.len()
10612 }
10613
10614 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10617 self.memtable.drain_sorted()
10618 }
10619
10620 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10621 self.runs_dir().join(format!("r-{run_id}.sr"))
10622 }
10623
10624 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10625 match self.runs_root.as_deref() {
10626 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10627 None => Ok(None),
10628 }
10629 }
10630
10631 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10632 match self.runs_root.as_deref() {
10633 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10634 None => Ok(None),
10635 }
10636 }
10637
10638 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10639 match self.runs_root.as_deref() {
10640 Some(root) => match root.remove_file(name) {
10641 Ok(()) => Ok(()),
10642 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10643 Err(error) => Err(error.into()),
10644 },
10645 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10646 Ok(()) => Ok(()),
10647 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10648 Err(error) => Err(error.into()),
10649 },
10650 }
10651 }
10652
10653 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10654 match self.runs_root.as_deref() {
10655 Some(root) => root
10656 .rename_file_new(source, destination)
10657 .map_err(Into::into),
10658 None => crate::durable_file::rename(
10659 &self.runs_dir().join(source),
10660 &self.runs_dir().join(destination),
10661 )
10662 .map_err(Into::into),
10663 }
10664 }
10665
10666 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10667 self.run_refs.iter().map(|run| run.run_id)
10668 }
10669
10670 pub(crate) fn table_dir(&self) -> &Path {
10671 &self.dir
10672 }
10673
10674 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10675 &self.schema
10676 }
10677
10678 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10679 let id = self.next_run_id;
10680 self.next_run_id = self
10681 .next_run_id
10682 .checked_add(1)
10683 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10684 Ok(id)
10685 }
10686
10687 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10688 self.run_refs.push(run_ref);
10689 }
10690
10691 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10701 self.retiring.push(crate::manifest::RetiredRun {
10702 run_id,
10703 retire_epoch,
10704 });
10705 }
10706
10707 pub(crate) fn reap_retiring(
10711 &mut self,
10712 min_active: Epoch,
10713 backup_pinned: &std::collections::HashSet<u128>,
10714 ) -> Result<usize> {
10715 if self.retiring.is_empty() {
10716 return Ok(0);
10717 }
10718 let mut reaped = 0;
10719 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
10720 for r in std::mem::take(&mut self.retiring) {
10726 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
10727 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
10728 reaped += 1;
10729 } else {
10730 kept.push(r);
10731 }
10732 }
10733 self.retiring = kept;
10734 if reaped > 0 {
10735 self.persist_manifest(self.current_epoch())?;
10736 }
10737 Ok(reaped)
10738 }
10739
10740 pub(crate) fn has_reapable_retiring(
10741 &self,
10742 min_active: Epoch,
10743 backup_pinned: &std::collections::HashSet<u128>,
10744 ) -> bool {
10745 self.retiring
10746 .iter()
10747 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
10748 }
10749
10750 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
10751 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
10752 return false;
10753 }
10754 self.live_count = self.live_count.saturating_add(run_ref.row_count);
10755 self.run_refs.push(run_ref);
10756 self.indexes_complete = false;
10757 true
10758 }
10759
10760 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
10761 self.kek.as_ref()
10762 }
10763
10764 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
10765 let mut reader = match self.runs_root.as_deref() {
10766 Some(root) => RunReader::open_file_with_cache(
10767 root.open_regular(format!("r-{run_id}.sr"))?,
10768 self.schema.clone(),
10769 self.kek.clone(),
10770 Some(self.page_cache.clone()),
10771 Some(self.decoded_cache.clone()),
10772 self.table_id,
10773 Some(&self.verified_runs),
10774 None,
10775 )?,
10776 None => RunReader::open_with_cache(
10777 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
10778 self.schema.clone(),
10779 self.kek.clone(),
10780 Some(self.page_cache.clone()),
10781 Some(self.decoded_cache.clone()),
10782 self.table_id,
10783 Some(&self.verified_runs),
10784 )?,
10785 };
10786 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
10790 reader.set_uniform_epoch(Epoch(rr.epoch_created));
10791 }
10792 Ok(reader)
10793 }
10794
10795 pub(crate) fn run_refs(&self) -> &[RunRef] {
10796 &self.run_refs
10797 }
10798
10799 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10800 self.retiring.iter().map(|run| run.run_id)
10801 }
10802
10803 pub(crate) fn runs_dir(&self) -> PathBuf {
10804 self.runs_root
10805 .as_deref()
10806 .and_then(|root| root.io_path().ok())
10807 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
10808 }
10809
10810 pub(crate) fn wal_dir(&self) -> PathBuf {
10811 self.dir.join(WAL_DIR)
10812 }
10813
10814 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
10815 self.run_refs = refs;
10816 }
10817
10818 pub(crate) fn compaction_zstd_level(&self) -> i32 {
10819 self.compaction_zstd_level
10820 }
10821
10822 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
10823 self.kek.clone()
10824 }
10825
10826 #[cfg(feature = "encryption")]
10830 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10831 self.kek.as_ref().map(|k| k.derive_idx_key())
10832 }
10833
10834 #[cfg(not(feature = "encryption"))]
10835 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10836 None
10837 }
10838
10839 #[cfg(feature = "encryption")]
10843 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10844 self.kek.as_ref().map(|k| *k.derive_meta_key())
10845 }
10846
10847 #[cfg(not(feature = "encryption"))]
10848 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10849 None
10850 }
10851
10852 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
10855 self.column_keys
10856 .iter()
10857 .map(|(&id, &(_, scheme))| (id, scheme))
10858 .collect()
10859 }
10860
10861 #[cfg(feature = "encryption")]
10866 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
10867 self.tokenize_value_enc(column_id, v)
10868 }
10869
10870 #[cfg(feature = "encryption")]
10871 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
10872 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10873 let (key, scheme) = self.column_keys.get(&column_id)?;
10874 let token: Vec<u8> = match (*scheme, v) {
10875 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
10876 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
10877 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
10878 _ => hmac_token(key, &v.encode_key()).to_vec(),
10879 };
10880 Some(Value::Bytes(token))
10881 }
10882
10883 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
10885 self.index_lookup_key_bytes(column_id, &v.encode_key())
10886 }
10887
10888 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
10891 #[cfg(feature = "encryption")]
10892 {
10893 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
10894 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
10895 if *scheme == SCHEME_HMAC_EQ {
10896 return hmac_token(key, encoded).to_vec();
10897 }
10898 }
10899 }
10900 let _ = column_id;
10901 encoded.to_vec()
10902 }
10903}
10904
10905fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
10906 let columnar::NativeColumn::Int64 { data, validity } = col else {
10907 return false;
10908 };
10909 if data.len() < n || !columnar::all_non_null(validity, n) {
10910 return false;
10911 }
10912 data.iter()
10913 .take(n)
10914 .zip(data.iter().skip(1))
10915 .all(|(a, b)| a < b)
10916}
10917
10918#[derive(Debug, Clone)]
10922pub struct ColumnStat {
10923 pub min: Option<Value>,
10924 pub max: Option<Value>,
10925 pub null_count: u64,
10926}
10927
10928#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10930pub enum NativeAgg {
10931 Count,
10932 Sum,
10933 Min,
10934 Max,
10935 Avg,
10936}
10937
10938#[derive(Debug, Clone, PartialEq)]
10940pub enum NativeAggResult {
10941 Count(u64),
10942 Int(i64),
10943 Float(f64),
10944 Null,
10946}
10947
10948#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10950pub enum ApproxAgg {
10951 Count,
10952 Sum,
10953 Avg,
10954}
10955
10956#[derive(Debug, Clone)]
10960pub struct ApproxResult {
10961 pub point: f64,
10963 pub ci_low: f64,
10965 pub ci_high: f64,
10967 pub n_population: u64,
10969 pub n_sample_live: usize,
10971 pub n_passing: usize,
10973}
10974
10975#[derive(Debug, Clone, PartialEq)]
10980pub enum AggState {
10981 Count(u64),
10983 SumI {
10985 sum: i128,
10986 count: u64,
10987 },
10988 SumF {
10990 sum: f64,
10991 count: u64,
10992 },
10993 AvgI {
10995 sum: i128,
10996 count: u64,
10997 },
10998 AvgF {
11000 sum: f64,
11001 count: u64,
11002 },
11003 MinI(i64),
11005 MaxI(i64),
11006 MinF(f64),
11008 MaxF(f64),
11009 Empty,
11011}
11012
11013impl AggState {
11014 pub fn merge(self, other: AggState) -> AggState {
11016 use AggState::*;
11017 match (self, other) {
11018 (Empty, x) | (x, Empty) => x,
11019 (Count(a), Count(b)) => Count(a + b),
11020 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11021 sum: sa + sb,
11022 count: ca + cb,
11023 },
11024 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11025 sum: sa + sb,
11026 count: ca + cb,
11027 },
11028 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11029 sum: sa + sb,
11030 count: ca + cb,
11031 },
11032 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11033 sum: sa + sb,
11034 count: ca + cb,
11035 },
11036 (MinI(a), MinI(b)) => MinI(a.min(b)),
11037 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11038 (MinF(a), MinF(b)) => MinF(a.min(b)),
11039 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11040 _ => Empty, }
11042 }
11043
11044 pub fn point(&self) -> Option<f64> {
11046 match self {
11047 AggState::Count(n) => Some(*n as f64),
11048 AggState::SumI { sum, .. } => Some(*sum as f64),
11049 AggState::SumF { sum, .. } => Some(*sum),
11050 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11051 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11052 AggState::MinI(n) => Some(*n as f64),
11053 AggState::MaxI(n) => Some(*n as f64),
11054 AggState::MinF(n) => Some(*n),
11055 AggState::MaxF(n) => Some(*n),
11056 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11057 }
11058 }
11059
11060 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11064 let is_float = matches!(ty, Some(TypeId::Float64));
11065 match (agg, result) {
11066 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11067 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11068 sum: x as i128,
11069 count: 1, },
11071 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11072 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11073 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11074 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11075 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11076 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11077 (NativeAgg::Count, _) => AggState::Empty,
11078 (_, NativeAggResult::Null) => AggState::Empty,
11079 _ => {
11080 let _ = is_float;
11081 AggState::Empty
11082 }
11083 }
11084 }
11085}
11086
11087#[derive(Debug, Clone)]
11090pub struct CachedAgg {
11091 pub state: AggState,
11092 pub watermark: u64,
11093 pub epoch: u64,
11094}
11095
11096#[derive(Debug, Clone)]
11098pub struct IncrementalAggResult {
11099 pub state: AggState,
11101 pub incremental: bool,
11104 pub delta_rows: u64,
11106}
11107
11108fn agg_state_from_rows(
11112 rows: &[Row],
11113 conditions: &[crate::query::Condition],
11114 index_sets: &[RowIdSet],
11115 column: Option<u16>,
11116 agg: NativeAgg,
11117 schema: &Schema,
11118 control: Option<&crate::ExecutionControl>,
11119) -> Result<AggState> {
11120 let mut count: u64 = 0;
11121 let mut sum_i: i128 = 0;
11122 let mut sum_f: f64 = 0.0;
11123 let mut mn_i: i64 = i64::MAX;
11124 let mut mx_i: i64 = i64::MIN;
11125 let mut mn_f: f64 = f64::INFINITY;
11126 let mut mx_f: f64 = f64::NEG_INFINITY;
11127 let mut saw_int = false;
11128 let mut saw_float = false;
11129 for (index, r) in rows.iter().enumerate() {
11130 execution_checkpoint(control, index)?;
11131 if !conditions
11132 .iter()
11133 .all(|c| condition_matches_row(c, r, schema))
11134 {
11135 continue;
11136 }
11137 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11138 continue;
11139 }
11140 match agg {
11141 NativeAgg::Count => match column {
11142 None => count += 1,
11144 Some(cid) => match r.columns.get(&cid) {
11147 None | Some(Value::Null) => {}
11148 Some(_) => count += 1,
11149 },
11150 },
11151 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11152 Some(Value::Int64(n)) => {
11153 count += 1;
11154 sum_i += *n as i128;
11155 mn_i = mn_i.min(*n);
11156 mx_i = mx_i.max(*n);
11157 saw_int = true;
11158 }
11159 Some(Value::Float64(f)) => {
11160 count += 1;
11161 sum_f += f;
11162 mn_f = mn_f.min(*f);
11163 mx_f = mx_f.max(*f);
11164 saw_float = true;
11165 }
11166 _ => {}
11167 },
11168 }
11169 }
11170 Ok(match agg {
11171 NativeAgg::Count => {
11172 if count == 0 {
11173 AggState::Empty
11174 } else {
11175 AggState::Count(count)
11176 }
11177 }
11178 NativeAgg::Sum => {
11179 if count == 0 {
11180 AggState::Empty
11181 } else if saw_int {
11182 AggState::SumI { sum: sum_i, count }
11183 } else {
11184 AggState::SumF { sum: sum_f, count }
11185 }
11186 }
11187 NativeAgg::Avg => {
11188 if count == 0 {
11189 AggState::Empty
11190 } else if saw_int {
11191 AggState::AvgI { sum: sum_i, count }
11192 } else {
11193 AggState::AvgF { sum: sum_f, count }
11194 }
11195 }
11196 NativeAgg::Min => {
11197 if !saw_int && !saw_float {
11198 AggState::Empty
11199 } else if saw_int {
11200 AggState::MinI(mn_i)
11201 } else {
11202 AggState::MinF(mn_f)
11203 }
11204 }
11205 NativeAgg::Max => {
11206 if !saw_int && !saw_float {
11207 AggState::Empty
11208 } else if saw_int {
11209 AggState::MaxI(mx_i)
11210 } else {
11211 AggState::MaxF(mx_f)
11212 }
11213 }
11214 })
11215}
11216
11217fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11221 use crate::query::Condition;
11222 match c {
11223 Condition::Pk(key) => match schema.primary_key() {
11224 Some(pk) => row
11225 .columns
11226 .get(&pk.id)
11227 .map(|v| v.encode_key() == *key)
11228 .unwrap_or(false),
11229 None => false,
11230 },
11231 Condition::BitmapEq { column_id, value } => row
11232 .columns
11233 .get(column_id)
11234 .map(|v| v.encode_key() == *value)
11235 .unwrap_or(false),
11236 Condition::BitmapIn { column_id, values } => {
11237 let key = row.columns.get(column_id).map(|v| v.encode_key());
11238 match key {
11239 Some(k) => values.contains(&k),
11240 None => false,
11241 }
11242 }
11243 Condition::BytesPrefix { column_id, prefix } => row
11244 .columns
11245 .get(column_id)
11246 .map(|v| v.encode_key().starts_with(prefix))
11247 .unwrap_or(false),
11248 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11249 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11250 _ => false,
11251 },
11252 Condition::RangeF64 {
11253 column_id,
11254 lo,
11255 lo_inclusive,
11256 hi,
11257 hi_inclusive,
11258 } => match row.columns.get(column_id) {
11259 Some(Value::Float64(n)) => {
11260 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11261 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11262 lo_ok && hi_ok
11263 }
11264 _ => false,
11265 },
11266 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11267 Some(Value::Bytes(b)) => {
11268 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11269 }
11270 _ => false,
11271 },
11272 Condition::FmContainsAll {
11273 column_id,
11274 patterns,
11275 } => match row.columns.get(column_id) {
11276 Some(Value::Bytes(b)) => patterns
11277 .iter()
11278 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11279 _ => false,
11280 },
11281 Condition::Ann { .. }
11282 | Condition::SparseMatch { .. }
11283 | Condition::MinHashSimilar { .. } => true,
11284 Condition::IsNull { column_id } => {
11285 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11286 }
11287 Condition::IsNotNull { column_id } => {
11288 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11289 }
11290 }
11291}
11292
11293fn as_f64(v: Option<&Value>) -> Option<f64> {
11295 match v {
11296 Some(Value::Int64(n)) => Some(*n as f64),
11297 Some(Value::Float64(f)) => Some(*f),
11298 _ => None,
11299 }
11300}
11301
11302fn accumulate_int(
11306 cursor: &mut dyn crate::cursor::Cursor,
11307 control: Option<&crate::ExecutionControl>,
11308) -> Result<(u64, i128, i64, i64)> {
11309 let mut count: u64 = 0;
11310 let mut sum: i128 = 0;
11311 let mut mn: i64 = i64::MAX;
11312 let mut mx: i64 = i64::MIN;
11313 while let Some(cols) = cursor.next_batch()? {
11314 execution_checkpoint(control, 0)?;
11315 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11316 if crate::columnar::all_non_null(validity, data.len()) {
11317 count += data.len() as u64;
11319 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11320 execution_checkpoint(control, chunk_index * 1024)?;
11321 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11322 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11323 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11324 }
11325 } else {
11326 for (i, &v) in data.iter().enumerate() {
11327 execution_checkpoint(control, i)?;
11328 if crate::columnar::validity_bit(validity, i) {
11329 count += 1;
11330 sum += v as i128;
11331 mn = mn.min(v);
11332 mx = mx.max(v);
11333 }
11334 }
11335 }
11336 }
11337 }
11338 Ok((count, sum, mn, mx))
11339}
11340
11341fn accumulate_float(
11343 cursor: &mut dyn crate::cursor::Cursor,
11344 control: Option<&crate::ExecutionControl>,
11345) -> Result<(u64, f64, f64, f64)> {
11346 let mut count: u64 = 0;
11347 let mut sum: f64 = 0.0;
11348 let mut mn: f64 = f64::INFINITY;
11349 let mut mx: f64 = f64::NEG_INFINITY;
11350 while let Some(cols) = cursor.next_batch()? {
11351 execution_checkpoint(control, 0)?;
11352 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11353 if crate::columnar::all_non_null(validity, data.len()) {
11354 count += data.len() as u64;
11355 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11356 execution_checkpoint(control, chunk_index * 1024)?;
11357 sum += chunk.iter().sum::<f64>();
11358 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11359 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11360 }
11361 } else {
11362 for (i, &v) in data.iter().enumerate() {
11363 execution_checkpoint(control, i)?;
11364 if crate::columnar::validity_bit(validity, i) {
11365 count += 1;
11366 sum += v;
11367 mn = mn.min(v);
11368 mx = mx.max(v);
11369 }
11370 }
11371 }
11372 }
11373 }
11374 Ok((count, sum, mn, mx))
11375}
11376
11377#[inline]
11378fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11379 if index.is_multiple_of(256) {
11380 control
11381 .map(crate::ExecutionControl::checkpoint)
11382 .transpose()?;
11383 }
11384 Ok(())
11385}
11386
11387fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11388 if count == 0 && !matches!(agg, NativeAgg::Count) {
11389 return NativeAggResult::Null;
11390 }
11391 match agg {
11392 NativeAgg::Count => NativeAggResult::Count(count),
11393 NativeAgg::Sum => match sum.try_into() {
11396 Ok(v) => NativeAggResult::Int(v),
11397 Err(_) => NativeAggResult::Null,
11398 },
11399 NativeAgg::Min => NativeAggResult::Int(mn),
11400 NativeAgg::Max => NativeAggResult::Int(mx),
11401 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11402 }
11403}
11404
11405fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11406 if count == 0 && !matches!(agg, NativeAgg::Count) {
11407 return NativeAggResult::Null;
11408 }
11409 match agg {
11410 NativeAgg::Count => NativeAggResult::Count(count),
11411 NativeAgg::Sum => NativeAggResult::Float(sum),
11412 NativeAgg::Min => NativeAggResult::Float(mn),
11413 NativeAgg::Max => NativeAggResult::Float(mx),
11414 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11415 }
11416}
11417
11418fn agg_int(
11421 stats: &[crate::page::PageStat],
11422 decode: fn(Option<&[u8]>) -> Option<i64>,
11423) -> Option<(Option<i64>, Option<i64>, u64)> {
11424 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11425 let mut any = false;
11426 for s in stats {
11427 if let Some(v) = decode(s.min.as_deref()) {
11428 mn = mn.min(v);
11429 any = true;
11430 }
11431 if let Some(v) = decode(s.max.as_deref()) {
11432 mx = mx.max(v);
11433 any = true;
11434 }
11435 nulls += s.null_count;
11436 }
11437 any.then_some((Some(mn), Some(mx), nulls))
11438}
11439
11440fn agg_float(
11442 stats: &[crate::page::PageStat],
11443 decode: fn(Option<&[u8]>) -> Option<f64>,
11444) -> Option<(Option<f64>, Option<f64>, u64)> {
11445 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11446 let mut any = false;
11447 for s in stats {
11448 if let Some(v) = decode(s.min.as_deref()) {
11449 mn = mn.min(v);
11450 any = true;
11451 }
11452 if let Some(v) = decode(s.max.as_deref()) {
11453 mx = mx.max(v);
11454 any = true;
11455 }
11456 nulls += s.null_count;
11457 }
11458 any.then_some((Some(mn), Some(mx), nulls))
11459}
11460
11461type SecondaryIndexes = (
11463 HashMap<u16, BitmapIndex>,
11464 HashMap<u16, AnnIndex>,
11465 HashMap<u16, FmIndex>,
11466 HashMap<u16, SparseIndex>,
11467 HashMap<u16, MinHashIndex>,
11468);
11469
11470fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11471 let mut bitmap = HashMap::new();
11472 let mut ann = HashMap::new();
11473 let mut fm = HashMap::new();
11474 let mut sparse = HashMap::new();
11475 let mut minhash = HashMap::new();
11476 for idef in &schema.indexes {
11477 match idef.kind {
11478 IndexKind::Bitmap => {
11479 bitmap.insert(idef.column_id, BitmapIndex::new());
11480 }
11481 IndexKind::Ann => {
11482 let dim = schema
11483 .columns
11484 .iter()
11485 .find(|c| c.id == idef.column_id)
11486 .and_then(|c| match c.ty {
11487 TypeId::Embedding { dim } => Some(dim as usize),
11488 _ => None,
11489 })
11490 .unwrap_or(0);
11491 let options = idef.options.ann.clone().unwrap_or_default();
11492 ann.insert(
11493 idef.column_id,
11494 AnnIndex::with_options(
11495 dim,
11496 options.m,
11497 options.ef_construction,
11498 options.ef_search,
11499 ),
11500 );
11501 }
11502 IndexKind::FmIndex => {
11503 fm.insert(idef.column_id, FmIndex::new());
11504 }
11505 IndexKind::Sparse => {
11506 sparse.insert(idef.column_id, SparseIndex::new());
11507 }
11508 IndexKind::MinHash => {
11509 let options = idef.options.minhash.clone().unwrap_or_default();
11510 minhash.insert(
11511 idef.column_id,
11512 MinHashIndex::with_options(options.permutations, options.bands),
11513 );
11514 }
11515 _ => {}
11516 }
11517 }
11518 (bitmap, ann, fm, sparse, minhash)
11519}
11520
11521const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11522 | ColumnFlags::AUTO_INCREMENT
11523 | ColumnFlags::ENCRYPTED
11524 | ColumnFlags::ENCRYPTED_INDEXABLE
11525 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11526
11527fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11528 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11529 return Err(MongrelError::Schema(
11530 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11531 ));
11532 }
11533 Ok(())
11534}
11535
11536fn validate_alter_column_type(
11537 schema: &Schema,
11538 old: &ColumnDef,
11539 next: &ColumnDef,
11540 has_stored_versions: bool,
11541) -> Result<()> {
11542 if old.ty == next.ty {
11543 return Ok(());
11544 }
11545 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11546 return Err(MongrelError::Schema(format!(
11547 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11548 old.name
11549 )));
11550 }
11551 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11552 return Ok(());
11553 }
11554 Err(MongrelError::Schema(format!(
11555 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11556 old.ty, next.ty
11557 )))
11558}
11559
11560fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11561 matches!(
11562 (old, new),
11563 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11564 )
11565}
11566
11567fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11573 let mut prev: Option<i64> = None;
11574 for r in rows {
11575 match r.columns.get(&pk_id) {
11576 Some(Value::Int64(v)) => {
11577 if prev.is_some_and(|p| p >= *v) {
11578 return false;
11579 }
11580 prev = Some(*v);
11581 }
11582 _ => return false,
11583 }
11584 }
11585 true
11586}
11587
11588#[allow(clippy::too_many_arguments)]
11589fn index_into(
11590 schema: &Schema,
11591 row: &Row,
11592 hot: &mut HotIndex,
11593 bitmap: &mut HashMap<u16, BitmapIndex>,
11594 ann: &mut HashMap<u16, AnnIndex>,
11595 fm: &mut HashMap<u16, FmIndex>,
11596 sparse: &mut HashMap<u16, SparseIndex>,
11597 minhash: &mut HashMap<u16, MinHashIndex>,
11598) {
11599 for idef in &schema.indexes {
11600 let Some(val) = row.columns.get(&idef.column_id) else {
11601 continue;
11602 };
11603 match idef.kind {
11604 IndexKind::Bitmap => {
11605 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11606 b.insert(val.encode_key(), row.row_id);
11607 }
11608 }
11609 IndexKind::Ann => {
11610 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11611 a.insert_validated(v, row.row_id);
11612 }
11613 }
11614 IndexKind::FmIndex => {
11615 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11616 f.insert(b.clone(), row.row_id);
11617 }
11618 }
11619 IndexKind::Sparse => {
11620 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11621 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11624 s.insert(&terms, row.row_id);
11625 }
11626 }
11627 }
11628 IndexKind::MinHash => {
11629 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11630 let tokens = crate::index::token_hashes_from_bytes(b);
11633 mh.insert(&tokens, row.row_id);
11634 }
11635 }
11636 _ => {}
11637 }
11638 }
11639 if let Some(pk_col) = schema.primary_key() {
11640 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11641 hot.insert(pk_val.encode_key(), row.row_id);
11642 }
11643 }
11644}
11645
11646#[allow(clippy::too_many_arguments)]
11649fn index_into_single(
11650 idef: &IndexDef,
11651 _schema: &Schema,
11652 row: &Row,
11653 _hot: &mut HotIndex,
11654 bitmap: &mut HashMap<u16, BitmapIndex>,
11655 ann: &mut HashMap<u16, AnnIndex>,
11656 fm: &mut HashMap<u16, FmIndex>,
11657 sparse: &mut HashMap<u16, SparseIndex>,
11658 minhash: &mut HashMap<u16, MinHashIndex>,
11659) {
11660 let Some(val) = row.columns.get(&idef.column_id) else {
11661 return;
11662 };
11663 match idef.kind {
11664 IndexKind::Bitmap => {
11665 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11666 b.insert(val.encode_key(), row.row_id);
11667 }
11668 }
11669 IndexKind::Ann => {
11670 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11671 a.insert_validated(v, row.row_id);
11672 }
11673 }
11674 IndexKind::FmIndex => {
11675 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11676 f.insert(b.clone(), row.row_id);
11677 }
11678 }
11679 IndexKind::Sparse => {
11680 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11681 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11682 s.insert(&terms, row.row_id);
11683 }
11684 }
11685 }
11686 IndexKind::MinHash => {
11687 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11688 let tokens = crate::index::token_hashes_from_bytes(b);
11689 mh.insert(&tokens, row.row_id);
11690 }
11691 }
11692 _ => {}
11693 }
11694}
11695
11696fn eval_partial_predicate(
11702 pred: &str,
11703 columns_map: &HashMap<u16, &Value>,
11704 name_to_id: &HashMap<&str, u16>,
11705) -> bool {
11706 let lower = pred.trim().to_ascii_lowercase();
11707 if let Some(rest) = lower.strip_suffix(" is not null") {
11709 let col_name = rest.trim();
11710 if let Some(col_id) = name_to_id.get(col_name) {
11711 return columns_map
11712 .get(col_id)
11713 .is_some_and(|v| !matches!(v, Value::Null));
11714 }
11715 }
11716 if let Some(rest) = lower.strip_suffix(" is null") {
11718 let col_name = rest.trim();
11719 if let Some(col_id) = name_to_id.get(col_name) {
11720 return columns_map
11721 .get(col_id)
11722 .is_none_or(|v| matches!(v, Value::Null));
11723 }
11724 }
11725 true
11728}
11729
11730#[allow(dead_code)]
11736fn bulk_index_key(
11737 column_keys: &HashMap<u16, ([u8; 32], u8)>,
11738 column_id: u16,
11739 ty: TypeId,
11740 col: &columnar::NativeColumn,
11741 i: usize,
11742) -> Option<Vec<u8>> {
11743 let encoded = columnar::encode_key_native(ty, col, i)?;
11744 #[cfg(feature = "encryption")]
11745 {
11746 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11747 if let Some((key, scheme)) = column_keys.get(&column_id) {
11748 return Some(match (*scheme, col) {
11749 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
11750 (_, columnar::NativeColumn::Int64 { data, .. }) => {
11751 ope_token_i64(key, data[i]).to_vec()
11752 }
11753 (_, columnar::NativeColumn::Float64 { data, .. }) => {
11754 ope_token_f64(key, data[i]).to_vec()
11755 }
11756 _ => hmac_token(key, &encoded).to_vec(),
11757 });
11758 }
11759 }
11760 #[cfg(not(feature = "encryption"))]
11761 {
11762 let _ = (column_id, column_keys, col);
11763 }
11764 Some(encoded)
11765}
11766
11767pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
11768 write_schema_with_after(dir, schema, || {})
11769}
11770
11771pub(crate) fn write_schema_durable(
11772 root: &crate::durable_file::DurableRoot,
11773 schema: &Schema,
11774) -> Result<()> {
11775 write_schema_durable_with_after(root, schema, || {})
11776}
11777
11778fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
11779where
11780 F: FnOnce(),
11781{
11782 let json = serde_json::to_string_pretty(schema)
11783 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
11784 crate::durable_file::write_atomic_with_after(
11785 &dir.join(SCHEMA_FILENAME),
11786 json.as_bytes(),
11787 after_publish,
11788 )?;
11789 Ok(())
11790}
11791
11792fn write_schema_durable_with_after<F>(
11793 root: &crate::durable_file::DurableRoot,
11794 schema: &Schema,
11795 after_publish: F,
11796) -> Result<()>
11797where
11798 F: FnOnce(),
11799{
11800 let json = serde_json::to_string_pretty(schema)
11801 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
11802 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
11803 Ok(())
11804}
11805
11806fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
11807 let mut schema_published = false;
11808 let schema_result = match table._root_guard.as_deref() {
11809 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
11810 schema_published = true;
11811 }),
11812 None => write_schema_with_after(&table.dir, &table.schema, || {
11813 schema_published = true;
11814 }),
11815 };
11816 if schema_result.is_err() && !schema_published {
11817 return schema_result;
11818 }
11819 match table.persist_manifest(table.current_epoch()) {
11820 Ok(()) => Ok(()),
11821 Err(manifest_error) => Err(match schema_result {
11822 Ok(()) => manifest_error,
11823 Err(schema_error) => MongrelError::Other(format!(
11824 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11825 )),
11826 }),
11827 }
11828}
11829
11830fn read_schema(dir: &Path) -> Result<Schema> {
11831 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
11832 read_schema_file(file)
11833}
11834
11835fn read_schema_file(file: std::fs::File) -> Result<Schema> {
11836 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
11837 use std::io::Read;
11838
11839 let length = file.metadata()?.len();
11840 if length > MAX_SCHEMA_BYTES {
11841 return Err(MongrelError::ResourceLimitExceeded {
11842 resource: "schema bytes",
11843 requested: usize::try_from(length).unwrap_or(usize::MAX),
11844 limit: MAX_SCHEMA_BYTES as usize,
11845 });
11846 }
11847 let mut bytes = Vec::with_capacity(length as usize);
11848 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
11849 if bytes.len() as u64 != length {
11850 return Err(MongrelError::Schema(
11851 "schema length changed while reading".into(),
11852 ));
11853 }
11854 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
11855}
11856
11857fn preflight_standalone_open(
11858 dir: &Path,
11859 runs_root: Option<&crate::durable_file::DurableRoot>,
11860 idx_root: Option<&crate::durable_file::DurableRoot>,
11861 manifest: &Manifest,
11862 schema: &Schema,
11863 records: &[crate::wal::Record],
11864 kek: Option<Arc<Kek>>,
11865) -> Result<()> {
11866 crate::wal::validate_shared_transaction_framing(records)?;
11867 if manifest.schema_id > schema.schema_id
11868 || manifest.flushed_epoch > manifest.current_epoch
11869 || manifest.global_idx_epoch > manifest.current_epoch
11870 || manifest.next_row_id == u64::MAX
11871 || manifest.auto_inc_next < 0
11872 || manifest.auto_inc_next == i64::MAX
11873 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
11874 {
11875 return Err(MongrelError::InvalidArgument(
11876 "manifest counters or schema identity are invalid".into(),
11877 ));
11878 }
11879 let mut run_ids = HashSet::new();
11880 let mut maximum_row_id = None::<u64>;
11881 for run in &manifest.runs {
11882 if run.run_id >= u64::MAX as u128
11883 || !run_ids.insert(run.run_id)
11884 || run.epoch_created > manifest.current_epoch
11885 {
11886 return Err(MongrelError::InvalidArgument(
11887 "manifest contains an invalid or duplicate active run".into(),
11888 ));
11889 }
11890 let mut reader = match runs_root {
11891 Some(root) => RunReader::open_file(
11892 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
11893 schema.clone(),
11894 kek.clone(),
11895 )?,
11896 None => RunReader::open(
11897 dir.join(RUNS_DIR)
11898 .join(format!("r-{}.sr", run.run_id as u64)),
11899 schema.clone(),
11900 kek.clone(),
11901 )?,
11902 };
11903 let header = reader.header();
11904 if header.run_id != run.run_id
11905 || header.level != run.level
11906 || header.row_count != run.row_count
11907 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
11908 || header.is_uniform_epoch() && header.epoch_created != 0
11909 || header.schema_id > schema.schema_id
11910 {
11911 return Err(MongrelError::InvalidArgument(format!(
11912 "run {} differs from its manifest",
11913 run.run_id
11914 )));
11915 }
11916 if header.row_count != 0 {
11917 maximum_row_id = Some(
11918 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
11919 );
11920 }
11921 reader.validate_all_pages()?;
11922 }
11923 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
11924 return Err(MongrelError::InvalidArgument(
11925 "manifest next_row_id does not advance beyond persisted rows".into(),
11926 ));
11927 }
11928 for run in &manifest.retiring {
11929 if run.run_id >= u64::MAX as u128
11930 || run.retire_epoch > manifest.current_epoch
11931 || !run_ids.insert(run.run_id)
11932 {
11933 return Err(MongrelError::InvalidArgument(
11934 "manifest contains an invalid or duplicate retired run".into(),
11935 ));
11936 }
11937 }
11938 #[cfg(feature = "encryption")]
11939 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
11940 #[cfg(not(feature = "encryption"))]
11941 let idx_dek: Option<Zeroizing<[u8; DEK_LEN]>> = None;
11942 match idx_root {
11943 Some(root) => {
11944 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
11945 }
11946 None => {
11947 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
11948 }
11949 }
11950
11951 let committed = records
11952 .iter()
11953 .filter_map(|record| match record.op {
11954 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
11955 _ => None,
11956 })
11957 .collect::<HashMap<_, _>>();
11958 for record in records {
11959 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
11960 continue;
11961 };
11962 match &record.op {
11963 Op::Put { table_id, rows } => {
11964 if *table_id != manifest.table_id {
11965 return Err(MongrelError::CorruptWal {
11966 offset: record.seq.0,
11967 reason: format!(
11968 "private WAL record references table {table_id}, expected {}",
11969 manifest.table_id
11970 ),
11971 });
11972 }
11973 let rows: Vec<Row> =
11974 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
11975 offset: record.seq.0,
11976 reason: format!("committed Put payload could not be decoded: {error}"),
11977 })?;
11978 for row in rows {
11979 if row.deleted || row.row_id.0 == u64::MAX {
11980 return Err(MongrelError::CorruptWal {
11981 offset: record.seq.0,
11982 reason: "committed Put contains an invalid row identity".into(),
11983 });
11984 }
11985 let cells = row.columns.into_iter().collect::<Vec<_>>();
11986 schema
11987 .validate_values(&cells)
11988 .map_err(|error| MongrelError::CorruptWal {
11989 offset: record.seq.0,
11990 reason: format!("committed Put violates table schema: {error}"),
11991 })?;
11992 if schema.auto_increment_column().is_some_and(|column| {
11993 matches!(
11994 cells.iter().find(|(id, _)| *id == column.id),
11995 Some((_, Value::Int64(value))) if *value == i64::MAX
11996 )
11997 }) {
11998 return Err(MongrelError::CorruptWal {
11999 offset: record.seq.0,
12000 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12001 });
12002 }
12003 }
12004 }
12005 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12006 if *table_id != manifest.table_id =>
12007 {
12008 return Err(MongrelError::CorruptWal {
12009 offset: record.seq.0,
12010 reason: format!(
12011 "private WAL record references table {table_id}, expected {}",
12012 manifest.table_id
12013 ),
12014 });
12015 }
12016 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12017 return Err(MongrelError::CorruptWal {
12018 offset: record.seq.0,
12019 reason: "private WAL contains shared spilled-run metadata".into(),
12020 });
12021 }
12022 _ => {}
12023 }
12024 }
12025 Ok(())
12026}
12027
12028fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12029 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12030}
12031
12032fn wal_segment_number(path: &Path) -> Option<u64> {
12033 path.file_stem()
12034 .and_then(|stem| stem.to_str())
12035 .and_then(|stem| stem.strip_prefix("seg-"))
12036 .and_then(|number| number.parse().ok())
12037}
12038
12039fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12040 let n = list_wal_numbers(wal_dir)?;
12041 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12042}
12043
12044fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12045 list_wal_numbers(wal_dir)?
12046 .map(|maximum| {
12047 maximum
12048 .checked_add(1)
12049 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12050 })
12051 .unwrap_or(Ok(0))
12052}
12053
12054fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12055 let mut max_n = None;
12056 let entries = match std::fs::read_dir(wal_dir) {
12057 Ok(entries) => entries,
12058 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12059 Err(error) => return Err(error.into()),
12060 };
12061 for entry in entries {
12062 let entry = entry?;
12063 let fname = entry.file_name();
12064 let Some(s) = fname.to_str() else {
12065 continue;
12066 };
12067 let Some(stripped) = s.strip_prefix("seg-") else {
12068 continue;
12069 };
12070 let Some(number) = stripped.strip_suffix(".wal") else {
12071 return Err(MongrelError::CorruptWal {
12072 offset: 0,
12073 reason: format!("malformed WAL segment name {s:?}"),
12074 });
12075 };
12076 let n = number
12077 .parse::<u32>()
12078 .map_err(|_| MongrelError::CorruptWal {
12079 offset: 0,
12080 reason: format!("malformed WAL segment name {s:?}"),
12081 })?;
12082 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12083 return Err(MongrelError::CorruptWal {
12084 offset: n as u64,
12085 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12086 });
12087 }
12088 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12089 }
12090 Ok(max_n)
12091}