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 }
2108 }
2109 if !checkpoint_valid {
2110 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2111 db.bitmap = bitmap;
2112 db.ann = ann;
2113 db.fm = fm;
2114 db.sparse = sparse;
2115 db.minhash = minhash;
2116 db.rebuild_indexes_from_runs()?;
2117 db.build_learned_ranges()?;
2118 }
2119
2120 for (epoch, group) in replayed_puts {
2125 let (losers, winner_pks) = db.partition_pk_winners(&group);
2126 for (key, &row_id) in &winner_pks {
2127 if let Some(old_rid) = db.hot.get(key) {
2128 if old_rid != row_id {
2129 db.tombstone_row(old_rid, epoch, false);
2130 }
2131 }
2132 }
2133 for &loser_rid in &losers {
2134 db.tombstone_row(loser_rid, epoch, false);
2135 }
2136 for (key, row_id) in winner_pks {
2137 db.insert_hot_pk(key, row_id);
2138 }
2139 if db.schema.primary_key().is_none() {
2140 for r in &group {
2141 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2142 }
2143 }
2144 for r in &group {
2145 if !losers.contains(&r.row_id) {
2146 db.index_row(r);
2147 }
2148 }
2149 }
2150 for (rid, epoch) in &replayed_deletes {
2154 db.remove_hot_for_row(*rid, *epoch);
2155 }
2156
2157 if recovered_manifest_dirty {
2158 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2159 db.live_count = rows.len() as u64;
2160 db.persist_manifest(Epoch(recovered_epoch))?;
2161 }
2162
2163 db.result_cache.lock().load_persistent();
2170 Ok(db)
2171 }
2172
2173 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2179 if self.reservoir_complete {
2180 return Ok(());
2181 }
2182 self.rebuild_reservoir()?;
2183 self.reservoir_complete = true;
2184 Ok(())
2185 }
2186
2187 fn rebuild_reservoir(&mut self) -> Result<()> {
2190 let snap = self.snapshot();
2191 let rows = self.visible_rows(snap)?;
2192 self.reservoir.reset();
2193 for r in rows {
2194 self.reservoir.offer(r.row_id.0);
2195 }
2196 Ok(())
2197 }
2198
2199 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2200 self.rebuild_indexes_from_runs_inner(None)
2201 }
2202
2203 fn rebuild_indexes_from_runs_inner(
2204 &mut self,
2205 control: Option<&crate::ExecutionControl>,
2206 ) -> Result<()> {
2207 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2210 crate::retention::PinSource::OnlineIndexBuild,
2211 self.current_epoch(),
2212 );
2213 self.hot = HotIndex::new();
2214 self.pk_by_row.clear();
2215 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2216 self.bitmap = bitmap;
2217 self.ann = ann;
2218 self.fm = fm;
2219 self.sparse = sparse;
2220 self.minhash = minhash;
2221 let snapshot = Epoch(u64::MAX);
2222 let ttl_now = unix_nanos_now();
2223 let mut scanned = 0_usize;
2224 for rr in self.run_refs.clone() {
2225 if let Some(control) = control {
2226 control.checkpoint()?;
2227 }
2228 let mut reader = self.open_reader(rr.run_id)?;
2229 for row in reader.visible_rows(snapshot)? {
2230 if scanned.is_multiple_of(256) {
2231 if let Some(control) = control {
2232 control.checkpoint()?;
2233 }
2234 }
2235 scanned += 1;
2236 if self.row_expired_at(&row, ttl_now) {
2237 continue;
2238 }
2239 let tok_row = self.tokenized_for_indexes(&row);
2240 index_into(
2241 &self.schema,
2242 &tok_row,
2243 &mut self.hot,
2244 &mut self.bitmap,
2245 &mut self.ann,
2246 &mut self.fm,
2247 &mut self.sparse,
2248 &mut self.minhash,
2249 );
2250 }
2251 }
2252 for row in self.mutable_run.visible_versions(snapshot) {
2253 if scanned.is_multiple_of(256) {
2254 if let Some(control) = control {
2255 control.checkpoint()?;
2256 }
2257 }
2258 scanned += 1;
2259 if row.deleted {
2260 self.remove_hot_for_row(row.row_id, snapshot);
2261 } else if !self.row_expired_at(&row, ttl_now) {
2262 self.index_row(&row);
2263 }
2264 }
2265 for row in self.memtable.visible_versions(snapshot) {
2266 if scanned.is_multiple_of(256) {
2267 if let Some(control) = control {
2268 control.checkpoint()?;
2269 }
2270 }
2271 scanned += 1;
2272 if row.deleted {
2273 self.remove_hot_for_row(row.row_id, snapshot);
2274 } else if !self.row_expired_at(&row, ttl_now) {
2275 self.index_row(&row);
2276 }
2277 }
2278 self.refresh_pk_by_row_from_hot();
2279 Ok(())
2280 }
2281
2282 fn refresh_pk_by_row_from_hot(&mut self) {
2283 self.pk_by_row_complete = true;
2284 if self.schema.primary_key().is_none() {
2285 self.pk_by_row.clear();
2286 return;
2287 }
2288 self.pk_by_row = ReversePkMap::from_entries(
2294 self.hot
2295 .entries()
2296 .into_iter()
2297 .map(|(key, row_id)| (row_id, key)),
2298 );
2299 }
2300
2301 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2302 if self.schema.primary_key().is_some() {
2303 self.pk_by_row.insert(row_id, key.clone());
2304 }
2305 self.hot.insert(key, row_id);
2306 }
2307
2308 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2312 self.build_learned_ranges_inner(None)
2313 }
2314
2315 fn build_learned_ranges_inner(
2316 &mut self,
2317 control: Option<&crate::ExecutionControl>,
2318 ) -> Result<()> {
2319 self.learned_range = Arc::new(HashMap::new());
2320 if self.run_refs.len() != 1 {
2321 return Ok(());
2322 }
2323 let cols: Vec<(u16, usize)> = self
2324 .schema
2325 .indexes
2326 .iter()
2327 .filter(|i| i.kind == IndexKind::LearnedRange)
2328 .map(|i| {
2329 (
2330 i.column_id,
2331 i.options
2332 .learned_range
2333 .as_ref()
2334 .map(|options| options.epsilon)
2335 .unwrap_or(16),
2336 )
2337 })
2338 .collect();
2339 if cols.is_empty() {
2340 return Ok(());
2341 }
2342 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2343 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2344 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2345 _ => return Ok(()),
2346 };
2347 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2348 if column_index % 256 == 0 {
2349 if let Some(control) = control {
2350 control.checkpoint()?;
2351 }
2352 }
2353 let ty = self
2354 .schema
2355 .columns
2356 .iter()
2357 .find(|c| c.id == cid)
2358 .map(|c| c.ty.clone())
2359 .unwrap_or(TypeId::Int64);
2360 match ty {
2361 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2362 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2363 let pairs: Vec<(i64, u64)> = data
2364 .iter()
2365 .zip(row_ids.iter())
2366 .map(|(v, r)| (*v, *r))
2367 .collect();
2368 Arc::make_mut(&mut self.learned_range).insert(
2369 cid,
2370 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2371 );
2372 }
2373 }
2374 TypeId::Float64 => {
2375 if let columnar::NativeColumn::Float64 { data, .. } =
2376 reader.column_native(cid)?
2377 {
2378 let pairs: Vec<(f64, u64)> = data
2379 .iter()
2380 .zip(row_ids.iter())
2381 .map(|(v, r)| (*v, *r))
2382 .collect();
2383 Arc::make_mut(&mut self.learned_range).insert(
2384 cid,
2385 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2386 );
2387 }
2388 }
2389 _ => {}
2390 }
2391 }
2392 Ok(())
2393 }
2394
2395 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2402 if self.indexes_complete {
2403 crate::trace::QueryTrace::record(|t| {
2404 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2405 });
2406 return Ok(());
2407 }
2408 crate::trace::QueryTrace::record(|t| {
2409 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2410 });
2411 self.rebuild_indexes_from_runs()?;
2412 self.build_learned_ranges()?;
2413 self.indexes_complete = true;
2414 let epoch = self.current_epoch();
2415 self.checkpoint_indexes(epoch);
2416 Ok(())
2417 }
2418
2419 #[doc(hidden)]
2422 pub fn ensure_indexes_complete_controlled<F>(
2423 &mut self,
2424 control: &crate::ExecutionControl,
2425 before_publish: F,
2426 ) -> Result<bool>
2427 where
2428 F: FnOnce() -> bool,
2429 {
2430 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2431 .map(|(changed, _)| changed)
2432 }
2433
2434 #[doc(hidden)]
2437 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2438 &mut self,
2439 control: &crate::ExecutionControl,
2440 before_publish: F,
2441 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2442 where
2443 F: FnOnce() -> bool,
2444 {
2445 if self.indexes_complete {
2446 crate::trace::QueryTrace::record(|trace| {
2447 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2448 });
2449 return Ok((false, None));
2450 }
2451 crate::trace::QueryTrace::record(|trace| {
2452 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2453 });
2454 control.checkpoint()?;
2455 let maintenance_epoch = self.current_epoch();
2456 self.rebuild_indexes_from_runs_inner(Some(control))?;
2457 self.build_learned_ranges_inner(Some(control))?;
2458 control.checkpoint()?;
2459 if !before_publish() {
2460 return Err(MongrelError::Cancelled);
2461 }
2462 self.indexes_complete = true;
2463 self.checkpoint_indexes(maintenance_epoch);
2464 Ok((
2465 true,
2466 Some(MaintenanceReceipt {
2467 epoch: maintenance_epoch,
2468 }),
2469 ))
2470 }
2471
2472 fn pending_epoch(&self) -> Epoch {
2473 Epoch(self.epoch.visible().0 + 1)
2474 }
2475
2476 fn is_shared(&self) -> bool {
2479 matches!(self.wal, WalSink::Shared(_))
2480 }
2481
2482 fn ensure_txn_id(&mut self) -> Result<u64> {
2486 if self.current_txn_id == 0 {
2487 let id = match &self.wal {
2488 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2489 WalSink::Private(_) => {
2490 return Err(MongrelError::Full(
2491 "standalone transaction id namespace exhausted".into(),
2492 ))
2493 }
2494 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2495 };
2496 self.current_txn_id = id;
2497 }
2498 Ok(self.current_txn_id)
2499 }
2500
2501 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2504 self.ensure_writable()?;
2505 let txn_id = self.ensure_txn_id()?;
2506 let table_id = self.table_id;
2507 match &mut self.wal {
2508 WalSink::Private(w) => {
2509 w.append_txn(txn_id, op)?;
2510 self.pending_private_mutations = true;
2511 }
2512 WalSink::Shared(s) => {
2513 s.wal.lock().append(txn_id, table_id, op)?;
2514 }
2515 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2516 }
2517 Ok(())
2518 }
2519
2520 fn ensure_writable(&self) -> Result<()> {
2521 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2522 return Err(MongrelError::ReadOnlyReplica);
2523 }
2524 if self.durable_commit_failed {
2525 return Err(MongrelError::Other(
2526 "table poisoned by post-commit failure; reopen required".into(),
2527 ));
2528 }
2529 Ok(())
2530 }
2531
2532 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2543 match &self.auth {
2544 Some(checker) => checker.check(&self.name, perm),
2545 None => Ok(()),
2546 }
2547 }
2548 pub fn require_select(&self) -> Result<()> {
2553 self.require(crate::auth_state::RequiredPermission::Select)
2554 }
2555 fn require_insert(&self) -> Result<()> {
2556 self.require(crate::auth_state::RequiredPermission::Insert)
2557 }
2558 #[allow(dead_code)]
2562 fn require_update(&self) -> Result<()> {
2563 self.require(crate::auth_state::RequiredPermission::Update)
2564 }
2565 fn require_delete(&self) -> Result<()> {
2566 self.require(crate::auth_state::RequiredPermission::Delete)
2567 }
2568
2569 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2572 self.require_insert()?;
2573 Ok(self.put_returning(columns)?.0)
2574 }
2575
2576 pub fn put_returning(
2581 &mut self,
2582 mut columns: Vec<(u16, Value)>,
2583 ) -> Result<(RowId, Option<i64>)> {
2584 self.require_insert()?;
2585 let assigned = self.fill_auto_inc(&mut columns)?;
2586 self.apply_defaults(&mut columns)?;
2587 self.schema.validate_values(&columns)?;
2588 let row_id = if self.schema.clustered {
2593 self.derive_clustered_row_id(&columns)?
2594 } else {
2595 self.allocator.alloc()?
2596 };
2597 let epoch = self.pending_epoch();
2598 let mut row = Row::new(row_id, epoch);
2599 for (col_id, val) in columns {
2600 row.columns.insert(col_id, val);
2601 }
2602 self.commit_rows(vec![row], assigned.is_some())?;
2603 Ok((row_id, assigned))
2604 }
2605
2606 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2609 self.require_insert()?;
2610 Ok(self
2611 .put_batch_returning(batch)?
2612 .into_iter()
2613 .map(|(r, _)| r)
2614 .collect())
2615 }
2616
2617 pub fn put_batch_returning(
2620 &mut self,
2621 batch: Vec<Vec<(u16, Value)>>,
2622 ) -> Result<Vec<(RowId, Option<i64>)>> {
2623 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2624 for mut cols in batch {
2625 let assigned = self.fill_auto_inc(&mut cols)?;
2626 self.apply_defaults(&mut cols)?;
2627 filled.push((cols, assigned));
2628 }
2629 for (cols, _) in &filled {
2630 self.schema.validate_values(cols)?;
2631 }
2632 let epoch = self.pending_epoch();
2633 let mut rows = Vec::with_capacity(filled.len());
2634 let mut ids = Vec::with_capacity(filled.len());
2635 let first_row_id = if self.schema.clustered {
2636 None
2637 } else {
2638 let count = u64::try_from(filled.len())
2639 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2640 Some(self.allocator.alloc_range(count)?.0)
2641 };
2642 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2643 let row_id = match first_row_id {
2644 Some(first) => RowId(first + row_index as u64),
2645 None => self.derive_clustered_row_id(&cols)?,
2646 };
2647 let mut row = Row::new(row_id, epoch);
2648 for (c, v) in cols {
2649 row.columns.insert(c, v);
2650 }
2651 ids.push((row_id, assigned));
2652 rows.push(row);
2653 }
2654 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2655 self.commit_rows(rows, all_auto_generated)?;
2656 Ok(ids)
2657 }
2658
2659 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2665 self.ensure_writable()?;
2666 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2667 return Ok(None);
2668 };
2669 let pos = columns.iter().position(|(c, _)| *c == cid);
2670 let assigned = match pos {
2671 Some(i) => match &columns[i].1 {
2672 Value::Null => {
2673 let next = self.alloc_auto_inc_value()?;
2674 columns[i].1 = Value::Int64(next);
2675 Some(next)
2676 }
2677 Value::Int64(n) => {
2678 self.advance_auto_inc_past(*n)?;
2679 None
2680 }
2681 other => {
2682 return Err(MongrelError::InvalidArgument(format!(
2683 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2684 other
2685 )))
2686 }
2687 },
2688 None => {
2689 let next = self.alloc_auto_inc_value()?;
2690 columns.push((cid, Value::Int64(next)));
2691 Some(next)
2692 }
2693 };
2694 Ok(assigned)
2695 }
2696
2697 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2703 for col in &self.schema.columns {
2704 let Some(expr) = &col.default_value else {
2705 continue;
2706 };
2707 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2709 continue;
2710 }
2711 let pos = columns.iter().position(|(c, _)| *c == col.id);
2712 let needs_default = match pos {
2713 None => true,
2714 Some(i) => matches!(columns[i].1, Value::Null),
2715 };
2716 if !needs_default {
2717 continue;
2718 }
2719 let v = match expr {
2720 crate::schema::DefaultExpr::Static(v) => v.clone(),
2721 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2722 crate::schema::DefaultExpr::Uuid => {
2723 let mut buf = [0u8; 16];
2724 getrandom::getrandom(&mut buf)
2725 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2726 Value::Uuid(buf)
2727 }
2728 };
2729 match pos {
2730 None => columns.push((col.id, v)),
2731 Some(i) => columns[i].1 = v,
2732 }
2733 }
2734 Ok(())
2735 }
2736
2737 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2739 self.ensure_auto_inc_seeded()?;
2740 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2742 let v = ai.next;
2743 ai.next = ai
2744 .next
2745 .checked_add(1)
2746 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2747 Ok(v)
2748 }
2749
2750 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2753 self.ensure_auto_inc_seeded()?;
2754 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2755 let floor = used
2756 .checked_add(1)
2757 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2758 .max(1);
2759 if ai.next < floor {
2760 ai.next = floor;
2761 }
2762 Ok(())
2763 }
2764
2765 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2770 let needs_seed = match self.auto_inc {
2771 Some(ai) => !ai.seeded,
2772 None => return Ok(()),
2773 };
2774 if !needs_seed {
2775 return Ok(());
2776 }
2777 if self.seed_empty_auto_inc() {
2778 return Ok(());
2779 }
2780 let cid = self
2781 .auto_inc
2782 .as_ref()
2783 .expect("auto-inc column present")
2784 .column_id;
2785 let max = self.scan_max_int64(cid)?;
2786 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2787 let floor = max
2788 .checked_add(1)
2789 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2790 .max(1);
2791 if ai.next < floor {
2792 ai.next = floor;
2793 }
2794 ai.seeded = true;
2795 Ok(())
2796 }
2797
2798 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2799 if n == 0 || self.auto_inc.is_none() {
2800 return Ok(None);
2801 }
2802 self.ensure_auto_inc_seeded()?;
2803 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2804 let start = ai.next;
2805 let count = i64::try_from(n)
2806 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2807 ai.next = ai
2808 .next
2809 .checked_add(count)
2810 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2811 Ok(Some(start))
2812 }
2813
2814 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2818 let mut max: i64 = 0;
2819 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2820 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2821 if *n > max {
2822 max = *n;
2823 }
2824 }
2825 }
2826 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2827 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2828 if *n > max {
2829 max = *n;
2830 }
2831 }
2832 }
2833 for rr in self.run_refs.clone() {
2834 let reader = self.open_reader(rr.run_id)?;
2835 if let Some(stats) = reader.column_page_stats(column_id) {
2836 for s in stats {
2837 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2838 if n > max {
2839 max = n;
2840 }
2841 }
2842 }
2843 } else if reader.has_column(column_id) {
2844 if let columnar::NativeColumn::Int64 { data, validity } =
2845 reader.column_native_shared(column_id)?
2846 {
2847 for (i, n) in data.iter().enumerate() {
2848 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2849 {
2850 max = *n;
2851 }
2852 }
2853 }
2854 }
2855 }
2856 Ok(max)
2857 }
2858
2859 fn seed_empty_auto_inc(&mut self) -> bool {
2860 let Some(ai) = self.auto_inc.as_mut() else {
2861 return false;
2862 };
2863 if ai.seeded || self.live_count != 0 {
2864 return false;
2865 }
2866 if ai.next < 1 {
2867 ai.next = 1;
2868 }
2869 ai.seeded = true;
2870 true
2871 }
2872
2873 fn advance_auto_inc_from_native_columns(
2874 &mut self,
2875 columns: &[(u16, columnar::NativeColumn)],
2876 n: usize,
2877 live_before: u64,
2878 ) -> Result<()> {
2879 let Some(ai) = self.auto_inc.as_mut() else {
2880 return Ok(());
2881 };
2882 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2883 return Ok(());
2884 };
2885 let columnar::NativeColumn::Int64 { data, validity } = col else {
2886 return Err(MongrelError::InvalidArgument(format!(
2887 "AUTO_INCREMENT column {} must be Int64",
2888 ai.column_id
2889 )));
2890 };
2891 let max = if native_int64_strictly_increasing(col, n) {
2892 data.get(n.saturating_sub(1)).copied()
2893 } else {
2894 data.iter()
2895 .take(n)
2896 .enumerate()
2897 .filter_map(|(i, v)| {
2898 if validity.is_empty() || columnar::validity_bit(validity, i) {
2899 Some(*v)
2900 } else {
2901 None
2902 }
2903 })
2904 .max()
2905 };
2906 if let Some(max) = max {
2907 let floor = max
2908 .checked_add(1)
2909 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2910 .max(1);
2911 if ai.next < floor {
2912 ai.next = floor;
2913 }
2914 if ai.seeded || live_before == 0 {
2915 ai.seeded = true;
2916 }
2917 }
2918 Ok(())
2919 }
2920
2921 fn fill_auto_inc_native_columns(
2922 &mut self,
2923 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2924 n: usize,
2925 ) -> Result<()> {
2926 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2927 return Ok(());
2928 };
2929 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2930 if let Some(start) = self.alloc_auto_inc_range(n)? {
2931 columns.push((
2932 cid,
2933 columnar::NativeColumn::Int64 {
2934 data: (start..start.saturating_add(n as i64)).collect(),
2935 validity: vec![0xFF; n.div_ceil(8)],
2936 },
2937 ));
2938 }
2939 return Ok(());
2940 };
2941
2942 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2943 return Err(MongrelError::InvalidArgument(format!(
2944 "AUTO_INCREMENT column {cid} must be Int64"
2945 )));
2946 };
2947 if data.len() < n {
2948 return Err(MongrelError::InvalidArgument(format!(
2949 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2950 data.len()
2951 )));
2952 }
2953 if columnar::all_non_null(validity, n) {
2954 return Ok(());
2955 }
2956 if validity.iter().all(|b| *b == 0) {
2957 if let Some(start) = self.alloc_auto_inc_range(n)? {
2958 for (i, slot) in data.iter_mut().take(n).enumerate() {
2959 *slot = start.saturating_add(i as i64);
2960 }
2961 *validity = vec![0xFF; n.div_ceil(8)];
2962 }
2963 return Ok(());
2964 }
2965
2966 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2967 for (i, slot) in data.iter_mut().enumerate().take(n) {
2968 if columnar::validity_bit(validity, i) {
2969 self.advance_auto_inc_past(*slot)?;
2970 } else {
2971 *slot = self.alloc_auto_inc_value()?;
2972 }
2973 }
2974 *validity = new_validity;
2975 Ok(())
2976 }
2977
2978 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2992 self.ensure_writable()?;
2993 if self.auto_inc.is_none() {
2994 return Ok(None);
2995 }
2996 Ok(Some(self.alloc_auto_inc_value()?))
2997 }
2998
2999 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3005 let payload = bincode::serialize(&rows)?;
3006 self.wal_append_data(Op::Put {
3007 table_id: self.table_id,
3008 rows: payload,
3009 })?;
3010 if self.is_shared() {
3011 self.pending_rows_auto_inc
3012 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3013 self.pending_rows.extend(rows);
3014 } else {
3015 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3016 }
3017 Ok(())
3018 }
3019
3020 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3023 self.ensure_indexes_complete()
3024 }
3025
3026 pub(crate) fn prepare_durable_publish_controlled(
3027 &mut self,
3028 control: &crate::ExecutionControl,
3029 ) -> Result<()> {
3030 self.ensure_indexes_complete_controlled(control, || true)?;
3031 Ok(())
3032 }
3033
3034 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3035 self.apply_put_rows_inner_prepared(rows, true);
3036 }
3037
3038 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3039 if check_existing_pk {
3040 self.ensure_indexes_complete()?;
3041 }
3042 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3043 Ok(())
3044 }
3045
3046 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3050 if rows.len() == 1 {
3054 let row = rows.into_iter().next().expect("len checked");
3055 self.apply_put_row_single(row, check_existing_pk);
3056 return;
3057 }
3058 let pk_id = self.schema.primary_key().map(|c| c.id);
3075 let probe = match pk_id {
3076 Some(pid) => {
3077 check_existing_pk
3078 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3079 }
3080 None => false,
3081 };
3082 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3085 for r in rows {
3086 for &cid in r.columns.keys() {
3087 self.pending_put_cols.insert(cid);
3088 }
3089 match pk_id {
3090 Some(pid) if probe || maintain_pk_by_row => {
3091 if let Some(pk_val) = r.columns.get(&pid) {
3092 let key = self.index_lookup_key(pid, pk_val);
3093 if probe {
3094 if let Some(old_rid) = self.hot.get(&key) {
3095 if old_rid != r.row_id {
3096 self.tombstone_row(old_rid, r.committed_epoch, true);
3097 }
3098 }
3099 }
3100 if maintain_pk_by_row {
3101 self.pk_by_row.insert(r.row_id, key);
3102 }
3103 }
3104 }
3105 Some(_) => {}
3106 None => {
3107 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3108 }
3109 }
3110 self.index_row(&r);
3111 self.reservoir.offer(r.row_id.0);
3112 self.memtable.upsert(r);
3113 self.live_count = self.live_count.saturating_add(1);
3116 }
3117 self.data_generation = self.data_generation.wrapping_add(1);
3118 }
3119
3120 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3124 for &cid in row.columns.keys() {
3125 self.pending_put_cols.insert(cid);
3126 }
3127 let epoch = row.committed_epoch;
3128 if let Some(pk_col) = self.schema.primary_key() {
3129 let pk_id = pk_col.id;
3130 if let Some(pk_val) = row.columns.get(&pk_id) {
3131 let maintain_pk_by_row = self.pk_by_row_complete;
3135 if check_existing_pk || maintain_pk_by_row {
3136 let key = self.index_lookup_key(pk_id, pk_val);
3137 if check_existing_pk {
3138 if let Some(old_rid) = self.hot.get(&key) {
3139 if old_rid != row.row_id {
3140 self.tombstone_row(old_rid, epoch, true);
3141 }
3142 }
3143 }
3144 if maintain_pk_by_row {
3145 self.pk_by_row.insert(row.row_id, key);
3146 }
3147 }
3148 }
3149 } else {
3150 self.hot
3151 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3152 }
3153 self.index_row(&row);
3154 self.reservoir.offer(row.row_id.0);
3155 self.memtable.upsert(row);
3156 self.live_count = self.live_count.saturating_add(1);
3157 self.data_generation = self.data_generation.wrapping_add(1);
3158 }
3159
3160 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3163 self.allocator.alloc()
3164 }
3165
3166 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3172 let pk = self.schema.primary_key().ok_or_else(|| {
3173 MongrelError::Schema("clustered table requires a single-column primary key".into())
3174 })?;
3175 let pk_val = columns
3176 .iter()
3177 .find(|(id, _)| *id == pk.id)
3178 .map(|(_, v)| v)
3179 .ok_or_else(|| {
3180 MongrelError::Schema(format!(
3181 "clustered table missing primary key column {} ({})",
3182 pk.id, pk.name
3183 ))
3184 })?;
3185 let key_bytes = pk_val.encode_key();
3186 let mut hash: u64 = 0xcbf29ce484222325;
3188 for &b in &key_bytes {
3189 hash ^= b as u64;
3190 hash = hash.wrapping_mul(0x100000001b3);
3191 }
3192 Ok(RowId(hash.max(1)))
3195 }
3196
3197 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3205 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3206 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3207 }
3208 let n = rows.len();
3209 for r in rows {
3210 for &cid in r.columns.keys() {
3211 self.pending_put_cols.insert(cid);
3212 }
3213 }
3214 let (losers, winner_pks) = self.partition_pk_winners(rows);
3215 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3216 for (key, &row_id) in &winner_pks {
3218 if let Some(old_rid) = self.hot.get(key) {
3219 if old_rid != row_id {
3220 self.tombstone_row(old_rid, epoch, true);
3221 }
3222 }
3223 }
3224 for &loser_rid in &losers {
3227 self.tombstone_row(loser_rid, epoch, false);
3228 }
3229 for (key, row_id) in winner_pks {
3231 self.insert_hot_pk(key, row_id);
3232 }
3233 if self.schema.primary_key().is_none() {
3234 for r in rows {
3235 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3236 }
3237 }
3238 for r in rows {
3239 self.allocator.advance_to(r.row_id)?;
3240 if !losers.contains(&r.row_id) {
3241 self.index_row(r);
3242 }
3243 }
3244 for r in rows {
3245 if !losers.contains(&r.row_id) {
3246 self.reservoir.offer(r.row_id.0);
3247 }
3248 }
3249 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3250 self.data_generation = self.data_generation.wrapping_add(1);
3251 Ok(())
3252 }
3253
3254 pub(crate) fn recover_apply(
3259 &mut self,
3260 rows: Vec<Row>,
3261 deletes: Vec<(RowId, Epoch)>,
3262 ) -> Result<()> {
3263 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3267 std::collections::BTreeMap::new();
3268 for row in rows {
3269 if row.row_id.0 >= u64::MAX - 1 {
3270 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3271 }
3272 self.allocator.advance_to(row.row_id)?;
3273 if let Some(ai) = self.auto_inc.as_mut() {
3278 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3279 let next = n.checked_add(1).ok_or_else(|| {
3280 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3281 })?;
3282 if next > ai.next {
3283 ai.next = next;
3284 }
3285 }
3286 }
3287 by_epoch.entry(row.committed_epoch).or_default().push(row);
3288 }
3289 for (epoch, group) in by_epoch {
3290 let (losers, winner_pks) = self.partition_pk_winners(&group);
3291 for (key, &row_id) in &winner_pks {
3293 if let Some(old_rid) = self.hot.get(key) {
3294 if old_rid != row_id {
3295 self.tombstone_row(old_rid, epoch, false);
3296 }
3297 }
3298 }
3299 for (key, row_id) in winner_pks {
3300 self.insert_hot_pk(key, row_id);
3301 }
3302 if self.schema.primary_key().is_none() {
3303 for r in &group {
3304 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3305 }
3306 }
3307 for r in &group {
3308 if !losers.contains(&r.row_id) {
3309 self.memtable.upsert(r.clone());
3310 self.index_row(r);
3311 }
3312 }
3313 }
3314 for (rid, epoch) in deletes {
3315 self.memtable.tombstone(rid, epoch);
3316 self.remove_hot_for_row(rid, epoch);
3317 }
3318 self.reservoir_complete = false;
3321 Ok(())
3322 }
3323
3324 pub(crate) fn flushed_epoch(&self) -> u64 {
3326 self.flushed_epoch
3327 }
3328
3329 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3330 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3331 }
3332
3333 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3335 self.schema.validate_values(cells)
3336 }
3337
3338 fn validate_columns_not_null(
3342 &self,
3343 columns: &[(u16, columnar::NativeColumn)],
3344 n: usize,
3345 ) -> Result<()> {
3346 let by_id: HashMap<u16, &columnar::NativeColumn> =
3347 columns.iter().map(|(id, c)| (*id, c)).collect();
3348 for col in &self.schema.columns {
3349 if !col.flags.contains(ColumnFlags::NULLABLE) {
3350 match by_id.get(&col.id) {
3351 None => {
3352 return Err(MongrelError::InvalidArgument(format!(
3353 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3354 col.name, col.id
3355 )));
3356 }
3357 Some(c) => {
3358 if c.null_count(n) != 0 {
3359 return Err(MongrelError::InvalidArgument(format!(
3360 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3361 col.name, col.id
3362 )));
3363 }
3364 }
3365 }
3366 }
3367 if let TypeId::Enum { variants } = &col.ty {
3368 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3369 if by_id.contains_key(&col.id) {
3370 return Err(MongrelError::InvalidArgument(format!(
3371 "column '{}' ({}) enum requires a bytes column",
3372 col.name, col.id
3373 )));
3374 }
3375 continue;
3376 };
3377 for index in 0..n {
3378 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3379 continue;
3380 };
3381 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3382 return Err(MongrelError::InvalidArgument(format!(
3383 "column '{}' ({}) enum value {:?} is not one of {:?}",
3384 col.name,
3385 col.id,
3386 String::from_utf8_lossy(value),
3387 variants
3388 )));
3389 }
3390 }
3391 }
3392 }
3393 Ok(())
3394 }
3395
3396 fn bulk_pk_winner_indices(
3401 &self,
3402 columns: &[(u16, columnar::NativeColumn)],
3403 n: usize,
3404 ) -> Option<Vec<usize>> {
3405 let pk_col = self.schema.primary_key()?;
3406 let pk_id = pk_col.id;
3407 let pk_ty = pk_col.ty.clone();
3408 let by_id: HashMap<u16, &columnar::NativeColumn> =
3409 columns.iter().map(|(id, c)| (*id, c)).collect();
3410 let pk_native = by_id.get(&pk_id)?;
3411 if native_int64_strictly_increasing(pk_native, n) {
3412 return None;
3413 }
3414 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3416 let mut null_pk_rows: Vec<usize> = Vec::new();
3417 for i in 0..n {
3418 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3419 Some(key) => {
3420 last.insert(key, i);
3421 }
3422 None => null_pk_rows.push(i),
3423 }
3424 }
3425 let mut winners: HashSet<usize> = last.values().copied().collect();
3426 for i in null_pk_rows {
3427 winners.insert(i);
3428 }
3429 Some((0..n).filter(|i| winners.contains(i)).collect())
3430 }
3431
3432 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3434 self.require_delete()?;
3435 let epoch = self.pending_epoch();
3436 self.wal_append_data(Op::Delete {
3437 table_id: self.table_id,
3438 row_ids: vec![row_id],
3439 })?;
3440 if self.is_shared() {
3441 self.pending_dels.push(row_id);
3442 } else {
3443 self.apply_delete(row_id, epoch);
3444 }
3445 Ok(())
3446 }
3447
3448 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3449 let pre = self.get(row_id, self.snapshot());
3450 self.delete(row_id)?;
3451 Ok(pre.map(|row| {
3452 let mut columns: Vec<_> = row.columns.into_iter().collect();
3453 columns.sort_by_key(|(id, _)| *id);
3454 OwnedRow { columns }
3455 }))
3456 }
3457
3458 pub fn truncate(&mut self) -> Result<()> {
3460 self.require_delete()?;
3461 let epoch = self.pending_epoch();
3462 self.wal_append_data(Op::TruncateTable {
3463 table_id: self.table_id,
3464 })?;
3465 self.pending_rows.clear();
3466 self.pending_rows_auto_inc.clear();
3467 self.pending_dels.clear();
3468 self.pending_truncate = Some(epoch);
3469 Ok(())
3470 }
3471
3472 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3474 self.run_refs.clear();
3479 self.retiring.clear();
3480 self.memtable = Memtable::new();
3481 self.mutable_run = MutableRun::new();
3482 self.hot = HotIndex::new();
3483 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3484 self.bitmap = bitmap;
3485 self.ann = ann;
3486 self.fm = fm;
3487 self.sparse = sparse;
3488 self.minhash = minhash;
3489 self.learned_range = Arc::new(HashMap::new());
3490 self.pk_by_row.clear();
3491 self.pk_by_row_complete = false;
3492 self.live_count = 0;
3493 self.reservoir = crate::reservoir::Reservoir::default();
3494 self.reservoir_complete = true;
3495 self.had_deletes = true;
3496 self.agg_cache = Arc::new(HashMap::new());
3497 self.global_idx_epoch = 0;
3498 self.indexes_complete = true;
3499 self.pending_delete_rids.clear();
3500 self.pending_put_cols.clear();
3501 self.pending_rows.clear();
3502 self.pending_rows_auto_inc.clear();
3503 self.pending_dels.clear();
3504 self.clear_result_cache();
3505 self.invalidate_index_checkpoint();
3506 self.data_generation = self.data_generation.wrapping_add(1);
3507 }
3508
3509 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3512 self.remove_hot_for_row(row_id, epoch);
3513 self.tombstone_row(row_id, epoch, true);
3514 self.data_generation = self.data_generation.wrapping_add(1);
3515 }
3516
3517 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3521 let tombstone = Row {
3522 row_id,
3523 committed_epoch: epoch,
3524 columns: std::collections::HashMap::new(),
3525 deleted: true,
3526 };
3527 self.memtable.upsert(tombstone);
3528 self.pk_by_row.remove(&row_id);
3529 if adjust_live_count {
3530 self.live_count = self.live_count.saturating_sub(1);
3531 }
3532 self.pending_delete_rids.insert(row_id.0 as u32);
3534 self.had_deletes = true;
3537 self.agg_cache = Arc::new(HashMap::new());
3538 }
3539
3540 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3544 let Some(pk_col) = self.schema.primary_key() else {
3545 return;
3546 };
3547 if self.pk_by_row_complete {
3550 if let Some(key) = self.pk_by_row.remove(&row_id) {
3551 if self.hot.get(&key) == Some(row_id) {
3552 self.hot.remove(&key);
3553 }
3554 }
3555 return;
3556 }
3557 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3576 if self.indexes_complete {
3577 let pk_val = self
3578 .memtable
3579 .get_version(row_id, lookup_epoch)
3580 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3581 .or_else(|| {
3582 self.mutable_run
3583 .get_version(row_id, lookup_epoch)
3584 .filter(|(_, r)| !r.deleted)
3585 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3586 })
3587 .or_else(|| {
3588 self.run_refs.iter().find_map(|rr| {
3589 let mut reader = self.open_reader(rr.run_id).ok()?;
3590 let (_, deleted, val) = reader
3591 .get_version_column(row_id, lookup_epoch, pk_col.id)
3592 .ok()??;
3593 if deleted {
3594 return None;
3595 }
3596 val
3597 })
3598 });
3599 if let Some(pk_val) = pk_val {
3600 let key = self.index_lookup_key(pk_col.id, &pk_val);
3601 if self.hot.get(&key) == Some(row_id) {
3602 self.hot.remove(&key);
3603 }
3604 return;
3605 }
3606 }
3607 self.refresh_pk_by_row_from_hot();
3612 if let Some(key) = self.pk_by_row.remove(&row_id) {
3613 if self.hot.get(&key) == Some(row_id) {
3614 self.hot.remove(&key);
3615 }
3616 }
3617 }
3618
3619 fn partition_pk_winners(
3624 &self,
3625 rows: &[Row],
3626 ) -> (
3627 std::collections::HashSet<RowId>,
3628 std::collections::HashMap<Vec<u8>, RowId>,
3629 ) {
3630 let mut losers = std::collections::HashSet::new();
3631 let Some(pk_col) = self.schema.primary_key() else {
3632 return (losers, std::collections::HashMap::new());
3633 };
3634 let pk_id = pk_col.id;
3635 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3636 std::collections::HashMap::new();
3637 for r in rows {
3638 let Some(pk_val) = r.columns.get(&pk_id) else {
3639 continue;
3640 };
3641 let key = self.index_lookup_key(pk_id, pk_val);
3642 if let Some(&old_rid) = winners.get(&key) {
3643 losers.insert(old_rid);
3644 }
3645 winners.insert(key, r.row_id);
3646 }
3647 (losers, winners)
3648 }
3649
3650 fn index_row(&mut self, row: &Row) {
3651 if row.deleted {
3652 return;
3653 }
3654 let any_predicate = self
3662 .schema
3663 .indexes
3664 .iter()
3665 .any(|idx| idx.predicate.is_some());
3666 if any_predicate {
3667 let columns_map: HashMap<u16, &Value> =
3668 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3669 let name_to_id: HashMap<&str, u16> = self
3670 .schema
3671 .columns
3672 .iter()
3673 .map(|c| (c.name.as_str(), c.id))
3674 .collect();
3675 for idx in &self.schema.indexes {
3676 if let Some(pred) = &idx.predicate {
3677 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3678 continue; }
3680 }
3681 index_into_single(
3683 idx,
3684 &self.schema,
3685 row,
3686 &mut self.hot,
3687 &mut self.bitmap,
3688 &mut self.ann,
3689 &mut self.fm,
3690 &mut self.sparse,
3691 &mut self.minhash,
3692 );
3693 }
3694 return;
3695 }
3696 if self.column_keys.is_empty() {
3700 index_into(
3701 &self.schema,
3702 row,
3703 &mut self.hot,
3704 &mut self.bitmap,
3705 &mut self.ann,
3706 &mut self.fm,
3707 &mut self.sparse,
3708 &mut self.minhash,
3709 );
3710 return;
3711 }
3712 let effective_row = self.tokenized_for_indexes(row);
3713 index_into(
3714 &self.schema,
3715 &effective_row,
3716 &mut self.hot,
3717 &mut self.bitmap,
3718 &mut self.ann,
3719 &mut self.fm,
3720 &mut self.sparse,
3721 &mut self.minhash,
3722 );
3723 }
3724
3725 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3731 if self.column_keys.is_empty() {
3732 return row.clone();
3733 }
3734 #[cfg(feature = "encryption")]
3735 {
3736 use crate::encryption::SCHEME_HMAC_EQ;
3737 let mut tok = row.clone();
3738 for (&cid, &(_, scheme)) in &self.column_keys {
3739 if scheme != SCHEME_HMAC_EQ {
3740 continue;
3741 }
3742 if let Some(v) = tok.columns.get(&cid).cloned() {
3743 if let Some(t) = self.tokenize_value(cid, &v) {
3744 tok.columns.insert(cid, t);
3745 }
3746 }
3747 }
3748 tok
3749 }
3750 #[cfg(not(feature = "encryption"))]
3751 {
3752 row.clone()
3753 }
3754 }
3755
3756 pub fn commit(&mut self) -> Result<Epoch> {
3761 self.commit_inner(None)
3762 }
3763
3764 #[doc(hidden)]
3767 pub fn commit_controlled<F>(
3768 &mut self,
3769 control: &crate::ExecutionControl,
3770 mut before_commit: F,
3771 ) -> Result<Epoch>
3772 where
3773 F: FnMut() -> Result<()>,
3774 {
3775 self.commit_inner(Some((control, &mut before_commit)))
3776 }
3777
3778 fn commit_inner(
3779 &mut self,
3780 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3781 ) -> Result<Epoch> {
3782 self.ensure_writable()?;
3783 if !self.has_pending_mutations() {
3784 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3785 return Err(MongrelError::Full(
3786 "standalone transaction id namespace exhausted".into(),
3787 ));
3788 }
3789 return Ok(self.epoch.visible());
3790 }
3791 self.commit_new_epoch_inner(controlled)
3792 }
3793
3794 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3798 self.commit_new_epoch_inner(None)
3799 }
3800
3801 fn commit_new_epoch_inner(
3802 &mut self,
3803 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3804 ) -> Result<Epoch> {
3805 self.ensure_writable()?;
3806 if self.is_shared() {
3807 self.commit_shared(controlled)
3808 } else {
3809 self.commit_private(controlled)
3810 }
3811 }
3812
3813 fn commit_private(
3815 &mut self,
3816 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3817 ) -> Result<Epoch> {
3818 let commit_lock = Arc::clone(&self.commit_lock);
3822 let _g = commit_lock.lock();
3823 let txn_id = self.ensure_txn_id()?;
3826 if let Some((control, before_commit)) = controlled {
3827 control.checkpoint()?;
3828 before_commit()?;
3829 }
3830 let new_epoch = self.epoch.bump_assigned();
3831 let epoch_authority = Arc::clone(&self.epoch);
3832 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3833 let wal_result = match &mut self.wal {
3837 WalSink::Private(w) => w
3838 .append_txn(
3839 txn_id,
3840 Op::TxnCommit {
3841 epoch: new_epoch.0,
3842 added_runs: Vec::new(),
3843 },
3844 )
3845 .and_then(|_| w.sync()),
3846 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3847 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
3848 };
3849 if let Err(error) = wal_result {
3850 self.durable_commit_failed = true;
3851 return Err(MongrelError::CommitOutcomeUnknown {
3852 epoch: new_epoch.0,
3853 message: error.to_string(),
3854 });
3855 }
3856 if let Some(epoch) = self.pending_truncate.take() {
3859 self.apply_truncate(epoch);
3860 }
3861 self.invalidate_pending_cache();
3862 let publish_result = self.persist_manifest(new_epoch);
3863 self.epoch.publish_in_order(new_epoch);
3867 epoch_guard.disarm();
3868 if let Err(error) = publish_result {
3869 self.durable_commit_failed = true;
3870 return Err(MongrelError::DurableCommit {
3871 epoch: new_epoch.0,
3872 message: error.to_string(),
3873 });
3874 }
3875 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
3876 self.pending_private_mutations = false;
3877 self.data_generation = self.data_generation.wrapping_add(1);
3878 Ok(new_epoch)
3879 }
3880
3881 fn commit_shared(
3887 &mut self,
3888 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3889 ) -> Result<Epoch> {
3890 use std::sync::atomic::Ordering;
3891 let s = match &self.wal {
3892 WalSink::Shared(s) => s.clone(),
3893 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3894 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3895 };
3896 if s.poisoned.load(Ordering::Relaxed) {
3897 return Err(MongrelError::Other(
3898 "database poisoned by fsync error".into(),
3899 ));
3900 }
3901 let commit_lock = Arc::clone(&self.commit_lock);
3908 let _g = commit_lock.lock();
3909 if !self.pending_rows.is_empty() {
3910 match controlled.as_ref() {
3911 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
3912 None => self.prepare_durable_publish()?,
3913 }
3914 }
3915 let txn_id = self.ensure_txn_id()?;
3918 let mut wal = s.wal.lock();
3919 if let Some((control, before_commit)) = controlled {
3920 control.checkpoint()?;
3921 before_commit()?;
3922 }
3923 let new_epoch = self.epoch.bump_assigned();
3924 let epoch_authority = Arc::clone(&self.epoch);
3925 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3926 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
3927 Ok(commit_seq) => commit_seq,
3928 Err(error) => {
3929 s.poisoned.store(true, Ordering::Relaxed);
3930 s.lifecycle.poison();
3931 return Err(MongrelError::CommitOutcomeUnknown {
3932 epoch: new_epoch.0,
3933 message: error.to_string(),
3934 });
3935 }
3936 };
3937 drop(wal);
3938 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
3939 s.poisoned.store(true, Ordering::Relaxed);
3940 s.lifecycle.poison();
3941 return Err(MongrelError::CommitOutcomeUnknown {
3942 epoch: new_epoch.0,
3943 message: error.to_string(),
3944 });
3945 }
3946
3947 if self.pending_truncate.take().is_some() {
3950 self.apply_truncate(new_epoch);
3951 }
3952 let mut rows = std::mem::take(&mut self.pending_rows);
3953 if !rows.is_empty() {
3954 for r in &mut rows {
3955 r.committed_epoch = new_epoch;
3956 }
3957 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3958 let all_auto_generated =
3959 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3960 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
3961 } else {
3962 self.pending_rows_auto_inc.clear();
3963 }
3964 let dels = std::mem::take(&mut self.pending_dels);
3965 for rid in dels {
3966 self.apply_delete(rid, new_epoch);
3967 }
3968
3969 self.invalidate_pending_cache();
3970 let publish_result = self.persist_manifest(new_epoch);
3971 self.epoch.publish_in_order(new_epoch);
3972 epoch_guard.disarm();
3973 let _ = s.change_wake.send(());
3974 if let Err(error) = publish_result {
3975 self.durable_commit_failed = true;
3976 s.poisoned.store(true, Ordering::Relaxed);
3977 s.lifecycle.poison();
3978 return Err(MongrelError::DurableCommit {
3979 epoch: new_epoch.0,
3980 message: error.to_string(),
3981 });
3982 }
3983 self.current_txn_id = 0;
3985 self.data_generation = self.data_generation.wrapping_add(1);
3986 Ok(new_epoch)
3987 }
3988
3989 pub fn flush(&mut self) -> Result<Epoch> {
3997 self.flush_with_outcome().map(|(epoch, _)| epoch)
3998 }
3999
4000 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4002 self.flush_with_outcome_inner(None)
4003 }
4004
4005 #[doc(hidden)]
4008 pub fn flush_with_outcome_controlled<F>(
4009 &mut self,
4010 control: &crate::ExecutionControl,
4011 mut before_commit: F,
4012 ) -> Result<(Epoch, bool)>
4013 where
4014 F: FnMut() -> Result<()>,
4015 {
4016 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4017 }
4018
4019 fn flush_with_outcome_inner(
4020 &mut self,
4021 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4022 ) -> Result<(Epoch, bool)> {
4023 match controlled.as_ref() {
4024 Some((control, _)) => {
4025 self.ensure_indexes_complete_controlled(control, || true)?;
4026 }
4027 None => self.ensure_indexes_complete()?,
4028 }
4029 let committed = self.has_pending_mutations();
4030 let epoch = self.commit_inner(controlled)?;
4031 let finish: Result<(Epoch, bool)> = (|| {
4032 let rows = self.memtable.drain_sorted();
4033 if !rows.is_empty() {
4034 self.mutable_run.insert_many(rows);
4035 }
4036 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4037 self.spill_mutable_run(epoch)?;
4038 self.mark_flushed(epoch)?;
4042 self.persist_manifest(epoch)?;
4043 self.build_learned_ranges()?;
4044 self.checkpoint_indexes(epoch);
4047 }
4048 Ok((epoch, committed))
4051 })();
4052 let outcome = match finish {
4053 Err(error) if committed => Err(MongrelError::DurableCommit {
4054 epoch: epoch.0,
4055 message: error.to_string(),
4056 }),
4057 result => result,
4058 };
4059 if outcome.is_ok() {
4060 let _ = self.publish_read_generation();
4066 }
4067 outcome
4068 }
4069
4070 fn has_pending_mutations(&self) -> bool {
4071 self.pending_private_mutations
4072 || !self.pending_rows.is_empty()
4073 || !self.pending_dels.is_empty()
4074 || self.pending_truncate.is_some()
4075 }
4076
4077 pub fn has_pending_writes(&self) -> bool {
4078 self.has_pending_mutations()
4079 }
4080
4081 pub fn force_flush(&mut self) -> Result<Epoch> {
4090 let saved = self.mutable_run_spill_bytes;
4091 self.mutable_run_spill_bytes = 1;
4092 let result = self.flush();
4093 self.mutable_run_spill_bytes = saved;
4094 result
4095 }
4096
4097 pub fn close(&mut self) -> Result<()> {
4104 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4105 self.force_flush()?;
4106 }
4107 Ok(())
4108 }
4109
4110 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4117 let op = Op::Flush {
4118 table_id: self.table_id,
4119 flushed_epoch: epoch.0,
4120 };
4121 match &mut self.wal {
4122 WalSink::Private(w) => {
4123 w.append_system(op)?;
4124 w.sync()?;
4125 }
4126 WalSink::Shared(s) => {
4127 s.wal.lock().append_system(op)?;
4132 }
4133 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4134 }
4135 self.flushed_epoch = epoch.0;
4136 if matches!(self.wal, WalSink::Private(_)) {
4137 self.rotate_wal(epoch)?;
4138 }
4139 Ok(())
4140 }
4141
4142 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4146 if self.mutable_run.is_empty() {
4147 return Ok(());
4148 }
4149 let run_id = self.alloc_run_id()?;
4150 let rows = self.mutable_run.drain_sorted();
4151 if rows.is_empty() {
4152 return Ok(());
4153 }
4154 let path = self.run_path(run_id);
4155 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4156 if let Some(kek) = &self.kek {
4157 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4158 }
4159 let header = match self.create_run_file(run_id)? {
4160 Some(file) => writer.write_file(file, &rows)?,
4161 None => writer.write(&path, &rows)?,
4162 };
4163 self.run_refs.push(RunRef {
4164 run_id: run_id as u128,
4165 level: 0,
4166 epoch_created: epoch.0,
4167 row_count: header.row_count,
4168 });
4169 Ok(())
4170 }
4171
4172 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4176 self.mutable_run_spill_bytes = bytes.max(1);
4177 }
4178
4179 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4183 self.compaction_zstd_level = level;
4184 }
4185
4186 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4190 self.result_cache.lock().set_max_bytes(max_bytes);
4191 }
4192
4193 pub(crate) fn clear_result_cache(&mut self) {
4197 self.result_cache.lock().clear();
4198 }
4199
4200 pub fn mutable_run_len(&self) -> usize {
4202 self.mutable_run.len()
4203 }
4204
4205 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4208 self.mutable_run.drain_sorted()
4209 }
4210
4211 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4213 let mut snapshot = self.mutable_run.clone();
4214 snapshot.drain_sorted()
4215 }
4216
4217 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4222 self.ensure_writable()?;
4223 let n = batch.len();
4224 if n == 0 {
4225 return Ok(self.current_epoch());
4226 }
4227 for row in &batch {
4228 self.schema.validate_values(row)?;
4229 }
4230 let epoch = self.commit_new_epoch()?;
4231 let live_before = self.live_count;
4232 self.spill_mutable_run(epoch)?;
4236 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4237 && self.indexes_complete
4238 && self.run_refs.is_empty()
4239 && self.memtable.is_empty()
4240 && self.mutable_run.is_empty();
4241 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4247 use rayon::prelude::*;
4248 self.schema
4249 .columns
4250 .par_iter()
4251 .map(|cdef| {
4252 (
4253 cdef.id,
4254 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4255 )
4256 })
4257 .collect::<Vec<_>>()
4258 };
4259 drop(batch);
4260 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4265 self.validate_columns_not_null(&user_columns, n)?;
4266 let winner_idx = self
4267 .bulk_pk_winner_indices(&user_columns, n)
4268 .filter(|idx| idx.len() != n);
4269 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4270 match winner_idx.as_deref() {
4271 Some(idx) => {
4272 let compacted = user_columns
4273 .iter()
4274 .map(|(id, c)| (*id, c.gather(idx)))
4275 .collect();
4276 (compacted, idx.len())
4277 }
4278 None => (std::mem::take(&mut user_columns), n),
4279 };
4280 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4281 let first = self.allocator.alloc_range(write_n as u64)?.0;
4282 for rid in first..first + write_n as u64 {
4283 self.reservoir.offer(rid);
4284 }
4285 let run_id = self.alloc_run_id()?;
4286 let path = self.run_path(run_id);
4287 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4288 .clean(true)
4289 .with_lz4()
4290 .with_native_endian();
4291 if let Some(kek) = &self.kek {
4292 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4293 }
4294 let header = match self.create_run_file(run_id)? {
4295 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4296 None => writer.write_native(&path, &write_columns, write_n, first)?,
4297 };
4298 self.run_refs.push(RunRef {
4299 run_id: run_id as u128,
4300 level: 0,
4301 epoch_created: epoch.0,
4302 row_count: header.row_count,
4303 });
4304 self.live_count = self.live_count.saturating_add(write_n as u64);
4305 if eager_index_build {
4306 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4307 self.index_columns_bulk(&write_columns, &row_ids);
4308 self.indexes_complete = true;
4309 self.build_learned_ranges()?;
4310 } else {
4311 self.indexes_complete = false;
4312 }
4313 self.mark_flushed(epoch)?;
4314 self.persist_manifest(epoch)?;
4315 if eager_index_build {
4316 self.checkpoint_indexes(epoch);
4317 }
4318 self.clear_result_cache();
4319 Ok(epoch)
4320 }
4321
4322 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4325 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4326 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4327 let segment_no = segment
4330 .file_stem()
4331 .and_then(|s| s.to_str())
4332 .and_then(|s| s.strip_prefix("seg-"))
4333 .and_then(|s| s.parse::<u64>().ok())
4334 .unwrap_or(0);
4335 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4336 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4337 wal.sync()?;
4338 self.wal = WalSink::Private(wal);
4339 Ok(())
4340 }
4341
4342 pub(crate) fn invalidate_pending_cache(&mut self) {
4347 self.result_cache
4348 .lock()
4349 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4350 self.pending_delete_rids.clear();
4351 self.pending_put_cols.clear();
4352 }
4353
4354 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4355 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4356 m.current_epoch = epoch.0;
4357 m.next_row_id = self.allocator.current().0;
4358 m.runs = self.run_refs.clone();
4359 m.live_count = self.live_count;
4360 m.global_idx_epoch = self.global_idx_epoch;
4361 m.flushed_epoch = self.flushed_epoch;
4362 m.retiring = self.retiring.clone();
4363 m.auto_inc_next = match self.auto_inc {
4367 Some(ai) if ai.seeded => ai.next,
4368 _ => 0,
4369 };
4370 m.ttl = self.ttl;
4371 let meta_dek = self.manifest_meta_dek();
4372 match self._root_guard.as_deref() {
4373 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4374 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4375 }
4376 Ok(())
4377 }
4378
4379 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4380 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4384 let live_count = u64::try_from(rows.len())
4385 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4386 let auto_inc = match self.auto_inc {
4387 Some(mut state) => {
4388 let maximum = self.scan_max_int64(state.column_id)?;
4389 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4390 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4391 })?;
4392 state.next = state.next.max(after_maximum).max(1);
4393 state.seeded = true;
4394 Some(state)
4395 }
4396 None => None,
4397 };
4398 Ok(RecoveryMetadataPlan {
4399 live_count,
4400 auto_inc,
4401 changed: live_count != self.live_count
4402 || auto_inc.is_some_and(|planned| {
4403 self.auto_inc.is_none_or(|current| {
4404 current.next != planned.next || current.seeded != planned.seeded
4405 })
4406 }),
4407 })
4408 }
4409
4410 pub(crate) fn apply_recovered_metadata(
4411 &mut self,
4412 plan: RecoveryMetadataPlan,
4413 epoch: Epoch,
4414 ) -> Result<()> {
4415 if !plan.changed {
4416 return Ok(());
4417 }
4418 self.live_count = plan.live_count;
4419 self.auto_inc = plan.auto_inc;
4420 self.persist_manifest(epoch)
4421 }
4422
4423 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4429 if !self.indexes_complete {
4432 return;
4433 }
4434 if crate::catalog::inject_hook("index.publish.before").is_err() {
4437 return;
4438 }
4439 if self.idx_root.is_none() {
4440 if let Some(root) = self._root_guard.as_ref() {
4441 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4442 return;
4443 };
4444 self.idx_root = Some(Arc::new(idx_root));
4445 }
4446 }
4447 let snap = global_idx::IndexSnapshot {
4448 hot: &self.hot,
4449 bitmap: &self.bitmap,
4450 ann: &self.ann,
4451 fm: &self.fm,
4452 sparse: &self.sparse,
4453 minhash: &self.minhash,
4454 learned_range: &self.learned_range,
4455 };
4456 let idx_dek = self.idx_dek();
4458 let written = match self.idx_root.as_deref() {
4459 Some(root) => global_idx::write_atomic_root(
4460 root,
4461 self.table_id,
4462 epoch.0,
4463 snap,
4464 idx_dek.as_deref(),
4465 ),
4466 None => global_idx::write_atomic(
4467 &self.dir,
4468 self.table_id,
4469 epoch.0,
4470 snap,
4471 idx_dek.as_deref(),
4472 ),
4473 };
4474 if written.is_ok() {
4475 self.global_idx_epoch = epoch.0;
4476 let _ = self.persist_manifest(epoch);
4477 let _ = crate::catalog::inject_hook("index.publish.after");
4479 }
4480 }
4481
4482 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4485 self.global_idx_epoch = 0;
4486 if let Some(root) = self.idx_root.as_deref() {
4487 let _ = root.remove_file(global_idx::IDX_FILENAME);
4488 } else {
4489 global_idx::remove(&self.dir);
4490 }
4491 let _ = self.persist_manifest(self.epoch.visible());
4492 }
4493
4494 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4499 self.indexes_complete = false;
4500 self.global_idx_epoch = 0;
4501 if let Some(root) = self.idx_root.as_deref() {
4502 let _ = root.remove_file(global_idx::IDX_FILENAME);
4503 } else {
4504 global_idx::remove(&self.dir);
4505 }
4506 }
4507
4508 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4509 self.indexes_complete = true;
4510 }
4511
4512 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4518 self.durable_commit_failed = true;
4519 if let WalSink::Shared(shared) = &self.wal {
4520 shared
4521 .poisoned
4522 .store(true, std::sync::atomic::Ordering::Relaxed);
4523 }
4524 }
4525
4526 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4530 self.durable_commit_failed = true;
4531 }
4532
4533 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4536 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4537 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4538 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4539 best = Some((epoch, row));
4540 }
4541 }
4542 for rr in &self.run_refs {
4543 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4544 continue;
4545 };
4546 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4547 continue;
4548 };
4549 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4550 best = Some((epoch, row));
4551 }
4552 }
4553 let now_nanos = unix_nanos_now();
4554 match best {
4555 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4556 Some((_, r)) => Some(r),
4557 None => None,
4558 }
4559 }
4560
4561 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4565 self.visible_rows_at_time(snapshot, unix_nanos_now())
4566 }
4567
4568 #[doc(hidden)]
4571 pub fn visible_rows_controlled(
4572 &self,
4573 snapshot: Snapshot,
4574 control: &crate::ExecutionControl,
4575 ) -> Result<Vec<Row>> {
4576 let mut out = Vec::new();
4577 self.for_each_visible_row_controlled(snapshot, control, |row| {
4578 out.push(row);
4579 Ok(())
4580 })?;
4581 Ok(out)
4582 }
4583
4584 #[doc(hidden)]
4587 pub fn for_each_visible_row_controlled<F>(
4588 &self,
4589 snapshot: Snapshot,
4590 control: &crate::ExecutionControl,
4591 visit: F,
4592 ) -> Result<()>
4593 where
4594 F: FnMut(Row) -> Result<()>,
4595 {
4596 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4597 control.checkpoint()?;
4598 let memtable = self.memtable.visible_versions(snapshot.epoch);
4599 if !memtable.is_empty() {
4600 sources.push(ControlledVisibleSource::memory(memtable));
4601 }
4602 control.checkpoint()?;
4603 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4604 if !mutable.is_empty() {
4605 sources.push(ControlledVisibleSource::memory(mutable));
4606 }
4607 for run in &self.run_refs {
4608 control.checkpoint()?;
4609 let reader = self.open_reader(run.run_id)?;
4610 sources.push(ControlledVisibleSource::run(
4611 reader.into_visible_version_cursor(snapshot.epoch)?,
4612 ));
4613 }
4614 let now_nanos = unix_nanos_now();
4615 merge_controlled_visible_sources(
4616 &mut sources,
4617 control,
4618 |row| self.row_expired_at(row, now_nanos),
4619 visit,
4620 )
4621 }
4622
4623 #[doc(hidden)]
4624 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4625 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4626 let mut fold = |row: Row| {
4627 best.entry(row.row_id.0)
4628 .and_modify(|e| {
4629 if row.committed_epoch > e.0 {
4630 *e = (row.committed_epoch, row.clone());
4631 }
4632 })
4633 .or_insert_with(|| (row.committed_epoch, row));
4634 };
4635 for row in self.memtable.visible_versions(snapshot.epoch) {
4636 fold(row);
4637 }
4638 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4639 fold(row);
4640 }
4641 for rr in &self.run_refs {
4642 let mut reader = self.open_reader(rr.run_id)?;
4643 for row in reader.visible_versions(snapshot.epoch)? {
4644 fold(row);
4645 }
4646 }
4647 let mut out: Vec<Row> = best
4648 .into_values()
4649 .filter_map(|(_, r)| {
4650 if r.deleted || self.row_expired_at(&r, now_nanos) {
4651 None
4652 } else {
4653 Some(r)
4654 }
4655 })
4656 .collect();
4657 out.sort_by_key(|r| r.row_id);
4658 Ok(out)
4659 }
4660
4661 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4668 if self.ttl.is_none()
4669 && self.memtable.is_empty()
4670 && self.mutable_run.is_empty()
4671 && self.run_refs.len() == 1
4672 {
4673 let rr = self.run_refs[0].clone();
4674 let mut reader = self.open_reader(rr.run_id)?;
4675 let idxs = reader.visible_indices(snapshot.epoch)?;
4676 let mut cols = Vec::with_capacity(self.schema.columns.len());
4677 for cdef in &self.schema.columns {
4678 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4679 }
4680 return Ok(cols);
4681 }
4682 let rows = self.visible_rows(snapshot)?;
4684 let mut cols: Vec<(u16, Vec<Value>)> = self
4685 .schema
4686 .columns
4687 .iter()
4688 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4689 .collect();
4690 for r in &rows {
4691 for (cid, vec) in cols.iter_mut() {
4692 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4693 }
4694 }
4695 Ok(cols)
4696 }
4697
4698 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4700 let row_id = self.hot.get(key)?;
4701 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4702 Some(row_id)
4703 } else {
4704 None
4705 }
4706 }
4707
4708 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4713 self.query_at_with_allowed(q, self.snapshot(), None)
4714 }
4715
4716 pub fn query_controlled(
4719 &mut self,
4720 q: &crate::query::Query,
4721 control: &crate::ExecutionControl,
4722 ) -> Result<Vec<Row>> {
4723 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4724 }
4725
4726 pub fn query_at_with_allowed(
4729 &mut self,
4730 q: &crate::query::Query,
4731 snapshot: Snapshot,
4732 allowed: Option<&std::collections::HashSet<RowId>>,
4733 ) -> Result<Vec<Row>> {
4734 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4735 }
4736
4737 #[doc(hidden)]
4738 pub fn query_at_with_allowed_controlled(
4739 &mut self,
4740 q: &crate::query::Query,
4741 snapshot: Snapshot,
4742 allowed: Option<&std::collections::HashSet<RowId>>,
4743 control: &crate::ExecutionControl,
4744 ) -> Result<Vec<Row>> {
4745 self.require_select()?;
4746 self.ensure_indexes_complete_controlled(control, || true)?;
4747 self.validate_native_query(q)?;
4748 self.query_conditions_at(
4749 &q.conditions,
4750 snapshot,
4751 allowed,
4752 q.limit,
4753 q.offset,
4754 None,
4755 unix_nanos_now(),
4756 Some(control),
4757 )
4758 }
4759
4760 #[doc(hidden)]
4761 pub fn query_at_with_allowed_after(
4762 &mut self,
4763 q: &crate::query::Query,
4764 snapshot: Snapshot,
4765 allowed: Option<&std::collections::HashSet<RowId>>,
4766 after_row_id: Option<RowId>,
4767 ) -> Result<Vec<Row>> {
4768 self.query_at_with_allowed_after_at_time(
4769 q,
4770 snapshot,
4771 allowed,
4772 after_row_id,
4773 unix_nanos_now(),
4774 )
4775 }
4776
4777 #[doc(hidden)]
4778 pub fn query_at_with_allowed_after_at_time(
4779 &mut self,
4780 q: &crate::query::Query,
4781 snapshot: Snapshot,
4782 allowed: Option<&std::collections::HashSet<RowId>>,
4783 after_row_id: Option<RowId>,
4784 query_time_nanos: i64,
4785 ) -> Result<Vec<Row>> {
4786 self.require_select()?;
4787 self.ensure_indexes_complete()?;
4788 self.validate_native_query(q)?;
4789 self.query_conditions_at(
4790 &q.conditions,
4791 snapshot,
4792 allowed,
4793 q.limit,
4794 q.offset,
4795 after_row_id,
4796 query_time_nanos,
4797 None,
4798 )
4799 }
4800
4801 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4802 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4803 return Err(MongrelError::InvalidArgument(format!(
4804 "query exceeds {} conditions",
4805 crate::query::MAX_HARD_CONDITIONS
4806 )));
4807 }
4808 if let Some(limit) = q.limit {
4809 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4810 return Err(MongrelError::InvalidArgument(format!(
4811 "query limit must be between 1 and {}",
4812 crate::query::MAX_FINAL_LIMIT
4813 )));
4814 }
4815 }
4816 if q.offset > crate::query::MAX_QUERY_OFFSET {
4817 return Err(MongrelError::InvalidArgument(format!(
4818 "query offset exceeds {}",
4819 crate::query::MAX_QUERY_OFFSET
4820 )));
4821 }
4822 Ok(())
4823 }
4824
4825 #[doc(hidden)]
4828 pub fn query_all_at(
4829 &mut self,
4830 conditions: &[crate::query::Condition],
4831 snapshot: Snapshot,
4832 ) -> Result<Vec<Row>> {
4833 self.require_select()?;
4834 self.ensure_indexes_complete()?;
4835 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4836 return Err(MongrelError::InvalidArgument(format!(
4837 "query exceeds {} conditions",
4838 crate::query::MAX_HARD_CONDITIONS
4839 )));
4840 }
4841 self.query_conditions_at(
4842 conditions,
4843 snapshot,
4844 None,
4845 None,
4846 0,
4847 None,
4848 unix_nanos_now(),
4849 None,
4850 )
4851 }
4852
4853 #[allow(clippy::too_many_arguments)]
4854 fn query_conditions_at(
4855 &self,
4856 conditions: &[crate::query::Condition],
4857 snapshot: Snapshot,
4858 allowed: Option<&std::collections::HashSet<RowId>>,
4859 limit: Option<usize>,
4860 offset: usize,
4861 after_row_id: Option<RowId>,
4862 query_time_nanos: i64,
4863 control: Option<&crate::ExecutionControl>,
4864 ) -> Result<Vec<Row>> {
4865 control
4866 .map(crate::ExecutionControl::checkpoint)
4867 .transpose()?;
4868 crate::trace::QueryTrace::record(|t| {
4869 t.run_count = self.run_refs.len();
4870 t.memtable_rows = self.memtable.len();
4871 t.mutable_run_rows = self.mutable_run.len();
4872 });
4873 if conditions.is_empty() {
4877 crate::trace::QueryTrace::record(|t| {
4878 t.scan_mode = crate::trace::ScanMode::Materialized;
4879 t.row_materialized = true;
4880 });
4881 let mut rows = match control {
4882 Some(control) => self.visible_rows_controlled(snapshot, control)?,
4883 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
4884 };
4885 if let Some(allowed) = allowed {
4886 let mut filtered = Vec::with_capacity(rows.len());
4887 for (index, row) in rows.into_iter().enumerate() {
4888 if index & 255 == 0 {
4889 control
4890 .map(crate::ExecutionControl::checkpoint)
4891 .transpose()?;
4892 }
4893 if allowed.contains(&row.row_id) {
4894 filtered.push(row);
4895 }
4896 }
4897 rows = filtered;
4898 }
4899 if let Some(after_row_id) = after_row_id {
4900 rows.retain(|row| row.row_id > after_row_id);
4901 }
4902 rows.drain(..offset.min(rows.len()));
4903 if let Some(limit) = limit {
4904 rows.truncate(limit);
4905 }
4906 return Ok(rows);
4907 }
4908 crate::trace::QueryTrace::record(|t| {
4909 t.conditions_pushed = conditions.len();
4910 t.scan_mode = crate::trace::ScanMode::Materialized;
4911 t.row_materialized = true;
4912 });
4913 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
4920 ordered.sort_by_key(|c| condition_cost_rank(c));
4921 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
4922 for c in &ordered {
4923 control
4924 .map(crate::ExecutionControl::checkpoint)
4925 .transpose()?;
4926 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
4927 let empty = s.is_empty();
4928 sets.push(s);
4929 if empty {
4930 break;
4931 }
4932 }
4933 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4934 if let Some(allowed) = allowed {
4935 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
4936 }
4937 if let Some(after_row_id) = after_row_id {
4938 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
4939 rids.drain(..first);
4940 }
4941 rids.drain(..offset.min(rids.len()));
4942 if let Some(limit) = limit {
4943 rids.truncate(limit);
4944 }
4945 control
4946 .map(crate::ExecutionControl::checkpoint)
4947 .transpose()?;
4948 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
4949 }
4950
4951 pub fn retrieve(
4953 &mut self,
4954 retriever: &crate::query::Retriever,
4955 ) -> Result<Vec<crate::query::RetrieverHit>> {
4956 self.retrieve_with_allowed(retriever, None)
4957 }
4958
4959 pub fn retrieve_at(
4960 &mut self,
4961 retriever: &crate::query::Retriever,
4962 snapshot: Snapshot,
4963 allowed: Option<&std::collections::HashSet<RowId>>,
4964 ) -> Result<Vec<crate::query::RetrieverHit>> {
4965 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
4966 }
4967
4968 pub fn retrieve_with_allowed(
4971 &mut self,
4972 retriever: &crate::query::Retriever,
4973 allowed: Option<&std::collections::HashSet<RowId>>,
4974 ) -> Result<Vec<crate::query::RetrieverHit>> {
4975 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
4976 }
4977
4978 pub fn retrieve_at_with_allowed(
4979 &mut self,
4980 retriever: &crate::query::Retriever,
4981 snapshot: Snapshot,
4982 allowed: Option<&std::collections::HashSet<RowId>>,
4983 ) -> Result<Vec<crate::query::RetrieverHit>> {
4984 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
4985 }
4986
4987 pub fn retrieve_at_with_allowed_and_context(
4988 &mut self,
4989 retriever: &crate::query::Retriever,
4990 snapshot: Snapshot,
4991 allowed: Option<&std::collections::HashSet<RowId>>,
4992 context: Option<&crate::query::AiExecutionContext>,
4993 ) -> Result<Vec<crate::query::RetrieverHit>> {
4994 self.require_select()?;
4995 self.ensure_indexes_complete()?;
4996 self.validate_retriever(retriever)?;
4997 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
4998 }
4999
5000 pub fn retrieve_at_with_candidate_authorization_and_context(
5001 &mut self,
5002 retriever: &crate::query::Retriever,
5003 snapshot: Snapshot,
5004 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5005 context: Option<&crate::query::AiExecutionContext>,
5006 ) -> Result<Vec<crate::query::RetrieverHit>> {
5007 self.require_select()?;
5008 self.ensure_indexes_complete()?;
5009 self.retrieve_at_with_candidate_authorization_on_generation(
5010 retriever,
5011 snapshot,
5012 authorization,
5013 context,
5014 )
5015 }
5016
5017 #[doc(hidden)]
5018 pub fn retrieve_at_with_candidate_authorization_on_generation(
5019 &self,
5020 retriever: &crate::query::Retriever,
5021 snapshot: Snapshot,
5022 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5023 context: Option<&crate::query::AiExecutionContext>,
5024 ) -> Result<Vec<crate::query::RetrieverHit>> {
5025 self.require_select()?;
5026 self.validate_retriever(retriever)?;
5027 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5028 }
5029
5030 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5031 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5032 let (column_id, k) = match retriever {
5033 Retriever::Ann {
5034 column_id,
5035 query,
5036 k,
5037 } => {
5038 let index = self.ann.get(column_id).ok_or_else(|| {
5039 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5040 })?;
5041 if query.len() != index.dim() {
5042 return Err(MongrelError::InvalidArgument(format!(
5043 "ANN query dimension must be {}, got {}",
5044 index.dim(),
5045 query.len()
5046 )));
5047 }
5048 if query.iter().any(|value| !value.is_finite()) {
5049 return Err(MongrelError::InvalidArgument(
5050 "ANN query values must be finite".into(),
5051 ));
5052 }
5053 (*column_id, *k)
5054 }
5055 Retriever::Sparse {
5056 column_id,
5057 query,
5058 k,
5059 } => {
5060 if !self.sparse.contains_key(column_id) {
5061 return Err(MongrelError::InvalidArgument(format!(
5062 "column {column_id} has no Sparse index"
5063 )));
5064 }
5065 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5066 return Err(MongrelError::InvalidArgument(
5067 "Sparse query must be non-empty with finite weights".into(),
5068 ));
5069 }
5070 if query.len() > MAX_SPARSE_TERMS {
5071 return Err(MongrelError::InvalidArgument(format!(
5072 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5073 )));
5074 }
5075 (*column_id, *k)
5076 }
5077 Retriever::MinHash {
5078 column_id,
5079 members,
5080 k,
5081 } => {
5082 if !self.minhash.contains_key(column_id) {
5083 return Err(MongrelError::InvalidArgument(format!(
5084 "column {column_id} has no MinHash index"
5085 )));
5086 }
5087 if members.is_empty() {
5088 return Err(MongrelError::InvalidArgument(
5089 "MinHash members must not be empty".into(),
5090 ));
5091 }
5092 if members.len() > MAX_SET_MEMBERS {
5093 return Err(MongrelError::InvalidArgument(format!(
5094 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5095 )));
5096 }
5097 let mut total_bytes = 0usize;
5098 for member in members {
5099 let bytes = member.encoded_len();
5100 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5101 return Err(MongrelError::InvalidArgument(format!(
5102 "MinHash member exceeds {} bytes",
5103 crate::query::MAX_SET_MEMBER_BYTES
5104 )));
5105 }
5106 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5107 MongrelError::InvalidArgument("MinHash input size overflow".into())
5108 })?;
5109 }
5110 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5111 return Err(MongrelError::InvalidArgument(format!(
5112 "MinHash input exceeds {} bytes",
5113 crate::query::MAX_SET_INPUT_BYTES
5114 )));
5115 }
5116 (*column_id, *k)
5117 }
5118 };
5119 if k == 0 {
5120 return Err(MongrelError::InvalidArgument(
5121 "retriever k must be > 0".into(),
5122 ));
5123 }
5124 if k > MAX_RETRIEVER_K {
5125 return Err(MongrelError::InvalidArgument(format!(
5126 "retriever k exceeds {MAX_RETRIEVER_K}"
5127 )));
5128 }
5129 debug_assert!(self
5130 .schema
5131 .columns
5132 .iter()
5133 .any(|column| column.id == column_id));
5134 Ok(())
5135 }
5136
5137 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5138 use crate::query::Condition;
5139 match condition {
5140 Condition::Ann {
5141 column_id,
5142 query,
5143 k,
5144 } => self.validate_retriever(&crate::query::Retriever::Ann {
5145 column_id: *column_id,
5146 query: query.clone(),
5147 k: *k,
5148 }),
5149 Condition::SparseMatch {
5150 column_id,
5151 query,
5152 k,
5153 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5154 column_id: *column_id,
5155 query: query.clone(),
5156 k: *k,
5157 }),
5158 Condition::MinHashSimilar {
5159 column_id,
5160 query,
5161 k,
5162 } => {
5163 if !self.minhash.contains_key(column_id) {
5164 return Err(MongrelError::InvalidArgument(format!(
5165 "column {column_id} has no MinHash index"
5166 )));
5167 }
5168 if query.is_empty() || *k == 0 {
5169 return Err(MongrelError::InvalidArgument(
5170 "MinHash query must be non-empty and k must be > 0".into(),
5171 ));
5172 }
5173 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5174 {
5175 return Err(MongrelError::InvalidArgument(format!(
5176 "MinHash query must have <= {} members and k <= {}",
5177 crate::query::MAX_SET_MEMBERS,
5178 crate::query::MAX_RETRIEVER_K
5179 )));
5180 }
5181 Ok(())
5182 }
5183 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5184 Err(MongrelError::InvalidArgument(format!(
5185 "bitmap IN exceeds {} values",
5186 crate::query::MAX_SET_MEMBERS
5187 )))
5188 }
5189 Condition::FmContainsAll { patterns, .. }
5190 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5191 {
5192 Err(MongrelError::InvalidArgument(format!(
5193 "FM query exceeds {} patterns",
5194 crate::query::MAX_HARD_CONDITIONS
5195 )))
5196 }
5197 _ => Ok(()),
5198 }
5199 }
5200
5201 fn retrieve_filtered(
5202 &self,
5203 retriever: &crate::query::Retriever,
5204 snapshot: Snapshot,
5205 hard_filter: Option<&RowIdSet>,
5206 allowed: Option<&std::collections::HashSet<RowId>>,
5207 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5208 context: Option<&crate::query::AiExecutionContext>,
5209 ) -> Result<Vec<crate::query::RetrieverHit>> {
5210 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5211 let started = std::time::Instant::now();
5212 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5213 Retriever::Ann {
5214 column_id,
5215 query,
5216 k,
5217 } => {
5218 let Some(index) = self.ann.get(column_id) else {
5219 return Ok(Vec::new());
5220 };
5221 let cap = ann_candidate_cap(index.len(), context);
5222 if cap == 0 {
5223 return Ok(Vec::new());
5224 }
5225 let mut breadth = (*k).max(1).min(cap);
5226 let mut eligibility = std::collections::HashMap::new();
5227 let mut filtered = loop {
5228 let mut seen = std::collections::HashSet::new();
5229 if let Some(context) = context {
5230 context.checkpoint()?;
5231 }
5232 let raw = index.search_with_context(query, breadth, context)?;
5233 let unchecked: Vec<_> = raw
5234 .iter()
5235 .map(|(row_id, _)| *row_id)
5236 .filter(|row_id| !eligibility.contains_key(row_id))
5237 .filter(|row_id| {
5238 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5239 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5240 })
5241 .collect();
5242 let eligible = self.eligible_and_authorized_candidate_ids(
5243 &unchecked,
5244 *column_id,
5245 snapshot,
5246 candidate_authorization,
5247 context,
5248 )?;
5249 for row_id in unchecked {
5250 eligibility.insert(row_id, eligible.contains(&row_id));
5251 }
5252 let filtered: Vec<_> = raw
5253 .into_iter()
5254 .filter(|(row_id, _)| {
5255 seen.insert(*row_id)
5256 && eligibility.get(row_id).copied().unwrap_or(false)
5257 })
5258 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
5259 .collect();
5260 if filtered.len() >= *k || breadth >= cap {
5261 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5262 crate::trace::QueryTrace::record(|trace| {
5263 trace.ann_candidate_cap_hit = true;
5264 });
5265 }
5266 break filtered;
5267 }
5268 breadth = breadth.saturating_mul(2).min(cap);
5269 };
5270 filtered.truncate(*k);
5271 filtered
5272 }
5273 Retriever::Sparse {
5274 column_id,
5275 query,
5276 k,
5277 } => self
5278 .sparse
5279 .get(column_id)
5280 .map(|index| -> Result<Vec<_>> {
5281 let mut breadth = (*k).max(1);
5282 let mut eligibility = std::collections::HashMap::new();
5283 loop {
5284 if let Some(context) = context {
5285 context.checkpoint()?;
5286 }
5287 let raw = index.search_with_context(query, breadth, context)?;
5288 let unchecked: Vec<_> = raw
5289 .iter()
5290 .map(|(row_id, _)| *row_id)
5291 .filter(|row_id| !eligibility.contains_key(row_id))
5292 .filter(|row_id| {
5293 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5294 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5295 })
5296 .collect();
5297 let eligible = self.eligible_and_authorized_candidate_ids(
5298 &unchecked,
5299 *column_id,
5300 snapshot,
5301 candidate_authorization,
5302 context,
5303 )?;
5304 for row_id in unchecked {
5305 eligibility.insert(row_id, eligible.contains(&row_id));
5306 }
5307 let filtered: Vec<_> = raw
5308 .iter()
5309 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5310 .take(*k)
5311 .map(|(row_id, score)| {
5312 (*row_id, RetrieverScore::SparseDotProduct(*score))
5313 })
5314 .collect();
5315 if filtered.len() >= *k || raw.len() < breadth {
5316 break Ok(filtered);
5317 }
5318 let next = breadth.saturating_mul(2);
5319 if next == breadth {
5320 break Ok(filtered);
5321 }
5322 breadth = next;
5323 }
5324 })
5325 .transpose()?
5326 .unwrap_or_default(),
5327 Retriever::MinHash {
5328 column_id,
5329 members,
5330 k,
5331 } => self
5332 .minhash
5333 .get(column_id)
5334 .map(|index| -> Result<Vec<_>> {
5335 let mut hashes = Vec::with_capacity(members.len());
5336 for member in members {
5337 if let Some(context) = context {
5338 context.consume(crate::query::work_units(
5339 member.encoded_len(),
5340 crate::query::PARSE_WORK_QUANTUM,
5341 ))?;
5342 }
5343 hashes.push(member.hash_v1());
5344 }
5345 let mut breadth = (*k).max(1);
5346 let mut eligibility = std::collections::HashMap::new();
5347 loop {
5348 if let Some(context) = context {
5349 context.checkpoint()?;
5350 }
5351 let raw = index.search_with_context(&hashes, breadth, context)?;
5352 let unchecked: Vec<_> = raw
5353 .iter()
5354 .map(|(row_id, _)| *row_id)
5355 .filter(|row_id| !eligibility.contains_key(row_id))
5356 .filter(|row_id| {
5357 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5358 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5359 })
5360 .collect();
5361 let eligible = self.eligible_and_authorized_candidate_ids(
5362 &unchecked,
5363 *column_id,
5364 snapshot,
5365 candidate_authorization,
5366 context,
5367 )?;
5368 for row_id in unchecked {
5369 eligibility.insert(row_id, eligible.contains(&row_id));
5370 }
5371 let filtered: Vec<_> = raw
5372 .iter()
5373 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5374 .take(*k)
5375 .map(|(row_id, score)| {
5376 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5377 })
5378 .collect();
5379 if filtered.len() >= *k || raw.len() < breadth {
5380 break Ok(filtered);
5381 }
5382 let next = breadth.saturating_mul(2);
5383 if next == breadth {
5384 break Ok(filtered);
5385 }
5386 breadth = next;
5387 }
5388 })
5389 .transpose()?
5390 .unwrap_or_default(),
5391 };
5392 let elapsed = started.elapsed().as_nanos() as u64;
5393 crate::trace::QueryTrace::record(|trace| {
5394 match retriever {
5395 Retriever::Ann { .. } => {
5396 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5397 }
5398 Retriever::Sparse { .. } => {
5399 trace.sparse_candidate_nanos =
5400 trace.sparse_candidate_nanos.saturating_add(elapsed)
5401 }
5402 Retriever::MinHash { .. } => {
5403 trace.minhash_candidate_nanos =
5404 trace.minhash_candidate_nanos.saturating_add(elapsed)
5405 }
5406 }
5407 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5408 });
5409 Ok(scored
5410 .into_iter()
5411 .enumerate()
5412 .map(|(rank, (row_id, score))| RetrieverHit {
5413 row_id,
5414 rank: rank + 1,
5415 score,
5416 })
5417 .collect())
5418 }
5419
5420 fn eligible_candidate_ids(
5421 &self,
5422 candidates: &[RowId],
5423 _column_id: u16,
5424 snapshot: Snapshot,
5425 context: Option<&crate::query::AiExecutionContext>,
5426 ) -> Result<std::collections::HashSet<RowId>> {
5427 if !self.had_deletes
5428 && self.ttl.is_none()
5429 && self.pending_put_cols.is_empty()
5430 && snapshot.epoch == self.snapshot().epoch
5431 {
5432 return Ok(candidates.iter().copied().collect());
5433 }
5434 let mut readers: Vec<_> = self
5435 .run_refs
5436 .iter()
5437 .map(|run| self.open_reader(run.run_id))
5438 .collect::<Result<_>>()?;
5439 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5440 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5441 for &row_id in candidates {
5442 if let Some(context) = context {
5443 context.consume(1)?;
5444 }
5445 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5446 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5447 let overlay = match (mem, mutable) {
5448 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5449 (Some(value), None) | (None, Some(value)) => Some(value),
5450 (None, None) => None,
5451 };
5452 if let Some((_, row)) = overlay {
5453 if !row.deleted && !self.row_expired_at(&row, now) {
5454 eligible.insert(row_id);
5455 }
5456 continue;
5457 }
5458 let mut best: Option<(Epoch, bool, usize)> = None;
5459 for (index, reader) in readers.iter_mut().enumerate() {
5460 if let Some((epoch, deleted)) =
5461 reader.get_version_visibility(row_id, snapshot.epoch)?
5462 {
5463 if best
5464 .as_ref()
5465 .map(|(best_epoch, ..)| epoch > *best_epoch)
5466 .unwrap_or(true)
5467 {
5468 best = Some((epoch, deleted, index));
5469 }
5470 }
5471 }
5472 let Some((_, false, reader_index)) = best else {
5473 continue;
5474 };
5475 if let Some(ttl) = self.ttl {
5476 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5477 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5478 {
5479 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5480 continue;
5481 }
5482 }
5483 }
5484 eligible.insert(row_id);
5485 }
5486 Ok(eligible)
5487 }
5488
5489 fn eligible_and_authorized_candidate_ids(
5490 &self,
5491 candidates: &[RowId],
5492 column_id: u16,
5493 snapshot: Snapshot,
5494 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5495 context: Option<&crate::query::AiExecutionContext>,
5496 ) -> Result<std::collections::HashSet<RowId>> {
5497 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5498 let Some(authorization) = authorization else {
5499 return Ok(eligible);
5500 };
5501 let candidates: Vec<_> = eligible.into_iter().collect();
5502 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5503 }
5504
5505 fn policy_allowed_candidate_ids(
5506 &self,
5507 candidates: &[RowId],
5508 snapshot: Snapshot,
5509 authorization: &crate::security::CandidateAuthorization<'_>,
5510 context: Option<&crate::query::AiExecutionContext>,
5511 ) -> Result<std::collections::HashSet<RowId>> {
5512 let started = std::time::Instant::now();
5513 if candidates.is_empty()
5514 || authorization.principal.is_admin
5515 || !authorization.security.rls_enabled(authorization.table)
5516 {
5517 return Ok(candidates.iter().copied().collect());
5518 }
5519 if let Some(context) = context {
5520 context.checkpoint()?;
5521 }
5522 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5523 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5524 .iter()
5525 .map(|row_id| {
5526 (
5527 *row_id,
5528 Row {
5529 row_id: *row_id,
5530 committed_epoch: snapshot.epoch,
5531 columns: std::collections::HashMap::new(),
5532 deleted: false,
5533 },
5534 )
5535 })
5536 .collect();
5537 let columns = authorization
5538 .security
5539 .select_policy_columns(authorization.table, authorization.principal);
5540 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5541 let mut decoded = 0usize;
5542 for column_id in &columns {
5543 if let Some(context) = context {
5544 context.checkpoint()?;
5545 }
5546 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5547 &row_ids, *column_id, snapshot, query_now, context,
5548 )? {
5549 if let Some(row) = rows.get_mut(&row_id) {
5550 row.columns.insert(*column_id, value);
5551 decoded = decoded.saturating_add(1);
5552 }
5553 }
5554 }
5555 if let Some(context) = context {
5556 context.consume(candidates.len().saturating_add(decoded))?;
5557 }
5558 let allowed = rows
5559 .into_values()
5560 .filter_map(|row| {
5561 authorization
5562 .security
5563 .row_allowed(
5564 authorization.table,
5565 crate::security::PolicyCommand::Select,
5566 &row,
5567 authorization.principal,
5568 false,
5569 )
5570 .then_some(row.row_id)
5571 })
5572 .collect();
5573 crate::trace::QueryTrace::record(|trace| {
5574 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5575 trace.rls_policy_columns_decoded =
5576 trace.rls_policy_columns_decoded.saturating_add(decoded);
5577 trace.authorization_nanos = trace
5578 .authorization_nanos
5579 .saturating_add(started.elapsed().as_nanos() as u64);
5580 });
5581 Ok(allowed)
5582 }
5583
5584 pub fn search(
5586 &mut self,
5587 request: &crate::query::SearchRequest,
5588 ) -> Result<Vec<crate::query::SearchHit>> {
5589 self.search_with_allowed(request, None)
5590 }
5591
5592 pub fn search_at(
5593 &mut self,
5594 request: &crate::query::SearchRequest,
5595 snapshot: Snapshot,
5596 authorized: Option<&std::collections::HashSet<RowId>>,
5597 ) -> Result<Vec<crate::query::SearchHit>> {
5598 self.search_at_with_allowed(request, snapshot, authorized)
5599 }
5600
5601 pub fn search_with_allowed(
5602 &mut self,
5603 request: &crate::query::SearchRequest,
5604 authorized: Option<&std::collections::HashSet<RowId>>,
5605 ) -> Result<Vec<crate::query::SearchHit>> {
5606 self.search_at_with_allowed(request, self.snapshot(), authorized)
5607 }
5608
5609 pub fn search_at_with_allowed(
5610 &mut self,
5611 request: &crate::query::SearchRequest,
5612 snapshot: Snapshot,
5613 authorized: Option<&std::collections::HashSet<RowId>>,
5614 ) -> Result<Vec<crate::query::SearchHit>> {
5615 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5616 }
5617
5618 pub fn search_at_with_allowed_and_context(
5619 &mut self,
5620 request: &crate::query::SearchRequest,
5621 snapshot: Snapshot,
5622 authorized: Option<&std::collections::HashSet<RowId>>,
5623 context: Option<&crate::query::AiExecutionContext>,
5624 ) -> Result<Vec<crate::query::SearchHit>> {
5625 self.ensure_indexes_complete()?;
5626 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5627 }
5628
5629 pub fn search_at_with_candidate_authorization_and_context(
5630 &mut self,
5631 request: &crate::query::SearchRequest,
5632 snapshot: Snapshot,
5633 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5634 context: Option<&crate::query::AiExecutionContext>,
5635 ) -> Result<Vec<crate::query::SearchHit>> {
5636 self.ensure_indexes_complete()?;
5637 self.search_at_with_filters_and_context(
5638 request,
5639 snapshot,
5640 None,
5641 authorization,
5642 context,
5643 None,
5644 )
5645 }
5646
5647 #[doc(hidden)]
5648 pub fn search_at_with_candidate_authorization_on_generation(
5649 &self,
5650 request: &crate::query::SearchRequest,
5651 snapshot: Snapshot,
5652 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5653 context: Option<&crate::query::AiExecutionContext>,
5654 ) -> Result<Vec<crate::query::SearchHit>> {
5655 self.search_at_with_filters_and_context(
5656 request,
5657 snapshot,
5658 None,
5659 authorization,
5660 context,
5661 None,
5662 )
5663 }
5664
5665 #[doc(hidden)]
5666 pub fn search_at_with_candidate_authorization_on_generation_after(
5667 &self,
5668 request: &crate::query::SearchRequest,
5669 snapshot: Snapshot,
5670 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5671 context: Option<&crate::query::AiExecutionContext>,
5672 after: Option<crate::query::SearchAfter>,
5673 ) -> Result<Vec<crate::query::SearchHit>> {
5674 self.search_at_with_filters_and_context(
5675 request,
5676 snapshot,
5677 None,
5678 authorization,
5679 context,
5680 after,
5681 )
5682 }
5683
5684 fn search_at_with_filters_and_context(
5685 &self,
5686 request: &crate::query::SearchRequest,
5687 snapshot: Snapshot,
5688 authorized: Option<&std::collections::HashSet<RowId>>,
5689 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5690 context: Option<&crate::query::AiExecutionContext>,
5691 after: Option<crate::query::SearchAfter>,
5692 ) -> Result<Vec<crate::query::SearchHit>> {
5693 use crate::query::{
5694 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5695 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5696 };
5697 let total_started = std::time::Instant::now();
5698 let rank_offset = after.map_or(0, |after| after.returned_count);
5699 self.require_select()?;
5700 if request.limit == 0 {
5701 return Err(MongrelError::InvalidArgument(
5702 "search limit must be > 0".into(),
5703 ));
5704 }
5705 if request.limit > MAX_FINAL_LIMIT {
5706 return Err(MongrelError::InvalidArgument(format!(
5707 "search limit exceeds {MAX_FINAL_LIMIT}"
5708 )));
5709 }
5710 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5711 return Err(MongrelError::InvalidArgument(
5712 "search-after score must be finite".into(),
5713 ));
5714 }
5715 if request.retrievers.is_empty() {
5716 return Err(MongrelError::InvalidArgument(
5717 "search requires at least one retriever".into(),
5718 ));
5719 }
5720 if request.retrievers.len() > MAX_RETRIEVERS {
5721 return Err(MongrelError::InvalidArgument(format!(
5722 "search exceeds {MAX_RETRIEVERS} retrievers"
5723 )));
5724 }
5725 if request.must.len() > MAX_HARD_CONDITIONS {
5726 return Err(MongrelError::InvalidArgument(format!(
5727 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5728 )));
5729 }
5730 for condition in &request.must {
5731 self.validate_condition(condition)?;
5732 }
5733 if request.must.iter().any(|condition| {
5734 matches!(
5735 condition,
5736 Condition::Ann { .. }
5737 | Condition::SparseMatch { .. }
5738 | Condition::MinHashSimilar { .. }
5739 )
5740 }) {
5741 return Err(MongrelError::InvalidArgument(
5742 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5743 .into(),
5744 ));
5745 }
5746 let mut names = std::collections::HashSet::new();
5747 for named in &request.retrievers {
5748 if named.name.is_empty()
5749 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5750 || !names.insert(named.name.as_str())
5751 {
5752 return Err(MongrelError::InvalidArgument(format!(
5753 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5754 crate::query::MAX_RETRIEVER_NAME_BYTES
5755 )));
5756 }
5757 if !named.weight.is_finite()
5758 || named.weight < 0.0
5759 || named.weight > MAX_RETRIEVER_WEIGHT
5760 {
5761 return Err(MongrelError::InvalidArgument(format!(
5762 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5763 )));
5764 }
5765 self.validate_retriever(&named.retriever)?;
5766 }
5767 let projection = request
5768 .projection
5769 .clone()
5770 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5771 if projection.len() > MAX_PROJECTION_COLUMNS {
5772 return Err(MongrelError::InvalidArgument(format!(
5773 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5774 )));
5775 }
5776 for column_id in &projection {
5777 if !self
5778 .schema
5779 .columns
5780 .iter()
5781 .any(|column| column.id == *column_id)
5782 {
5783 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5784 }
5785 }
5786 if let Some(crate::query::Rerank::ExactVector {
5787 embedding_column,
5788 query,
5789 candidate_limit,
5790 weight,
5791 ..
5792 }) = &request.rerank
5793 {
5794 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5795 {
5796 return Err(MongrelError::InvalidArgument(format!(
5797 "rerank candidate_limit must be between search limit and {}",
5798 crate::query::MAX_RETRIEVER_K
5799 )));
5800 }
5801 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5802 return Err(MongrelError::InvalidArgument(format!(
5803 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5804 )));
5805 }
5806 let column = self
5807 .schema
5808 .columns
5809 .iter()
5810 .find(|column| column.id == *embedding_column)
5811 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5812 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5813 return Err(MongrelError::InvalidArgument(format!(
5814 "rerank column {embedding_column} is not an embedding"
5815 )));
5816 };
5817 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5818 return Err(MongrelError::InvalidArgument(format!(
5819 "rerank query must contain {dim} finite values"
5820 )));
5821 }
5822 }
5823
5824 let hard_filter_started = std::time::Instant::now();
5825 let hard_filter = if request.must.is_empty() {
5826 None
5827 } else {
5828 let mut sets = Vec::with_capacity(request.must.len());
5829 for condition in &request.must {
5830 if let Some(context) = context {
5831 context.checkpoint()?;
5832 }
5833 sets.push(self.resolve_condition(condition, snapshot)?);
5834 }
5835 Some(RowIdSet::intersect_many(sets))
5836 };
5837 crate::trace::QueryTrace::record(|trace| {
5838 trace.hard_filter_nanos = trace
5839 .hard_filter_nanos
5840 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
5841 });
5842 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
5843 return Ok(Vec::new());
5844 }
5845
5846 let constant = match request.fusion {
5847 Fusion::ReciprocalRank { constant } => constant,
5848 };
5849 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
5850 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
5851 let mut fusion_nanos = 0u64;
5852 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
5853 std::collections::HashMap::new();
5854 for named in retrievers {
5855 if named.weight == 0.0 {
5856 continue;
5857 }
5858 if let Some(context) = context {
5859 context.checkpoint()?;
5860 }
5861 let hits = self.retrieve_filtered(
5862 &named.retriever,
5863 snapshot,
5864 hard_filter.as_ref(),
5865 authorized,
5866 candidate_authorization,
5867 context,
5868 )?;
5869 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
5870 let fusion_started = std::time::Instant::now();
5871 for hit in hits {
5872 if let Some(context) = context {
5873 context.consume(1)?;
5874 }
5875 let contribution = named.weight / (constant as f64 + hit.rank as f64);
5876 if !contribution.is_finite() {
5877 return Err(MongrelError::InvalidArgument(
5878 "retriever contribution must be finite".into(),
5879 ));
5880 }
5881 let max_fused_candidates = context.map_or(
5882 crate::query::MAX_FUSED_CANDIDATES,
5883 crate::query::AiExecutionContext::max_fused_candidates,
5884 );
5885 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
5886 return Err(MongrelError::WorkBudgetExceeded);
5887 }
5888 let entry = fused.entry(hit.row_id).or_default();
5889 entry.0 += contribution;
5890 if !entry.0.is_finite() {
5891 return Err(MongrelError::InvalidArgument(
5892 "fused score must be finite".into(),
5893 ));
5894 }
5895 entry.1.push(ComponentScore {
5896 retriever_name: retriever_name.clone(),
5897 rank: hit.rank,
5898 raw_score: hit.score,
5899 contribution,
5900 });
5901 }
5902 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
5903 }
5904 let union_size = fused.len();
5905 let mut ranked: Vec<_> = fused
5906 .into_iter()
5907 .map(|(row_id, (fused_score, components))| {
5908 (row_id, fused_score, components, None, fused_score)
5909 })
5910 .collect();
5911 let order = |(a_row, _, _, _, a_score): &(
5912 RowId,
5913 f64,
5914 Vec<ComponentScore>,
5915 Option<f32>,
5916 f64,
5917 ),
5918 (b_row, _, _, _, b_score): &(
5919 RowId,
5920 f64,
5921 Vec<ComponentScore>,
5922 Option<f32>,
5923 f64,
5924 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
5925 if let Some(crate::query::Rerank::ExactVector {
5926 embedding_column,
5927 query,
5928 metric,
5929 candidate_limit,
5930 weight,
5931 }) = &request.rerank
5932 {
5933 let fused_order = |(a_row, a_score, ..): &(
5934 RowId,
5935 f64,
5936 Vec<ComponentScore>,
5937 Option<f32>,
5938 f64,
5939 ),
5940 (b_row, b_score, ..): &(
5941 RowId,
5942 f64,
5943 Vec<ComponentScore>,
5944 Option<f32>,
5945 f64,
5946 )| {
5947 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
5948 };
5949 let selection_started = std::time::Instant::now();
5950 if let Some(context) = context {
5951 context.consume(ranked.len())?;
5952 }
5953 if ranked.len() > *candidate_limit {
5954 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
5955 ranked.truncate(*candidate_limit);
5956 }
5957 ranked.sort_by(fused_order);
5958 fusion_nanos =
5959 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5960 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5961 if let Some(context) = context {
5962 context.consume(row_ids.len())?;
5963 }
5964 let query_now =
5965 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5966 let gather_started = std::time::Instant::now();
5967 let vectors = self.values_for_rids_batch_at_with_context(
5968 &row_ids,
5969 *embedding_column,
5970 snapshot,
5971 query_now,
5972 context,
5973 )?;
5974 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5975 let vector_work =
5976 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
5977 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
5978 if let Some(context) = context {
5979 context.consume(vector_work)?;
5980 }
5981 query
5982 .iter()
5983 .map(|value| f64::from(*value).powi(2))
5984 .sum::<f64>()
5985 .sqrt()
5986 } else {
5987 0.0
5988 };
5989 let score_started = std::time::Instant::now();
5990 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
5991 for (row_id, value) in vectors {
5992 let Value::Embedding(vector) = value else {
5993 continue;
5994 };
5995 let score = match metric {
5996 crate::query::VectorMetric::DotProduct => {
5997 if let Some(context) = context {
5998 context.consume(vector_work)?;
5999 }
6000 query
6001 .iter()
6002 .zip(&vector)
6003 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6004 .sum::<f64>()
6005 }
6006 crate::query::VectorMetric::Cosine => {
6007 if let Some(context) = context {
6008 context.consume(vector_work.saturating_mul(2))?;
6009 }
6010 let dot = query
6011 .iter()
6012 .zip(&vector)
6013 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6014 .sum::<f64>();
6015 let norm = vector
6016 .iter()
6017 .map(|value| f64::from(*value).powi(2))
6018 .sum::<f64>()
6019 .sqrt();
6020 if query_norm == 0.0 || norm == 0.0 {
6021 0.0
6022 } else {
6023 dot / (query_norm * norm)
6024 }
6025 }
6026 crate::query::VectorMetric::Euclidean => {
6027 if let Some(context) = context {
6028 context.consume(vector_work)?;
6029 }
6030 query
6031 .iter()
6032 .zip(&vector)
6033 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6034 .sum::<f64>()
6035 .sqrt()
6036 }
6037 };
6038 if !score.is_finite() {
6039 return Err(MongrelError::InvalidArgument(
6040 "exact rerank score must be finite".into(),
6041 ));
6042 }
6043 scores.insert(row_id, score as f32);
6044 }
6045 let mut reranked = Vec::with_capacity(ranked.len());
6046 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6047 let Some(score) = scores.get(&row_id).copied() else {
6048 continue;
6049 };
6050 let ordering_score = match metric {
6051 crate::query::VectorMetric::Euclidean => -f64::from(score),
6052 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6053 f64::from(score)
6054 }
6055 };
6056 let final_score = fused_score + *weight * ordering_score;
6057 if !final_score.is_finite() {
6058 return Err(MongrelError::InvalidArgument(
6059 "final rerank score must be finite".into(),
6060 ));
6061 }
6062 reranked.push((row_id, fused_score, components, Some(score), final_score));
6063 }
6064 ranked = reranked;
6065 ranked.sort_by(order);
6066 crate::trace::QueryTrace::record(|trace| {
6067 trace.exact_vector_gather_nanos =
6068 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6069 trace.exact_vector_score_nanos = trace
6070 .exact_vector_score_nanos
6071 .saturating_add(score_started.elapsed().as_nanos() as u64);
6072 });
6073 }
6074 if let Some(after) = after {
6075 ranked.retain(|(row_id, _, _, _, final_score)| {
6076 final_score.total_cmp(&after.final_score).is_lt()
6077 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6078 });
6079 }
6080 let projection_started = std::time::Instant::now();
6081 let sentinel = projection
6082 .first()
6083 .copied()
6084 .or_else(|| self.schema.columns.first().map(|column| column.id));
6085 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6086 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6087 let mut projection_rows = 0usize;
6088 let mut projection_cells = 0usize;
6089 while out.len() < request.limit && !ranked.is_empty() {
6090 if let Some(context) = context {
6091 context.checkpoint()?;
6092 context.consume(ranked.len())?;
6093 }
6094 let needed = request.limit - out.len();
6095 let window_size = ranked
6096 .len()
6097 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6098 let selection_started = std::time::Instant::now();
6099 let mut remainder = if ranked.len() > window_size {
6100 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6101 ranked.split_off(window_size)
6102 } else {
6103 Vec::new()
6104 };
6105 ranked.sort_by(order);
6106 fusion_nanos =
6107 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6108 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6109 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6110 if let Some(context) = context {
6111 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6112 }
6113 projection_rows = projection_rows.saturating_add(row_ids.len());
6114 projection_cells =
6115 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6116 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6117 std::collections::HashMap::new();
6118 if let Some(column_id) = sentinel {
6119 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6120 &row_ids, column_id, snapshot, query_now, context,
6121 )? {
6122 cells.entry(row_id).or_default().insert(column_id, value);
6123 }
6124 }
6125 for &column_id in &projection {
6126 if Some(column_id) == sentinel {
6127 continue;
6128 }
6129 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6130 &row_ids, column_id, snapshot, query_now, context,
6131 )? {
6132 cells.entry(row_id).or_default().insert(column_id, value);
6133 }
6134 }
6135 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6136 ranked.drain(..)
6137 {
6138 let Some(row_cells) = cells.remove(&row_id) else {
6139 continue;
6140 };
6141 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6142 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6143 out.push(SearchHit {
6144 row_id,
6145 cells: projection
6146 .iter()
6147 .filter_map(|column_id| {
6148 row_cells
6149 .get(column_id)
6150 .cloned()
6151 .map(|value| (*column_id, value))
6152 })
6153 .collect(),
6154 components,
6155 fused_score,
6156 exact_rerank_score,
6157 final_score,
6158 final_rank,
6159 });
6160 if out.len() == request.limit {
6161 break;
6162 }
6163 }
6164 ranked.append(&mut remainder);
6165 }
6166 crate::trace::QueryTrace::record(|trace| {
6167 trace.union_size = union_size;
6168 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6169 trace.projection_nanos = trace
6170 .projection_nanos
6171 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6172 trace.total_nanos = trace
6173 .total_nanos
6174 .saturating_add(total_started.elapsed().as_nanos() as u64);
6175 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6176 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6177 if let Some(context) = context {
6178 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6179 }
6180 });
6181 Ok(out)
6182 }
6183
6184 pub fn set_similarity(
6187 &mut self,
6188 request: &crate::query::SetSimilarityRequest,
6189 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6190 self.set_similarity_with_allowed(request, None)
6191 }
6192
6193 pub fn set_similarity_at(
6194 &mut self,
6195 request: &crate::query::SetSimilarityRequest,
6196 snapshot: Snapshot,
6197 allowed: Option<&std::collections::HashSet<RowId>>,
6198 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6199 self.set_similarity_explained_at(request, snapshot, allowed)
6200 .map(|(hits, _)| hits)
6201 }
6202
6203 pub fn ann_rerank(
6205 &mut self,
6206 request: &crate::query::AnnRerankRequest,
6207 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6208 self.ann_rerank_with_allowed(request, None)
6209 }
6210
6211 pub fn ann_rerank_with_allowed(
6212 &mut self,
6213 request: &crate::query::AnnRerankRequest,
6214 allowed: Option<&std::collections::HashSet<RowId>>,
6215 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6216 self.ann_rerank_at(request, self.snapshot(), allowed)
6217 }
6218
6219 pub fn ann_rerank_at(
6220 &mut self,
6221 request: &crate::query::AnnRerankRequest,
6222 snapshot: Snapshot,
6223 allowed: Option<&std::collections::HashSet<RowId>>,
6224 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6225 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6226 }
6227
6228 pub fn ann_rerank_at_with_context(
6229 &mut self,
6230 request: &crate::query::AnnRerankRequest,
6231 snapshot: Snapshot,
6232 allowed: Option<&std::collections::HashSet<RowId>>,
6233 context: Option<&crate::query::AiExecutionContext>,
6234 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6235 self.ensure_indexes_complete()?;
6236 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6237 }
6238
6239 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6240 &mut self,
6241 request: &crate::query::AnnRerankRequest,
6242 snapshot: Snapshot,
6243 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6244 context: Option<&crate::query::AiExecutionContext>,
6245 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6246 self.ensure_indexes_complete()?;
6247 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6248 }
6249
6250 #[doc(hidden)]
6251 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6252 &self,
6253 request: &crate::query::AnnRerankRequest,
6254 snapshot: Snapshot,
6255 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6256 context: Option<&crate::query::AiExecutionContext>,
6257 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6258 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6259 }
6260
6261 fn ann_rerank_at_with_filters_and_context(
6262 &self,
6263 request: &crate::query::AnnRerankRequest,
6264 snapshot: Snapshot,
6265 allowed: Option<&std::collections::HashSet<RowId>>,
6266 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6267 context: Option<&crate::query::AiExecutionContext>,
6268 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6269 use crate::query::{
6270 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6271 };
6272 if request.candidate_k == 0 || request.limit == 0 {
6273 return Err(MongrelError::InvalidArgument(
6274 "candidate_k and limit must be > 0".into(),
6275 ));
6276 }
6277 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6278 return Err(MongrelError::InvalidArgument(format!(
6279 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6280 )));
6281 }
6282 let retriever = Retriever::Ann {
6283 column_id: request.column_id,
6284 query: request.query.clone(),
6285 k: request.candidate_k,
6286 };
6287 self.require_select()?;
6288 self.validate_retriever(&retriever)?;
6289 let hits = self.retrieve_filtered(
6290 &retriever,
6291 snapshot,
6292 None,
6293 allowed,
6294 candidate_authorization,
6295 context,
6296 )?;
6297 let distances: std::collections::HashMap<_, _> = hits
6298 .iter()
6299 .filter_map(|hit| match hit.score {
6300 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
6301 _ => None,
6302 })
6303 .collect();
6304 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6305 if let Some(context) = context {
6306 context.consume(row_ids.len())?;
6307 }
6308 let gather_started = std::time::Instant::now();
6309 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6310 let values = self.values_for_rids_batch_at_with_context(
6311 &row_ids,
6312 request.column_id,
6313 snapshot,
6314 query_now,
6315 context,
6316 )?;
6317 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6318 let score_started = std::time::Instant::now();
6319 let vector_work =
6320 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6321 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6322 if let Some(context) = context {
6323 context.consume(vector_work)?;
6324 }
6325 request
6326 .query
6327 .iter()
6328 .map(|value| f64::from(*value).powi(2))
6329 .sum::<f64>()
6330 .sqrt()
6331 } else {
6332 0.0
6333 };
6334 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6335 for (row_id, value) in values {
6336 let Value::Embedding(vector) = value else {
6337 continue;
6338 };
6339 let exact_score = match request.metric {
6340 VectorMetric::DotProduct => {
6341 if let Some(context) = context {
6342 context.consume(vector_work)?;
6343 }
6344 request
6345 .query
6346 .iter()
6347 .zip(&vector)
6348 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6349 .sum::<f64>()
6350 }
6351 VectorMetric::Cosine => {
6352 if let Some(context) = context {
6353 context.consume(vector_work.saturating_mul(2))?;
6354 }
6355 let dot = request
6356 .query
6357 .iter()
6358 .zip(&vector)
6359 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6360 .sum::<f64>();
6361 let norm = vector
6362 .iter()
6363 .map(|value| f64::from(*value).powi(2))
6364 .sum::<f64>()
6365 .sqrt();
6366 if query_norm == 0.0 || norm == 0.0 {
6367 0.0
6368 } else {
6369 dot / (query_norm * norm)
6370 }
6371 }
6372 VectorMetric::Euclidean => {
6373 if let Some(context) = context {
6374 context.consume(vector_work)?;
6375 }
6376 request
6377 .query
6378 .iter()
6379 .zip(&vector)
6380 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6381 .sum::<f64>()
6382 .sqrt()
6383 }
6384 };
6385 let exact_score = exact_score as f32;
6386 if !exact_score.is_finite() {
6387 return Err(MongrelError::InvalidArgument(
6388 "exact ANN score must be finite".into(),
6389 ));
6390 }
6391 reranked.push(AnnRerankHit {
6392 row_id,
6393 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
6394 exact_score,
6395 });
6396 }
6397 reranked.sort_by(|left, right| {
6398 let score = match request.metric {
6399 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6400 VectorMetric::Cosine | VectorMetric::DotProduct => {
6401 right.exact_score.total_cmp(&left.exact_score)
6402 }
6403 };
6404 score.then_with(|| left.row_id.cmp(&right.row_id))
6405 });
6406 reranked.truncate(request.limit);
6407 crate::trace::QueryTrace::record(|trace| {
6408 trace.exact_vector_gather_nanos =
6409 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6410 trace.exact_vector_score_nanos = trace
6411 .exact_vector_score_nanos
6412 .saturating_add(score_started.elapsed().as_nanos() as u64);
6413 });
6414 Ok(reranked)
6415 }
6416
6417 pub fn set_similarity_with_allowed(
6418 &mut self,
6419 request: &crate::query::SetSimilarityRequest,
6420 allowed: Option<&std::collections::HashSet<RowId>>,
6421 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6422 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6423 .map(|(hits, _)| hits)
6424 }
6425
6426 pub fn set_similarity_explained(
6427 &mut self,
6428 request: &crate::query::SetSimilarityRequest,
6429 ) -> Result<(
6430 Vec<crate::query::SetSimilarityHit>,
6431 crate::query::SetSimilarityTrace,
6432 )> {
6433 self.set_similarity_explained_at(request, self.snapshot(), None)
6434 }
6435
6436 fn set_similarity_explained_at(
6437 &mut self,
6438 request: &crate::query::SetSimilarityRequest,
6439 snapshot: Snapshot,
6440 allowed: Option<&std::collections::HashSet<RowId>>,
6441 ) -> Result<(
6442 Vec<crate::query::SetSimilarityHit>,
6443 crate::query::SetSimilarityTrace,
6444 )> {
6445 self.ensure_indexes_complete()?;
6446 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6447 }
6448
6449 pub fn set_similarity_at_with_context(
6450 &mut self,
6451 request: &crate::query::SetSimilarityRequest,
6452 snapshot: Snapshot,
6453 allowed: Option<&std::collections::HashSet<RowId>>,
6454 context: Option<&crate::query::AiExecutionContext>,
6455 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6456 self.ensure_indexes_complete()?;
6457 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6458 .map(|(hits, _)| hits)
6459 }
6460
6461 pub fn set_similarity_at_with_candidate_authorization_and_context(
6462 &mut self,
6463 request: &crate::query::SetSimilarityRequest,
6464 snapshot: Snapshot,
6465 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6466 context: Option<&crate::query::AiExecutionContext>,
6467 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6468 self.ensure_indexes_complete()?;
6469 self.set_similarity_explained_at_with_context(
6470 request,
6471 snapshot,
6472 None,
6473 authorization,
6474 context,
6475 )
6476 .map(|(hits, _)| hits)
6477 }
6478
6479 #[doc(hidden)]
6480 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6481 &self,
6482 request: &crate::query::SetSimilarityRequest,
6483 snapshot: Snapshot,
6484 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6485 context: Option<&crate::query::AiExecutionContext>,
6486 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6487 self.set_similarity_explained_at_with_context(
6488 request,
6489 snapshot,
6490 None,
6491 authorization,
6492 context,
6493 )
6494 .map(|(hits, _)| hits)
6495 }
6496
6497 fn set_similarity_explained_at_with_context(
6498 &self,
6499 request: &crate::query::SetSimilarityRequest,
6500 snapshot: Snapshot,
6501 allowed: Option<&std::collections::HashSet<RowId>>,
6502 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6503 context: Option<&crate::query::AiExecutionContext>,
6504 ) -> Result<(
6505 Vec<crate::query::SetSimilarityHit>,
6506 crate::query::SetSimilarityTrace,
6507 )> {
6508 use crate::query::{
6509 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6510 MAX_SET_MEMBERS,
6511 };
6512 let mut trace = crate::query::SetSimilarityTrace::default();
6513 if request.members.is_empty() {
6514 return Ok((Vec::new(), trace));
6515 }
6516 if request.candidate_k == 0 || request.limit == 0 {
6517 return Err(MongrelError::InvalidArgument(
6518 "candidate_k and limit must be > 0".into(),
6519 ));
6520 }
6521 if request.candidate_k > MAX_RETRIEVER_K
6522 || request.limit > MAX_FINAL_LIMIT
6523 || request.members.len() > MAX_SET_MEMBERS
6524 {
6525 return Err(MongrelError::InvalidArgument(format!(
6526 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6527 )));
6528 }
6529 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6530 return Err(MongrelError::InvalidArgument(
6531 "min_jaccard must be finite and between 0 and 1".into(),
6532 ));
6533 }
6534 let started = std::time::Instant::now();
6535 let retriever = Retriever::MinHash {
6536 column_id: request.column_id,
6537 members: request.members.clone(),
6538 k: request.candidate_k,
6539 };
6540 self.require_select()?;
6541 self.validate_retriever(&retriever)?;
6542 let hits = self.retrieve_filtered(
6543 &retriever,
6544 snapshot,
6545 None,
6546 allowed,
6547 candidate_authorization,
6548 context,
6549 )?;
6550 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6551 trace.candidate_count = hits.len();
6552 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6553 if let Some(context) = context {
6554 context.consume(row_ids.len())?;
6555 }
6556 let started = std::time::Instant::now();
6557 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6558 let values = self.values_for_rids_batch_at_with_context(
6559 &row_ids,
6560 request.column_id,
6561 snapshot,
6562 query_now,
6563 context,
6564 )?;
6565 trace.gather_us = started.elapsed().as_micros() as u64;
6566 if let Some(context) = context {
6567 context.consume(request.members.len())?;
6568 }
6569 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6570 let estimates: std::collections::HashMap<_, _> = hits
6571 .into_iter()
6572 .filter_map(|hit| match hit.score {
6573 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6574 _ => None,
6575 })
6576 .collect();
6577 let started = std::time::Instant::now();
6578 let mut parsed = Vec::with_capacity(values.len());
6579 for (row_id, value) in values {
6580 let Value::Bytes(bytes) = value else {
6581 continue;
6582 };
6583 if let Some(context) = context {
6584 context.consume(crate::query::work_units(
6585 bytes.len(),
6586 crate::query::PARSE_WORK_QUANTUM,
6587 ))?;
6588 }
6589 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6590 continue;
6591 };
6592 if let Some(context) = context {
6593 context.consume(members.len())?;
6594 }
6595 let stored = members
6596 .into_iter()
6597 .filter_map(|member| match member {
6598 serde_json::Value::String(value) => {
6599 Some(crate::query::SetMember::String(value))
6600 }
6601 serde_json::Value::Number(value) => {
6602 Some(crate::query::SetMember::Number(value))
6603 }
6604 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6605 _ => None,
6606 })
6607 .collect::<std::collections::HashSet<_>>();
6608 parsed.push((row_id, stored));
6609 }
6610 trace.parse_us = started.elapsed().as_micros() as u64;
6611 trace.verified_count = parsed.len();
6612 let started = std::time::Instant::now();
6613 let mut exact = Vec::new();
6614 for (row_id, stored) in parsed {
6615 if let Some(context) = context {
6616 context.consume(query.len().saturating_add(stored.len()))?;
6617 }
6618 let union = query.union(&stored).count();
6619 let score = if union == 0 {
6620 1.0
6621 } else {
6622 query.intersection(&stored).count() as f32 / union as f32
6623 };
6624 if score >= request.min_jaccard {
6625 exact.push(SetSimilarityHit {
6626 row_id,
6627 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6628 exact_jaccard: score,
6629 });
6630 }
6631 }
6632 exact.sort_by(|a, b| {
6633 b.exact_jaccard
6634 .total_cmp(&a.exact_jaccard)
6635 .then_with(|| a.row_id.cmp(&b.row_id))
6636 });
6637 exact.truncate(request.limit);
6638 trace.score_us = started.elapsed().as_micros() as u64;
6639 crate::trace::QueryTrace::record(|query_trace| {
6640 query_trace.exact_set_gather_nanos = query_trace
6641 .exact_set_gather_nanos
6642 .saturating_add(trace.gather_us.saturating_mul(1_000));
6643 query_trace.exact_set_parse_nanos = query_trace
6644 .exact_set_parse_nanos
6645 .saturating_add(trace.parse_us.saturating_mul(1_000));
6646 query_trace.exact_set_score_nanos = query_trace
6647 .exact_set_score_nanos
6648 .saturating_add(trace.score_us.saturating_mul(1_000));
6649 });
6650 Ok((exact, trace))
6651 }
6652
6653 fn values_for_rids_batch_at(
6655 &self,
6656 row_ids: &[u64],
6657 column_id: u16,
6658 snapshot: Snapshot,
6659 now: i64,
6660 ) -> Result<Vec<(RowId, Value)>> {
6661 if self.ttl.is_none()
6662 && self.memtable.is_empty()
6663 && self.mutable_run.is_empty()
6664 && self.run_refs.len() == 1
6665 {
6666 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6667 if row_ids.len().saturating_mul(24) < reader.row_count() {
6672 let mut values = Vec::with_capacity(row_ids.len());
6673 for &raw_row_id in row_ids {
6674 let row_id = RowId(raw_row_id);
6675 if let Some((_, false, Some(value))) =
6676 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6677 {
6678 values.push((row_id, value));
6679 }
6680 }
6681 return Ok(values);
6682 }
6683 let (positions, visible_row_ids) =
6684 reader.visible_positions_with_rids(snapshot.epoch)?;
6685 let requested: Vec<(RowId, usize)> = row_ids
6686 .iter()
6687 .filter_map(|raw| {
6688 visible_row_ids
6689 .binary_search(&(*raw as i64))
6690 .ok()
6691 .map(|index| (RowId(*raw), positions[index]))
6692 })
6693 .collect();
6694 let values = reader.gather_column(
6695 column_id,
6696 &requested
6697 .iter()
6698 .map(|(_, position)| *position)
6699 .collect::<Vec<_>>(),
6700 )?;
6701 return Ok(requested
6702 .into_iter()
6703 .zip(values)
6704 .map(|((row_id, _), value)| (row_id, value))
6705 .collect());
6706 }
6707 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6708 }
6709
6710 fn values_for_rids_batch_at_with_context(
6711 &self,
6712 row_ids: &[u64],
6713 column_id: u16,
6714 snapshot: Snapshot,
6715 now: i64,
6716 context: Option<&crate::query::AiExecutionContext>,
6717 ) -> Result<Vec<(RowId, Value)>> {
6718 let Some(context) = context else {
6719 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6720 };
6721 let mut values = Vec::with_capacity(row_ids.len());
6722 for chunk in row_ids.chunks(256) {
6723 context.checkpoint()?;
6724 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6725 }
6726 Ok(values)
6727 }
6728
6729 fn values_for_rids_at(
6731 &self,
6732 row_ids: &[u64],
6733 column_id: u16,
6734 snapshot: Snapshot,
6735 now: i64,
6736 ) -> Result<Vec<(RowId, Value)>> {
6737 let mut readers: Vec<_> = self
6738 .run_refs
6739 .iter()
6740 .map(|run| self.open_reader(run.run_id))
6741 .collect::<Result<_>>()?;
6742 let mut out = Vec::with_capacity(row_ids.len());
6743 for &raw_row_id in row_ids {
6744 let row_id = RowId(raw_row_id);
6745 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6746 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6747 let overlay = match (mem, mutable) {
6748 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6749 (a_epoch, a)
6750 } else {
6751 (b_epoch, b)
6752 }),
6753 (Some(value), None) | (None, Some(value)) => Some(value),
6754 (None, None) => None,
6755 };
6756 if let Some((_, row)) = overlay {
6757 if !row.deleted && !self.row_expired_at(&row, now) {
6758 if let Some(value) = row.columns.get(&column_id) {
6759 out.push((row_id, value.clone()));
6760 }
6761 }
6762 continue;
6763 }
6764
6765 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6766 for (index, reader) in readers.iter_mut().enumerate() {
6767 if let Some((epoch, deleted, value)) =
6768 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6769 {
6770 if best
6771 .as_ref()
6772 .map(|(best_epoch, ..)| epoch > *best_epoch)
6773 .unwrap_or(true)
6774 {
6775 best = Some((epoch, deleted, value, index));
6776 }
6777 }
6778 }
6779 let Some((_, false, Some(value), reader_index)) = best else {
6780 continue;
6781 };
6782 if let Some(ttl) = self.ttl {
6783 if ttl.column_id != column_id {
6784 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6785 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6786 {
6787 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6788 continue;
6789 }
6790 }
6791 } else if let Value::Int64(timestamp) = value {
6792 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6793 continue;
6794 }
6795 }
6796 }
6797 out.push((row_id, value));
6798 }
6799 Ok(out)
6800 }
6801
6802 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6807 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6808 }
6809
6810 pub fn rows_for_rids_with_context(
6811 &self,
6812 rids: &[u64],
6813 snapshot: Snapshot,
6814 context: &crate::query::AiExecutionContext,
6815 ) -> Result<Vec<Row>> {
6816 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6817 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
6818 }
6819
6820 fn rows_for_rids_at_time(
6821 &self,
6822 rids: &[u64],
6823 snapshot: Snapshot,
6824 ttl_now: i64,
6825 control: Option<&crate::ExecutionControl>,
6826 ) -> Result<Vec<Row>> {
6827 use std::collections::HashMap;
6828 let mut rows = Vec::with_capacity(rids.len());
6829 let tier_size = self.memtable.len() + self.mutable_run.len();
6846 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
6847 if rids.len().saturating_mul(24) < tier_size {
6848 for &rid in rids {
6849 if overlay.len() & 255 == 0 {
6850 control
6851 .map(crate::ExecutionControl::checkpoint)
6852 .transpose()?;
6853 }
6854 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
6855 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
6856 let newest = match (mem, mrun) {
6857 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
6858 (Some((_, mr)), None) => Some(mr),
6859 (None, Some((_, rr))) => Some(rr),
6860 (None, None) => None,
6861 };
6862 if let Some(row) = newest {
6863 overlay.insert(rid, row);
6864 }
6865 }
6866 } else {
6867 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
6868 overlay
6869 .entry(row.row_id.0)
6870 .and_modify(|e| {
6871 if row.committed_epoch > e.committed_epoch {
6872 *e = row.clone();
6873 }
6874 })
6875 .or_insert(row);
6876 };
6877 for (index, row) in self
6878 .memtable
6879 .visible_versions(snapshot.epoch)
6880 .into_iter()
6881 .enumerate()
6882 {
6883 if index & 255 == 0 {
6884 control
6885 .map(crate::ExecutionControl::checkpoint)
6886 .transpose()?;
6887 }
6888 fold_newest(row, &mut overlay);
6889 }
6890 for (index, row) in self
6891 .mutable_run
6892 .visible_versions(snapshot.epoch)
6893 .into_iter()
6894 .enumerate()
6895 {
6896 if index & 255 == 0 {
6897 control
6898 .map(crate::ExecutionControl::checkpoint)
6899 .transpose()?;
6900 }
6901 fold_newest(row, &mut overlay);
6902 }
6903 }
6904 if self.run_refs.len() == 1 {
6905 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6906 if rids.len().saturating_mul(24) < reader.row_count() {
6914 for (index, &rid) in rids.iter().enumerate() {
6915 if index & 255 == 0 {
6916 control
6917 .map(crate::ExecutionControl::checkpoint)
6918 .transpose()?;
6919 }
6920 if let Some(r) = overlay.get(&rid) {
6921 if !r.deleted {
6922 rows.push(r.clone());
6923 }
6924 continue;
6925 }
6926 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
6927 if !row.deleted {
6928 rows.push(row);
6929 }
6930 }
6931 }
6932 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6933 return Ok(rows);
6934 }
6935 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6944 enum Src {
6947 Overlay,
6948 Run,
6949 }
6950 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
6951 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
6952 for (index, rid) in rids.iter().enumerate() {
6953 if index & 255 == 0 {
6954 control
6955 .map(crate::ExecutionControl::checkpoint)
6956 .transpose()?;
6957 }
6958 if overlay.contains_key(rid) {
6959 plan.push(Src::Overlay);
6960 continue;
6961 }
6962 match vis_rids.binary_search(&(*rid as i64)) {
6963 Ok(i) => {
6964 plan.push(Src::Run);
6965 fetch.push(positions[i]);
6966 }
6967 Err(_) => { }
6968 }
6969 }
6970 let fetched = reader.materialize_batch(&fetch)?;
6971 let mut fetched_iter = fetched.into_iter();
6972 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
6973 if index & 255 == 0 {
6974 control
6975 .map(crate::ExecutionControl::checkpoint)
6976 .transpose()?;
6977 }
6978 match src {
6979 Src::Overlay => {
6980 if let Some(r) = overlay.get(rid) {
6981 if !r.deleted {
6982 rows.push(r.clone());
6983 }
6984 }
6985 }
6986 Src::Run => {
6987 if let Some(row) = fetched_iter.next() {
6988 if !row.deleted {
6989 rows.push(row);
6990 }
6991 }
6992 }
6993 }
6994 }
6995 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6996 return Ok(rows);
6997 }
6998 let mut readers: Vec<_> = self
7002 .run_refs
7003 .iter()
7004 .map(|rr| self.open_reader(rr.run_id))
7005 .collect::<Result<Vec<_>>>()?;
7006 for (index, rid) in rids.iter().enumerate() {
7007 if index & 255 == 0 {
7008 control
7009 .map(crate::ExecutionControl::checkpoint)
7010 .transpose()?;
7011 }
7012 if let Some(r) = overlay.get(rid) {
7013 if !r.deleted {
7014 rows.push(r.clone());
7015 }
7016 continue;
7017 }
7018 let mut best: Option<(Epoch, Row)> = None;
7019 for reader in readers.iter_mut() {
7020 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7021 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7022 best = Some((epoch, row));
7023 }
7024 }
7025 }
7026 if let Some((_, r)) = best {
7027 if !r.deleted {
7028 rows.push(r);
7029 }
7030 }
7031 }
7032 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7033 Ok(rows)
7034 }
7035
7036 pub fn indexes_complete(&self) -> bool {
7046 self.indexes_complete
7047 }
7048
7049 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7051 self.index_build_policy
7052 }
7053
7054 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7058 self.index_build_policy = policy;
7059 }
7060
7061 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7066 if !self.indexes_complete {
7070 return None;
7071 }
7072 let b = self.bitmap.get(&column_id)?;
7073 let result: Vec<Vec<u8>> = b
7074 .keys()
7075 .into_iter()
7076 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7077 .collect();
7078 Some(result)
7079 }
7080
7081 pub fn fk_join_row_ids(
7082 &self,
7083 fk_column_id: u16,
7084 pk_values: &[Vec<u8>],
7085 fk_conditions: &[crate::query::Condition],
7086 snapshot: Snapshot,
7087 ) -> Result<Vec<u64>> {
7088 let Some(b) = self.bitmap.get(&fk_column_id) else {
7089 return Ok(Vec::new());
7090 };
7091 let mut join_set = {
7092 let mut acc = roaring::RoaringBitmap::new();
7093 for v in pk_values {
7094 acc |= b.get(v);
7095 }
7096 RowIdSet::from_roaring(acc)
7097 };
7098 if !fk_conditions.is_empty() {
7099 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7100 sets.push(join_set);
7101 for c in fk_conditions {
7102 sets.push(self.resolve_condition(c, snapshot)?);
7103 }
7104 join_set = RowIdSet::intersect_many(sets);
7105 }
7106 Ok(join_set.into_sorted_vec())
7107 }
7108
7109 pub fn fk_join_count(
7115 &self,
7116 fk_column_id: u16,
7117 pk_values: &[Vec<u8>],
7118 fk_conditions: &[crate::query::Condition],
7119 snapshot: Snapshot,
7120 ) -> Result<u64> {
7121 let Some(b) = self.bitmap.get(&fk_column_id) else {
7122 return Ok(0);
7123 };
7124 let mut acc = roaring::RoaringBitmap::new();
7125 for v in pk_values {
7126 acc |= b.get(v);
7127 }
7128 if fk_conditions.is_empty() {
7129 return Ok(acc.len());
7130 }
7131 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7132 sets.push(RowIdSet::from_roaring(acc));
7133 for c in fk_conditions {
7134 sets.push(self.resolve_condition(c, snapshot)?);
7135 }
7136 Ok(RowIdSet::intersect_many(sets).len() as u64)
7137 }
7138
7139 fn resolve_condition(
7144 &self,
7145 c: &crate::query::Condition,
7146 snapshot: Snapshot,
7147 ) -> Result<RowIdSet> {
7148 self.resolve_condition_with_allowed(c, snapshot, None)
7149 }
7150
7151 fn resolve_condition_with_allowed(
7152 &self,
7153 c: &crate::query::Condition,
7154 snapshot: Snapshot,
7155 allowed: Option<&std::collections::HashSet<RowId>>,
7156 ) -> Result<RowIdSet> {
7157 use crate::query::Condition;
7158 self.validate_condition(c)?;
7159 Ok(match c {
7160 Condition::Pk(key) => {
7161 let lookup = self
7162 .schema
7163 .primary_key()
7164 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7165 .unwrap_or_else(|| key.clone());
7166 self.hot
7167 .get(&lookup)
7168 .map(|r| RowIdSet::one(r.0))
7169 .unwrap_or_else(RowIdSet::empty)
7170 }
7171 Condition::BitmapEq { column_id, value } => {
7172 let lookup = self.index_lookup_key_bytes(*column_id, value);
7173 self.bitmap
7174 .get(column_id)
7175 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7176 .unwrap_or_else(RowIdSet::empty)
7177 }
7178 Condition::BitmapIn { column_id, values } => {
7179 let bm = self.bitmap.get(column_id);
7180 let mut acc = roaring::RoaringBitmap::new();
7181 if let Some(b) = bm {
7182 for v in values {
7183 let lookup = self.index_lookup_key_bytes(*column_id, v);
7184 acc |= b.get(&lookup);
7185 }
7186 }
7187 RowIdSet::from_roaring(acc)
7188 }
7189 Condition::BytesPrefix { column_id, prefix } => {
7190 if let Some(b) = self.bitmap.get(column_id) {
7195 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7196 let mut acc = roaring::RoaringBitmap::new();
7197 for key in b.keys() {
7198 if key.starts_with(&lookup_prefix) {
7199 acc |= b.get(&key);
7200 }
7201 }
7202 RowIdSet::from_roaring(acc)
7203 } else {
7204 RowIdSet::empty()
7205 }
7206 }
7207 Condition::FmContains { column_id, pattern } => self
7208 .fm
7209 .get(column_id)
7210 .map(|f| {
7211 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7212 })
7213 .unwrap_or_else(RowIdSet::empty),
7214 Condition::FmContainsAll {
7215 column_id,
7216 patterns,
7217 } => {
7218 if let Some(f) = self.fm.get(column_id) {
7221 let sets: Vec<RowIdSet> = patterns
7222 .iter()
7223 .map(|pat| {
7224 RowIdSet::from_unsorted(
7225 f.locate(pat).into_iter().map(|r| r.0).collect(),
7226 )
7227 })
7228 .collect();
7229 RowIdSet::intersect_many(sets)
7230 } else {
7231 RowIdSet::empty()
7232 }
7233 }
7234 Condition::Ann {
7235 column_id,
7236 query,
7237 k,
7238 } => RowIdSet::from_unsorted(
7239 self.retrieve_filtered(
7240 &crate::query::Retriever::Ann {
7241 column_id: *column_id,
7242 query: query.clone(),
7243 k: *k,
7244 },
7245 snapshot,
7246 None,
7247 allowed,
7248 None,
7249 None,
7250 )?
7251 .into_iter()
7252 .map(|hit| hit.row_id.0)
7253 .collect(),
7254 ),
7255 Condition::SparseMatch {
7256 column_id,
7257 query,
7258 k,
7259 } => RowIdSet::from_unsorted(
7260 self.retrieve_filtered(
7261 &crate::query::Retriever::Sparse {
7262 column_id: *column_id,
7263 query: query.clone(),
7264 k: *k,
7265 },
7266 snapshot,
7267 None,
7268 allowed,
7269 None,
7270 None,
7271 )?
7272 .into_iter()
7273 .map(|hit| hit.row_id.0)
7274 .collect(),
7275 ),
7276 Condition::MinHashSimilar {
7277 column_id,
7278 query,
7279 k,
7280 } => match self.minhash.get(column_id) {
7281 Some(index) => {
7282 let candidates = index.candidate_row_ids(query);
7283 let eligible =
7284 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7285 RowIdSet::from_unsorted(
7286 index
7287 .search_filtered(query, *k, |row_id| {
7288 eligible.contains(&row_id)
7289 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7290 })
7291 .into_iter()
7292 .map(|(row_id, _)| row_id.0)
7293 .collect(),
7294 )
7295 }
7296 None => RowIdSet::empty(),
7297 },
7298 Condition::Range { column_id, lo, hi } => {
7299 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7308 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7309 } else if self.run_refs.len() == 1 {
7310 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7311 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7312 } else {
7313 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7314 };
7315 set.remove_many(self.overlay_rid_set(snapshot));
7316 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7317 set
7318 }
7319 Condition::RangeF64 {
7320 column_id,
7321 lo,
7322 lo_inclusive,
7323 hi,
7324 hi_inclusive,
7325 } => {
7326 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7329 RowIdSet::from_unsorted(
7330 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7331 .into_iter()
7332 .collect(),
7333 )
7334 } else if self.run_refs.len() == 1 {
7335 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7336 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7337 } else {
7338 return self.range_scan_f64(
7339 *column_id,
7340 *lo,
7341 *lo_inclusive,
7342 *hi,
7343 *hi_inclusive,
7344 snapshot,
7345 );
7346 };
7347 set.remove_many(self.overlay_rid_set(snapshot));
7348 self.range_scan_overlay_f64(
7349 &mut set,
7350 *column_id,
7351 *lo,
7352 *lo_inclusive,
7353 *hi,
7354 *hi_inclusive,
7355 snapshot,
7356 );
7357 set
7358 }
7359 Condition::IsNull { column_id } => {
7360 let mut set = if self.run_refs.len() == 1 {
7361 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7362 r.null_row_id_set(*column_id, true)?
7363 } else {
7364 return self.null_scan(*column_id, true, snapshot);
7365 };
7366 set.remove_many(self.overlay_rid_set(snapshot));
7367 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7368 set
7369 }
7370 Condition::IsNotNull { column_id } => {
7371 let mut set = if self.run_refs.len() == 1 {
7372 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7373 r.null_row_id_set(*column_id, false)?
7374 } else {
7375 return self.null_scan(*column_id, false, snapshot);
7376 };
7377 set.remove_many(self.overlay_rid_set(snapshot));
7378 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7379 set
7380 }
7381 })
7382 }
7383
7384 fn range_scan_i64(
7392 &self,
7393 column_id: u16,
7394 lo: i64,
7395 hi: i64,
7396 snapshot: Snapshot,
7397 ) -> Result<RowIdSet> {
7398 let mut row_ids = Vec::new();
7399 let overlay_rids = self.overlay_rid_set(snapshot);
7400 for rr in &self.run_refs {
7401 let mut reader = self.open_reader(rr.run_id)?;
7402 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7403 for rid in matched {
7404 if !overlay_rids.contains(&rid) {
7405 row_ids.push(rid);
7406 }
7407 }
7408 }
7409 let mut s = RowIdSet::from_unsorted(row_ids);
7410 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7411 Ok(s)
7412 }
7413
7414 fn range_scan_f64(
7417 &self,
7418 column_id: u16,
7419 lo: f64,
7420 lo_inclusive: bool,
7421 hi: f64,
7422 hi_inclusive: bool,
7423 snapshot: Snapshot,
7424 ) -> Result<RowIdSet> {
7425 let mut row_ids = Vec::new();
7426 let overlay_rids = self.overlay_rid_set(snapshot);
7427 for rr in &self.run_refs {
7428 let mut reader = self.open_reader(rr.run_id)?;
7429 let matched = reader.range_row_ids_visible_f64(
7430 column_id,
7431 lo,
7432 lo_inclusive,
7433 hi,
7434 hi_inclusive,
7435 snapshot.epoch,
7436 )?;
7437 for rid in matched {
7438 if !overlay_rids.contains(&rid) {
7439 row_ids.push(rid);
7440 }
7441 }
7442 }
7443 let mut s = RowIdSet::from_unsorted(row_ids);
7444 self.range_scan_overlay_f64(
7445 &mut s,
7446 column_id,
7447 lo,
7448 lo_inclusive,
7449 hi,
7450 hi_inclusive,
7451 snapshot,
7452 );
7453 Ok(s)
7454 }
7455
7456 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7458 let mut s = HashSet::new();
7459 for row in self.memtable.visible_versions(snapshot.epoch) {
7460 s.insert(row.row_id.0);
7461 }
7462 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7463 s.insert(row.row_id.0);
7464 }
7465 s
7466 }
7467
7468 fn range_scan_overlay_i64(
7469 &self,
7470 s: &mut RowIdSet,
7471 column_id: u16,
7472 lo: i64,
7473 hi: i64,
7474 snapshot: Snapshot,
7475 ) {
7476 let mut newest: HashMap<u64, &Row> = HashMap::new();
7481 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7482 let memtable = self.memtable.visible_versions(snapshot.epoch);
7483 for r in &mutable {
7484 newest.entry(r.row_id.0).or_insert(r);
7485 }
7486 for r in &memtable {
7487 newest.insert(r.row_id.0, r);
7488 }
7489 for row in newest.values() {
7490 if !row.deleted {
7491 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7492 if *v >= lo && *v <= hi {
7493 s.insert(row.row_id.0);
7494 }
7495 }
7496 }
7497 }
7498 }
7499
7500 #[allow(clippy::too_many_arguments)]
7501 fn range_scan_overlay_f64(
7502 &self,
7503 s: &mut RowIdSet,
7504 column_id: u16,
7505 lo: f64,
7506 lo_inclusive: bool,
7507 hi: f64,
7508 hi_inclusive: bool,
7509 snapshot: Snapshot,
7510 ) {
7511 let mut newest: HashMap<u64, &Row> = HashMap::new();
7514 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7515 let memtable = self.memtable.visible_versions(snapshot.epoch);
7516 for r in &mutable {
7517 newest.entry(r.row_id.0).or_insert(r);
7518 }
7519 for r in &memtable {
7520 newest.insert(r.row_id.0, r);
7521 }
7522 for row in newest.values() {
7523 if !row.deleted {
7524 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7525 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7526 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7527 if ok_lo && ok_hi {
7528 s.insert(row.row_id.0);
7529 }
7530 }
7531 }
7532 }
7533 }
7534
7535 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7538 let mut row_ids = Vec::new();
7539 let overlay_rids = self.overlay_rid_set(snapshot);
7540 for rr in &self.run_refs {
7541 let mut reader = self.open_reader(rr.run_id)?;
7542 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7543 for rid in matched {
7544 if !overlay_rids.contains(&rid) {
7545 row_ids.push(rid);
7546 }
7547 }
7548 }
7549 let mut s = RowIdSet::from_unsorted(row_ids);
7550 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7551 Ok(s)
7552 }
7553
7554 fn null_scan_overlay(
7558 &self,
7559 s: &mut RowIdSet,
7560 column_id: u16,
7561 want_nulls: bool,
7562 snapshot: Snapshot,
7563 ) {
7564 let mut newest: HashMap<u64, &Row> = HashMap::new();
7565 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7566 let memtable = self.memtable.visible_versions(snapshot.epoch);
7567 for r in &mutable {
7568 newest.entry(r.row_id.0).or_insert(r);
7569 }
7570 for r in &memtable {
7571 newest.insert(r.row_id.0, r);
7572 }
7573 for row in newest.values() {
7574 if row.deleted {
7575 continue;
7576 }
7577 let is_null = !row.columns.contains_key(&column_id)
7578 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7579 if is_null == want_nulls {
7580 s.insert(row.row_id.0);
7581 }
7582 }
7583 }
7584
7585 pub fn snapshot(&self) -> Snapshot {
7586 Snapshot::at(self.epoch.visible())
7587 }
7588
7589 pub fn data_generation(&self) -> u64 {
7591 self.data_generation
7592 }
7593
7594 pub(crate) fn bump_data_generation(&mut self) {
7595 self.data_generation = self.data_generation.wrapping_add(1);
7596 }
7597
7598 pub fn table_id(&self) -> u64 {
7600 self.table_id
7601 }
7602
7603 fn seal_generations(&mut self) {
7609 self.memtable.seal();
7610 self.mutable_run.seal();
7611 self.hot.seal();
7612 for index in self.bitmap.values_mut() {
7613 index.seal();
7614 }
7615 for index in self.ann.values_mut() {
7616 index.seal();
7617 }
7618 for index in self.fm.values_mut() {
7619 index.seal();
7620 }
7621 for index in self.sparse.values_mut() {
7622 index.seal();
7623 }
7624 for index in self.minhash.values_mut() {
7625 index.seal();
7626 }
7627 self.pk_by_row.seal();
7628 }
7629
7630 fn capture_read_generation(&self) -> ReadGeneration {
7635 let visible_through = self.current_epoch();
7636 ReadGeneration {
7637 schema: Arc::new(self.schema.clone()),
7638 base_runs: Arc::new(self.run_refs.clone()),
7639 deltas: TableDeltas {
7640 memtable: self.memtable.clone(),
7641 mutable_run: self.mutable_run.clone(),
7642 hot: self.hot.clone(),
7643 pk_by_row: self.pk_by_row.clone(),
7644 },
7645 indexes: Arc::new(IndexGeneration::capture(
7646 &self.bitmap,
7647 &self.learned_range,
7648 &self.fm,
7649 &self.ann,
7650 &self.sparse,
7651 &self.minhash,
7652 visible_through,
7653 )),
7654 visible_through,
7655 }
7656 }
7657
7658 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
7663 self.ensure_indexes_complete()?;
7664 self.seal_generations();
7665 let view = Arc::new(self.capture_read_generation());
7666 self.published.store(Arc::clone(&view));
7667 Ok(view)
7668 }
7669
7670 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
7676 self.published.load_full()
7677 }
7678
7679 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
7681 &self.pins
7682 }
7683
7684 pub fn version_gc_floor(&self) -> Epoch {
7689 self.min_active_snapshot()
7690 .unwrap_or_else(|| self.current_epoch())
7691 }
7692
7693 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
7700 let mut report = self.pins.report();
7701 let transaction_floor = [
7702 self.pinned.keys().next().copied(),
7703 self.snapshots.min_pinned(),
7704 ]
7705 .into_iter()
7706 .flatten()
7707 .min();
7708 if let Some(epoch) = transaction_floor {
7709 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
7710 }
7711 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
7712 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
7713 }
7714 report
7715 }
7716
7717 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7718 self.publish_read_generation()?;
7719 let mut generation = self.clone();
7720 generation.read_only = true;
7721 generation.wal = WalSink::ReadOnly;
7722 generation.pending_delete_rids.clear();
7723 generation.pending_put_cols.clear();
7724 generation.pending_rows.clear();
7725 generation.pending_rows_auto_inc.clear();
7726 generation.pending_dels.clear();
7727 generation.pending_truncate = None;
7728 generation.agg_cache = Arc::new(HashMap::new());
7729 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
7732 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
7735 crate::retention::PinSource::ReadGeneration,
7736 self.current_epoch(),
7737 )));
7738 Ok(generation)
7739 }
7740
7741 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7742 (std::mem::size_of::<Self>() as u64)
7743 .saturating_add(self.memtable.approx_bytes())
7744 .saturating_add(self.mutable_run.approx_bytes())
7745 .saturating_add(self.live_count.saturating_mul(64))
7746 }
7747
7748 pub fn pin_snapshot(&mut self) -> Snapshot {
7751 let e = self.epoch.visible();
7752 *self.pinned.entry(e).or_insert(0) += 1;
7753 Snapshot::at(e)
7754 }
7755
7756 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7758 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7759 *count -= 1;
7760 if *count == 0 {
7761 self.pinned.remove(&snap.epoch);
7762 }
7763 }
7764 }
7765
7766 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7778 let local = self.pinned.keys().next().copied();
7779 let global = self.snapshots.min_pinned();
7780 let history = self.snapshots.history_floor(self.current_epoch());
7781 let pinned = self.pins.oldest_pinned();
7782 [local, global, history, pinned].into_iter().flatten().min()
7783 }
7784
7785 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7789 self.ensure_writable()?;
7790 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7791 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7792 }
7793
7794 pub fn clear_ttl(&mut self) -> Result<()> {
7795 self.ensure_writable()?;
7796 self.apply_ttl_policy_at(None, self.current_epoch())
7797 }
7798
7799 pub fn ttl(&self) -> Option<TtlPolicy> {
7800 self.ttl
7801 }
7802
7803 pub(crate) fn prepare_ttl_policy(
7804 &self,
7805 column_name: &str,
7806 duration_nanos: u64,
7807 ) -> Result<TtlPolicy> {
7808 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7809 return Err(MongrelError::InvalidArgument(
7810 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7811 ));
7812 }
7813 let column = self
7814 .schema
7815 .columns
7816 .iter()
7817 .find(|column| column.name == column_name)
7818 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7819 if column.ty != TypeId::TimestampNanos {
7820 return Err(MongrelError::Schema(format!(
7821 "TTL column {column_name} must be TimestampNanos, is {:?}",
7822 column.ty
7823 )));
7824 }
7825 Ok(TtlPolicy {
7826 column_id: column.id,
7827 duration_nanos,
7828 })
7829 }
7830
7831 pub(crate) fn apply_ttl_policy_at(
7832 &mut self,
7833 policy: Option<TtlPolicy>,
7834 epoch: Epoch,
7835 ) -> Result<()> {
7836 if let Some(policy) = policy {
7837 let column = self
7838 .schema
7839 .columns
7840 .iter()
7841 .find(|column| column.id == policy.column_id)
7842 .ok_or_else(|| {
7843 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
7844 })?;
7845 if column.ty != TypeId::TimestampNanos
7846 || policy.duration_nanos == 0
7847 || policy.duration_nanos > i64::MAX as u64
7848 {
7849 return Err(MongrelError::Schema("invalid TTL policy".into()));
7850 }
7851 }
7852 self.ttl = policy;
7853 self.agg_cache = Arc::new(HashMap::new());
7854 self.clear_result_cache();
7855 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7856 self.persist_manifest(epoch)
7857 }
7858
7859 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
7860 let Some(policy) = self.ttl else {
7861 return false;
7862 };
7863 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
7864 return false;
7865 };
7866 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
7867 }
7868
7869 pub fn current_epoch(&self) -> Epoch {
7870 self.epoch.visible()
7871 }
7872
7873 pub fn memtable_len(&self) -> usize {
7874 self.memtable.len()
7875 }
7876
7877 pub fn count(&self) -> u64 {
7880 if self.ttl.is_none()
7881 && self.pending_put_cols.is_empty()
7882 && self.pending_delete_rids.is_empty()
7883 && self.pending_rows.is_empty()
7884 && self.pending_dels.is_empty()
7885 && self.pending_truncate.is_none()
7886 {
7887 self.live_count
7888 } else {
7889 self.visible_rows(self.snapshot())
7890 .map(|rows| rows.len() as u64)
7891 .unwrap_or(self.live_count)
7892 }
7893 }
7894
7895 pub fn count_conditions(
7899 &mut self,
7900 conditions: &[crate::query::Condition],
7901 snapshot: Snapshot,
7902 ) -> Result<Option<u64>> {
7903 use crate::query::Condition;
7904 if self.ttl.is_some() {
7905 if conditions.is_empty() {
7906 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
7907 }
7908 let mut sets = Vec::with_capacity(conditions.len());
7909 for condition in conditions {
7910 sets.push(self.resolve_condition(condition, snapshot)?);
7911 }
7912 let survivors = RowIdSet::intersect_many(sets);
7913 let rows = self.visible_rows(snapshot)?;
7914 return Ok(Some(
7915 rows.into_iter()
7916 .filter(|row| survivors.contains(row.row_id.0))
7917 .count() as u64,
7918 ));
7919 }
7920 if conditions.is_empty() {
7921 return Ok(Some(self.count()));
7922 }
7923 let served = |c: &Condition| {
7924 matches!(
7925 c,
7926 Condition::Pk(_)
7927 | Condition::BitmapEq { .. }
7928 | Condition::BitmapIn { .. }
7929 | Condition::BytesPrefix { .. }
7930 | Condition::FmContains { .. }
7931 | Condition::FmContainsAll { .. }
7932 | Condition::Ann { .. }
7933 | Condition::Range { .. }
7934 | Condition::RangeF64 { .. }
7935 | Condition::SparseMatch { .. }
7936 | Condition::MinHashSimilar { .. }
7937 | Condition::IsNull { .. }
7938 | Condition::IsNotNull { .. }
7939 )
7940 };
7941 if !conditions.iter().all(served) {
7942 return Ok(None);
7943 }
7944 self.ensure_indexes_complete()?;
7945 if !self.pending_put_cols.is_empty()
7946 || !self.pending_delete_rids.is_empty()
7947 || !self.pending_rows.is_empty()
7948 || !self.pending_dels.is_empty()
7949 || self.pending_truncate.is_some()
7950 {
7951 let mut sets = Vec::with_capacity(conditions.len());
7952 for condition in conditions {
7953 sets.push(self.resolve_condition(condition, snapshot)?);
7954 }
7955 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7956 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
7957 }
7958 let mut sets = Vec::with_capacity(conditions.len());
7959 for condition in conditions {
7960 sets.push(self.resolve_condition(condition, snapshot)?);
7961 }
7962 let mut rids = RowIdSet::intersect_many(sets);
7963 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7973 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
7974 }
7975 let count = rids.len() as u64;
7976 crate::trace::QueryTrace::record(|t| {
7977 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
7978 t.survivor_count = Some(count as usize);
7979 t.conditions_pushed = conditions.len();
7980 });
7981 Ok(Some(count))
7982 }
7983
7984 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
7989 let mut out = Vec::new();
7990 for row in self.memtable.visible_versions(snapshot.epoch) {
7991 if row.deleted {
7992 out.push(row.row_id.0);
7993 }
7994 }
7995 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7996 if row.deleted {
7997 out.push(row.row_id.0);
7998 }
7999 }
8000 out
8001 }
8002
8003 pub fn bulk_load_columns(
8012 &mut self,
8013 user_columns: Vec<(u16, columnar::NativeColumn)>,
8014 ) -> Result<Epoch> {
8015 self.bulk_load_columns_with(user_columns, 3, false, true)
8016 }
8017
8018 pub fn bulk_load_fast(
8025 &mut self,
8026 user_columns: Vec<(u16, columnar::NativeColumn)>,
8027 ) -> Result<Epoch> {
8028 self.bulk_load_columns_with(user_columns, -1, true, false)
8029 }
8030
8031 fn bulk_load_columns_with(
8032 &mut self,
8033 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8034 zstd_level: i32,
8035 force_plain: bool,
8036 lz4: bool,
8037 ) -> Result<Epoch> {
8038 self.ensure_writable()?;
8039 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8040 if n == 0 {
8041 return Ok(self.current_epoch());
8042 }
8043 let epoch = self.commit_new_epoch()?;
8044 let live_before = self.live_count;
8045 self.spill_mutable_run(epoch)?;
8047 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8048 && self.indexes_complete
8049 && self.run_refs.is_empty()
8050 && self.memtable.is_empty()
8051 && self.mutable_run.is_empty();
8052 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8055 self.validate_columns_not_null(&user_columns, n)?;
8056 let winner_idx = self
8057 .bulk_pk_winner_indices(&user_columns, n)
8058 .filter(|idx| idx.len() != n);
8059 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8060 match winner_idx.as_deref() {
8061 Some(idx) => {
8062 let compacted = user_columns
8063 .iter()
8064 .map(|(id, c)| (*id, c.gather(idx)))
8065 .collect();
8066 (compacted, idx.len())
8067 }
8068 None => (user_columns, n),
8069 };
8070 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8071 let first = self.allocator.alloc_range(write_n as u64)?.0;
8072 for rid in first..first + write_n as u64 {
8073 self.reservoir.offer(rid);
8074 }
8075 let run_id = self.alloc_run_id()?;
8076 let path = self.run_path(run_id);
8077 let mut writer =
8078 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8079 if force_plain {
8080 writer = writer.with_plain();
8081 } else if lz4 {
8082 writer = writer.with_lz4();
8085 } else {
8086 writer = writer.with_zstd_level(zstd_level);
8087 }
8088 if let Some(kek) = &self.kek {
8089 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8090 }
8091 let header = match self.create_run_file(run_id)? {
8092 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8093 None => writer.write_native(&path, &write_columns, write_n, first)?,
8094 };
8095 self.run_refs.push(RunRef {
8096 run_id: run_id as u128,
8097 level: 0,
8098 epoch_created: epoch.0,
8099 row_count: header.row_count,
8100 });
8101 self.live_count = self.live_count.saturating_add(write_n as u64);
8102 if eager_index_build {
8103 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8104 self.index_columns_bulk(&write_columns, &row_ids);
8105 self.indexes_complete = true;
8106 self.build_learned_ranges()?;
8107 } else {
8108 self.indexes_complete = false;
8112 }
8113 self.mark_flushed(epoch)?;
8114 self.persist_manifest(epoch)?;
8115 if eager_index_build {
8116 self.checkpoint_indexes(epoch);
8117 }
8118 self.clear_result_cache();
8119 self.data_generation = self.data_generation.wrapping_add(1);
8120 Ok(epoch)
8121 }
8122
8123 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8141 let n = row_ids.len();
8142 if n == 0 {
8143 return;
8144 }
8145 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8146 columns.iter().map(|(id, c)| (*id, c)).collect();
8147 let ty_of: std::collections::HashMap<u16, TypeId> = self
8148 .schema
8149 .columns
8150 .iter()
8151 .map(|c| (c.id, c.ty.clone()))
8152 .collect();
8153 let pk_id = self.schema.primary_key().map(|c| c.id);
8154
8155 for (i, &rid) in row_ids.iter().enumerate() {
8156 let row_id = RowId(rid);
8157 if let Some(pid) = pk_id {
8158 if let Some(col) = by_id.get(&pid) {
8159 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8160 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8161 self.insert_hot_pk(key, row_id);
8162 }
8163 }
8164 }
8165 for idef in &self.schema.indexes {
8166 let Some(col) = by_id.get(&idef.column_id) else {
8167 continue;
8168 };
8169 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8170 match idef.kind {
8171 IndexKind::Bitmap => {
8172 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8173 if let Some(key) =
8174 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8175 {
8176 b.insert(key, row_id);
8177 }
8178 }
8179 }
8180 IndexKind::FmIndex => {
8181 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8182 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8183 f.insert(bytes.to_vec(), row_id);
8184 }
8185 }
8186 }
8187 IndexKind::Sparse => {
8188 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8189 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8190 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8191 s.insert(&terms, row_id);
8192 }
8193 }
8194 }
8195 }
8196 IndexKind::MinHash => {
8197 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8198 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8199 let tokens = crate::index::token_hashes_from_bytes(bytes);
8200 mh.insert(&tokens, row_id);
8201 }
8202 }
8203 }
8204 _ => {}
8205 }
8206 }
8207 }
8208 }
8209
8210 pub fn visible_columns_native(
8215 &self,
8216 snapshot: Snapshot,
8217 projection: Option<&[u16]>,
8218 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8219 self.visible_columns_native_inner(snapshot, projection, None)
8220 }
8221
8222 pub fn visible_columns_native_with_control(
8223 &self,
8224 snapshot: Snapshot,
8225 projection: Option<&[u16]>,
8226 control: &crate::ExecutionControl,
8227 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8228 self.visible_columns_native_inner(snapshot, projection, Some(control))
8229 }
8230
8231 fn visible_columns_native_inner(
8232 &self,
8233 snapshot: Snapshot,
8234 projection: Option<&[u16]>,
8235 control: Option<&crate::ExecutionControl>,
8236 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8237 execution_checkpoint(control, 0)?;
8238 let wanted: Vec<u16> = match projection {
8239 Some(p) => p.to_vec(),
8240 None => self.schema.columns.iter().map(|c| c.id).collect(),
8241 };
8242 if self.ttl.is_none()
8243 && self.memtable.is_empty()
8244 && self.mutable_run.is_empty()
8245 && self.run_refs.len() == 1
8246 {
8247 let rr = self.run_refs[0].clone();
8248 let mut reader = self.open_reader(rr.run_id)?;
8249 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8250 execution_checkpoint(control, 0)?;
8251 let all_visible = idxs.len() == reader.row_count();
8252 if reader.has_mmap() && control.is_none() {
8258 use rayon::prelude::*;
8259 let valid: Vec<u16> = wanted
8262 .iter()
8263 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8264 .copied()
8265 .collect();
8266 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8268 .par_iter()
8269 .filter_map(|cid| {
8270 reader
8271 .column_native_shared(*cid)
8272 .ok()
8273 .map(|col| (*cid, col))
8274 })
8275 .collect();
8276 let cols = decoded
8277 .into_iter()
8278 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8279 .collect();
8280 return Ok(cols);
8281 }
8282 let mut cols = Vec::with_capacity(wanted.len());
8283 for (index, cid) in wanted.iter().enumerate() {
8284 execution_checkpoint(control, index)?;
8285 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8286 Some(c) => c,
8287 None => continue,
8288 };
8289 let col = reader.column_native(cdef.id)?;
8290 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8291 }
8292 return Ok(cols);
8293 }
8294 let vcols = self.visible_columns(snapshot)?;
8295 execution_checkpoint(control, 0)?;
8296 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8297 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8298 .into_iter()
8299 .filter(|(id, _)| want_set.contains(id))
8300 .map(|(id, vals)| {
8301 let ty = self
8302 .schema
8303 .columns
8304 .iter()
8305 .find(|c| c.id == id)
8306 .map(|c| c.ty.clone())
8307 .unwrap_or(TypeId::Bytes);
8308 (id, columnar::values_to_native(ty, &vals))
8309 })
8310 .collect();
8311 Ok(out)
8312 }
8313
8314 pub fn run_count(&self) -> usize {
8315 self.run_refs.len()
8316 }
8317
8318 pub fn memtable_is_empty(&self) -> bool {
8320 self.memtable.is_empty()
8321 }
8322
8323 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8327 self.page_cache.stats()
8328 }
8329
8330 pub fn reset_page_cache_stats(&self) {
8332 self.page_cache.reset_stats();
8333 }
8334
8335 pub fn run_ids(&self) -> Vec<u128> {
8338 self.run_refs.iter().map(|r| r.run_id).collect()
8339 }
8340
8341 pub fn single_run_is_clean(&self) -> bool {
8345 if self.ttl.is_some() || self.run_refs.len() != 1 {
8346 return false;
8347 }
8348 self.open_reader(self.run_refs[0].run_id)
8349 .map(|r| r.is_clean())
8350 .unwrap_or(false)
8351 }
8352
8353 fn resolve_footprint(
8360 &self,
8361 conditions: &[crate::query::Condition],
8362 snapshot: Snapshot,
8363 ) -> roaring::RoaringBitmap {
8364 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8365 return roaring::RoaringBitmap::new();
8366 }
8367 if self.run_refs.is_empty() {
8368 return roaring::RoaringBitmap::new();
8369 }
8370 if self.run_refs.len() == 1 {
8372 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8373 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8374 return rids.to_roaring_lossy();
8375 }
8376 }
8377 }
8378 roaring::RoaringBitmap::new()
8379 }
8380
8381 pub fn query_columns_native_cached(
8392 &mut self,
8393 conditions: &[crate::query::Condition],
8394 projection: Option<&[u16]>,
8395 snapshot: Snapshot,
8396 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8397 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8398 }
8399
8400 pub fn query_columns_native_cached_with_control(
8401 &mut self,
8402 conditions: &[crate::query::Condition],
8403 projection: Option<&[u16]>,
8404 snapshot: Snapshot,
8405 control: &crate::ExecutionControl,
8406 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8407 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8408 }
8409
8410 fn query_columns_native_cached_inner(
8411 &mut self,
8412 conditions: &[crate::query::Condition],
8413 projection: Option<&[u16]>,
8414 snapshot: Snapshot,
8415 control: Option<&crate::ExecutionControl>,
8416 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8417 execution_checkpoint(control, 0)?;
8418 if self.ttl.is_some() {
8421 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8422 }
8423 if conditions.is_empty() {
8424 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8425 }
8426 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8430 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8431 crate::trace::QueryTrace::record(|t| {
8432 t.result_cache_hit = true;
8433 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8434 });
8435 return Ok(Some((*hit).clone()));
8436 }
8437 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8438 execution_checkpoint(control, 0)?;
8439 if let Some(cols) = &res {
8440 let footprint = self.resolve_footprint(conditions, snapshot);
8441 let condition_cols = crate::query::condition_columns(conditions);
8442 execution_checkpoint(control, 0)?;
8443 self.result_cache.lock().insert(
8444 key,
8445 CachedEntry {
8446 data: CachedData::Columns(Arc::new(cols.clone())),
8447 footprint,
8448 condition_cols,
8449 },
8450 );
8451 }
8452 Ok(res)
8453 }
8454
8455 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8460 if self.ttl.is_some() {
8461 return self.query(q);
8462 }
8463 if q.conditions.is_empty() {
8464 return self.query(q);
8465 }
8466 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8467 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8468 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8469 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8470 crate::trace::QueryTrace::record(|t| {
8471 t.result_cache_hit = true;
8472 t.scan_mode = crate::trace::ScanMode::Materialized;
8473 });
8474 return Ok((*hit).clone());
8475 }
8476 let rows = self.query(q)?;
8477 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8478 let condition_cols = crate::query::condition_columns(&q.conditions);
8479 self.result_cache.lock().insert(
8480 key,
8481 CachedEntry {
8482 data: CachedData::Rows(Arc::new(rows.clone())),
8483 footprint,
8484 condition_cols,
8485 },
8486 );
8487 Ok(rows)
8488 }
8489
8490 #[allow(clippy::type_complexity)]
8505 pub fn query_columns_native_traced(
8506 &mut self,
8507 conditions: &[crate::query::Condition],
8508 projection: Option<&[u16]>,
8509 snapshot: Snapshot,
8510 ) -> Result<(
8511 Option<Vec<(u16, columnar::NativeColumn)>>,
8512 crate::trace::QueryTrace,
8513 )> {
8514 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8515 self.query_columns_native(conditions, projection, snapshot)
8516 });
8517 Ok((result?, trace))
8518 }
8519
8520 #[allow(clippy::type_complexity)]
8523 pub fn query_columns_native_cached_traced(
8524 &mut self,
8525 conditions: &[crate::query::Condition],
8526 projection: Option<&[u16]>,
8527 snapshot: Snapshot,
8528 ) -> Result<(
8529 Option<Vec<(u16, columnar::NativeColumn)>>,
8530 crate::trace::QueryTrace,
8531 )> {
8532 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8533 self.query_columns_native_cached(conditions, projection, snapshot)
8534 });
8535 Ok((result?, trace))
8536 }
8537
8538 pub fn native_page_cursor_traced(
8540 &self,
8541 snapshot: Snapshot,
8542 projection: Vec<(u16, TypeId)>,
8543 conditions: &[crate::query::Condition],
8544 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8545 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8546 self.native_page_cursor(snapshot, projection, conditions)
8547 });
8548 Ok((result?, trace))
8549 }
8550
8551 pub fn native_multi_run_cursor_traced(
8553 &self,
8554 snapshot: Snapshot,
8555 projection: Vec<(u16, TypeId)>,
8556 conditions: &[crate::query::Condition],
8557 ) -> Result<(
8558 Option<crate::cursor::MultiRunCursor>,
8559 crate::trace::QueryTrace,
8560 )> {
8561 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8562 self.native_multi_run_cursor(snapshot, projection, conditions)
8563 });
8564 Ok((result?, trace))
8565 }
8566
8567 pub fn count_conditions_traced(
8569 &mut self,
8570 conditions: &[crate::query::Condition],
8571 snapshot: Snapshot,
8572 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8573 let (result, trace) =
8574 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8575 Ok((result?, trace))
8576 }
8577
8578 pub fn query_traced(
8580 &mut self,
8581 q: &crate::query::Query,
8582 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8583 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8584 Ok((result?, trace))
8585 }
8586
8587 pub fn query_columns_native(
8592 &mut self,
8593 conditions: &[crate::query::Condition],
8594 projection: Option<&[u16]>,
8595 snapshot: Snapshot,
8596 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8597 self.query_columns_native_inner(conditions, projection, snapshot, None)
8598 }
8599
8600 pub fn query_columns_native_with_control(
8601 &mut self,
8602 conditions: &[crate::query::Condition],
8603 projection: Option<&[u16]>,
8604 snapshot: Snapshot,
8605 control: &crate::ExecutionControl,
8606 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8607 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8608 }
8609
8610 fn query_columns_native_inner(
8611 &mut self,
8612 conditions: &[crate::query::Condition],
8613 projection: Option<&[u16]>,
8614 snapshot: Snapshot,
8615 control: Option<&crate::ExecutionControl>,
8616 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8617 use crate::query::Condition;
8618 execution_checkpoint(control, 0)?;
8619 if self.ttl.is_some() {
8622 return Ok(None);
8623 }
8624 if conditions.is_empty() {
8625 return Ok(None);
8626 }
8627 self.ensure_indexes_complete()?;
8628
8629 let served = |c: &Condition| {
8634 matches!(
8635 c,
8636 Condition::Pk(_)
8637 | Condition::BitmapEq { .. }
8638 | Condition::BitmapIn { .. }
8639 | Condition::BytesPrefix { .. }
8640 | Condition::FmContains { .. }
8641 | Condition::FmContainsAll { .. }
8642 | Condition::Ann { .. }
8643 | Condition::Range { .. }
8644 | Condition::RangeF64 { .. }
8645 | Condition::SparseMatch { .. }
8646 | Condition::MinHashSimilar { .. }
8647 | Condition::IsNull { .. }
8648 | Condition::IsNotNull { .. }
8649 )
8650 };
8651 if !conditions.iter().all(served) {
8652 return Ok(None);
8653 }
8654 let fast_path =
8655 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8656 crate::trace::QueryTrace::record(|t| {
8657 t.run_count = self.run_refs.len();
8658 t.memtable_rows = self.memtable.len();
8659 t.mutable_run_rows = self.mutable_run.len();
8660 t.conditions_pushed = conditions.len();
8661 t.learned_range_used = conditions.iter().any(|c| match c {
8662 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8663 self.learned_range.contains_key(column_id)
8664 }
8665 _ => false,
8666 });
8667 });
8668 let col_ids: Vec<u16> = projection
8670 .map(|p| p.to_vec())
8671 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8672 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8673 .iter()
8674 .map(|&cid| {
8675 let ty = self
8676 .schema
8677 .columns
8678 .iter()
8679 .find(|c| c.id == cid)
8680 .map(|c| c.ty.clone())
8681 .unwrap_or(TypeId::Bytes);
8682 (cid, ty)
8683 })
8684 .collect();
8685
8686 if fast_path {
8692 let needs_column = conditions.iter().any(|c| match c {
8695 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8696 Condition::RangeF64 { column_id, .. } => {
8697 !self.learned_range.contains_key(column_id)
8698 }
8699 _ => false,
8700 });
8701 let mut reader_opt: Option<RunReader> = if needs_column {
8702 Some(self.open_reader(self.run_refs[0].run_id)?)
8703 } else {
8704 None
8705 };
8706 let mut sets: Vec<RowIdSet> = Vec::new();
8707 for (index, c) in conditions.iter().enumerate() {
8708 execution_checkpoint(control, index)?;
8709 let s = match c {
8710 Condition::Range { column_id, lo, hi }
8711 if !self.learned_range.contains_key(column_id) =>
8712 {
8713 if reader_opt.is_none() {
8714 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8715 }
8716 reader_opt
8717 .as_mut()
8718 .expect("reader opened for range")
8719 .range_row_id_set_i64(*column_id, *lo, *hi)?
8720 }
8721 Condition::RangeF64 {
8722 column_id,
8723 lo,
8724 lo_inclusive,
8725 hi,
8726 hi_inclusive,
8727 } if !self.learned_range.contains_key(column_id) => {
8728 if reader_opt.is_none() {
8729 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8730 }
8731 reader_opt
8732 .as_mut()
8733 .expect("reader opened for range")
8734 .range_row_id_set_f64(
8735 *column_id,
8736 *lo,
8737 *lo_inclusive,
8738 *hi,
8739 *hi_inclusive,
8740 )?
8741 }
8742 _ => self.resolve_condition(c, snapshot)?,
8743 };
8744 sets.push(s);
8745 }
8746 let candidates = RowIdSet::intersect_many(sets);
8747 crate::trace::QueryTrace::record(|t| {
8748 t.survivor_count = Some(candidates.len());
8749 });
8750 if candidates.is_empty() {
8751 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8752 .iter()
8753 .map(|&id| {
8754 (
8755 id,
8756 columnar::null_native(
8757 proj_pairs
8758 .iter()
8759 .find(|(c, _)| c == &id)
8760 .map(|(_, t)| t.clone())
8761 .unwrap_or(TypeId::Bytes),
8762 0,
8763 ),
8764 )
8765 })
8766 .collect();
8767 return Ok(Some(cols));
8768 }
8769 let mut reader = match reader_opt.take() {
8770 Some(r) => r,
8771 None => self.open_reader(self.run_refs[0].run_id)?,
8772 };
8773 let candidate_ids = candidates.into_sorted_vec();
8774 let (positions, fast_rid) = if let Some(positions) =
8775 reader.positions_for_row_ids_fast(&candidate_ids)
8776 {
8777 (positions, true)
8778 } else {
8779 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8780 match col {
8781 columnar::NativeColumn::Int64 { data, .. } => {
8782 let mut p = Vec::with_capacity(candidate_ids.len());
8783 for (index, rid) in candidate_ids.iter().enumerate() {
8784 execution_checkpoint(control, index)?;
8785 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8786 p.push(position);
8787 }
8788 }
8789 p.sort_unstable();
8790 (p, false)
8791 }
8792 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8793 }
8794 };
8795 crate::trace::QueryTrace::record(|t| {
8796 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8797 t.fast_row_id_map = fast_rid;
8798 });
8799 let mut cols = Vec::with_capacity(col_ids.len());
8800 for (index, cid) in col_ids.iter().enumerate() {
8801 execution_checkpoint(control, index)?;
8802 let col = reader.column_native(*cid)?;
8803 cols.push((*cid, col.gather(&positions)));
8804 }
8805 return Ok(Some(cols));
8806 }
8807
8808 if !self.run_refs.is_empty() {
8821 use crate::cursor::{
8822 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8823 };
8824 let remaining: usize;
8825 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
8826 let c = self
8827 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
8828 .expect("single-run cursor should build when run_refs.len() == 1");
8829 remaining = c.remaining_rows();
8830 Box::new(c)
8831 } else {
8832 let c = self
8833 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
8834 .expect("multi-run cursor should build when run_refs.len() >= 1");
8835 remaining = c.remaining_rows();
8836 Box::new(c)
8837 };
8838 crate::trace::QueryTrace::record(|t| {
8839 if t.survivor_count.is_none() {
8840 t.survivor_count = Some(remaining);
8841 }
8842 });
8843 let cols = match control {
8844 Some(control) => {
8845 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
8846 }
8847 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
8848 };
8849 return Ok(Some(cols));
8850 }
8851
8852 crate::trace::QueryTrace::record(|t| {
8857 t.scan_mode = crate::trace::ScanMode::Materialized;
8858 t.row_materialized = true;
8859 });
8860 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8861 for (index, c) in conditions.iter().enumerate() {
8862 execution_checkpoint(control, index)?;
8863 sets.push(self.resolve_condition(c, snapshot)?);
8864 }
8865 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8866 let rows = self.rows_for_rids(&rids, snapshot)?;
8867 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
8868 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
8869 execution_checkpoint(control, index)?;
8870 let vals: Vec<Value> = rows
8871 .iter()
8872 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8873 .collect();
8874 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
8875 }
8876 Ok(Some(cols))
8877 }
8878
8879 pub fn native_page_cursor(
8894 &self,
8895 snapshot: Snapshot,
8896 projection: Vec<(u16, TypeId)>,
8897 conditions: &[crate::query::Condition],
8898 ) -> Result<Option<NativePageCursor>> {
8899 use crate::cursor::build_page_plans;
8900 if self.ttl.is_some() {
8901 return Ok(None);
8902 }
8903 if !conditions.is_empty() && !self.indexes_complete {
8906 return Ok(None);
8907 }
8908 if self.run_refs.len() != 1 {
8909 return Ok(None);
8910 }
8911 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8912 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8913
8914 let overlay_rids: HashSet<u64> = {
8917 let mut s = HashSet::new();
8918 for row in self.memtable.visible_versions(snapshot.epoch) {
8919 s.insert(row.row_id.0);
8920 }
8921 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8922 s.insert(row.row_id.0);
8923 }
8924 s
8925 };
8926
8927 let survivors = if conditions.is_empty() {
8931 None
8932 } else {
8933 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
8934 };
8935
8936 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
8943 survivors.clone()
8944 } else if let Some(s) = &survivors {
8945 let mut run_set = s.clone();
8946 run_set.remove_many(overlay_rids.iter().copied());
8947 Some(run_set)
8948 } else {
8949 Some(RowIdSet::from_unsorted(
8950 rids.iter()
8951 .map(|&r| r as u64)
8952 .filter(|r| !overlay_rids.contains(r))
8953 .collect(),
8954 ))
8955 };
8956
8957 let overlay_rows = if overlay_rids.is_empty() {
8958 Vec::new()
8959 } else {
8960 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8961 self.overlay_visible_rows(snapshot, bound)
8962 };
8963
8964 let plans = if positions.is_empty() {
8966 Vec::new()
8967 } else {
8968 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
8969 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
8970 };
8971
8972 let overlay = if overlay_rows.is_empty() {
8974 None
8975 } else {
8976 let filtered =
8977 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8978 if filtered.is_empty() {
8979 None
8980 } else {
8981 Some(self.materialize_overlay(&filtered, &projection))
8982 }
8983 };
8984
8985 let overlay_row_count = overlay
8986 .as_ref()
8987 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8988 .unwrap_or(0);
8989 crate::trace::QueryTrace::record(|t| {
8990 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
8991 t.run_count = self.run_refs.len();
8992 t.memtable_rows = self.memtable.len();
8993 t.mutable_run_rows = self.mutable_run.len();
8994 t.overlay_rows = overlay_row_count;
8995 t.conditions_pushed = conditions.len();
8996 t.pages_decoded = plans
8997 .iter()
8998 .map(|p| p.positions.len())
8999 .sum::<usize>()
9000 .min(1);
9001 });
9002
9003 Ok(Some(NativePageCursor::new_with_overlay(
9004 reader, projection, plans, overlay,
9005 )))
9006 }
9007 #[allow(clippy::type_complexity)]
9017 pub fn native_multi_run_cursor(
9018 &self,
9019 snapshot: Snapshot,
9020 projection: Vec<(u16, TypeId)>,
9021 conditions: &[crate::query::Condition],
9022 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9023 use crate::cursor::{MultiRunCursor, RunStream};
9024 use crate::sorted_run::SYS_ROW_ID;
9025 use std::collections::{BinaryHeap, HashMap, HashSet};
9026 if self.ttl.is_some() {
9027 return Ok(None);
9028 }
9029 if !conditions.is_empty() && !self.indexes_complete {
9032 return Ok(None);
9033 }
9034 if self.run_refs.is_empty() {
9035 return Ok(None);
9036 }
9037
9038 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9040 Vec::with_capacity(self.run_refs.len());
9041 for rr in &self.run_refs {
9042 let mut reader = self.open_reader(rr.run_id)?;
9043 let (rids, eps, del) = reader.system_columns_native()?;
9044 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9045 run_meta.push((reader, rids, eps, del, page_rows));
9046 }
9047
9048 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9052 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9053 for i in 0..rids.len() {
9054 let rid = rids[i] as u64;
9055 let e = eps[i] as u64;
9056 if e > snapshot.epoch.0 {
9057 continue;
9058 }
9059 let is_del = del[i] != 0;
9060 best.entry(rid)
9061 .and_modify(|cur| {
9062 if e > cur.0 {
9063 *cur = (e, run_idx, i, is_del);
9064 }
9065 })
9066 .or_insert((e, run_idx, i, is_del));
9067 }
9068 }
9069
9070 let overlay_rids: HashSet<u64> = {
9072 let mut s = HashSet::new();
9073 for row in self.memtable.visible_versions(snapshot.epoch) {
9074 s.insert(row.row_id.0);
9075 }
9076 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9077 s.insert(row.row_id.0);
9078 }
9079 s
9080 };
9081
9082 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9084 None
9085 } else {
9086 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9087 for c in conditions {
9088 sets.push(self.resolve_condition(c, snapshot)?);
9089 }
9090 Some(RowIdSet::intersect_many(sets))
9091 };
9092
9093 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9097 for (rid, (_, run_idx, pos, deleted)) in &best {
9098 if *deleted {
9099 continue;
9100 }
9101 if overlay_rids.contains(rid) {
9102 continue;
9103 }
9104 if let Some(s) = &survivors {
9105 if !s.contains(*rid) {
9106 continue;
9107 }
9108 }
9109 per_run[*run_idx].push((*rid, *pos));
9110 }
9111 for v in per_run.iter_mut() {
9112 v.sort_unstable_by_key(|&(rid, _)| rid);
9113 }
9114
9115 let mut streams = Vec::with_capacity(run_meta.len());
9117 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9118 let mut total = 0usize;
9119 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9120 let mut starts = Vec::with_capacity(page_rows.len());
9121 let mut acc = 0usize;
9122 for &r in &page_rows {
9123 starts.push(acc);
9124 acc += r;
9125 }
9126 let mut survivors_vec: Vec<(u64, usize, usize)> =
9127 Vec::with_capacity(per_run[run_idx].len());
9128 for &(rid, pos) in &per_run[run_idx] {
9129 let page_seq = match starts.partition_point(|&s| s <= pos) {
9130 0 => continue,
9131 p => p - 1,
9132 };
9133 let within = pos - starts[page_seq];
9134 survivors_vec.push((rid, page_seq, within));
9135 }
9136 total += survivors_vec.len();
9137 if let Some(&(rid, _, _)) = survivors_vec.first() {
9138 heap.push(std::cmp::Reverse((rid, run_idx)));
9139 }
9140 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9141 }
9142
9143 let overlay_rows = if overlay_rids.is_empty() {
9145 Vec::new()
9146 } else {
9147 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9148 self.overlay_visible_rows(snapshot, bound)
9149 };
9150 let overlay = if overlay_rows.is_empty() {
9151 None
9152 } else {
9153 let filtered =
9154 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9155 if filtered.is_empty() {
9156 None
9157 } else {
9158 Some(self.materialize_overlay(&filtered, &projection))
9159 }
9160 };
9161
9162 let overlay_row_count = overlay
9163 .as_ref()
9164 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9165 .unwrap_or(0);
9166 crate::trace::QueryTrace::record(|t| {
9167 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9168 t.run_count = self.run_refs.len();
9169 t.memtable_rows = self.memtable.len();
9170 t.mutable_run_rows = self.mutable_run.len();
9171 t.overlay_rows = overlay_row_count;
9172 t.conditions_pushed = conditions.len();
9173 t.survivor_count = Some(total);
9174 });
9175
9176 Ok(Some(MultiRunCursor::new(
9177 streams, projection, heap, total, overlay,
9178 )))
9179 }
9180
9181 fn overlay_materialization_bound<'a>(
9193 conditions: &[crate::query::Condition],
9194 survivors: &'a Option<RowIdSet>,
9195 ) -> Option<&'a RowIdSet> {
9196 use crate::query::Condition;
9197 let has_range = conditions
9198 .iter()
9199 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9200 if has_range {
9201 None
9202 } else {
9203 survivors.as_ref()
9204 }
9205 }
9206
9207 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9219 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9220 let mut fold = |row: Row| {
9221 if let Some(b) = bound {
9222 if !b.contains(row.row_id.0) {
9223 return;
9224 }
9225 }
9226 best.entry(row.row_id.0)
9227 .and_modify(|(be, br)| {
9228 if row.committed_epoch > *be {
9229 *be = row.committed_epoch;
9230 *br = row.clone();
9231 }
9232 })
9233 .or_insert_with(|| (row.committed_epoch, row));
9234 };
9235 for row in self.memtable.visible_versions(snapshot.epoch) {
9236 fold(row);
9237 }
9238 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9239 fold(row);
9240 }
9241 let mut out: Vec<Row> = best
9242 .into_values()
9243 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9244 .collect();
9245 out.sort_by_key(|r| r.row_id);
9246 out
9247 }
9248
9249 fn filter_overlay_rows(
9257 &self,
9258 rows: Vec<Row>,
9259 conditions: &[crate::query::Condition],
9260 survivors: Option<&RowIdSet>,
9261 snapshot: Snapshot,
9262 ) -> Result<Vec<Row>> {
9263 if conditions.is_empty() {
9264 return Ok(rows);
9265 }
9266 use crate::query::Condition;
9267 let all_index_served = !conditions
9271 .iter()
9272 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9273 if all_index_served {
9274 return Ok(rows
9275 .into_iter()
9276 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9277 .collect());
9278 }
9279 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9282 for c in conditions {
9283 let s = match c {
9284 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9285 _ => self.resolve_condition(c, snapshot)?,
9286 };
9287 per_cond_sets.push(s);
9288 }
9289 Ok(rows
9290 .into_iter()
9291 .filter(|row| {
9292 conditions.iter().enumerate().all(|(i, c)| match c {
9293 Condition::Range { column_id, lo, hi } => {
9294 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9295 }
9296 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9297 match row.columns.get(column_id) {
9298 Some(Value::Float64(v)) => {
9299 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9300 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9301 lo_ok && hi_ok
9302 }
9303 _ => false,
9304 }
9305 }
9306 _ => per_cond_sets[i].contains(row.row_id.0),
9307 })
9308 })
9309 .collect())
9310 }
9311
9312 fn materialize_overlay(
9315 &self,
9316 rows: &[Row],
9317 projection: &[(u16, TypeId)],
9318 ) -> Vec<columnar::NativeColumn> {
9319 if projection.is_empty() {
9320 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9321 }
9322 let mut cols = Vec::with_capacity(projection.len());
9323 for (cid, ty) in projection {
9324 let vals: Vec<Value> = rows
9325 .iter()
9326 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9327 .collect();
9328 cols.push(columnar::values_to_native(ty.clone(), &vals));
9329 }
9330 cols
9331 }
9332
9333 fn resolve_survivor_rids(
9338 &self,
9339 conditions: &[crate::query::Condition],
9340 reader: &mut RunReader,
9341 snapshot: Snapshot,
9342 ) -> Result<RowIdSet> {
9343 use crate::query::Condition;
9344 let mut sets: Vec<RowIdSet> = Vec::new();
9345 for c in conditions {
9346 self.validate_condition(c)?;
9347 let s: RowIdSet = match c {
9348 Condition::Pk(key) => {
9349 let lookup = self
9350 .schema
9351 .primary_key()
9352 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9353 .unwrap_or_else(|| key.clone());
9354 self.hot
9355 .get(&lookup)
9356 .map(|r| RowIdSet::one(r.0))
9357 .unwrap_or_else(RowIdSet::empty)
9358 }
9359 Condition::BitmapEq { column_id, value } => {
9360 let lookup = self.index_lookup_key_bytes(*column_id, value);
9361 self.bitmap
9362 .get(column_id)
9363 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9364 .unwrap_or_else(RowIdSet::empty)
9365 }
9366 Condition::BitmapIn { column_id, values } => {
9367 let bm = self.bitmap.get(column_id);
9368 let mut acc = roaring::RoaringBitmap::new();
9369 if let Some(b) = bm {
9370 for v in values {
9371 let lookup = self.index_lookup_key_bytes(*column_id, v);
9372 acc |= b.get(&lookup);
9373 }
9374 }
9375 RowIdSet::from_roaring(acc)
9376 }
9377 Condition::BytesPrefix { column_id, prefix } => {
9378 if let Some(b) = self.bitmap.get(column_id) {
9379 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9380 let mut acc = roaring::RoaringBitmap::new();
9381 for key in b.keys() {
9382 if key.starts_with(&lookup_prefix) {
9383 acc |= b.get(&key);
9384 }
9385 }
9386 RowIdSet::from_roaring(acc)
9387 } else {
9388 RowIdSet::empty()
9389 }
9390 }
9391 Condition::FmContains { column_id, pattern } => self
9392 .fm
9393 .get(column_id)
9394 .map(|f| {
9395 RowIdSet::from_unsorted(
9396 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9397 )
9398 })
9399 .unwrap_or_else(RowIdSet::empty),
9400 Condition::FmContainsAll {
9401 column_id,
9402 patterns,
9403 } => {
9404 if let Some(f) = self.fm.get(column_id) {
9405 let sets: Vec<RowIdSet> = patterns
9406 .iter()
9407 .map(|pat| {
9408 RowIdSet::from_unsorted(
9409 f.locate(pat).into_iter().map(|r| r.0).collect(),
9410 )
9411 })
9412 .collect();
9413 RowIdSet::intersect_many(sets)
9414 } else {
9415 RowIdSet::empty()
9416 }
9417 }
9418 Condition::Ann {
9419 column_id,
9420 query,
9421 k,
9422 } => RowIdSet::from_unsorted(
9423 self.retrieve_filtered(
9424 &crate::query::Retriever::Ann {
9425 column_id: *column_id,
9426 query: query.clone(),
9427 k: *k,
9428 },
9429 snapshot,
9430 None,
9431 None,
9432 None,
9433 None,
9434 )?
9435 .into_iter()
9436 .map(|hit| hit.row_id.0)
9437 .collect(),
9438 ),
9439 Condition::SparseMatch {
9440 column_id,
9441 query,
9442 k,
9443 } => RowIdSet::from_unsorted(
9444 self.retrieve_filtered(
9445 &crate::query::Retriever::Sparse {
9446 column_id: *column_id,
9447 query: query.clone(),
9448 k: *k,
9449 },
9450 snapshot,
9451 None,
9452 None,
9453 None,
9454 None,
9455 )?
9456 .into_iter()
9457 .map(|hit| hit.row_id.0)
9458 .collect(),
9459 ),
9460 Condition::MinHashSimilar {
9461 column_id,
9462 query,
9463 k,
9464 } => match self.minhash.get(column_id) {
9465 Some(index) => {
9466 let candidates = index.candidate_row_ids(query);
9467 let eligible =
9468 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9469 RowIdSet::from_unsorted(
9470 index
9471 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9472 .into_iter()
9473 .map(|(row_id, _)| row_id.0)
9474 .collect(),
9475 )
9476 }
9477 None => RowIdSet::empty(),
9478 },
9479 Condition::Range { column_id, lo, hi } => {
9480 if let Some(li) = self.learned_range.get(column_id) {
9481 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9482 } else {
9483 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9484 }
9485 }
9486 Condition::RangeF64 {
9487 column_id,
9488 lo,
9489 lo_inclusive,
9490 hi,
9491 hi_inclusive,
9492 } => {
9493 if let Some(li) = self.learned_range.get(column_id) {
9494 RowIdSet::from_unsorted(
9495 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9496 .into_iter()
9497 .collect(),
9498 )
9499 } else {
9500 reader.range_row_id_set_f64(
9501 *column_id,
9502 *lo,
9503 *lo_inclusive,
9504 *hi,
9505 *hi_inclusive,
9506 )?
9507 }
9508 }
9509 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9510 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9511 };
9512 sets.push(s);
9513 }
9514 Ok(RowIdSet::intersect_many(sets))
9515 }
9516
9517 pub fn scan_cursor(
9538 &self,
9539 snapshot: Snapshot,
9540 projection: Vec<(u16, TypeId)>,
9541 conditions: &[crate::query::Condition],
9542 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9543 if self.ttl.is_some() {
9544 return Ok(None);
9545 }
9546 if !conditions.is_empty() && !self.indexes_complete {
9552 return Ok(None);
9553 }
9554 if self.run_refs.len() == 1 {
9555 Ok(self
9556 .native_page_cursor(snapshot, projection, conditions)?
9557 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9558 } else {
9559 Ok(self
9560 .native_multi_run_cursor(snapshot, projection, conditions)?
9561 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9562 }
9563 }
9564
9565 pub fn aggregate_native(
9579 &self,
9580 snapshot: Snapshot,
9581 column: Option<u16>,
9582 conditions: &[crate::query::Condition],
9583 agg: NativeAgg,
9584 ) -> Result<Option<NativeAggResult>> {
9585 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9586 }
9587
9588 pub fn aggregate_native_with_control(
9589 &self,
9590 snapshot: Snapshot,
9591 column: Option<u16>,
9592 conditions: &[crate::query::Condition],
9593 agg: NativeAgg,
9594 control: &crate::ExecutionControl,
9595 ) -> Result<Option<NativeAggResult>> {
9596 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9597 }
9598
9599 fn aggregate_native_inner(
9600 &self,
9601 snapshot: Snapshot,
9602 column: Option<u16>,
9603 conditions: &[crate::query::Condition],
9604 agg: NativeAgg,
9605 control: Option<&crate::ExecutionControl>,
9606 ) -> Result<Option<NativeAggResult>> {
9607 execution_checkpoint(control, 0)?;
9608 if self.ttl.is_some() {
9609 return Ok(None);
9610 }
9611 if self.run_refs.len() == 1 && conditions.is_empty() {
9613 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9614 return Ok(Some(res));
9615 }
9616 }
9617 if matches!(agg, NativeAgg::Count) && column.is_none() {
9621 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
9622 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
9623 }
9624 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9625 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
9626 }
9627 let cid = match column {
9630 Some(c) => c,
9631 None => return Ok(None),
9632 };
9633 let ty = self.column_type(cid);
9634 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
9635 execution_checkpoint(control, 0)?;
9636 return match ty {
9637 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9638 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9639 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9640 }
9641 TypeId::Float64 => {
9642 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9643 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9644 }
9645 _ => Ok(None),
9646 };
9647 }
9648 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9650 execution_checkpoint(control, 0)?;
9651 match ty {
9652 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9653 let mut count = 0u64;
9654 let mut sum = 0i128;
9655 let mut mn = i64::MAX;
9656 let mut mx = i64::MIN;
9657 for row in &rows {
9658 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
9659 count += 1;
9660 sum += i128::from(*v);
9661 mn = mn.min(*v);
9662 mx = mx.max(*v);
9663 }
9664 }
9665 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9666 }
9667 TypeId::Float64 => {
9668 let mut count = 0u64;
9669 let mut sum = 0.0f64;
9670 let mut mn = f64::INFINITY;
9671 let mut mx = f64::NEG_INFINITY;
9672 for row in &rows {
9673 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
9674 count += 1;
9675 sum += *v;
9676 mn = mn.min(*v);
9677 mx = mx.max(*v);
9678 }
9679 }
9680 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9681 }
9682 _ => Ok(None),
9683 }
9684 }
9685
9686 fn visible_rows_filtered(
9688 &self,
9689 snapshot: Snapshot,
9690 conditions: &[crate::query::Condition],
9691 control: Option<&crate::ExecutionControl>,
9692 ) -> Result<Vec<Row>> {
9693 let rows = if let Some(control) = control {
9694 self.visible_rows_controlled(snapshot, control)?
9695 } else {
9696 self.visible_rows(snapshot)?
9697 };
9698 if conditions.is_empty() {
9699 return Ok(rows);
9700 }
9701 Ok(rows
9702 .into_iter()
9703 .filter(|row| {
9704 conditions
9705 .iter()
9706 .all(|cond| condition_matches_row(cond, row, &self.schema))
9707 })
9708 .collect())
9709 }
9710
9711 fn aggregate_from_stats(
9719 &self,
9720 snapshot: Snapshot,
9721 column: Option<u16>,
9722 agg: NativeAgg,
9723 ) -> Result<Option<NativeAggResult>> {
9724 let cid = match (agg, column) {
9725 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9726 _ => return Ok(None), };
9728 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9729 return Ok(None);
9730 };
9731 let Some(cs) = stats.get(&cid) else {
9732 return Ok(None);
9733 };
9734 match agg {
9735 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9737 self.live_count.saturating_sub(cs.null_count),
9738 ))),
9739 NativeAgg::Min | NativeAgg::Max => {
9740 let bound = if agg == NativeAgg::Min {
9741 &cs.min
9742 } else {
9743 &cs.max
9744 };
9745 match bound {
9746 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9747 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9748 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9753 None => Ok(None),
9754 }
9755 }
9756 _ => Ok(None),
9757 }
9758 }
9759
9760 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9769 if self.ttl.is_some() {
9770 return Ok(None);
9771 }
9772 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9773 return Ok(None);
9774 }
9775 self.ensure_indexes_complete()?;
9778 let reader = self.open_reader(self.run_refs[0].run_id)?;
9779 if self.live_count != reader.row_count() as u64 {
9780 return Ok(None);
9781 }
9782 let Some(bm) = self.bitmap.get(&column_id) else {
9783 return Ok(None); };
9785 let mut distinct = bm.value_count() as u64;
9786 if !bm.get(&Value::Null.encode_key()).is_empty() {
9789 distinct = distinct.saturating_sub(1);
9790 }
9791 Ok(Some(distinct))
9792 }
9793
9794 pub fn aggregate_incremental(
9806 &mut self,
9807 cache_key: u64,
9808 conditions: &[crate::query::Condition],
9809 column: Option<u16>,
9810 agg: NativeAgg,
9811 ) -> Result<IncrementalAggResult> {
9812 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9813 }
9814
9815 pub fn aggregate_incremental_with_control(
9816 &mut self,
9817 cache_key: u64,
9818 conditions: &[crate::query::Condition],
9819 column: Option<u16>,
9820 agg: NativeAgg,
9821 control: &crate::ExecutionControl,
9822 ) -> Result<IncrementalAggResult> {
9823 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9824 }
9825
9826 fn aggregate_incremental_inner(
9827 &mut self,
9828 cache_key: u64,
9829 conditions: &[crate::query::Condition],
9830 column: Option<u16>,
9831 agg: NativeAgg,
9832 control: Option<&crate::ExecutionControl>,
9833 ) -> Result<IncrementalAggResult> {
9834 execution_checkpoint(control, 0)?;
9835 let snap = self.snapshot();
9836 let cur_wm = self.allocator.current().0;
9837 let cur_epoch = snap.epoch.0;
9838 let incremental_ok = self.ttl.is_none()
9845 && !self.had_deletes
9846 && self.memtable.is_empty()
9847 && self.mutable_run.is_empty();
9848
9849 if incremental_ok {
9852 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
9853 if cached.epoch == cur_epoch {
9854 return Ok(IncrementalAggResult {
9855 state: cached.state,
9856 incremental: true,
9857 delta_rows: 0,
9858 });
9859 }
9860 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
9861 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
9862 let mut delta_rids = Vec::with_capacity(delta_len);
9863 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
9864 execution_checkpoint(control, index)?;
9865 delta_rids.push(row_id);
9866 }
9867 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
9868 execution_checkpoint(control, 0)?;
9869 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9870 let delta_state = agg_state_from_rows(
9871 &delta_rows,
9872 conditions,
9873 &index_sets,
9874 column,
9875 agg,
9876 &self.schema,
9877 control,
9878 )?;
9879 let merged = cached.state.merge(delta_state);
9880 let delta_n = delta_rids.len() as u64;
9881 Arc::make_mut(&mut self.agg_cache).insert(
9882 cache_key,
9883 CachedAgg {
9884 state: merged.clone(),
9885 watermark: cur_wm,
9886 epoch: cur_epoch,
9887 },
9888 );
9889 return Ok(IncrementalAggResult {
9890 state: merged,
9891 incremental: true,
9892 delta_rows: delta_n,
9893 });
9894 }
9895 }
9896 }
9897
9898 let cursor_ok =
9903 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9904 let state = if cursor_ok && agg != NativeAgg::Avg {
9905 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
9906 Some(result) => {
9907 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
9908 }
9909 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
9910 }
9911 } else {
9912 self.agg_state_full_scan(conditions, column, agg, snap, control)?
9913 };
9914 if incremental_ok {
9916 Arc::make_mut(&mut self.agg_cache).insert(
9917 cache_key,
9918 CachedAgg {
9919 state: state.clone(),
9920 watermark: cur_wm,
9921 epoch: cur_epoch,
9922 },
9923 );
9924 }
9925 Ok(IncrementalAggResult {
9926 state,
9927 incremental: false,
9928 delta_rows: 0,
9929 })
9930 }
9931
9932 fn agg_state_full_scan(
9935 &self,
9936 conditions: &[crate::query::Condition],
9937 column: Option<u16>,
9938 agg: NativeAgg,
9939 snap: Snapshot,
9940 control: Option<&crate::ExecutionControl>,
9941 ) -> Result<AggState> {
9942 execution_checkpoint(control, 0)?;
9943 let rows = self.visible_rows(snap)?;
9944 execution_checkpoint(control, 0)?;
9945 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9946 agg_state_from_rows(
9947 &rows,
9948 conditions,
9949 &index_sets,
9950 column,
9951 agg,
9952 &self.schema,
9953 control,
9954 )
9955 }
9956
9957 fn resolve_index_conditions(
9960 &self,
9961 conditions: &[crate::query::Condition],
9962 snapshot: Snapshot,
9963 ) -> Result<Vec<RowIdSet>> {
9964 use crate::query::Condition;
9965 let mut sets = Vec::new();
9966 for c in conditions {
9967 if matches!(
9968 c,
9969 Condition::Ann { .. }
9970 | Condition::SparseMatch { .. }
9971 | Condition::MinHashSimilar { .. }
9972 ) {
9973 sets.push(self.resolve_condition(c, snapshot)?);
9974 }
9975 }
9976 Ok(sets)
9977 }
9978
9979 fn column_type(&self, cid: u16) -> TypeId {
9980 self.schema
9981 .columns
9982 .iter()
9983 .find(|c| c.id == cid)
9984 .map(|c| c.ty.clone())
9985 .unwrap_or(TypeId::Bytes)
9986 }
9987
9988 pub fn approx_aggregate(
9997 &mut self,
9998 conditions: &[crate::query::Condition],
9999 column: Option<u16>,
10000 agg: ApproxAgg,
10001 z: f64,
10002 ) -> Result<Option<ApproxResult>> {
10003 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10004 }
10005
10006 pub fn approx_aggregate_with_candidate_authorization(
10009 &mut self,
10010 conditions: &[crate::query::Condition],
10011 column: Option<u16>,
10012 agg: ApproxAgg,
10013 z: f64,
10014 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10015 ) -> Result<Option<ApproxResult>> {
10016 use crate::query::Condition;
10017 self.ensure_reservoir_complete()?;
10018 let snapshot = self.snapshot();
10019 let n_pop = self.count();
10020 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10021 if sample_rids.is_empty() {
10022 return Ok(None);
10023 }
10024 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10026 let s = live_sample.len();
10027 if s == 0 {
10028 return Ok(None);
10029 }
10030 let authorized = authorization
10031 .map(|authorization| {
10032 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10033 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10034 })
10035 .transpose()?;
10036
10037 let mut index_sets: Vec<RowIdSet> = Vec::new();
10040 for c in conditions {
10041 if matches!(
10042 c,
10043 Condition::Ann { .. }
10044 | Condition::SparseMatch { .. }
10045 | Condition::MinHashSimilar { .. }
10046 ) {
10047 index_sets.push(self.resolve_condition(c, snapshot)?);
10048 }
10049 }
10050
10051 let cid = match (agg, column) {
10053 (ApproxAgg::Count, _) => None,
10054 (_, Some(c)) => Some(c),
10055 _ => return Ok(None),
10056 };
10057 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10058 for r in &live_sample {
10059 if authorized
10060 .as_ref()
10061 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10062 {
10063 continue;
10064 }
10065 if !conditions
10067 .iter()
10068 .all(|c| condition_matches_row(c, r, &self.schema))
10069 {
10070 continue;
10071 }
10072 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10074 continue;
10075 }
10076 if let Some(cid) = cid {
10077 let mut cells = r
10078 .columns
10079 .get(&cid)
10080 .cloned()
10081 .map(|value| vec![(cid, value)])
10082 .unwrap_or_default();
10083 if let Some(authorization) = authorization {
10084 authorization.security.apply_masks_to_cells(
10085 authorization.table,
10086 &mut cells,
10087 authorization.principal,
10088 );
10089 }
10090 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10091 passing_vals.push(v);
10092 } } else {
10094 passing_vals.push(0.0); }
10096 }
10097 let m = passing_vals.len();
10098
10099 let (point, half) = match agg {
10100 ApproxAgg::Count => {
10101 let p = m as f64 / s as f64;
10103 let point = n_pop as f64 * p;
10104 let var = if s > 1 {
10105 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10106 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10107 } else {
10108 0.0
10109 };
10110 (point, z * var.sqrt())
10111 }
10112 ApproxAgg::Sum => {
10113 let y: Vec<f64> = live_sample
10115 .iter()
10116 .map(|r| {
10117 let passes_row = authorized
10118 .as_ref()
10119 .is_none_or(|authorized| authorized.contains(&r.row_id))
10120 && conditions
10121 .iter()
10122 .all(|c| condition_matches_row(c, r, &self.schema))
10123 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10124 if passes_row {
10125 cid.and_then(|cid| {
10126 let mut cells = r
10127 .columns
10128 .get(&cid)
10129 .cloned()
10130 .map(|value| vec![(cid, value)])
10131 .unwrap_or_default();
10132 if let Some(authorization) = authorization {
10133 authorization.security.apply_masks_to_cells(
10134 authorization.table,
10135 &mut cells,
10136 authorization.principal,
10137 );
10138 }
10139 as_f64(cells.first().map(|(_, value)| value))
10140 })
10141 .unwrap_or(0.0)
10142 } else {
10143 0.0
10144 }
10145 })
10146 .collect();
10147 let mean_y = y.iter().sum::<f64>() / s as f64;
10148 let point = n_pop as f64 * mean_y;
10149 let var = if s > 1 {
10150 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10151 let var_y = ss / (s - 1) as f64;
10152 n_pop as f64 * n_pop as f64 * var_y / s as f64
10153 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10154 } else {
10155 0.0
10156 };
10157 (point, z * var.sqrt())
10158 }
10159 ApproxAgg::Avg => {
10160 if m == 0 {
10161 return Ok(Some(ApproxResult {
10162 point: 0.0,
10163 ci_low: 0.0,
10164 ci_high: 0.0,
10165 n_population: n_pop,
10166 n_sample_live: s,
10167 n_passing: 0,
10168 }));
10169 }
10170 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10171 let half = if m > 1 {
10172 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10173 let sd = (ss / (m - 1) as f64).sqrt();
10174 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10175 z * sd / (m as f64).sqrt() * fpc.sqrt()
10176 } else {
10177 0.0
10178 };
10179 (mean, half)
10180 }
10181 };
10182
10183 Ok(Some(ApproxResult {
10184 point,
10185 ci_low: point - half,
10186 ci_high: point + half,
10187 n_population: n_pop,
10188 n_sample_live: s,
10189 n_passing: m,
10190 }))
10191 }
10192
10193 pub fn exact_column_stats(
10201 &self,
10202 _snapshot: Snapshot,
10203 projection: &[u16],
10204 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10205 if self.ttl.is_some()
10206 || !(self.memtable.is_empty()
10207 && self.mutable_run.is_empty()
10208 && self.run_refs.len() == 1)
10209 {
10210 return Ok(None);
10211 }
10212 let reader = self.open_reader(self.run_refs[0].run_id)?;
10213 if self.live_count != reader.row_count() as u64 {
10214 return Ok(None);
10215 }
10216 let mut out = HashMap::new();
10217 for &cid in projection {
10218 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10219 Some(c) => c,
10220 None => continue,
10221 };
10222 let Some(stats) = reader.column_page_stats(cid) else {
10224 out.insert(
10225 cid,
10226 ColumnStat {
10227 min: None,
10228 max: None,
10229 null_count: self.live_count,
10230 },
10231 );
10232 continue;
10233 };
10234 let stat = match cdef.ty {
10235 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10236 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10237 min: mn.map(Value::Int64),
10238 max: mx.map(Value::Int64),
10239 null_count: n,
10240 })
10241 }
10242 TypeId::Float64 => {
10243 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10244 min: mn.map(Value::Float64),
10245 max: mx.map(Value::Float64),
10246 null_count: n,
10247 })
10248 }
10249 _ => None,
10250 };
10251 if let Some(s) = stat {
10252 out.insert(cid, s);
10253 }
10254 }
10255 Ok(Some(out))
10256 }
10257
10258 pub fn dir(&self) -> &Path {
10259 &self.dir
10260 }
10261
10262 pub fn schema(&self) -> &Schema {
10263 &self.schema
10264 }
10265
10266 pub(crate) fn set_catalog_name(&mut self, name: String) {
10267 self.name = name;
10268 }
10269
10270 pub(crate) fn prepare_alter_column(
10271 &mut self,
10272 column_name: &str,
10273 change: &AlterColumn,
10274 ) -> Result<(ColumnDef, Option<Schema>)> {
10275 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10276 return Err(MongrelError::InvalidArgument(
10277 "ALTER COLUMN requires committing staged writes first".into(),
10278 ));
10279 }
10280 let old = self
10281 .schema
10282 .columns
10283 .iter()
10284 .find(|c| c.name == column_name)
10285 .cloned()
10286 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10287 let mut next = old.clone();
10288
10289 if let Some(name) = &change.name {
10290 let trimmed = name.trim();
10291 if trimmed.is_empty() {
10292 return Err(MongrelError::InvalidArgument(
10293 "ALTER COLUMN name must not be empty".into(),
10294 ));
10295 }
10296 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10297 return Err(MongrelError::Schema(format!(
10298 "column {trimmed} already exists"
10299 )));
10300 }
10301 next.name = trimmed.to_string();
10302 }
10303
10304 if let Some(ty) = &change.ty {
10305 next.ty = ty.clone();
10306 }
10307 if let Some(flags) = change.flags {
10308 validate_alter_column_flags(old.flags, flags)?;
10309 next.flags = flags;
10310 }
10311
10312 if let Some(default_change) = &change.default_value {
10313 next.default_value = default_change.clone();
10314 }
10315 if let Some(source_change) = &change.embedding_source {
10316 next.embedding_source = source_change.clone();
10317 }
10318
10319 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10320 if old.flags.contains(ColumnFlags::NULLABLE)
10321 && !next.flags.contains(ColumnFlags::NULLABLE)
10322 && self.column_has_nulls(old.id)?
10323 {
10324 return Err(MongrelError::InvalidArgument(format!(
10325 "column '{}' contains NULL values",
10326 old.name
10327 )));
10328 }
10329 if next == old {
10330 return Ok((next, None));
10331 }
10332 let mut schema = self.schema.clone();
10333 let index = schema
10334 .columns
10335 .iter()
10336 .position(|column| column.id == next.id)
10337 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10338 schema.columns[index] = next.clone();
10339 schema.schema_id = schema
10340 .schema_id
10341 .checked_add(1)
10342 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10343 schema.validate_auto_increment()?;
10344 schema.validate_defaults()?;
10345 Ok((next, Some(schema)))
10346 }
10347
10348 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10349 self.schema = schema;
10350 self.auto_inc = resolve_auto_inc(&self.schema);
10351 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10352 self.clear_result_cache();
10353 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10354 }
10355
10356 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10357 checkpoint_current_schema(self)
10358 }
10359
10360 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10361 self.ensure_writable()?;
10362 let previous_schema = self.schema.clone();
10363 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10364 if let Some(schema) = schema {
10365 self.apply_altered_schema_prepared(schema);
10366 self.checkpoint_standalone_schema_change(previous_schema)?;
10367 }
10368 Ok(column)
10369 }
10370
10371 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10372 if self.live_count == 0 {
10373 return Ok(false);
10374 }
10375 let snap = self.snapshot();
10376 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10377 Ok(columns
10378 .first()
10379 .map(|(_, col)| col.null_count(col.len()) != 0)
10380 .unwrap_or(true))
10381 }
10382
10383 fn has_stored_versions(&self) -> bool {
10384 !self.memtable.is_empty()
10385 || !self.mutable_run.is_empty()
10386 || self.run_refs.iter().any(|r| r.row_count > 0)
10387 || !self.retiring.is_empty()
10388 }
10389
10390 pub fn add_column(
10395 &mut self,
10396 name: &str,
10397 ty: TypeId,
10398 flags: ColumnFlags,
10399 default_value: Option<crate::schema::DefaultExpr>,
10400 ) -> Result<u16> {
10401 self.add_column_with_id(name, ty, flags, default_value, None)
10402 }
10403
10404 pub fn add_column_with_id(
10405 &mut self,
10406 name: &str,
10407 ty: TypeId,
10408 flags: ColumnFlags,
10409 default_value: Option<crate::schema::DefaultExpr>,
10410 requested_id: Option<u16>,
10411 ) -> Result<u16> {
10412 self.ensure_writable()?;
10413 if self.schema.columns.iter().any(|c| c.name == name) {
10414 return Err(MongrelError::Schema(format!(
10415 "column {name} already exists"
10416 )));
10417 }
10418 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10419 if self.schema.columns.iter().any(|c| c.id == id) {
10420 return Err(MongrelError::Schema(format!(
10421 "column id {id} already exists"
10422 )));
10423 }
10424 id
10425 } else {
10426 self.schema
10427 .columns
10428 .iter()
10429 .map(|c| c.id)
10430 .max()
10431 .unwrap_or(0)
10432 .checked_add(1)
10433 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10434 };
10435 let previous_schema = self.schema.clone();
10436 let mut next_schema = previous_schema.clone();
10437 next_schema.columns.push(ColumnDef {
10438 id,
10439 name: name.to_string(),
10440 ty,
10441 flags,
10442 default_value,
10443 embedding_source: None,
10444 });
10445 next_schema.schema_id = next_schema
10446 .schema_id
10447 .checked_add(1)
10448 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10449 next_schema.validate_auto_increment()?;
10450 next_schema.validate_defaults()?;
10451 self.apply_altered_schema_prepared(next_schema);
10452 self.checkpoint_standalone_schema_change(previous_schema)?;
10453 Ok(id)
10454 }
10455
10456 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10465 self.ensure_writable()?;
10466 let cid = self
10467 .schema
10468 .columns
10469 .iter()
10470 .find(|c| c.name == column_name)
10471 .map(|c| c.id)
10472 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10473 let ty = self
10474 .schema
10475 .columns
10476 .iter()
10477 .find(|c| c.id == cid)
10478 .map(|c| c.ty.clone())
10479 .unwrap_or(TypeId::Int64);
10480 if !matches!(
10481 ty,
10482 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10483 ) {
10484 return Err(MongrelError::Schema(format!(
10485 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10486 )));
10487 }
10488 if self
10489 .schema
10490 .indexes
10491 .iter()
10492 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10493 {
10494 return Ok(()); }
10496 let previous_schema = self.schema.clone();
10497 let previous_learned_range = Arc::clone(&self.learned_range);
10498 let mut next_schema = previous_schema.clone();
10499 next_schema.indexes.push(IndexDef {
10500 name: format!("{}_learned_range", column_name),
10501 column_id: cid,
10502 kind: IndexKind::LearnedRange,
10503 predicate: None,
10504 options: Default::default(),
10505 });
10506 next_schema.schema_id = next_schema
10507 .schema_id
10508 .checked_add(1)
10509 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10510 self.apply_altered_schema_prepared(next_schema);
10511 if let Err(error) = self.build_learned_ranges() {
10512 self.apply_altered_schema_prepared(previous_schema);
10513 self.learned_range = previous_learned_range;
10514 return Err(error);
10515 }
10516 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10517 if !matches!(
10518 &error,
10519 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10520 ) {
10521 self.learned_range = previous_learned_range;
10522 }
10523 return Err(error);
10524 }
10525 Ok(())
10526 }
10527
10528 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10529 let mut schema_published = false;
10530 let schema_result = match self._root_guard.as_deref() {
10531 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10532 schema_published = true;
10533 }),
10534 None => write_schema_with_after(&self.dir, &self.schema, || {
10535 schema_published = true;
10536 }),
10537 };
10538 if schema_result.is_err() && !schema_published {
10539 self.apply_altered_schema_prepared(previous_schema);
10540 return schema_result;
10541 }
10542
10543 let manifest_result = self.persist_manifest(self.current_epoch());
10544 match (schema_result, manifest_result) {
10545 (_, Ok(())) => Ok(()),
10546 (Ok(()), Err(error)) => {
10547 self.poison_after_maintenance_publish_failure();
10548 Err(MongrelError::DurableCommit {
10549 epoch: self.current_epoch().0,
10550 message: format!(
10551 "schema is durable but matching manifest publication failed: {error}"
10552 ),
10553 })
10554 }
10555 (Err(schema_error), Err(manifest_error)) => {
10556 self.poison_after_maintenance_publish_failure();
10557 Err(MongrelError::CommitOutcomeUnknown {
10558 epoch: self.current_epoch().0,
10559 message: format!(
10560 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10561 ),
10562 })
10563 }
10564 }
10565 }
10566
10567 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10570 self.sync_byte_threshold = threshold;
10571 if let WalSink::Private(w) = &mut self.wal {
10572 w.set_sync_byte_threshold(threshold);
10573 }
10574 }
10575
10576 pub fn page_cache_flush(&self) {
10580 self.page_cache.flush_to_disk();
10581 }
10582
10583 pub fn page_cache_len(&self) -> usize {
10585 self.page_cache.len()
10586 }
10587
10588 pub fn decoded_cache_len(&self) -> usize {
10591 self.decoded_cache.len()
10592 }
10593
10594 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10597 self.memtable.drain_sorted()
10598 }
10599
10600 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10601 self.runs_dir().join(format!("r-{run_id}.sr"))
10602 }
10603
10604 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10605 match self.runs_root.as_deref() {
10606 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10607 None => Ok(None),
10608 }
10609 }
10610
10611 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10612 match self.runs_root.as_deref() {
10613 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10614 None => Ok(None),
10615 }
10616 }
10617
10618 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10619 match self.runs_root.as_deref() {
10620 Some(root) => match root.remove_file(name) {
10621 Ok(()) => Ok(()),
10622 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10623 Err(error) => Err(error.into()),
10624 },
10625 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10626 Ok(()) => Ok(()),
10627 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10628 Err(error) => Err(error.into()),
10629 },
10630 }
10631 }
10632
10633 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10634 match self.runs_root.as_deref() {
10635 Some(root) => root
10636 .rename_file_new(source, destination)
10637 .map_err(Into::into),
10638 None => crate::durable_file::rename(
10639 &self.runs_dir().join(source),
10640 &self.runs_dir().join(destination),
10641 )
10642 .map_err(Into::into),
10643 }
10644 }
10645
10646 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10647 self.run_refs.iter().map(|run| run.run_id)
10648 }
10649
10650 pub(crate) fn table_dir(&self) -> &Path {
10651 &self.dir
10652 }
10653
10654 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10655 &self.schema
10656 }
10657
10658 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10659 let id = self.next_run_id;
10660 self.next_run_id = self
10661 .next_run_id
10662 .checked_add(1)
10663 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10664 Ok(id)
10665 }
10666
10667 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10668 self.run_refs.push(run_ref);
10669 }
10670
10671 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10681 self.retiring.push(crate::manifest::RetiredRun {
10682 run_id,
10683 retire_epoch,
10684 });
10685 }
10686
10687 pub(crate) fn reap_retiring(
10691 &mut self,
10692 min_active: Epoch,
10693 backup_pinned: &std::collections::HashSet<u128>,
10694 ) -> Result<usize> {
10695 if self.retiring.is_empty() {
10696 return Ok(0);
10697 }
10698 let mut reaped = 0;
10699 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
10700 for r in std::mem::take(&mut self.retiring) {
10706 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
10707 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
10708 reaped += 1;
10709 } else {
10710 kept.push(r);
10711 }
10712 }
10713 self.retiring = kept;
10714 if reaped > 0 {
10715 self.persist_manifest(self.current_epoch())?;
10716 }
10717 Ok(reaped)
10718 }
10719
10720 pub(crate) fn has_reapable_retiring(
10721 &self,
10722 min_active: Epoch,
10723 backup_pinned: &std::collections::HashSet<u128>,
10724 ) -> bool {
10725 self.retiring
10726 .iter()
10727 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
10728 }
10729
10730 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
10731 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
10732 return false;
10733 }
10734 self.live_count = self.live_count.saturating_add(run_ref.row_count);
10735 self.run_refs.push(run_ref);
10736 self.indexes_complete = false;
10737 true
10738 }
10739
10740 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
10741 self.kek.as_ref()
10742 }
10743
10744 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
10745 let mut reader = match self.runs_root.as_deref() {
10746 Some(root) => RunReader::open_file_with_cache(
10747 root.open_regular(format!("r-{run_id}.sr"))?,
10748 self.schema.clone(),
10749 self.kek.clone(),
10750 Some(self.page_cache.clone()),
10751 Some(self.decoded_cache.clone()),
10752 self.table_id,
10753 Some(&self.verified_runs),
10754 None,
10755 )?,
10756 None => RunReader::open_with_cache(
10757 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
10758 self.schema.clone(),
10759 self.kek.clone(),
10760 Some(self.page_cache.clone()),
10761 Some(self.decoded_cache.clone()),
10762 self.table_id,
10763 Some(&self.verified_runs),
10764 )?,
10765 };
10766 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
10770 reader.set_uniform_epoch(Epoch(rr.epoch_created));
10771 }
10772 Ok(reader)
10773 }
10774
10775 pub(crate) fn run_refs(&self) -> &[RunRef] {
10776 &self.run_refs
10777 }
10778
10779 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10780 self.retiring.iter().map(|run| run.run_id)
10781 }
10782
10783 pub(crate) fn runs_dir(&self) -> PathBuf {
10784 self.runs_root
10785 .as_deref()
10786 .and_then(|root| root.io_path().ok())
10787 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
10788 }
10789
10790 pub(crate) fn wal_dir(&self) -> PathBuf {
10791 self.dir.join(WAL_DIR)
10792 }
10793
10794 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
10795 self.run_refs = refs;
10796 }
10797
10798 pub(crate) fn compaction_zstd_level(&self) -> i32 {
10799 self.compaction_zstd_level
10800 }
10801
10802 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
10803 self.kek.clone()
10804 }
10805
10806 #[cfg(feature = "encryption")]
10810 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10811 self.kek.as_ref().map(|k| k.derive_idx_key())
10812 }
10813
10814 #[cfg(not(feature = "encryption"))]
10815 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10816 None
10817 }
10818
10819 #[cfg(feature = "encryption")]
10823 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10824 self.kek.as_ref().map(|k| *k.derive_meta_key())
10825 }
10826
10827 #[cfg(not(feature = "encryption"))]
10828 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10829 None
10830 }
10831
10832 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
10835 self.column_keys
10836 .iter()
10837 .map(|(&id, &(_, scheme))| (id, scheme))
10838 .collect()
10839 }
10840
10841 #[cfg(feature = "encryption")]
10846 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
10847 self.tokenize_value_enc(column_id, v)
10848 }
10849
10850 #[cfg(feature = "encryption")]
10851 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
10852 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10853 let (key, scheme) = self.column_keys.get(&column_id)?;
10854 let token: Vec<u8> = match (*scheme, v) {
10855 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
10856 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
10857 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
10858 _ => hmac_token(key, &v.encode_key()).to_vec(),
10859 };
10860 Some(Value::Bytes(token))
10861 }
10862
10863 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
10865 self.index_lookup_key_bytes(column_id, &v.encode_key())
10866 }
10867
10868 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
10871 #[cfg(feature = "encryption")]
10872 {
10873 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
10874 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
10875 if *scheme == SCHEME_HMAC_EQ {
10876 return hmac_token(key, encoded).to_vec();
10877 }
10878 }
10879 }
10880 let _ = column_id;
10881 encoded.to_vec()
10882 }
10883}
10884
10885fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
10886 let columnar::NativeColumn::Int64 { data, validity } = col else {
10887 return false;
10888 };
10889 if data.len() < n || !columnar::all_non_null(validity, n) {
10890 return false;
10891 }
10892 data.iter()
10893 .take(n)
10894 .zip(data.iter().skip(1))
10895 .all(|(a, b)| a < b)
10896}
10897
10898#[derive(Debug, Clone)]
10902pub struct ColumnStat {
10903 pub min: Option<Value>,
10904 pub max: Option<Value>,
10905 pub null_count: u64,
10906}
10907
10908#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10910pub enum NativeAgg {
10911 Count,
10912 Sum,
10913 Min,
10914 Max,
10915 Avg,
10916}
10917
10918#[derive(Debug, Clone, PartialEq)]
10920pub enum NativeAggResult {
10921 Count(u64),
10922 Int(i64),
10923 Float(f64),
10924 Null,
10926}
10927
10928#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10930pub enum ApproxAgg {
10931 Count,
10932 Sum,
10933 Avg,
10934}
10935
10936#[derive(Debug, Clone)]
10940pub struct ApproxResult {
10941 pub point: f64,
10943 pub ci_low: f64,
10945 pub ci_high: f64,
10947 pub n_population: u64,
10949 pub n_sample_live: usize,
10951 pub n_passing: usize,
10953}
10954
10955#[derive(Debug, Clone, PartialEq)]
10960pub enum AggState {
10961 Count(u64),
10963 SumI {
10965 sum: i128,
10966 count: u64,
10967 },
10968 SumF {
10970 sum: f64,
10971 count: u64,
10972 },
10973 AvgI {
10975 sum: i128,
10976 count: u64,
10977 },
10978 AvgF {
10980 sum: f64,
10981 count: u64,
10982 },
10983 MinI(i64),
10985 MaxI(i64),
10986 MinF(f64),
10988 MaxF(f64),
10989 Empty,
10991}
10992
10993impl AggState {
10994 pub fn merge(self, other: AggState) -> AggState {
10996 use AggState::*;
10997 match (self, other) {
10998 (Empty, x) | (x, Empty) => x,
10999 (Count(a), Count(b)) => Count(a + b),
11000 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11001 sum: sa + sb,
11002 count: ca + cb,
11003 },
11004 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11005 sum: sa + sb,
11006 count: ca + cb,
11007 },
11008 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11009 sum: sa + sb,
11010 count: ca + cb,
11011 },
11012 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11013 sum: sa + sb,
11014 count: ca + cb,
11015 },
11016 (MinI(a), MinI(b)) => MinI(a.min(b)),
11017 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11018 (MinF(a), MinF(b)) => MinF(a.min(b)),
11019 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11020 _ => Empty, }
11022 }
11023
11024 pub fn point(&self) -> Option<f64> {
11026 match self {
11027 AggState::Count(n) => Some(*n as f64),
11028 AggState::SumI { sum, .. } => Some(*sum as f64),
11029 AggState::SumF { sum, .. } => Some(*sum),
11030 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11031 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11032 AggState::MinI(n) => Some(*n as f64),
11033 AggState::MaxI(n) => Some(*n as f64),
11034 AggState::MinF(n) => Some(*n),
11035 AggState::MaxF(n) => Some(*n),
11036 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11037 }
11038 }
11039
11040 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11044 let is_float = matches!(ty, Some(TypeId::Float64));
11045 match (agg, result) {
11046 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11047 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11048 sum: x as i128,
11049 count: 1, },
11051 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11052 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11053 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11054 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11055 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11056 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11057 (NativeAgg::Count, _) => AggState::Empty,
11058 (_, NativeAggResult::Null) => AggState::Empty,
11059 _ => {
11060 let _ = is_float;
11061 AggState::Empty
11062 }
11063 }
11064 }
11065}
11066
11067#[derive(Debug, Clone)]
11070pub struct CachedAgg {
11071 pub state: AggState,
11072 pub watermark: u64,
11073 pub epoch: u64,
11074}
11075
11076#[derive(Debug, Clone)]
11078pub struct IncrementalAggResult {
11079 pub state: AggState,
11081 pub incremental: bool,
11084 pub delta_rows: u64,
11086}
11087
11088fn agg_state_from_rows(
11092 rows: &[Row],
11093 conditions: &[crate::query::Condition],
11094 index_sets: &[RowIdSet],
11095 column: Option<u16>,
11096 agg: NativeAgg,
11097 schema: &Schema,
11098 control: Option<&crate::ExecutionControl>,
11099) -> Result<AggState> {
11100 let mut count: u64 = 0;
11101 let mut sum_i: i128 = 0;
11102 let mut sum_f: f64 = 0.0;
11103 let mut mn_i: i64 = i64::MAX;
11104 let mut mx_i: i64 = i64::MIN;
11105 let mut mn_f: f64 = f64::INFINITY;
11106 let mut mx_f: f64 = f64::NEG_INFINITY;
11107 let mut saw_int = false;
11108 let mut saw_float = false;
11109 for (index, r) in rows.iter().enumerate() {
11110 execution_checkpoint(control, index)?;
11111 if !conditions
11112 .iter()
11113 .all(|c| condition_matches_row(c, r, schema))
11114 {
11115 continue;
11116 }
11117 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11118 continue;
11119 }
11120 match agg {
11121 NativeAgg::Count => match column {
11122 None => count += 1,
11124 Some(cid) => match r.columns.get(&cid) {
11127 None | Some(Value::Null) => {}
11128 Some(_) => count += 1,
11129 },
11130 },
11131 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11132 Some(Value::Int64(n)) => {
11133 count += 1;
11134 sum_i += *n as i128;
11135 mn_i = mn_i.min(*n);
11136 mx_i = mx_i.max(*n);
11137 saw_int = true;
11138 }
11139 Some(Value::Float64(f)) => {
11140 count += 1;
11141 sum_f += f;
11142 mn_f = mn_f.min(*f);
11143 mx_f = mx_f.max(*f);
11144 saw_float = true;
11145 }
11146 _ => {}
11147 },
11148 }
11149 }
11150 Ok(match agg {
11151 NativeAgg::Count => {
11152 if count == 0 {
11153 AggState::Empty
11154 } else {
11155 AggState::Count(count)
11156 }
11157 }
11158 NativeAgg::Sum => {
11159 if count == 0 {
11160 AggState::Empty
11161 } else if saw_int {
11162 AggState::SumI { sum: sum_i, count }
11163 } else {
11164 AggState::SumF { sum: sum_f, count }
11165 }
11166 }
11167 NativeAgg::Avg => {
11168 if count == 0 {
11169 AggState::Empty
11170 } else if saw_int {
11171 AggState::AvgI { sum: sum_i, count }
11172 } else {
11173 AggState::AvgF { sum: sum_f, count }
11174 }
11175 }
11176 NativeAgg::Min => {
11177 if !saw_int && !saw_float {
11178 AggState::Empty
11179 } else if saw_int {
11180 AggState::MinI(mn_i)
11181 } else {
11182 AggState::MinF(mn_f)
11183 }
11184 }
11185 NativeAgg::Max => {
11186 if !saw_int && !saw_float {
11187 AggState::Empty
11188 } else if saw_int {
11189 AggState::MaxI(mx_i)
11190 } else {
11191 AggState::MaxF(mx_f)
11192 }
11193 }
11194 })
11195}
11196
11197fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11201 use crate::query::Condition;
11202 match c {
11203 Condition::Pk(key) => match schema.primary_key() {
11204 Some(pk) => row
11205 .columns
11206 .get(&pk.id)
11207 .map(|v| v.encode_key() == *key)
11208 .unwrap_or(false),
11209 None => false,
11210 },
11211 Condition::BitmapEq { column_id, value } => row
11212 .columns
11213 .get(column_id)
11214 .map(|v| v.encode_key() == *value)
11215 .unwrap_or(false),
11216 Condition::BitmapIn { column_id, values } => {
11217 let key = row.columns.get(column_id).map(|v| v.encode_key());
11218 match key {
11219 Some(k) => values.contains(&k),
11220 None => false,
11221 }
11222 }
11223 Condition::BytesPrefix { column_id, prefix } => row
11224 .columns
11225 .get(column_id)
11226 .map(|v| v.encode_key().starts_with(prefix))
11227 .unwrap_or(false),
11228 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11229 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11230 _ => false,
11231 },
11232 Condition::RangeF64 {
11233 column_id,
11234 lo,
11235 lo_inclusive,
11236 hi,
11237 hi_inclusive,
11238 } => match row.columns.get(column_id) {
11239 Some(Value::Float64(n)) => {
11240 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11241 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11242 lo_ok && hi_ok
11243 }
11244 _ => false,
11245 },
11246 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11247 Some(Value::Bytes(b)) => {
11248 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11249 }
11250 _ => false,
11251 },
11252 Condition::FmContainsAll {
11253 column_id,
11254 patterns,
11255 } => match row.columns.get(column_id) {
11256 Some(Value::Bytes(b)) => patterns
11257 .iter()
11258 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11259 _ => false,
11260 },
11261 Condition::Ann { .. }
11262 | Condition::SparseMatch { .. }
11263 | Condition::MinHashSimilar { .. } => true,
11264 Condition::IsNull { column_id } => {
11265 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11266 }
11267 Condition::IsNotNull { column_id } => {
11268 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11269 }
11270 }
11271}
11272
11273fn as_f64(v: Option<&Value>) -> Option<f64> {
11275 match v {
11276 Some(Value::Int64(n)) => Some(*n as f64),
11277 Some(Value::Float64(f)) => Some(*f),
11278 _ => None,
11279 }
11280}
11281
11282fn accumulate_int(
11286 cursor: &mut dyn crate::cursor::Cursor,
11287 control: Option<&crate::ExecutionControl>,
11288) -> Result<(u64, i128, i64, i64)> {
11289 let mut count: u64 = 0;
11290 let mut sum: i128 = 0;
11291 let mut mn: i64 = i64::MAX;
11292 let mut mx: i64 = i64::MIN;
11293 while let Some(cols) = cursor.next_batch()? {
11294 execution_checkpoint(control, 0)?;
11295 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11296 if crate::columnar::all_non_null(validity, data.len()) {
11297 count += data.len() as u64;
11299 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11300 execution_checkpoint(control, chunk_index * 1024)?;
11301 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11302 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11303 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11304 }
11305 } else {
11306 for (i, &v) in data.iter().enumerate() {
11307 execution_checkpoint(control, i)?;
11308 if crate::columnar::validity_bit(validity, i) {
11309 count += 1;
11310 sum += v as i128;
11311 mn = mn.min(v);
11312 mx = mx.max(v);
11313 }
11314 }
11315 }
11316 }
11317 }
11318 Ok((count, sum, mn, mx))
11319}
11320
11321fn accumulate_float(
11323 cursor: &mut dyn crate::cursor::Cursor,
11324 control: Option<&crate::ExecutionControl>,
11325) -> Result<(u64, f64, f64, f64)> {
11326 let mut count: u64 = 0;
11327 let mut sum: f64 = 0.0;
11328 let mut mn: f64 = f64::INFINITY;
11329 let mut mx: f64 = f64::NEG_INFINITY;
11330 while let Some(cols) = cursor.next_batch()? {
11331 execution_checkpoint(control, 0)?;
11332 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11333 if crate::columnar::all_non_null(validity, data.len()) {
11334 count += data.len() as u64;
11335 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11336 execution_checkpoint(control, chunk_index * 1024)?;
11337 sum += chunk.iter().sum::<f64>();
11338 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11339 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11340 }
11341 } else {
11342 for (i, &v) in data.iter().enumerate() {
11343 execution_checkpoint(control, i)?;
11344 if crate::columnar::validity_bit(validity, i) {
11345 count += 1;
11346 sum += v;
11347 mn = mn.min(v);
11348 mx = mx.max(v);
11349 }
11350 }
11351 }
11352 }
11353 }
11354 Ok((count, sum, mn, mx))
11355}
11356
11357#[inline]
11358fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11359 if index.is_multiple_of(256) {
11360 control
11361 .map(crate::ExecutionControl::checkpoint)
11362 .transpose()?;
11363 }
11364 Ok(())
11365}
11366
11367fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11368 if count == 0 && !matches!(agg, NativeAgg::Count) {
11369 return NativeAggResult::Null;
11370 }
11371 match agg {
11372 NativeAgg::Count => NativeAggResult::Count(count),
11373 NativeAgg::Sum => match sum.try_into() {
11376 Ok(v) => NativeAggResult::Int(v),
11377 Err(_) => NativeAggResult::Null,
11378 },
11379 NativeAgg::Min => NativeAggResult::Int(mn),
11380 NativeAgg::Max => NativeAggResult::Int(mx),
11381 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11382 }
11383}
11384
11385fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11386 if count == 0 && !matches!(agg, NativeAgg::Count) {
11387 return NativeAggResult::Null;
11388 }
11389 match agg {
11390 NativeAgg::Count => NativeAggResult::Count(count),
11391 NativeAgg::Sum => NativeAggResult::Float(sum),
11392 NativeAgg::Min => NativeAggResult::Float(mn),
11393 NativeAgg::Max => NativeAggResult::Float(mx),
11394 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11395 }
11396}
11397
11398fn agg_int(
11401 stats: &[crate::page::PageStat],
11402 decode: fn(Option<&[u8]>) -> Option<i64>,
11403) -> Option<(Option<i64>, Option<i64>, u64)> {
11404 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11405 let mut any = false;
11406 for s in stats {
11407 if let Some(v) = decode(s.min.as_deref()) {
11408 mn = mn.min(v);
11409 any = true;
11410 }
11411 if let Some(v) = decode(s.max.as_deref()) {
11412 mx = mx.max(v);
11413 any = true;
11414 }
11415 nulls += s.null_count;
11416 }
11417 any.then_some((Some(mn), Some(mx), nulls))
11418}
11419
11420fn agg_float(
11422 stats: &[crate::page::PageStat],
11423 decode: fn(Option<&[u8]>) -> Option<f64>,
11424) -> Option<(Option<f64>, Option<f64>, u64)> {
11425 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11426 let mut any = false;
11427 for s in stats {
11428 if let Some(v) = decode(s.min.as_deref()) {
11429 mn = mn.min(v);
11430 any = true;
11431 }
11432 if let Some(v) = decode(s.max.as_deref()) {
11433 mx = mx.max(v);
11434 any = true;
11435 }
11436 nulls += s.null_count;
11437 }
11438 any.then_some((Some(mn), Some(mx), nulls))
11439}
11440
11441type SecondaryIndexes = (
11443 HashMap<u16, BitmapIndex>,
11444 HashMap<u16, AnnIndex>,
11445 HashMap<u16, FmIndex>,
11446 HashMap<u16, SparseIndex>,
11447 HashMap<u16, MinHashIndex>,
11448);
11449
11450fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11451 let mut bitmap = HashMap::new();
11452 let mut ann = HashMap::new();
11453 let mut fm = HashMap::new();
11454 let mut sparse = HashMap::new();
11455 let mut minhash = HashMap::new();
11456 for idef in &schema.indexes {
11457 match idef.kind {
11458 IndexKind::Bitmap => {
11459 bitmap.insert(idef.column_id, BitmapIndex::new());
11460 }
11461 IndexKind::Ann => {
11462 let dim = schema
11463 .columns
11464 .iter()
11465 .find(|c| c.id == idef.column_id)
11466 .and_then(|c| match c.ty {
11467 TypeId::Embedding { dim } => Some(dim as usize),
11468 _ => None,
11469 })
11470 .unwrap_or(0);
11471 let options = idef.options.ann.clone().unwrap_or_default();
11472 ann.insert(
11473 idef.column_id,
11474 AnnIndex::with_options(
11475 dim,
11476 options.m,
11477 options.ef_construction,
11478 options.ef_search,
11479 ),
11480 );
11481 }
11482 IndexKind::FmIndex => {
11483 fm.insert(idef.column_id, FmIndex::new());
11484 }
11485 IndexKind::Sparse => {
11486 sparse.insert(idef.column_id, SparseIndex::new());
11487 }
11488 IndexKind::MinHash => {
11489 let options = idef.options.minhash.clone().unwrap_or_default();
11490 minhash.insert(
11491 idef.column_id,
11492 MinHashIndex::with_options(options.permutations, options.bands),
11493 );
11494 }
11495 _ => {}
11496 }
11497 }
11498 (bitmap, ann, fm, sparse, minhash)
11499}
11500
11501const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11502 | ColumnFlags::AUTO_INCREMENT
11503 | ColumnFlags::ENCRYPTED
11504 | ColumnFlags::ENCRYPTED_INDEXABLE
11505 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11506
11507fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11508 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11509 return Err(MongrelError::Schema(
11510 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11511 ));
11512 }
11513 Ok(())
11514}
11515
11516fn validate_alter_column_type(
11517 schema: &Schema,
11518 old: &ColumnDef,
11519 next: &ColumnDef,
11520 has_stored_versions: bool,
11521) -> Result<()> {
11522 if old.ty == next.ty {
11523 return Ok(());
11524 }
11525 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11526 return Err(MongrelError::Schema(format!(
11527 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11528 old.name
11529 )));
11530 }
11531 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11532 return Ok(());
11533 }
11534 Err(MongrelError::Schema(format!(
11535 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11536 old.ty, next.ty
11537 )))
11538}
11539
11540fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11541 matches!(
11542 (old, new),
11543 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11544 )
11545}
11546
11547fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11553 let mut prev: Option<i64> = None;
11554 for r in rows {
11555 match r.columns.get(&pk_id) {
11556 Some(Value::Int64(v)) => {
11557 if prev.is_some_and(|p| p >= *v) {
11558 return false;
11559 }
11560 prev = Some(*v);
11561 }
11562 _ => return false,
11563 }
11564 }
11565 true
11566}
11567
11568#[allow(clippy::too_many_arguments)]
11569fn index_into(
11570 schema: &Schema,
11571 row: &Row,
11572 hot: &mut HotIndex,
11573 bitmap: &mut HashMap<u16, BitmapIndex>,
11574 ann: &mut HashMap<u16, AnnIndex>,
11575 fm: &mut HashMap<u16, FmIndex>,
11576 sparse: &mut HashMap<u16, SparseIndex>,
11577 minhash: &mut HashMap<u16, MinHashIndex>,
11578) {
11579 for idef in &schema.indexes {
11580 let Some(val) = row.columns.get(&idef.column_id) else {
11581 continue;
11582 };
11583 match idef.kind {
11584 IndexKind::Bitmap => {
11585 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11586 b.insert(val.encode_key(), row.row_id);
11587 }
11588 }
11589 IndexKind::Ann => {
11590 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11591 a.insert_validated(v, row.row_id);
11592 }
11593 }
11594 IndexKind::FmIndex => {
11595 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11596 f.insert(b.clone(), row.row_id);
11597 }
11598 }
11599 IndexKind::Sparse => {
11600 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11601 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11604 s.insert(&terms, row.row_id);
11605 }
11606 }
11607 }
11608 IndexKind::MinHash => {
11609 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11610 let tokens = crate::index::token_hashes_from_bytes(b);
11613 mh.insert(&tokens, row.row_id);
11614 }
11615 }
11616 _ => {}
11617 }
11618 }
11619 if let Some(pk_col) = schema.primary_key() {
11620 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11621 hot.insert(pk_val.encode_key(), row.row_id);
11622 }
11623 }
11624}
11625
11626#[allow(clippy::too_many_arguments)]
11629fn index_into_single(
11630 idef: &IndexDef,
11631 _schema: &Schema,
11632 row: &Row,
11633 _hot: &mut HotIndex,
11634 bitmap: &mut HashMap<u16, BitmapIndex>,
11635 ann: &mut HashMap<u16, AnnIndex>,
11636 fm: &mut HashMap<u16, FmIndex>,
11637 sparse: &mut HashMap<u16, SparseIndex>,
11638 minhash: &mut HashMap<u16, MinHashIndex>,
11639) {
11640 let Some(val) = row.columns.get(&idef.column_id) else {
11641 return;
11642 };
11643 match idef.kind {
11644 IndexKind::Bitmap => {
11645 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11646 b.insert(val.encode_key(), row.row_id);
11647 }
11648 }
11649 IndexKind::Ann => {
11650 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11651 a.insert_validated(v, row.row_id);
11652 }
11653 }
11654 IndexKind::FmIndex => {
11655 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11656 f.insert(b.clone(), row.row_id);
11657 }
11658 }
11659 IndexKind::Sparse => {
11660 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11661 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11662 s.insert(&terms, row.row_id);
11663 }
11664 }
11665 }
11666 IndexKind::MinHash => {
11667 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11668 let tokens = crate::index::token_hashes_from_bytes(b);
11669 mh.insert(&tokens, row.row_id);
11670 }
11671 }
11672 _ => {}
11673 }
11674}
11675
11676fn eval_partial_predicate(
11682 pred: &str,
11683 columns_map: &HashMap<u16, &Value>,
11684 name_to_id: &HashMap<&str, u16>,
11685) -> bool {
11686 let lower = pred.trim().to_ascii_lowercase();
11687 if let Some(rest) = lower.strip_suffix(" is not null") {
11689 let col_name = rest.trim();
11690 if let Some(col_id) = name_to_id.get(col_name) {
11691 return columns_map
11692 .get(col_id)
11693 .is_some_and(|v| !matches!(v, Value::Null));
11694 }
11695 }
11696 if let Some(rest) = lower.strip_suffix(" is null") {
11698 let col_name = rest.trim();
11699 if let Some(col_id) = name_to_id.get(col_name) {
11700 return columns_map
11701 .get(col_id)
11702 .is_none_or(|v| matches!(v, Value::Null));
11703 }
11704 }
11705 true
11708}
11709
11710#[allow(dead_code)]
11716fn bulk_index_key(
11717 column_keys: &HashMap<u16, ([u8; 32], u8)>,
11718 column_id: u16,
11719 ty: TypeId,
11720 col: &columnar::NativeColumn,
11721 i: usize,
11722) -> Option<Vec<u8>> {
11723 let encoded = columnar::encode_key_native(ty, col, i)?;
11724 #[cfg(feature = "encryption")]
11725 {
11726 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11727 if let Some((key, scheme)) = column_keys.get(&column_id) {
11728 return Some(match (*scheme, col) {
11729 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
11730 (_, columnar::NativeColumn::Int64 { data, .. }) => {
11731 ope_token_i64(key, data[i]).to_vec()
11732 }
11733 (_, columnar::NativeColumn::Float64 { data, .. }) => {
11734 ope_token_f64(key, data[i]).to_vec()
11735 }
11736 _ => hmac_token(key, &encoded).to_vec(),
11737 });
11738 }
11739 }
11740 #[cfg(not(feature = "encryption"))]
11741 {
11742 let _ = (column_id, column_keys, col);
11743 }
11744 Some(encoded)
11745}
11746
11747pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
11748 write_schema_with_after(dir, schema, || {})
11749}
11750
11751pub(crate) fn write_schema_durable(
11752 root: &crate::durable_file::DurableRoot,
11753 schema: &Schema,
11754) -> Result<()> {
11755 write_schema_durable_with_after(root, schema, || {})
11756}
11757
11758fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
11759where
11760 F: FnOnce(),
11761{
11762 let json = serde_json::to_string_pretty(schema)
11763 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
11764 crate::durable_file::write_atomic_with_after(
11765 &dir.join(SCHEMA_FILENAME),
11766 json.as_bytes(),
11767 after_publish,
11768 )?;
11769 Ok(())
11770}
11771
11772fn write_schema_durable_with_after<F>(
11773 root: &crate::durable_file::DurableRoot,
11774 schema: &Schema,
11775 after_publish: F,
11776) -> Result<()>
11777where
11778 F: FnOnce(),
11779{
11780 let json = serde_json::to_string_pretty(schema)
11781 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
11782 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
11783 Ok(())
11784}
11785
11786fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
11787 let mut schema_published = false;
11788 let schema_result = match table._root_guard.as_deref() {
11789 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
11790 schema_published = true;
11791 }),
11792 None => write_schema_with_after(&table.dir, &table.schema, || {
11793 schema_published = true;
11794 }),
11795 };
11796 if schema_result.is_err() && !schema_published {
11797 return schema_result;
11798 }
11799 match table.persist_manifest(table.current_epoch()) {
11800 Ok(()) => Ok(()),
11801 Err(manifest_error) => Err(match schema_result {
11802 Ok(()) => manifest_error,
11803 Err(schema_error) => MongrelError::Other(format!(
11804 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11805 )),
11806 }),
11807 }
11808}
11809
11810fn read_schema(dir: &Path) -> Result<Schema> {
11811 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
11812 read_schema_file(file)
11813}
11814
11815fn read_schema_file(file: std::fs::File) -> Result<Schema> {
11816 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
11817 use std::io::Read;
11818
11819 let length = file.metadata()?.len();
11820 if length > MAX_SCHEMA_BYTES {
11821 return Err(MongrelError::ResourceLimitExceeded {
11822 resource: "schema bytes",
11823 requested: usize::try_from(length).unwrap_or(usize::MAX),
11824 limit: MAX_SCHEMA_BYTES as usize,
11825 });
11826 }
11827 let mut bytes = Vec::with_capacity(length as usize);
11828 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
11829 if bytes.len() as u64 != length {
11830 return Err(MongrelError::Schema(
11831 "schema length changed while reading".into(),
11832 ));
11833 }
11834 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
11835}
11836
11837fn preflight_standalone_open(
11838 dir: &Path,
11839 runs_root: Option<&crate::durable_file::DurableRoot>,
11840 idx_root: Option<&crate::durable_file::DurableRoot>,
11841 manifest: &Manifest,
11842 schema: &Schema,
11843 records: &[crate::wal::Record],
11844 kek: Option<Arc<Kek>>,
11845) -> Result<()> {
11846 crate::wal::validate_shared_transaction_framing(records)?;
11847 if manifest.schema_id > schema.schema_id
11848 || manifest.flushed_epoch > manifest.current_epoch
11849 || manifest.global_idx_epoch > manifest.current_epoch
11850 || manifest.next_row_id == u64::MAX
11851 || manifest.auto_inc_next < 0
11852 || manifest.auto_inc_next == i64::MAX
11853 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
11854 {
11855 return Err(MongrelError::InvalidArgument(
11856 "manifest counters or schema identity are invalid".into(),
11857 ));
11858 }
11859 let mut run_ids = HashSet::new();
11860 let mut maximum_row_id = None::<u64>;
11861 for run in &manifest.runs {
11862 if run.run_id >= u64::MAX as u128
11863 || !run_ids.insert(run.run_id)
11864 || run.epoch_created > manifest.current_epoch
11865 {
11866 return Err(MongrelError::InvalidArgument(
11867 "manifest contains an invalid or duplicate active run".into(),
11868 ));
11869 }
11870 let mut reader = match runs_root {
11871 Some(root) => RunReader::open_file(
11872 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
11873 schema.clone(),
11874 kek.clone(),
11875 )?,
11876 None => RunReader::open(
11877 dir.join(RUNS_DIR)
11878 .join(format!("r-{}.sr", run.run_id as u64)),
11879 schema.clone(),
11880 kek.clone(),
11881 )?,
11882 };
11883 let header = reader.header();
11884 if header.run_id != run.run_id
11885 || header.level != run.level
11886 || header.row_count != run.row_count
11887 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
11888 || header.is_uniform_epoch() && header.epoch_created != 0
11889 || header.schema_id > schema.schema_id
11890 {
11891 return Err(MongrelError::InvalidArgument(format!(
11892 "run {} differs from its manifest",
11893 run.run_id
11894 )));
11895 }
11896 if header.row_count != 0 {
11897 maximum_row_id = Some(
11898 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
11899 );
11900 }
11901 reader.validate_all_pages()?;
11902 }
11903 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
11904 return Err(MongrelError::InvalidArgument(
11905 "manifest next_row_id does not advance beyond persisted rows".into(),
11906 ));
11907 }
11908 for run in &manifest.retiring {
11909 if run.run_id >= u64::MAX as u128
11910 || run.retire_epoch > manifest.current_epoch
11911 || !run_ids.insert(run.run_id)
11912 {
11913 return Err(MongrelError::InvalidArgument(
11914 "manifest contains an invalid or duplicate retired run".into(),
11915 ));
11916 }
11917 }
11918 #[cfg(feature = "encryption")]
11919 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
11920 #[cfg(not(feature = "encryption"))]
11921 let idx_dek: Option<Zeroizing<[u8; DEK_LEN]>> = None;
11922 match idx_root {
11923 Some(root) => {
11924 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
11925 }
11926 None => {
11927 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
11928 }
11929 }
11930
11931 let committed = records
11932 .iter()
11933 .filter_map(|record| match record.op {
11934 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
11935 _ => None,
11936 })
11937 .collect::<HashMap<_, _>>();
11938 for record in records {
11939 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
11940 continue;
11941 };
11942 match &record.op {
11943 Op::Put { table_id, rows } => {
11944 if *table_id != manifest.table_id {
11945 return Err(MongrelError::CorruptWal {
11946 offset: record.seq.0,
11947 reason: format!(
11948 "private WAL record references table {table_id}, expected {}",
11949 manifest.table_id
11950 ),
11951 });
11952 }
11953 let rows: Vec<Row> =
11954 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
11955 offset: record.seq.0,
11956 reason: format!("committed Put payload could not be decoded: {error}"),
11957 })?;
11958 for row in rows {
11959 if row.deleted || row.row_id.0 == u64::MAX {
11960 return Err(MongrelError::CorruptWal {
11961 offset: record.seq.0,
11962 reason: "committed Put contains an invalid row identity".into(),
11963 });
11964 }
11965 let cells = row.columns.into_iter().collect::<Vec<_>>();
11966 schema
11967 .validate_values(&cells)
11968 .map_err(|error| MongrelError::CorruptWal {
11969 offset: record.seq.0,
11970 reason: format!("committed Put violates table schema: {error}"),
11971 })?;
11972 if schema.auto_increment_column().is_some_and(|column| {
11973 matches!(
11974 cells.iter().find(|(id, _)| *id == column.id),
11975 Some((_, Value::Int64(value))) if *value == i64::MAX
11976 )
11977 }) {
11978 return Err(MongrelError::CorruptWal {
11979 offset: record.seq.0,
11980 reason: "committed Put exhausts AUTO_INCREMENT".into(),
11981 });
11982 }
11983 }
11984 }
11985 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
11986 if *table_id != manifest.table_id =>
11987 {
11988 return Err(MongrelError::CorruptWal {
11989 offset: record.seq.0,
11990 reason: format!(
11991 "private WAL record references table {table_id}, expected {}",
11992 manifest.table_id
11993 ),
11994 });
11995 }
11996 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
11997 return Err(MongrelError::CorruptWal {
11998 offset: record.seq.0,
11999 reason: "private WAL contains shared spilled-run metadata".into(),
12000 });
12001 }
12002 _ => {}
12003 }
12004 }
12005 Ok(())
12006}
12007
12008fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12009 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12010}
12011
12012fn wal_segment_number(path: &Path) -> Option<u64> {
12013 path.file_stem()
12014 .and_then(|stem| stem.to_str())
12015 .and_then(|stem| stem.strip_prefix("seg-"))
12016 .and_then(|number| number.parse().ok())
12017}
12018
12019fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12020 let n = list_wal_numbers(wal_dir)?;
12021 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12022}
12023
12024fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12025 list_wal_numbers(wal_dir)?
12026 .map(|maximum| {
12027 maximum
12028 .checked_add(1)
12029 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12030 })
12031 .unwrap_or(Ok(0))
12032}
12033
12034fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12035 let mut max_n = None;
12036 let entries = match std::fs::read_dir(wal_dir) {
12037 Ok(entries) => entries,
12038 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12039 Err(error) => return Err(error.into()),
12040 };
12041 for entry in entries {
12042 let entry = entry?;
12043 let fname = entry.file_name();
12044 let Some(s) = fname.to_str() else {
12045 continue;
12046 };
12047 let Some(stripped) = s.strip_prefix("seg-") else {
12048 continue;
12049 };
12050 let Some(number) = stripped.strip_suffix(".wal") else {
12051 return Err(MongrelError::CorruptWal {
12052 offset: 0,
12053 reason: format!("malformed WAL segment name {s:?}"),
12054 });
12055 };
12056 let n = number
12057 .parse::<u32>()
12058 .map_err(|_| MongrelError::CorruptWal {
12059 offset: 0,
12060 reason: format!("malformed WAL segment name {s:?}"),
12061 })?;
12062 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12063 return Err(MongrelError::CorruptWal {
12064 offset: n as u64,
12065 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12066 });
12067 }
12068 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12069 }
12070 Ok(max_n)
12071}