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 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1067 use crate::encryption::Cipher;
1068 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1069 let mut nonce = [0u8; 12];
1070 crate::encryption::fill_random(&mut nonce).ok()?;
1071 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1072 let mut out = Vec::with_capacity(12 + ct.len());
1073 out.extend_from_slice(&nonce);
1074 out.extend_from_slice(&ct);
1075 Some(out)
1076 }
1077
1078 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1080 use crate::encryption::Cipher;
1081 if bytes.len() < 28 {
1082 return None;
1083 }
1084 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1085 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1086 let ct = &bytes[12..];
1087 cipher.decrypt_page(&nonce, ct).ok()
1088 }
1089
1090 fn load_persistent(&mut self) {
1093 let Some(dir) = self.dir.as_ref().cloned() else {
1094 return;
1095 };
1096 let entries = match std::fs::read_dir(&dir) {
1097 Ok(e) => e,
1098 Err(_) => return,
1099 };
1100 for entry in entries.flatten() {
1101 let path = entry.path();
1102 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1104 let _ = std::fs::remove_file(&path);
1105 continue;
1106 }
1107 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1108 continue;
1109 }
1110 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1111 Some(s) => s,
1112 None => continue,
1113 };
1114 let key = match u64::from_str_radix(stem, 16) {
1115 Ok(k) => k,
1116 Err(_) => continue,
1117 };
1118 let bytes = match std::fs::read(&path) {
1119 Ok(b) => b,
1120 Err(_) => continue,
1121 };
1122 let plaintext = if let Some(dek) = &self.cache_dek {
1124 match self.decrypt_cache(&bytes, dek) {
1125 Some(p) => p,
1126 None => {
1127 let _ = std::fs::remove_file(&path);
1128 continue;
1129 }
1130 }
1131 } else {
1132 bytes
1133 };
1134 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1135 Ok(serialized) => {
1136 if let Some(entry) = serialized.into_entry() {
1137 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1138 self.entries.insert(key, entry);
1139 self.order.push_back(key);
1140 } else {
1141 let _ = std::fs::remove_file(&path);
1142 }
1143 }
1144 Err(_) => {
1145 let _ = std::fs::remove_file(&path);
1146 }
1147 }
1148 }
1149 self.evict();
1150 }
1151
1152 fn set_max_bytes(&mut self, max_bytes: u64) {
1153 self.max_bytes = max_bytes;
1154 self.evict();
1155 }
1156
1157 fn touch(&mut self, key: u64) {
1159 self.order.retain(|k| *k != key);
1160 self.order.push_back(key);
1161 }
1162
1163 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1164 let res = self.entries.get(&key).and_then(|e| match &e.data {
1165 CachedData::Rows(r) => Some(r.clone()),
1166 CachedData::Columns(_) => None,
1167 });
1168 if res.is_some() {
1169 self.touch(key);
1170 return res;
1171 }
1172 if let Some(entry) = self.load_from_disk(key) {
1174 let res = match &entry.data {
1175 CachedData::Rows(r) => Some(r.clone()),
1176 CachedData::Columns(_) => None,
1177 };
1178 if res.is_some() {
1179 let approx = entry.data.approx_bytes();
1180 self.bytes = self.bytes.saturating_add(approx);
1181 self.entries.insert(key, entry);
1182 self.order.push_back(key);
1183 self.evict();
1184 return res;
1185 }
1186 }
1187 None
1188 }
1189
1190 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1191 let res = self.entries.get(&key).and_then(|e| match &e.data {
1192 CachedData::Columns(c) => Some(c.clone()),
1193 CachedData::Rows(_) => None,
1194 });
1195 if res.is_some() {
1196 self.touch(key);
1197 return res;
1198 }
1199 if let Some(entry) = self.load_from_disk(key) {
1201 let res = match &entry.data {
1202 CachedData::Columns(c) => Some(c.clone()),
1203 CachedData::Rows(_) => None,
1204 };
1205 if res.is_some() {
1206 let approx = entry.data.approx_bytes();
1207 self.bytes = self.bytes.saturating_add(approx);
1208 self.entries.insert(key, entry);
1209 self.order.push_back(key);
1210 self.evict();
1211 return res;
1212 }
1213 }
1214 None
1215 }
1216
1217 fn insert(&mut self, key: u64, entry: CachedEntry) {
1218 let approx = entry.data.approx_bytes();
1219 if self.entries.remove(&key).is_some() {
1220 self.order.retain(|k| *k != key);
1221 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1222 }
1223 self.store_to_disk(key, &entry);
1225 self.bytes = self.bytes.saturating_add(approx);
1226 self.entries.insert(key, entry);
1227 self.order.push_back(key);
1228 self.evict();
1229 }
1230
1231 fn invalidate(
1240 &mut self,
1241 delete_rids: &roaring::RoaringBitmap,
1242 put_cols: &std::collections::HashSet<u16>,
1243 ) {
1244 if self.entries.is_empty() {
1245 return;
1246 }
1247 let has_deletes = !delete_rids.is_empty();
1248 let to_remove: std::collections::HashSet<u64> = self
1249 .entries
1250 .iter()
1251 .filter(|(_, e)| {
1252 let delete_hit = if e.footprint.is_empty() {
1253 has_deletes
1254 } else {
1255 e.footprint.intersection_len(delete_rids) > 0
1256 };
1257 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1258 delete_hit || col_hit
1259 })
1260 .map(|(&k, _)| k)
1261 .collect();
1262 for key in &to_remove {
1263 if let Some(e) = self.entries.remove(key) {
1264 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1265 }
1266 self.remove_from_disk(*key);
1267 }
1268 if !to_remove.is_empty() {
1269 self.order.retain(|k| !to_remove.contains(k));
1270 }
1271 }
1272
1273 fn clear(&mut self) {
1274 if let Some(dir) = &self.dir {
1276 if let Ok(entries) = std::fs::read_dir(dir) {
1277 for entry in entries.flatten() {
1278 let path = entry.path();
1279 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1280 let _ = std::fs::remove_file(&path);
1281 }
1282 }
1283 }
1284 }
1285 self.entries.clear();
1286 self.order.clear();
1287 self.bytes = 0;
1288 }
1289
1290 fn evict(&mut self) {
1291 while self.bytes > self.max_bytes {
1292 let Some(k) = self.order.pop_front() else {
1293 break;
1294 };
1295 if let Some(e) = self.entries.remove(&k) {
1296 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1297 self.remove_from_disk(k);
1301 }
1302 }
1303 }
1304}
1305
1306type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1313
1314fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1315 let _ = kek;
1316 {
1317 if let Some(k) = kek {
1318 return (
1319 Some(k.derive_table_wal_key(_table_id)),
1320 Some(k.derive_cache_key()),
1321 );
1322 }
1323 }
1324 (None, None)
1325}
1326
1327fn read_table_encryption_salt_root(
1328 root: &crate::durable_file::DurableRoot,
1329) -> Result<[u8; crate::encryption::SALT_LEN]> {
1330 use std::io::Read;
1331
1332 let mut file = root
1333 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1334 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1335 let length = file.metadata()?.len();
1336 if length != crate::encryption::SALT_LEN as u64 {
1337 return Err(MongrelError::InvalidArgument(format!(
1338 "salt file is {length} bytes, expected {}",
1339 crate::encryption::SALT_LEN
1340 )));
1341 }
1342 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1343 file.read_exact(&mut salt)?;
1344 Ok(salt)
1345}
1346
1347fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1349 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1350}
1351
1352fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1353 let Some(kek) = kek else {
1354 return HashMap::new();
1355 };
1356 {
1357 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1358 schema
1359 .columns
1360 .iter()
1361 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1362 .map(|c| {
1363 let scheme = if schema
1364 .indexes
1365 .iter()
1366 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1367 {
1368 SCHEME_OPE_RANGE
1369 } else {
1370 SCHEME_HMAC_EQ
1371 };
1372 let key: [u8; 32] = *kek.derive_column_key(c.id);
1373 (c.id, (key, scheme))
1374 })
1375 .collect()
1376 }
1377}
1378
1379pub(crate) struct SharedCtx {
1384 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1385 pub epoch: Arc<EpochAuthority>,
1386 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1387 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1388 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1389 pub kek: Option<Arc<Kek>>,
1390 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1396 pub shared: Option<SharedWalCtx>,
1400 pub table_name: Option<String>,
1403 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1406 pub read_only: bool,
1408}
1409
1410#[derive(Clone)]
1416pub(crate) struct SharedWalCtx {
1417 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1418 pub group: Arc<GroupCommit>,
1419 pub poisoned: Arc<AtomicBool>,
1420 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1421 pub change_wake: tokio::sync::broadcast::Sender<()>,
1422 pub lifecycle: Arc<crate::core::LifecycleController>,
1425}
1426
1427enum WalSink {
1430 Private(Wal),
1431 Shared(SharedWalCtx),
1432 ReadOnly,
1433}
1434
1435impl Clone for WalSink {
1436 fn clone(&self) -> Self {
1437 match self {
1438 Self::Shared(shared) => Self::Shared(shared.clone()),
1439 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1440 }
1441 }
1442}
1443
1444impl SharedCtx {
1445 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1449 let n_shards = if cache_dir.is_some() {
1453 1
1454 } else {
1455 crate::cache::CACHE_SHARDS
1456 };
1457 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1458 let page_cache = if let Some(d) = cache_dir {
1459 Arc::new(crate::cache::Sharded::new(1, || {
1460 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1461 }))
1462 } else {
1463 Arc::new(crate::cache::Sharded::new(n_shards, || {
1464 crate::cache::PageCache::new(per_shard)
1465 }))
1466 };
1467 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1468 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1469 crate::cache::CACHE_SHARDS,
1470 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1471 ));
1472 Self {
1473 root_guard: None,
1474 epoch: Arc::new(EpochAuthority::new(0)),
1475 page_cache,
1476 decoded_cache,
1477 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1478 kek,
1479 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1480 shared: None,
1481 table_name: None,
1482 auth: None,
1483 read_only: false,
1484 }
1485 }
1486}
1487
1488fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1492 use crate::query::Condition;
1493 match c {
1494 Condition::Pk(_)
1496 | Condition::BitmapEq { .. }
1497 | Condition::BitmapIn { .. }
1498 | Condition::BytesPrefix { .. }
1499 | Condition::IsNull { .. }
1500 | Condition::IsNotNull { .. } => 0,
1501 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1503 1
1504 }
1505 Condition::FmContains { .. }
1507 | Condition::FmContainsAll { .. }
1508 | Condition::Ann { .. }
1509 | Condition::SparseMatch { .. } => 2,
1510 }
1511}
1512
1513impl Table {
1514 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1515 let dir = dir.as_ref().to_path_buf();
1516 crate::durable_file::create_directory_all(&dir)?;
1517 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1518 let pinned = root.io_path()?;
1519 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1520 ctx.root_guard = Some(root);
1521 Self::create_in(&pinned, schema, table_id, ctx)
1522 }
1523
1524 pub fn create_encrypted(
1535 dir: impl AsRef<Path>,
1536 schema: Schema,
1537 table_id: u64,
1538 passphrase: &str,
1539 ) -> Result<Self> {
1540 let dir = dir.as_ref().to_path_buf();
1541 crate::durable_file::create_directory_all(&dir)?;
1542 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1543 root.create_directory_all(META_DIR)?;
1544 let salt = crate::encryption::random_salt()?;
1545 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1546 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1547 let pinned = root.io_path()?;
1548 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1549 ctx.root_guard = Some(root);
1550 Self::create_in(&pinned, schema, table_id, ctx)
1551 }
1552
1553 pub fn create_with_key(
1558 dir: impl AsRef<Path>,
1559 schema: Schema,
1560 table_id: u64,
1561 key: &[u8],
1562 ) -> Result<Self> {
1563 let dir = dir.as_ref().to_path_buf();
1564 crate::durable_file::create_directory_all(&dir)?;
1565 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1566 root.create_directory_all(META_DIR)?;
1567 let salt = crate::encryption::random_salt()?;
1568 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1569 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1570 let pinned = root.io_path()?;
1571 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1572 ctx.root_guard = Some(root);
1573 Self::create_in(&pinned, schema, table_id, ctx)
1574 }
1575
1576 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1578 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1579 let salt = read_table_encryption_salt_root(&root)?;
1580 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1581 let pinned = root.io_path()?;
1582 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1583 ctx.root_guard = Some(root);
1584 Self::open_in(&pinned, ctx)
1585 }
1586
1587 pub(crate) fn create_in(
1588 dir: impl AsRef<Path>,
1589 schema: Schema,
1590 table_id: u64,
1591 ctx: SharedCtx,
1592 ) -> Result<Self> {
1593 schema.validate_auto_increment()?;
1594 schema.validate_defaults()?;
1595 schema.validate_ai()?;
1596 for index in &schema.indexes {
1597 index.validate_options()?;
1598 }
1599 let dir = dir.as_ref().to_path_buf();
1600 let runs_root = match ctx.root_guard.as_ref() {
1601 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1602 None => {
1603 crate::durable_file::create_directory_all(&dir)?;
1604 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1605 None
1606 }
1607 };
1608 match ctx.root_guard.as_deref() {
1609 Some(root) => write_schema_durable(root, &schema)?,
1610 None => write_schema(&dir, &schema)?,
1611 }
1612 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1613 let (wal, current_txn_id) = match ctx.shared.clone() {
1616 Some(s) => (WalSink::Shared(s), 0),
1617 None => {
1618 let pinned_wal_root = match ctx.root_guard.as_deref() {
1619 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1620 None => None,
1621 };
1622 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1623 root.io_path()?
1624 } else {
1625 let wal_dir = dir.join(WAL_DIR);
1626 crate::durable_file::create_directory_all(&wal_dir)?;
1627 wal_dir
1628 };
1629 let mut w = if let Some(ref dk) = wal_dek {
1630 Wal::create_with_cipher(
1631 wal_dir.join("seg-000000.wal"),
1632 Epoch(0),
1633 Some(make_cipher(dk)),
1634 0,
1635 )?
1636 } else {
1637 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1638 };
1639 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1640 (WalSink::Private(w), 1)
1641 }
1642 };
1643 let mut manifest = Manifest::new(table_id, schema.schema_id);
1644 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1649 match ctx.root_guard.as_deref() {
1650 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1651 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1652 }
1653 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1654 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1655 let auto_inc = resolve_auto_inc(&schema);
1656 let rcache_dir = dir.join(RCACHE_DIR);
1657 let initial_view = ReadGeneration::empty(&schema);
1658 Ok(Self {
1659 dir,
1660 _root_guard: ctx.root_guard,
1661 runs_root,
1662 idx_root: None,
1663 table_id,
1664 name: ctx.table_name.unwrap_or_default(),
1665 auth: ctx.auth,
1666 read_only: ctx.read_only,
1667 durable_commit_failed: false,
1668 wal,
1669 memtable: Memtable::new(),
1670 mutable_run: MutableRun::new(),
1671 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1672 compaction_zstd_level: 3,
1673 allocator: RowIdAllocator::new(0),
1674 epoch: ctx.epoch,
1675 data_generation: 0,
1676 schema,
1677 hot: HotIndex::new(),
1678 kek: ctx.kek,
1679 column_keys,
1680 run_refs: Vec::new(),
1681 retiring: Vec::new(),
1682 next_run_id: 1,
1683 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1684 current_txn_id,
1685 pending_private_mutations: false,
1686 bitmap,
1687 ann,
1688 fm,
1689 sparse,
1690 minhash,
1691 learned_range: Arc::new(HashMap::new()),
1692 pk_by_row: ReversePkMap::new(),
1693 pinned: BTreeMap::new(),
1694 live_count: 0,
1695 reservoir: crate::reservoir::Reservoir::default(),
1696 reservoir_complete: true,
1697 had_deletes: false,
1698 agg_cache: Arc::new(HashMap::new()),
1699 global_idx_epoch: 0,
1700 indexes_complete: true,
1701 index_build_policy: IndexBuildPolicy::default(),
1702 pk_by_row_complete: false,
1703 flushed_epoch: 0,
1704 page_cache: ctx.page_cache,
1705 decoded_cache: ctx.decoded_cache,
1706 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1707 snapshots: ctx.snapshots,
1708 commit_lock: ctx.commit_lock,
1709 result_cache: Arc::new(parking_lot::Mutex::new(
1710 ResultCache::new()
1711 .with_dir(rcache_dir)
1712 .with_cache_dek(cache_dek.clone()),
1713 )),
1714 pending_delete_rids: roaring::RoaringBitmap::new(),
1715 pending_put_cols: std::collections::HashSet::new(),
1716 pending_rows: Vec::new(),
1717 pending_rows_auto_inc: Vec::new(),
1718 pending_dels: Vec::new(),
1719 pending_truncate: None,
1720 wal_dek,
1721 auto_inc,
1722 ttl: None,
1723 pins: Arc::new(crate::retention::PinRegistry::new()),
1724 published: Arc::new(ArcSwap::from_pointee(initial_view)),
1725 read_generation_pin: None,
1726 })
1727 }
1728
1729 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1733 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1734 let pinned = root.io_path()?;
1735 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1736 ctx.root_guard = Some(root);
1737 Self::open_in(&pinned, ctx)
1738 }
1739
1740 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1743 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1744 let salt = read_table_encryption_salt_root(&root)?;
1745 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1746 let pinned = root.io_path()?;
1747 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1748 ctx.root_guard = Some(root);
1749 let t = Self::open_in(&pinned, ctx)?;
1750 Ok(t)
1751 }
1752
1753 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1754 let dir = dir.as_ref().to_path_buf();
1755 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1756 let mut manifest = match ctx.root_guard.as_ref() {
1757 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1758 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1759 };
1760 let schema: Schema = match ctx.root_guard.as_ref() {
1761 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1762 None => read_schema(&dir)?,
1763 };
1764 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1770 let runs_root = match ctx.root_guard.as_ref() {
1771 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1772 None => None,
1773 };
1774 let idx_root = match ctx.root_guard.as_ref() {
1775 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1776 Ok(root) => Some(Arc::new(root)),
1777 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1778 Err(error) => return Err(error.into()),
1779 },
1780 None => None,
1781 };
1782 schema.validate_auto_increment()?;
1783 schema.validate_defaults()?;
1784 schema.validate_ai()?;
1785 for index in &schema.indexes {
1786 index.validate_options()?;
1787 }
1788 let replay_epoch = Epoch(manifest.current_epoch);
1789 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1790 let private_replayed = if ctx.shared.is_none() {
1791 match latest_wal_segment(&dir.join(WAL_DIR))? {
1792 Some(path) => {
1793 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1794 crate::wal::replay_with_cipher(path, cipher)?
1795 }
1796 None => Vec::new(),
1797 }
1798 } else {
1799 Vec::new()
1800 };
1801 if ctx.shared.is_none() {
1802 preflight_standalone_open(
1803 &dir,
1804 runs_root.as_deref(),
1805 idx_root.as_deref(),
1806 &manifest,
1807 &schema,
1808 &private_replayed,
1809 ctx.kek.clone(),
1810 )?;
1811 }
1812 let next_run_id = derive_next_run_id(
1813 &dir,
1814 runs_root.as_deref(),
1815 &manifest.runs,
1816 &manifest.retiring,
1817 )?;
1818 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1822 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1823 None => {
1824 let replayed = private_replayed;
1825 let wal_dir = dir.join(WAL_DIR);
1831 crate::durable_file::create_directory_all(&wal_dir)?;
1832 let segment = next_wal_segment(&wal_dir)?;
1833 let segment_no = wal_segment_number(&segment).unwrap_or(0);
1834 let temporary = wal_dir.join(format!(
1835 ".recovery-{}-{}-{segment_no:06}.tmp",
1836 std::process::id(),
1837 std::time::SystemTime::now()
1838 .duration_since(std::time::UNIX_EPOCH)
1839 .unwrap_or_default()
1840 .as_nanos()
1841 ));
1842 let mut w = Wal::create_with_cipher(
1843 &temporary,
1844 replay_epoch,
1845 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1846 segment_no,
1847 )?;
1848 for record in &replayed {
1849 w.append_txn(record.txn_id, record.op.clone())?;
1850 }
1851 let mut w = w.publish_as(segment)?;
1852 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1853 let next_txn_id = replayed
1854 .iter()
1855 .map(|record| record.txn_id)
1856 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
1857 .max()
1858 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
1859 .unwrap_or(1);
1860 (WalSink::Private(w), replayed, next_txn_id)
1861 }
1862 };
1863
1864 let mut memtable = Memtable::new();
1865 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1866 let persisted_epoch = manifest.current_epoch;
1867 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1874 s.next = manifest.auto_inc_next;
1875 s.seeded = manifest.auto_inc_next > 0;
1876 s
1877 });
1878
1879 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1886 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1887 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
1888 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1889 std::collections::BTreeMap::new();
1890 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1891 let mut recovered_epoch = manifest.current_epoch;
1892 let mut recovered_manifest_dirty = schema_manifest_repair;
1893 let mut saw_delete = false;
1894 for record in replayed {
1895 let txn_id = record.txn_id;
1896 match record.op {
1897 Op::Put { rows, .. } => {
1898 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1899 for row in &rows {
1900 allocator.advance_to(row.row_id)?;
1901 if let Some(ai) = auto_inc.as_mut() {
1902 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1903 let next = n.checked_add(1).ok_or_else(|| {
1904 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
1905 })?;
1906 if next > ai.next {
1907 ai.next = next;
1908 }
1909 }
1910 }
1911 }
1912 staged_puts.entry(txn_id).or_default().extend(rows);
1913 }
1914 Op::Delete { row_ids, .. } => {
1915 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1916 }
1917 Op::TxnCommit { epoch, .. } => {
1918 let commit_epoch = Epoch(epoch);
1919 recovered_epoch = recovered_epoch.max(epoch);
1920 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
1921 memtable = Memtable::new();
1922 replayed_puts.clear();
1923 replayed_deletes.clear();
1924 manifest.runs.clear();
1925 manifest.retiring.clear();
1926 manifest.live_count = 0;
1927 manifest.global_idx_epoch = 0;
1928 manifest.current_epoch = manifest.current_epoch.max(epoch);
1929 recovered_manifest_dirty = true;
1930 saw_delete = true;
1931 }
1932 if let Some(puts) = staged_puts.remove(&txn_id) {
1933 if commit_epoch.0 > manifest.flushed_epoch {
1934 for row in &puts {
1935 memtable.upsert(row.clone());
1936 }
1937 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1938 }
1939 }
1940 if let Some(dels) = staged_deletes.remove(&txn_id) {
1941 saw_delete = true;
1942 if commit_epoch.0 > manifest.flushed_epoch {
1943 for rid in dels {
1944 memtable.tombstone(rid, commit_epoch);
1945 replayed_deletes.push((rid, commit_epoch));
1946 }
1947 }
1948 }
1949 }
1950 Op::TxnAbort => {
1951 staged_puts.remove(&txn_id);
1952 staged_deletes.remove(&txn_id);
1953 staged_truncates.remove(&txn_id);
1954 }
1955 Op::TruncateTable { .. } => {
1956 staged_puts.remove(&txn_id);
1957 staged_deletes.remove(&txn_id);
1958 staged_truncates.insert(txn_id);
1959 }
1960 Op::ExternalTableState { .. }
1961 | Op::Flush { .. }
1962 | Op::Ddl(_)
1963 | Op::BeforeImage { .. }
1964 | Op::CommitTimestamp { .. }
1965 | Op::SpilledRows { .. } => {}
1966 }
1967 }
1968
1969 let rcache_dir = dir.join(RCACHE_DIR);
1970 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1971 let initial_view = ReadGeneration::empty(&schema);
1972 let mut db = Self {
1973 dir,
1974 _root_guard: ctx.root_guard,
1975 runs_root,
1976 idx_root,
1977 table_id: manifest.table_id,
1978 name: ctx.table_name.unwrap_or_default(),
1979 auth: ctx.auth,
1980 read_only: ctx.read_only,
1981 durable_commit_failed: false,
1982 wal,
1983 memtable,
1984 mutable_run: MutableRun::new(),
1985 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1986 compaction_zstd_level: 3,
1987 allocator,
1988 epoch: ctx.epoch,
1989 data_generation: persisted_epoch,
1990 schema,
1991 hot: HotIndex::new(),
1992 kek: ctx.kek,
1993 column_keys,
1994 run_refs: manifest.runs.clone(),
1995 retiring: manifest.retiring.clone(),
1996 next_run_id,
1997 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1998 current_txn_id,
1999 pending_private_mutations: false,
2000 bitmap: HashMap::new(),
2001 ann: HashMap::new(),
2002 fm: HashMap::new(),
2003 sparse: HashMap::new(),
2004 minhash: HashMap::new(),
2005 learned_range: Arc::new(HashMap::new()),
2006 pk_by_row: ReversePkMap::new(),
2007 pinned: BTreeMap::new(),
2008 live_count: manifest.live_count,
2009 reservoir: crate::reservoir::Reservoir::default(),
2010 reservoir_complete: false,
2011 had_deletes: saw_delete
2012 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2013 != manifest.live_count,
2014 agg_cache: Arc::new(HashMap::new()),
2015 global_idx_epoch: manifest.global_idx_epoch,
2016 indexes_complete: true,
2017 index_build_policy: IndexBuildPolicy::default(),
2018 pk_by_row_complete: false,
2019 flushed_epoch: manifest.flushed_epoch,
2020 page_cache: ctx.page_cache,
2021 decoded_cache: ctx.decoded_cache,
2022 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2023 snapshots: ctx.snapshots,
2024 commit_lock: ctx.commit_lock,
2025 result_cache: Arc::new(parking_lot::Mutex::new(
2026 ResultCache::new()
2027 .with_dir(rcache_dir)
2028 .with_cache_dek(cache_dek.clone()),
2029 )),
2030 pending_delete_rids: roaring::RoaringBitmap::new(),
2031 pending_put_cols: std::collections::HashSet::new(),
2032 pending_rows: Vec::new(),
2033 pending_rows_auto_inc: Vec::new(),
2034 pending_dels: Vec::new(),
2035 pending_truncate: None,
2036 wal_dek,
2037 auto_inc,
2038 ttl: manifest.ttl,
2039 pins: Arc::new(crate::retention::PinRegistry::new()),
2040 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2041 read_generation_pin: None,
2042 };
2043
2044 db.epoch.advance_recovered(Epoch(recovered_epoch));
2047
2048 let checkpoint = match db.idx_root.as_deref() {
2053 Some(root) => {
2054 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2055 }
2056 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2057 };
2058 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2059 c.epoch_built == manifest.global_idx_epoch
2060 && manifest.global_idx_epoch > 0
2061 && manifest
2062 .runs
2063 .iter()
2064 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2065 });
2066 if let Some(loaded) = checkpoint {
2067 if checkpoint_valid {
2068 db.hot = loaded.hot;
2069 db.bitmap = loaded.bitmap;
2070 db.ann = loaded.ann;
2071 db.fm = loaded.fm;
2072 db.sparse = loaded.sparse;
2073 db.minhash = loaded.minhash;
2074 db.learned_range = Arc::new(loaded.learned_range);
2075 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2080 for (cid, idx) in bitmap0 {
2081 db.bitmap.entry(cid).or_insert(idx);
2082 }
2083 for (cid, idx) in ann0 {
2084 db.ann.entry(cid).or_insert(idx);
2085 }
2086 for (cid, idx) in fm0 {
2087 db.fm.entry(cid).or_insert(idx);
2088 }
2089 for (cid, idx) in sparse0 {
2090 db.sparse.entry(cid).or_insert(idx);
2091 }
2092 for (cid, idx) in minhash0 {
2093 db.minhash.entry(cid).or_insert(idx);
2094 }
2095 }
2098 }
2099 if !checkpoint_valid {
2100 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2101 db.bitmap = bitmap;
2102 db.ann = ann;
2103 db.fm = fm;
2104 db.sparse = sparse;
2105 db.minhash = minhash;
2106 db.rebuild_indexes_from_runs()?;
2107 db.build_learned_ranges()?;
2108 }
2109
2110 for (epoch, group) in replayed_puts {
2115 let (losers, winner_pks) = db.partition_pk_winners(&group);
2116 for (key, &row_id) in &winner_pks {
2117 if let Some(old_rid) = db.hot.get(key) {
2118 if old_rid != row_id {
2119 db.tombstone_row(old_rid, epoch, false);
2120 }
2121 }
2122 }
2123 for &loser_rid in &losers {
2124 db.tombstone_row(loser_rid, epoch, false);
2125 }
2126 for (key, row_id) in winner_pks {
2127 db.insert_hot_pk(key, row_id);
2128 }
2129 if db.schema.primary_key().is_none() {
2130 for r in &group {
2131 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2132 }
2133 }
2134 for r in &group {
2135 if !losers.contains(&r.row_id) {
2136 db.index_row(r);
2137 }
2138 }
2139 }
2140 for (rid, epoch) in &replayed_deletes {
2144 db.remove_hot_for_row(*rid, *epoch);
2145 }
2146
2147 if recovered_manifest_dirty {
2148 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2149 db.live_count = rows.len() as u64;
2150 db.persist_manifest(Epoch(recovered_epoch))?;
2151 }
2152
2153 db.result_cache.lock().load_persistent();
2160 Ok(db)
2161 }
2162
2163 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2169 if self.reservoir_complete {
2170 return Ok(());
2171 }
2172 self.rebuild_reservoir()?;
2173 self.reservoir_complete = true;
2174 Ok(())
2175 }
2176
2177 fn rebuild_reservoir(&mut self) -> Result<()> {
2180 let snap = self.snapshot();
2181 let rows = self.visible_rows(snap)?;
2182 self.reservoir.reset();
2183 for r in rows {
2184 self.reservoir.offer(r.row_id.0);
2185 }
2186 Ok(())
2187 }
2188
2189 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2190 self.rebuild_indexes_from_runs_inner(None)
2191 }
2192
2193 fn rebuild_indexes_from_runs_inner(
2194 &mut self,
2195 control: Option<&crate::ExecutionControl>,
2196 ) -> Result<()> {
2197 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2200 crate::retention::PinSource::OnlineIndexBuild,
2201 self.current_epoch(),
2202 );
2203 self.hot = HotIndex::new();
2204 self.pk_by_row.clear();
2205 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2206 self.bitmap = bitmap;
2207 self.ann = ann;
2208 self.fm = fm;
2209 self.sparse = sparse;
2210 self.minhash = minhash;
2211 let snapshot = Epoch(u64::MAX);
2212 let ttl_now = unix_nanos_now();
2213 let mut scanned = 0_usize;
2214 for rr in self.run_refs.clone() {
2215 if let Some(control) = control {
2216 control.checkpoint()?;
2217 }
2218 let mut reader = self.open_reader(rr.run_id)?;
2219 for row in reader.visible_rows(snapshot)? {
2220 if scanned.is_multiple_of(256) {
2221 if let Some(control) = control {
2222 control.checkpoint()?;
2223 }
2224 }
2225 scanned += 1;
2226 if self.row_expired_at(&row, ttl_now) {
2227 continue;
2228 }
2229 let tok_row = self.tokenized_for_indexes(&row);
2230 index_into(
2231 &self.schema,
2232 &tok_row,
2233 &mut self.hot,
2234 &mut self.bitmap,
2235 &mut self.ann,
2236 &mut self.fm,
2237 &mut self.sparse,
2238 &mut self.minhash,
2239 );
2240 }
2241 }
2242 for row in self.mutable_run.visible_versions(snapshot) {
2243 if scanned.is_multiple_of(256) {
2244 if let Some(control) = control {
2245 control.checkpoint()?;
2246 }
2247 }
2248 scanned += 1;
2249 if row.deleted {
2250 self.remove_hot_for_row(row.row_id, snapshot);
2251 } else if !self.row_expired_at(&row, ttl_now) {
2252 self.index_row(&row);
2253 }
2254 }
2255 for row in self.memtable.visible_versions(snapshot) {
2256 if scanned.is_multiple_of(256) {
2257 if let Some(control) = control {
2258 control.checkpoint()?;
2259 }
2260 }
2261 scanned += 1;
2262 if row.deleted {
2263 self.remove_hot_for_row(row.row_id, snapshot);
2264 } else if !self.row_expired_at(&row, ttl_now) {
2265 self.index_row(&row);
2266 }
2267 }
2268 self.refresh_pk_by_row_from_hot();
2269 Ok(())
2270 }
2271
2272 fn refresh_pk_by_row_from_hot(&mut self) {
2273 self.pk_by_row_complete = true;
2274 if self.schema.primary_key().is_none() {
2275 self.pk_by_row.clear();
2276 return;
2277 }
2278 self.pk_by_row = ReversePkMap::from_entries(
2284 self.hot
2285 .entries()
2286 .into_iter()
2287 .map(|(key, row_id)| (row_id, key)),
2288 );
2289 }
2290
2291 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2292 if self.schema.primary_key().is_some() {
2293 self.pk_by_row.insert(row_id, key.clone());
2294 }
2295 self.hot.insert(key, row_id);
2296 }
2297
2298 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2302 self.build_learned_ranges_inner(None)
2303 }
2304
2305 fn build_learned_ranges_inner(
2306 &mut self,
2307 control: Option<&crate::ExecutionControl>,
2308 ) -> Result<()> {
2309 self.learned_range = Arc::new(HashMap::new());
2310 if self.run_refs.len() != 1 {
2311 return Ok(());
2312 }
2313 let cols: Vec<(u16, usize)> = self
2314 .schema
2315 .indexes
2316 .iter()
2317 .filter(|i| i.kind == IndexKind::LearnedRange)
2318 .map(|i| {
2319 (
2320 i.column_id,
2321 i.options
2322 .learned_range
2323 .as_ref()
2324 .map(|options| options.epsilon)
2325 .unwrap_or(16),
2326 )
2327 })
2328 .collect();
2329 if cols.is_empty() {
2330 return Ok(());
2331 }
2332 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2333 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2334 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2335 _ => return Ok(()),
2336 };
2337 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2338 if column_index % 256 == 0 {
2339 if let Some(control) = control {
2340 control.checkpoint()?;
2341 }
2342 }
2343 let ty = self
2344 .schema
2345 .columns
2346 .iter()
2347 .find(|c| c.id == cid)
2348 .map(|c| c.ty.clone())
2349 .unwrap_or(TypeId::Int64);
2350 match ty {
2351 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2352 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2353 let pairs: Vec<(i64, u64)> = data
2354 .iter()
2355 .zip(row_ids.iter())
2356 .map(|(v, r)| (*v, *r))
2357 .collect();
2358 Arc::make_mut(&mut self.learned_range).insert(
2359 cid,
2360 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2361 );
2362 }
2363 }
2364 TypeId::Float64 => {
2365 if let columnar::NativeColumn::Float64 { data, .. } =
2366 reader.column_native(cid)?
2367 {
2368 let pairs: Vec<(f64, u64)> = data
2369 .iter()
2370 .zip(row_ids.iter())
2371 .map(|(v, r)| (*v, *r))
2372 .collect();
2373 Arc::make_mut(&mut self.learned_range).insert(
2374 cid,
2375 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2376 );
2377 }
2378 }
2379 _ => {}
2380 }
2381 }
2382 Ok(())
2383 }
2384
2385 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2392 if self.indexes_complete {
2393 crate::trace::QueryTrace::record(|t| {
2394 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2395 });
2396 return Ok(());
2397 }
2398 crate::trace::QueryTrace::record(|t| {
2399 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2400 });
2401 self.rebuild_indexes_from_runs()?;
2402 self.build_learned_ranges()?;
2403 self.indexes_complete = true;
2404 let epoch = self.current_epoch();
2405 self.checkpoint_indexes(epoch);
2406 Ok(())
2407 }
2408
2409 #[doc(hidden)]
2412 pub fn ensure_indexes_complete_controlled<F>(
2413 &mut self,
2414 control: &crate::ExecutionControl,
2415 before_publish: F,
2416 ) -> Result<bool>
2417 where
2418 F: FnOnce() -> bool,
2419 {
2420 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2421 .map(|(changed, _)| changed)
2422 }
2423
2424 #[doc(hidden)]
2427 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2428 &mut self,
2429 control: &crate::ExecutionControl,
2430 before_publish: F,
2431 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2432 where
2433 F: FnOnce() -> bool,
2434 {
2435 if self.indexes_complete {
2436 crate::trace::QueryTrace::record(|trace| {
2437 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2438 });
2439 return Ok((false, None));
2440 }
2441 crate::trace::QueryTrace::record(|trace| {
2442 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2443 });
2444 control.checkpoint()?;
2445 let maintenance_epoch = self.current_epoch();
2446 self.rebuild_indexes_from_runs_inner(Some(control))?;
2447 self.build_learned_ranges_inner(Some(control))?;
2448 control.checkpoint()?;
2449 if !before_publish() {
2450 return Err(MongrelError::Cancelled);
2451 }
2452 self.indexes_complete = true;
2453 self.checkpoint_indexes(maintenance_epoch);
2454 Ok((
2455 true,
2456 Some(MaintenanceReceipt {
2457 epoch: maintenance_epoch,
2458 }),
2459 ))
2460 }
2461
2462 fn pending_epoch(&self) -> Epoch {
2463 Epoch(self.epoch.visible().0 + 1)
2464 }
2465
2466 fn is_shared(&self) -> bool {
2469 matches!(self.wal, WalSink::Shared(_))
2470 }
2471
2472 fn ensure_txn_id(&mut self) -> Result<u64> {
2476 if self.current_txn_id == 0 {
2477 let id = match &self.wal {
2478 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2479 WalSink::Private(_) => {
2480 return Err(MongrelError::Full(
2481 "standalone transaction id namespace exhausted".into(),
2482 ))
2483 }
2484 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2485 };
2486 self.current_txn_id = id;
2487 }
2488 Ok(self.current_txn_id)
2489 }
2490
2491 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2494 self.ensure_writable()?;
2495 let txn_id = self.ensure_txn_id()?;
2496 let table_id = self.table_id;
2497 match &mut self.wal {
2498 WalSink::Private(w) => {
2499 w.append_txn(txn_id, op)?;
2500 self.pending_private_mutations = true;
2501 }
2502 WalSink::Shared(s) => {
2503 s.wal.lock().append(txn_id, table_id, op)?;
2504 }
2505 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2506 }
2507 Ok(())
2508 }
2509
2510 fn ensure_writable(&self) -> Result<()> {
2511 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2512 return Err(MongrelError::ReadOnlyReplica);
2513 }
2514 if self.durable_commit_failed {
2515 return Err(MongrelError::Other(
2516 "table poisoned by post-commit failure; reopen required".into(),
2517 ));
2518 }
2519 Ok(())
2520 }
2521
2522 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2533 match &self.auth {
2534 Some(checker) => checker.check(&self.name, perm),
2535 None => Ok(()),
2536 }
2537 }
2538 pub fn require_select(&self) -> Result<()> {
2543 self.require(crate::auth_state::RequiredPermission::Select)
2544 }
2545 fn require_insert(&self) -> Result<()> {
2546 self.require(crate::auth_state::RequiredPermission::Insert)
2547 }
2548 #[allow(dead_code)]
2552 fn require_update(&self) -> Result<()> {
2553 self.require(crate::auth_state::RequiredPermission::Update)
2554 }
2555 fn require_delete(&self) -> Result<()> {
2556 self.require(crate::auth_state::RequiredPermission::Delete)
2557 }
2558
2559 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2562 self.require_insert()?;
2563 Ok(self.put_returning(columns)?.0)
2564 }
2565
2566 pub fn put_returning(
2571 &mut self,
2572 mut columns: Vec<(u16, Value)>,
2573 ) -> Result<(RowId, Option<i64>)> {
2574 self.require_insert()?;
2575 let assigned = self.fill_auto_inc(&mut columns)?;
2576 self.apply_defaults(&mut columns)?;
2577 self.schema.validate_values(&columns)?;
2578 let row_id = if self.schema.clustered {
2583 self.derive_clustered_row_id(&columns)?
2584 } else {
2585 self.allocator.alloc()?
2586 };
2587 let epoch = self.pending_epoch();
2588 let mut row = Row::new(row_id, epoch);
2589 for (col_id, val) in columns {
2590 row.columns.insert(col_id, val);
2591 }
2592 self.commit_rows(vec![row], assigned.is_some())?;
2593 Ok((row_id, assigned))
2594 }
2595
2596 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2599 self.require_insert()?;
2600 Ok(self
2601 .put_batch_returning(batch)?
2602 .into_iter()
2603 .map(|(r, _)| r)
2604 .collect())
2605 }
2606
2607 pub fn put_batch_returning(
2610 &mut self,
2611 batch: Vec<Vec<(u16, Value)>>,
2612 ) -> Result<Vec<(RowId, Option<i64>)>> {
2613 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2614 for mut cols in batch {
2615 let assigned = self.fill_auto_inc(&mut cols)?;
2616 self.apply_defaults(&mut cols)?;
2617 filled.push((cols, assigned));
2618 }
2619 for (cols, _) in &filled {
2620 self.schema.validate_values(cols)?;
2621 }
2622 let epoch = self.pending_epoch();
2623 let mut rows = Vec::with_capacity(filled.len());
2624 let mut ids = Vec::with_capacity(filled.len());
2625 let first_row_id = if self.schema.clustered {
2626 None
2627 } else {
2628 let count = u64::try_from(filled.len())
2629 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2630 Some(self.allocator.alloc_range(count)?.0)
2631 };
2632 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2633 let row_id = match first_row_id {
2634 Some(first) => RowId(first + row_index as u64),
2635 None => self.derive_clustered_row_id(&cols)?,
2636 };
2637 let mut row = Row::new(row_id, epoch);
2638 for (c, v) in cols {
2639 row.columns.insert(c, v);
2640 }
2641 ids.push((row_id, assigned));
2642 rows.push(row);
2643 }
2644 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2645 self.commit_rows(rows, all_auto_generated)?;
2646 Ok(ids)
2647 }
2648
2649 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2655 self.ensure_writable()?;
2656 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2657 return Ok(None);
2658 };
2659 let pos = columns.iter().position(|(c, _)| *c == cid);
2660 let assigned = match pos {
2661 Some(i) => match &columns[i].1 {
2662 Value::Null => {
2663 let next = self.alloc_auto_inc_value()?;
2664 columns[i].1 = Value::Int64(next);
2665 Some(next)
2666 }
2667 Value::Int64(n) => {
2668 self.advance_auto_inc_past(*n)?;
2669 None
2670 }
2671 other => {
2672 return Err(MongrelError::InvalidArgument(format!(
2673 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2674 other
2675 )))
2676 }
2677 },
2678 None => {
2679 let next = self.alloc_auto_inc_value()?;
2680 columns.push((cid, Value::Int64(next)));
2681 Some(next)
2682 }
2683 };
2684 Ok(assigned)
2685 }
2686
2687 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2693 for col in &self.schema.columns {
2694 let Some(expr) = &col.default_value else {
2695 continue;
2696 };
2697 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2699 continue;
2700 }
2701 let pos = columns.iter().position(|(c, _)| *c == col.id);
2702 let needs_default = match pos {
2703 None => true,
2704 Some(i) => matches!(columns[i].1, Value::Null),
2705 };
2706 if !needs_default {
2707 continue;
2708 }
2709 let v = match expr {
2710 crate::schema::DefaultExpr::Static(v) => v.clone(),
2711 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2712 crate::schema::DefaultExpr::Uuid => {
2713 let mut buf = [0u8; 16];
2714 getrandom::getrandom(&mut buf)
2715 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2716 Value::Uuid(buf)
2717 }
2718 };
2719 match pos {
2720 None => columns.push((col.id, v)),
2721 Some(i) => columns[i].1 = v,
2722 }
2723 }
2724 Ok(())
2725 }
2726
2727 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2729 self.ensure_auto_inc_seeded()?;
2730 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2732 let v = ai.next;
2733 ai.next = ai
2734 .next
2735 .checked_add(1)
2736 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2737 Ok(v)
2738 }
2739
2740 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2743 self.ensure_auto_inc_seeded()?;
2744 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2745 let floor = used
2746 .checked_add(1)
2747 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2748 .max(1);
2749 if ai.next < floor {
2750 ai.next = floor;
2751 }
2752 Ok(())
2753 }
2754
2755 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2760 let needs_seed = match self.auto_inc {
2761 Some(ai) => !ai.seeded,
2762 None => return Ok(()),
2763 };
2764 if !needs_seed {
2765 return Ok(());
2766 }
2767 if self.seed_empty_auto_inc() {
2768 return Ok(());
2769 }
2770 let cid = self
2771 .auto_inc
2772 .as_ref()
2773 .expect("auto-inc column present")
2774 .column_id;
2775 let max = self.scan_max_int64(cid)?;
2776 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2777 let floor = max
2778 .checked_add(1)
2779 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2780 .max(1);
2781 if ai.next < floor {
2782 ai.next = floor;
2783 }
2784 ai.seeded = true;
2785 Ok(())
2786 }
2787
2788 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2789 if n == 0 || self.auto_inc.is_none() {
2790 return Ok(None);
2791 }
2792 self.ensure_auto_inc_seeded()?;
2793 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2794 let start = ai.next;
2795 let count = i64::try_from(n)
2796 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2797 ai.next = ai
2798 .next
2799 .checked_add(count)
2800 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2801 Ok(Some(start))
2802 }
2803
2804 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2808 let mut max: i64 = 0;
2809 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2810 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2811 if *n > max {
2812 max = *n;
2813 }
2814 }
2815 }
2816 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2817 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2818 if *n > max {
2819 max = *n;
2820 }
2821 }
2822 }
2823 for rr in self.run_refs.clone() {
2824 let reader = self.open_reader(rr.run_id)?;
2825 if let Some(stats) = reader.column_page_stats(column_id) {
2826 for s in stats {
2827 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2828 if n > max {
2829 max = n;
2830 }
2831 }
2832 }
2833 } else if reader.has_column(column_id) {
2834 if let columnar::NativeColumn::Int64 { data, validity } =
2835 reader.column_native_shared(column_id)?
2836 {
2837 for (i, n) in data.iter().enumerate() {
2838 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2839 {
2840 max = *n;
2841 }
2842 }
2843 }
2844 }
2845 }
2846 Ok(max)
2847 }
2848
2849 fn seed_empty_auto_inc(&mut self) -> bool {
2850 let Some(ai) = self.auto_inc.as_mut() else {
2851 return false;
2852 };
2853 if ai.seeded || self.live_count != 0 {
2854 return false;
2855 }
2856 if ai.next < 1 {
2857 ai.next = 1;
2858 }
2859 ai.seeded = true;
2860 true
2861 }
2862
2863 fn advance_auto_inc_from_native_columns(
2864 &mut self,
2865 columns: &[(u16, columnar::NativeColumn)],
2866 n: usize,
2867 live_before: u64,
2868 ) -> Result<()> {
2869 let Some(ai) = self.auto_inc.as_mut() else {
2870 return Ok(());
2871 };
2872 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2873 return Ok(());
2874 };
2875 let columnar::NativeColumn::Int64 { data, validity } = col else {
2876 return Err(MongrelError::InvalidArgument(format!(
2877 "AUTO_INCREMENT column {} must be Int64",
2878 ai.column_id
2879 )));
2880 };
2881 let max = if native_int64_strictly_increasing(col, n) {
2882 data.get(n.saturating_sub(1)).copied()
2883 } else {
2884 data.iter()
2885 .take(n)
2886 .enumerate()
2887 .filter_map(|(i, v)| {
2888 if validity.is_empty() || columnar::validity_bit(validity, i) {
2889 Some(*v)
2890 } else {
2891 None
2892 }
2893 })
2894 .max()
2895 };
2896 if let Some(max) = max {
2897 let floor = max
2898 .checked_add(1)
2899 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2900 .max(1);
2901 if ai.next < floor {
2902 ai.next = floor;
2903 }
2904 if ai.seeded || live_before == 0 {
2905 ai.seeded = true;
2906 }
2907 }
2908 Ok(())
2909 }
2910
2911 fn fill_auto_inc_native_columns(
2912 &mut self,
2913 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2914 n: usize,
2915 ) -> Result<()> {
2916 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2917 return Ok(());
2918 };
2919 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2920 if let Some(start) = self.alloc_auto_inc_range(n)? {
2921 columns.push((
2922 cid,
2923 columnar::NativeColumn::Int64 {
2924 data: (start..start.saturating_add(n as i64)).collect(),
2925 validity: vec![0xFF; n.div_ceil(8)],
2926 },
2927 ));
2928 }
2929 return Ok(());
2930 };
2931
2932 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2933 return Err(MongrelError::InvalidArgument(format!(
2934 "AUTO_INCREMENT column {cid} must be Int64"
2935 )));
2936 };
2937 if data.len() < n {
2938 return Err(MongrelError::InvalidArgument(format!(
2939 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2940 data.len()
2941 )));
2942 }
2943 if columnar::all_non_null(validity, n) {
2944 return Ok(());
2945 }
2946 if validity.iter().all(|b| *b == 0) {
2947 if let Some(start) = self.alloc_auto_inc_range(n)? {
2948 for (i, slot) in data.iter_mut().take(n).enumerate() {
2949 *slot = start.saturating_add(i as i64);
2950 }
2951 *validity = vec![0xFF; n.div_ceil(8)];
2952 }
2953 return Ok(());
2954 }
2955
2956 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2957 for (i, slot) in data.iter_mut().enumerate().take(n) {
2958 if columnar::validity_bit(validity, i) {
2959 self.advance_auto_inc_past(*slot)?;
2960 } else {
2961 *slot = self.alloc_auto_inc_value()?;
2962 }
2963 }
2964 *validity = new_validity;
2965 Ok(())
2966 }
2967
2968 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2982 self.ensure_writable()?;
2983 if self.auto_inc.is_none() {
2984 return Ok(None);
2985 }
2986 Ok(Some(self.alloc_auto_inc_value()?))
2987 }
2988
2989 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2995 let payload = bincode::serialize(&rows)?;
2996 self.wal_append_data(Op::Put {
2997 table_id: self.table_id,
2998 rows: payload,
2999 })?;
3000 if self.is_shared() {
3001 self.pending_rows_auto_inc
3002 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3003 self.pending_rows.extend(rows);
3004 } else {
3005 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3006 }
3007 Ok(())
3008 }
3009
3010 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3013 self.ensure_indexes_complete()
3014 }
3015
3016 pub(crate) fn prepare_durable_publish_controlled(
3017 &mut self,
3018 control: &crate::ExecutionControl,
3019 ) -> Result<()> {
3020 self.ensure_indexes_complete_controlled(control, || true)?;
3021 Ok(())
3022 }
3023
3024 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3025 self.apply_put_rows_inner_prepared(rows, true);
3026 }
3027
3028 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3029 if check_existing_pk {
3030 self.ensure_indexes_complete()?;
3031 }
3032 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3033 Ok(())
3034 }
3035
3036 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3040 if rows.len() == 1 {
3044 let row = rows.into_iter().next().expect("len checked");
3045 self.apply_put_row_single(row, check_existing_pk);
3046 return;
3047 }
3048 let pk_id = self.schema.primary_key().map(|c| c.id);
3065 let probe = match pk_id {
3066 Some(pid) => {
3067 check_existing_pk
3068 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3069 }
3070 None => false,
3071 };
3072 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3075 for r in rows {
3076 for &cid in r.columns.keys() {
3077 self.pending_put_cols.insert(cid);
3078 }
3079 match pk_id {
3080 Some(pid) if probe || maintain_pk_by_row => {
3081 if let Some(pk_val) = r.columns.get(&pid) {
3082 let key = self.index_lookup_key(pid, pk_val);
3083 if probe {
3084 if let Some(old_rid) = self.hot.get(&key) {
3085 if old_rid != r.row_id {
3086 self.tombstone_row(old_rid, r.committed_epoch, true);
3087 }
3088 }
3089 }
3090 if maintain_pk_by_row {
3091 self.pk_by_row.insert(r.row_id, key);
3092 }
3093 }
3094 }
3095 Some(_) => {}
3096 None => {
3097 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3098 }
3099 }
3100 self.index_row(&r);
3101 self.reservoir.offer(r.row_id.0);
3102 self.memtable.upsert(r);
3103 self.live_count = self.live_count.saturating_add(1);
3106 }
3107 self.data_generation = self.data_generation.wrapping_add(1);
3108 }
3109
3110 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3114 for &cid in row.columns.keys() {
3115 self.pending_put_cols.insert(cid);
3116 }
3117 let epoch = row.committed_epoch;
3118 if let Some(pk_col) = self.schema.primary_key() {
3119 let pk_id = pk_col.id;
3120 if let Some(pk_val) = row.columns.get(&pk_id) {
3121 let maintain_pk_by_row = self.pk_by_row_complete;
3125 if check_existing_pk || maintain_pk_by_row {
3126 let key = self.index_lookup_key(pk_id, pk_val);
3127 if check_existing_pk {
3128 if let Some(old_rid) = self.hot.get(&key) {
3129 if old_rid != row.row_id {
3130 self.tombstone_row(old_rid, epoch, true);
3131 }
3132 }
3133 }
3134 if maintain_pk_by_row {
3135 self.pk_by_row.insert(row.row_id, key);
3136 }
3137 }
3138 }
3139 } else {
3140 self.hot
3141 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3142 }
3143 self.index_row(&row);
3144 self.reservoir.offer(row.row_id.0);
3145 self.memtable.upsert(row);
3146 self.live_count = self.live_count.saturating_add(1);
3147 self.data_generation = self.data_generation.wrapping_add(1);
3148 }
3149
3150 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3153 self.allocator.alloc()
3154 }
3155
3156 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3162 let pk = self.schema.primary_key().ok_or_else(|| {
3163 MongrelError::Schema("clustered table requires a single-column primary key".into())
3164 })?;
3165 let pk_val = columns
3166 .iter()
3167 .find(|(id, _)| *id == pk.id)
3168 .map(|(_, v)| v)
3169 .ok_or_else(|| {
3170 MongrelError::Schema(format!(
3171 "clustered table missing primary key column {} ({})",
3172 pk.id, pk.name
3173 ))
3174 })?;
3175 let key_bytes = pk_val.encode_key();
3176 let mut hash: u64 = 0xcbf29ce484222325;
3178 for &b in &key_bytes {
3179 hash ^= b as u64;
3180 hash = hash.wrapping_mul(0x100000001b3);
3181 }
3182 Ok(RowId(hash.max(1)))
3185 }
3186
3187 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3195 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3196 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3197 }
3198 let n = rows.len();
3199 for r in rows {
3200 for &cid in r.columns.keys() {
3201 self.pending_put_cols.insert(cid);
3202 }
3203 }
3204 let (losers, winner_pks) = self.partition_pk_winners(rows);
3205 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3206 for (key, &row_id) in &winner_pks {
3208 if let Some(old_rid) = self.hot.get(key) {
3209 if old_rid != row_id {
3210 self.tombstone_row(old_rid, epoch, true);
3211 }
3212 }
3213 }
3214 for &loser_rid in &losers {
3217 self.tombstone_row(loser_rid, epoch, false);
3218 }
3219 for (key, row_id) in winner_pks {
3221 self.insert_hot_pk(key, row_id);
3222 }
3223 if self.schema.primary_key().is_none() {
3224 for r in rows {
3225 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3226 }
3227 }
3228 for r in rows {
3229 self.allocator.advance_to(r.row_id)?;
3230 if !losers.contains(&r.row_id) {
3231 self.index_row(r);
3232 }
3233 }
3234 for r in rows {
3235 if !losers.contains(&r.row_id) {
3236 self.reservoir.offer(r.row_id.0);
3237 }
3238 }
3239 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3240 self.data_generation = self.data_generation.wrapping_add(1);
3241 Ok(())
3242 }
3243
3244 pub(crate) fn recover_apply(
3249 &mut self,
3250 rows: Vec<Row>,
3251 deletes: Vec<(RowId, Epoch)>,
3252 ) -> Result<()> {
3253 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3257 std::collections::BTreeMap::new();
3258 for row in rows {
3259 if row.row_id.0 >= u64::MAX - 1 {
3260 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3261 }
3262 self.allocator.advance_to(row.row_id)?;
3263 if let Some(ai) = self.auto_inc.as_mut() {
3268 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3269 let next = n.checked_add(1).ok_or_else(|| {
3270 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3271 })?;
3272 if next > ai.next {
3273 ai.next = next;
3274 }
3275 }
3276 }
3277 by_epoch.entry(row.committed_epoch).or_default().push(row);
3278 }
3279 for (epoch, group) in by_epoch {
3280 let (losers, winner_pks) = self.partition_pk_winners(&group);
3281 for (key, &row_id) in &winner_pks {
3283 if let Some(old_rid) = self.hot.get(key) {
3284 if old_rid != row_id {
3285 self.tombstone_row(old_rid, epoch, false);
3286 }
3287 }
3288 }
3289 for (key, row_id) in winner_pks {
3290 self.insert_hot_pk(key, row_id);
3291 }
3292 if self.schema.primary_key().is_none() {
3293 for r in &group {
3294 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3295 }
3296 }
3297 for r in &group {
3298 if !losers.contains(&r.row_id) {
3299 self.memtable.upsert(r.clone());
3300 self.index_row(r);
3301 }
3302 }
3303 }
3304 for (rid, epoch) in deletes {
3305 self.memtable.tombstone(rid, epoch);
3306 self.remove_hot_for_row(rid, epoch);
3307 }
3308 self.reservoir_complete = false;
3311 Ok(())
3312 }
3313
3314 pub(crate) fn flushed_epoch(&self) -> u64 {
3316 self.flushed_epoch
3317 }
3318
3319 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3320 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3321 }
3322
3323 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3325 self.schema.validate_values(cells)
3326 }
3327
3328 fn validate_columns_not_null(
3332 &self,
3333 columns: &[(u16, columnar::NativeColumn)],
3334 n: usize,
3335 ) -> Result<()> {
3336 let by_id: HashMap<u16, &columnar::NativeColumn> =
3337 columns.iter().map(|(id, c)| (*id, c)).collect();
3338 for col in &self.schema.columns {
3339 if !col.flags.contains(ColumnFlags::NULLABLE) {
3340 match by_id.get(&col.id) {
3341 None => {
3342 return Err(MongrelError::InvalidArgument(format!(
3343 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3344 col.name, col.id
3345 )));
3346 }
3347 Some(c) => {
3348 if c.null_count(n) != 0 {
3349 return Err(MongrelError::InvalidArgument(format!(
3350 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3351 col.name, col.id
3352 )));
3353 }
3354 }
3355 }
3356 }
3357 if let TypeId::Enum { variants } = &col.ty {
3358 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3359 if by_id.contains_key(&col.id) {
3360 return Err(MongrelError::InvalidArgument(format!(
3361 "column '{}' ({}) enum requires a bytes column",
3362 col.name, col.id
3363 )));
3364 }
3365 continue;
3366 };
3367 for index in 0..n {
3368 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3369 continue;
3370 };
3371 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3372 return Err(MongrelError::InvalidArgument(format!(
3373 "column '{}' ({}) enum value {:?} is not one of {:?}",
3374 col.name,
3375 col.id,
3376 String::from_utf8_lossy(value),
3377 variants
3378 )));
3379 }
3380 }
3381 }
3382 }
3383 Ok(())
3384 }
3385
3386 fn bulk_pk_winner_indices(
3391 &self,
3392 columns: &[(u16, columnar::NativeColumn)],
3393 n: usize,
3394 ) -> Option<Vec<usize>> {
3395 let pk_col = self.schema.primary_key()?;
3396 let pk_id = pk_col.id;
3397 let pk_ty = pk_col.ty.clone();
3398 let by_id: HashMap<u16, &columnar::NativeColumn> =
3399 columns.iter().map(|(id, c)| (*id, c)).collect();
3400 let pk_native = by_id.get(&pk_id)?;
3401 if native_int64_strictly_increasing(pk_native, n) {
3402 return None;
3403 }
3404 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3406 let mut null_pk_rows: Vec<usize> = Vec::new();
3407 for i in 0..n {
3408 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3409 Some(key) => {
3410 last.insert(key, i);
3411 }
3412 None => null_pk_rows.push(i),
3413 }
3414 }
3415 let mut winners: HashSet<usize> = last.values().copied().collect();
3416 for i in null_pk_rows {
3417 winners.insert(i);
3418 }
3419 Some((0..n).filter(|i| winners.contains(i)).collect())
3420 }
3421
3422 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3424 self.require_delete()?;
3425 let epoch = self.pending_epoch();
3426 self.wal_append_data(Op::Delete {
3427 table_id: self.table_id,
3428 row_ids: vec![row_id],
3429 })?;
3430 if self.is_shared() {
3431 self.pending_dels.push(row_id);
3432 } else {
3433 self.apply_delete(row_id, epoch);
3434 }
3435 Ok(())
3436 }
3437
3438 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3439 let pre = self.get(row_id, self.snapshot());
3440 self.delete(row_id)?;
3441 Ok(pre.map(|row| {
3442 let mut columns: Vec<_> = row.columns.into_iter().collect();
3443 columns.sort_by_key(|(id, _)| *id);
3444 OwnedRow { columns }
3445 }))
3446 }
3447
3448 pub fn truncate(&mut self) -> Result<()> {
3450 self.require_delete()?;
3451 let epoch = self.pending_epoch();
3452 self.wal_append_data(Op::TruncateTable {
3453 table_id: self.table_id,
3454 })?;
3455 self.pending_rows.clear();
3456 self.pending_rows_auto_inc.clear();
3457 self.pending_dels.clear();
3458 self.pending_truncate = Some(epoch);
3459 Ok(())
3460 }
3461
3462 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3464 self.run_refs.clear();
3469 self.retiring.clear();
3470 self.memtable = Memtable::new();
3471 self.mutable_run = MutableRun::new();
3472 self.hot = HotIndex::new();
3473 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3474 self.bitmap = bitmap;
3475 self.ann = ann;
3476 self.fm = fm;
3477 self.sparse = sparse;
3478 self.minhash = minhash;
3479 self.learned_range = Arc::new(HashMap::new());
3480 self.pk_by_row.clear();
3481 self.pk_by_row_complete = false;
3482 self.live_count = 0;
3483 self.reservoir = crate::reservoir::Reservoir::default();
3484 self.reservoir_complete = true;
3485 self.had_deletes = true;
3486 self.agg_cache = Arc::new(HashMap::new());
3487 self.global_idx_epoch = 0;
3488 self.indexes_complete = true;
3489 self.pending_delete_rids.clear();
3490 self.pending_put_cols.clear();
3491 self.pending_rows.clear();
3492 self.pending_rows_auto_inc.clear();
3493 self.pending_dels.clear();
3494 self.clear_result_cache();
3495 self.invalidate_index_checkpoint();
3496 self.data_generation = self.data_generation.wrapping_add(1);
3497 }
3498
3499 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3502 self.remove_hot_for_row(row_id, epoch);
3503 self.tombstone_row(row_id, epoch, true);
3504 self.data_generation = self.data_generation.wrapping_add(1);
3505 }
3506
3507 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3511 let tombstone = Row {
3512 row_id,
3513 committed_epoch: epoch,
3514 columns: std::collections::HashMap::new(),
3515 deleted: true,
3516 };
3517 self.memtable.upsert(tombstone);
3518 self.pk_by_row.remove(&row_id);
3519 if adjust_live_count {
3520 self.live_count = self.live_count.saturating_sub(1);
3521 }
3522 self.pending_delete_rids.insert(row_id.0 as u32);
3524 self.had_deletes = true;
3527 self.agg_cache = Arc::new(HashMap::new());
3528 }
3529
3530 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3534 let Some(pk_col) = self.schema.primary_key() else {
3535 return;
3536 };
3537 if self.pk_by_row_complete {
3540 if let Some(key) = self.pk_by_row.remove(&row_id) {
3541 if self.hot.get(&key) == Some(row_id) {
3542 self.hot.remove(&key);
3543 }
3544 }
3545 return;
3546 }
3547 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3566 if self.indexes_complete {
3567 let pk_val = self
3568 .memtable
3569 .get_version(row_id, lookup_epoch)
3570 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3571 .or_else(|| {
3572 self.mutable_run
3573 .get_version(row_id, lookup_epoch)
3574 .filter(|(_, r)| !r.deleted)
3575 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3576 })
3577 .or_else(|| {
3578 self.run_refs.iter().find_map(|rr| {
3579 let mut reader = self.open_reader(rr.run_id).ok()?;
3580 let (_, deleted, val) = reader
3581 .get_version_column(row_id, lookup_epoch, pk_col.id)
3582 .ok()??;
3583 if deleted {
3584 return None;
3585 }
3586 val
3587 })
3588 });
3589 if let Some(pk_val) = pk_val {
3590 let key = self.index_lookup_key(pk_col.id, &pk_val);
3591 if self.hot.get(&key) == Some(row_id) {
3592 self.hot.remove(&key);
3593 }
3594 return;
3595 }
3596 }
3597 self.refresh_pk_by_row_from_hot();
3602 if let Some(key) = self.pk_by_row.remove(&row_id) {
3603 if self.hot.get(&key) == Some(row_id) {
3604 self.hot.remove(&key);
3605 }
3606 }
3607 }
3608
3609 fn partition_pk_winners(
3614 &self,
3615 rows: &[Row],
3616 ) -> (
3617 std::collections::HashSet<RowId>,
3618 std::collections::HashMap<Vec<u8>, RowId>,
3619 ) {
3620 let mut losers = std::collections::HashSet::new();
3621 let Some(pk_col) = self.schema.primary_key() else {
3622 return (losers, std::collections::HashMap::new());
3623 };
3624 let pk_id = pk_col.id;
3625 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3626 std::collections::HashMap::new();
3627 for r in rows {
3628 let Some(pk_val) = r.columns.get(&pk_id) else {
3629 continue;
3630 };
3631 let key = self.index_lookup_key(pk_id, pk_val);
3632 if let Some(&old_rid) = winners.get(&key) {
3633 losers.insert(old_rid);
3634 }
3635 winners.insert(key, r.row_id);
3636 }
3637 (losers, winners)
3638 }
3639
3640 fn index_row(&mut self, row: &Row) {
3641 if row.deleted {
3642 return;
3643 }
3644 let any_predicate = self
3652 .schema
3653 .indexes
3654 .iter()
3655 .any(|idx| idx.predicate.is_some());
3656 if any_predicate {
3657 let columns_map: HashMap<u16, &Value> =
3658 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3659 let name_to_id: HashMap<&str, u16> = self
3660 .schema
3661 .columns
3662 .iter()
3663 .map(|c| (c.name.as_str(), c.id))
3664 .collect();
3665 for idx in &self.schema.indexes {
3666 if let Some(pred) = &idx.predicate {
3667 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3668 continue; }
3670 }
3671 index_into_single(
3673 idx,
3674 &self.schema,
3675 row,
3676 &mut self.hot,
3677 &mut self.bitmap,
3678 &mut self.ann,
3679 &mut self.fm,
3680 &mut self.sparse,
3681 &mut self.minhash,
3682 );
3683 }
3684 return;
3685 }
3686 if self.column_keys.is_empty() {
3690 index_into(
3691 &self.schema,
3692 row,
3693 &mut self.hot,
3694 &mut self.bitmap,
3695 &mut self.ann,
3696 &mut self.fm,
3697 &mut self.sparse,
3698 &mut self.minhash,
3699 );
3700 return;
3701 }
3702 let effective_row = self.tokenized_for_indexes(row);
3703 index_into(
3704 &self.schema,
3705 &effective_row,
3706 &mut self.hot,
3707 &mut self.bitmap,
3708 &mut self.ann,
3709 &mut self.fm,
3710 &mut self.sparse,
3711 &mut self.minhash,
3712 );
3713 }
3714
3715 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3721 if self.column_keys.is_empty() {
3722 return row.clone();
3723 }
3724 {
3725 use crate::encryption::SCHEME_HMAC_EQ;
3726 let mut tok = row.clone();
3727 for (&cid, &(_, scheme)) in &self.column_keys {
3728 if scheme != SCHEME_HMAC_EQ {
3729 continue;
3730 }
3731 if let Some(v) = tok.columns.get(&cid).cloned() {
3732 if let Some(t) = self.tokenize_value(cid, &v) {
3733 tok.columns.insert(cid, t);
3734 }
3735 }
3736 }
3737 tok
3738 }
3739 }
3740
3741 pub fn commit(&mut self) -> Result<Epoch> {
3746 self.commit_inner(None)
3747 }
3748
3749 #[doc(hidden)]
3752 pub fn commit_controlled<F>(
3753 &mut self,
3754 control: &crate::ExecutionControl,
3755 mut before_commit: F,
3756 ) -> Result<Epoch>
3757 where
3758 F: FnMut() -> Result<()>,
3759 {
3760 self.commit_inner(Some((control, &mut before_commit)))
3761 }
3762
3763 fn commit_inner(
3764 &mut self,
3765 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3766 ) -> Result<Epoch> {
3767 self.ensure_writable()?;
3768 if !self.has_pending_mutations() {
3769 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3770 return Err(MongrelError::Full(
3771 "standalone transaction id namespace exhausted".into(),
3772 ));
3773 }
3774 return Ok(self.epoch.visible());
3775 }
3776 self.commit_new_epoch_inner(controlled)
3777 }
3778
3779 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3783 self.commit_new_epoch_inner(None)
3784 }
3785
3786 fn commit_new_epoch_inner(
3787 &mut self,
3788 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3789 ) -> Result<Epoch> {
3790 self.ensure_writable()?;
3791 if self.is_shared() {
3792 self.commit_shared(controlled)
3793 } else {
3794 self.commit_private(controlled)
3795 }
3796 }
3797
3798 fn commit_private(
3800 &mut self,
3801 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3802 ) -> Result<Epoch> {
3803 let commit_lock = Arc::clone(&self.commit_lock);
3807 let _g = commit_lock.lock();
3808 let txn_id = self.ensure_txn_id()?;
3811 if let Some((control, before_commit)) = controlled {
3812 control.checkpoint()?;
3813 before_commit()?;
3814 }
3815 let new_epoch = self.epoch.bump_assigned();
3816 let epoch_authority = Arc::clone(&self.epoch);
3817 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3818 let wal_result = match &mut self.wal {
3822 WalSink::Private(w) => w
3823 .append_txn(
3824 txn_id,
3825 Op::TxnCommit {
3826 epoch: new_epoch.0,
3827 added_runs: Vec::new(),
3828 },
3829 )
3830 .and_then(|_| w.sync()),
3831 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3832 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
3833 };
3834 if let Err(error) = wal_result {
3835 self.durable_commit_failed = true;
3836 return Err(MongrelError::CommitOutcomeUnknown {
3837 epoch: new_epoch.0,
3838 message: error.to_string(),
3839 });
3840 }
3841 if let Some(epoch) = self.pending_truncate.take() {
3844 self.apply_truncate(epoch);
3845 }
3846 self.invalidate_pending_cache();
3847 let publish_result = self.persist_manifest(new_epoch);
3848 self.epoch.publish_in_order(new_epoch);
3852 epoch_guard.disarm();
3853 if let Err(error) = publish_result {
3854 self.durable_commit_failed = true;
3855 return Err(MongrelError::DurableCommit {
3856 epoch: new_epoch.0,
3857 message: error.to_string(),
3858 });
3859 }
3860 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
3861 self.pending_private_mutations = false;
3862 self.data_generation = self.data_generation.wrapping_add(1);
3863 Ok(new_epoch)
3864 }
3865
3866 fn commit_shared(
3872 &mut self,
3873 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3874 ) -> Result<Epoch> {
3875 use std::sync::atomic::Ordering;
3876 let s = match &self.wal {
3877 WalSink::Shared(s) => s.clone(),
3878 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3879 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3880 };
3881 if s.poisoned.load(Ordering::Relaxed) {
3882 return Err(MongrelError::Other(
3883 "database poisoned by fsync error".into(),
3884 ));
3885 }
3886 let commit_lock = Arc::clone(&self.commit_lock);
3893 let _g = commit_lock.lock();
3894 if !self.pending_rows.is_empty() {
3895 match controlled.as_ref() {
3896 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
3897 None => self.prepare_durable_publish()?,
3898 }
3899 }
3900 let txn_id = self.ensure_txn_id()?;
3903 let mut wal = s.wal.lock();
3904 if let Some((control, before_commit)) = controlled {
3905 control.checkpoint()?;
3906 before_commit()?;
3907 }
3908 let new_epoch = self.epoch.bump_assigned();
3909 let epoch_authority = Arc::clone(&self.epoch);
3910 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3911 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
3912 Ok(commit_seq) => commit_seq,
3913 Err(error) => {
3914 s.poisoned.store(true, Ordering::Relaxed);
3915 s.lifecycle.poison();
3916 return Err(MongrelError::CommitOutcomeUnknown {
3917 epoch: new_epoch.0,
3918 message: error.to_string(),
3919 });
3920 }
3921 };
3922 drop(wal);
3923 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
3924 s.poisoned.store(true, Ordering::Relaxed);
3925 s.lifecycle.poison();
3926 return Err(MongrelError::CommitOutcomeUnknown {
3927 epoch: new_epoch.0,
3928 message: error.to_string(),
3929 });
3930 }
3931
3932 if self.pending_truncate.take().is_some() {
3935 self.apply_truncate(new_epoch);
3936 }
3937 let mut rows = std::mem::take(&mut self.pending_rows);
3938 if !rows.is_empty() {
3939 for r in &mut rows {
3940 r.committed_epoch = new_epoch;
3941 }
3942 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3943 let all_auto_generated =
3944 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3945 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
3946 } else {
3947 self.pending_rows_auto_inc.clear();
3948 }
3949 let dels = std::mem::take(&mut self.pending_dels);
3950 for rid in dels {
3951 self.apply_delete(rid, new_epoch);
3952 }
3953
3954 self.invalidate_pending_cache();
3955 let publish_result = self.persist_manifest(new_epoch);
3956 self.epoch.publish_in_order(new_epoch);
3957 epoch_guard.disarm();
3958 let _ = s.change_wake.send(());
3959 if let Err(error) = publish_result {
3960 self.durable_commit_failed = true;
3961 s.poisoned.store(true, Ordering::Relaxed);
3962 s.lifecycle.poison();
3963 return Err(MongrelError::DurableCommit {
3964 epoch: new_epoch.0,
3965 message: error.to_string(),
3966 });
3967 }
3968 self.current_txn_id = 0;
3970 self.data_generation = self.data_generation.wrapping_add(1);
3971 Ok(new_epoch)
3972 }
3973
3974 pub fn flush(&mut self) -> Result<Epoch> {
3982 self.flush_with_outcome().map(|(epoch, _)| epoch)
3983 }
3984
3985 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
3987 self.flush_with_outcome_inner(None)
3988 }
3989
3990 #[doc(hidden)]
3993 pub fn flush_with_outcome_controlled<F>(
3994 &mut self,
3995 control: &crate::ExecutionControl,
3996 mut before_commit: F,
3997 ) -> Result<(Epoch, bool)>
3998 where
3999 F: FnMut() -> Result<()>,
4000 {
4001 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4002 }
4003
4004 fn flush_with_outcome_inner(
4005 &mut self,
4006 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4007 ) -> Result<(Epoch, bool)> {
4008 match controlled.as_ref() {
4009 Some((control, _)) => {
4010 self.ensure_indexes_complete_controlled(control, || true)?;
4011 }
4012 None => self.ensure_indexes_complete()?,
4013 }
4014 let committed = self.has_pending_mutations();
4015 let epoch = self.commit_inner(controlled)?;
4016 let finish: Result<(Epoch, bool)> = (|| {
4017 let rows = self.memtable.drain_sorted();
4018 if !rows.is_empty() {
4019 self.mutable_run.insert_many(rows);
4020 }
4021 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4022 self.spill_mutable_run(epoch)?;
4023 self.mark_flushed(epoch)?;
4027 self.persist_manifest(epoch)?;
4028 self.build_learned_ranges()?;
4029 self.checkpoint_indexes(epoch);
4032 }
4033 Ok((epoch, committed))
4036 })();
4037 let outcome = match finish {
4038 Err(error) if committed => Err(MongrelError::DurableCommit {
4039 epoch: epoch.0,
4040 message: error.to_string(),
4041 }),
4042 result => result,
4043 };
4044 if outcome.is_ok() {
4045 let _ = self.publish_read_generation();
4051 }
4052 outcome
4053 }
4054
4055 fn has_pending_mutations(&self) -> bool {
4056 self.pending_private_mutations
4057 || !self.pending_rows.is_empty()
4058 || !self.pending_dels.is_empty()
4059 || self.pending_truncate.is_some()
4060 }
4061
4062 pub fn has_pending_writes(&self) -> bool {
4063 self.has_pending_mutations()
4064 }
4065
4066 pub fn force_flush(&mut self) -> Result<Epoch> {
4075 let saved = self.mutable_run_spill_bytes;
4076 self.mutable_run_spill_bytes = 1;
4077 let result = self.flush();
4078 self.mutable_run_spill_bytes = saved;
4079 result
4080 }
4081
4082 pub fn close(&mut self) -> Result<()> {
4089 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4090 self.force_flush()?;
4091 }
4092 Ok(())
4093 }
4094
4095 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4102 let op = Op::Flush {
4103 table_id: self.table_id,
4104 flushed_epoch: epoch.0,
4105 };
4106 match &mut self.wal {
4107 WalSink::Private(w) => {
4108 w.append_system(op)?;
4109 w.sync()?;
4110 }
4111 WalSink::Shared(s) => {
4112 s.wal.lock().append_system(op)?;
4117 }
4118 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4119 }
4120 self.flushed_epoch = epoch.0;
4121 if matches!(self.wal, WalSink::Private(_)) {
4122 self.rotate_wal(epoch)?;
4123 }
4124 Ok(())
4125 }
4126
4127 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4131 if self.mutable_run.is_empty() {
4132 return Ok(());
4133 }
4134 let run_id = self.alloc_run_id()?;
4135 let rows = self.mutable_run.drain_sorted();
4136 if rows.is_empty() {
4137 return Ok(());
4138 }
4139 let path = self.run_path(run_id);
4140 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4141 if let Some(kek) = &self.kek {
4142 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4143 }
4144 let header = match self.create_run_file(run_id)? {
4145 Some(file) => writer.write_file(file, &rows)?,
4146 None => writer.write(&path, &rows)?,
4147 };
4148 self.run_refs.push(RunRef {
4149 run_id: run_id as u128,
4150 level: 0,
4151 epoch_created: epoch.0,
4152 row_count: header.row_count,
4153 });
4154 Ok(())
4155 }
4156
4157 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4161 self.mutable_run_spill_bytes = bytes.max(1);
4162 }
4163
4164 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4168 self.compaction_zstd_level = level;
4169 }
4170
4171 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4175 self.result_cache.lock().set_max_bytes(max_bytes);
4176 }
4177
4178 pub(crate) fn clear_result_cache(&mut self) {
4182 self.result_cache.lock().clear();
4183 }
4184
4185 pub fn mutable_run_len(&self) -> usize {
4187 self.mutable_run.len()
4188 }
4189
4190 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4193 self.mutable_run.drain_sorted()
4194 }
4195
4196 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4198 let mut snapshot = self.mutable_run.clone();
4199 snapshot.drain_sorted()
4200 }
4201
4202 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4207 self.ensure_writable()?;
4208 let n = batch.len();
4209 if n == 0 {
4210 return Ok(self.current_epoch());
4211 }
4212 for row in &batch {
4213 self.schema.validate_values(row)?;
4214 }
4215 let epoch = self.commit_new_epoch()?;
4216 let live_before = self.live_count;
4217 self.spill_mutable_run(epoch)?;
4221 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4222 && self.indexes_complete
4223 && self.run_refs.is_empty()
4224 && self.memtable.is_empty()
4225 && self.mutable_run.is_empty();
4226 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4232 use rayon::prelude::*;
4233 self.schema
4234 .columns
4235 .par_iter()
4236 .map(|cdef| {
4237 (
4238 cdef.id,
4239 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4240 )
4241 })
4242 .collect::<Vec<_>>()
4243 };
4244 drop(batch);
4245 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4250 self.validate_columns_not_null(&user_columns, n)?;
4251 let winner_idx = self
4252 .bulk_pk_winner_indices(&user_columns, n)
4253 .filter(|idx| idx.len() != n);
4254 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4255 match winner_idx.as_deref() {
4256 Some(idx) => {
4257 let compacted = user_columns
4258 .iter()
4259 .map(|(id, c)| (*id, c.gather(idx)))
4260 .collect();
4261 (compacted, idx.len())
4262 }
4263 None => (std::mem::take(&mut user_columns), n),
4264 };
4265 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4266 let first = self.allocator.alloc_range(write_n as u64)?.0;
4267 for rid in first..first + write_n as u64 {
4268 self.reservoir.offer(rid);
4269 }
4270 let run_id = self.alloc_run_id()?;
4271 let path = self.run_path(run_id);
4272 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4273 .clean(true)
4274 .with_lz4()
4275 .with_native_endian();
4276 if let Some(kek) = &self.kek {
4277 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4278 }
4279 let header = match self.create_run_file(run_id)? {
4280 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4281 None => writer.write_native(&path, &write_columns, write_n, first)?,
4282 };
4283 self.run_refs.push(RunRef {
4284 run_id: run_id as u128,
4285 level: 0,
4286 epoch_created: epoch.0,
4287 row_count: header.row_count,
4288 });
4289 self.live_count = self.live_count.saturating_add(write_n as u64);
4290 if eager_index_build {
4291 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4292 self.index_columns_bulk(&write_columns, &row_ids);
4293 self.indexes_complete = true;
4294 self.build_learned_ranges()?;
4295 } else {
4296 self.indexes_complete = false;
4297 }
4298 self.mark_flushed(epoch)?;
4299 self.persist_manifest(epoch)?;
4300 if eager_index_build {
4301 self.checkpoint_indexes(epoch);
4302 }
4303 self.clear_result_cache();
4304 Ok(epoch)
4305 }
4306
4307 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4310 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4311 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4312 let segment_no = segment
4315 .file_stem()
4316 .and_then(|s| s.to_str())
4317 .and_then(|s| s.strip_prefix("seg-"))
4318 .and_then(|s| s.parse::<u64>().ok())
4319 .unwrap_or(0);
4320 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4321 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4322 wal.sync()?;
4323 self.wal = WalSink::Private(wal);
4324 Ok(())
4325 }
4326
4327 pub(crate) fn invalidate_pending_cache(&mut self) {
4332 self.result_cache
4333 .lock()
4334 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4335 self.pending_delete_rids.clear();
4336 self.pending_put_cols.clear();
4337 }
4338
4339 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4340 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4341 m.current_epoch = epoch.0;
4342 m.next_row_id = self.allocator.current().0;
4343 m.runs = self.run_refs.clone();
4344 m.live_count = self.live_count;
4345 m.global_idx_epoch = self.global_idx_epoch;
4346 m.flushed_epoch = self.flushed_epoch;
4347 m.retiring = self.retiring.clone();
4348 m.auto_inc_next = match self.auto_inc {
4352 Some(ai) if ai.seeded => ai.next,
4353 _ => 0,
4354 };
4355 m.ttl = self.ttl;
4356 let meta_dek = self.manifest_meta_dek();
4357 match self._root_guard.as_deref() {
4358 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4359 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4360 }
4361 Ok(())
4362 }
4363
4364 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4365 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4369 let live_count = u64::try_from(rows.len())
4370 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4371 let auto_inc = match self.auto_inc {
4372 Some(mut state) => {
4373 let maximum = self.scan_max_int64(state.column_id)?;
4374 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4375 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4376 })?;
4377 state.next = state.next.max(after_maximum).max(1);
4378 state.seeded = true;
4379 Some(state)
4380 }
4381 None => None,
4382 };
4383 Ok(RecoveryMetadataPlan {
4384 live_count,
4385 auto_inc,
4386 changed: live_count != self.live_count
4387 || auto_inc.is_some_and(|planned| {
4388 self.auto_inc.is_none_or(|current| {
4389 current.next != planned.next || current.seeded != planned.seeded
4390 })
4391 }),
4392 })
4393 }
4394
4395 pub(crate) fn apply_recovered_metadata(
4396 &mut self,
4397 plan: RecoveryMetadataPlan,
4398 epoch: Epoch,
4399 ) -> Result<()> {
4400 if !plan.changed {
4401 return Ok(());
4402 }
4403 self.live_count = plan.live_count;
4404 self.auto_inc = plan.auto_inc;
4405 self.persist_manifest(epoch)
4406 }
4407
4408 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4414 if !self.indexes_complete {
4417 return;
4418 }
4419 if crate::catalog::inject_hook("index.publish.before").is_err() {
4422 return;
4423 }
4424 if self.idx_root.is_none() {
4425 if let Some(root) = self._root_guard.as_ref() {
4426 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4427 return;
4428 };
4429 self.idx_root = Some(Arc::new(idx_root));
4430 }
4431 }
4432 let snap = global_idx::IndexSnapshot {
4433 hot: &self.hot,
4434 bitmap: &self.bitmap,
4435 ann: &self.ann,
4436 fm: &self.fm,
4437 sparse: &self.sparse,
4438 minhash: &self.minhash,
4439 learned_range: &self.learned_range,
4440 };
4441 let idx_dek = self.idx_dek();
4443 let written = match self.idx_root.as_deref() {
4444 Some(root) => global_idx::write_atomic_root(
4445 root,
4446 self.table_id,
4447 epoch.0,
4448 snap,
4449 idx_dek.as_deref(),
4450 ),
4451 None => global_idx::write_atomic(
4452 &self.dir,
4453 self.table_id,
4454 epoch.0,
4455 snap,
4456 idx_dek.as_deref(),
4457 ),
4458 };
4459 if written.is_ok() {
4460 self.global_idx_epoch = epoch.0;
4461 let _ = self.persist_manifest(epoch);
4462 let _ = crate::catalog::inject_hook("index.publish.after");
4464 }
4465 }
4466
4467 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4470 self.global_idx_epoch = 0;
4471 if let Some(root) = self.idx_root.as_deref() {
4472 let _ = root.remove_file(global_idx::IDX_FILENAME);
4473 } else {
4474 global_idx::remove(&self.dir);
4475 }
4476 let _ = self.persist_manifest(self.epoch.visible());
4477 }
4478
4479 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4484 self.indexes_complete = false;
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 }
4492
4493 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4494 self.indexes_complete = true;
4495 }
4496
4497 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4503 self.durable_commit_failed = true;
4504 if let WalSink::Shared(shared) = &self.wal {
4505 shared
4506 .poisoned
4507 .store(true, std::sync::atomic::Ordering::Relaxed);
4508 }
4509 }
4510
4511 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4515 self.durable_commit_failed = true;
4516 }
4517
4518 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4521 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4522 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4523 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4524 best = Some((epoch, row));
4525 }
4526 }
4527 for rr in &self.run_refs {
4528 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4529 continue;
4530 };
4531 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4532 continue;
4533 };
4534 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4535 best = Some((epoch, row));
4536 }
4537 }
4538 let now_nanos = unix_nanos_now();
4539 match best {
4540 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4541 Some((_, r)) => Some(r),
4542 None => None,
4543 }
4544 }
4545
4546 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4550 self.visible_rows_at_time(snapshot, unix_nanos_now())
4551 }
4552
4553 #[doc(hidden)]
4556 pub fn visible_rows_controlled(
4557 &self,
4558 snapshot: Snapshot,
4559 control: &crate::ExecutionControl,
4560 ) -> Result<Vec<Row>> {
4561 let mut out = Vec::new();
4562 self.for_each_visible_row_controlled(snapshot, control, |row| {
4563 out.push(row);
4564 Ok(())
4565 })?;
4566 Ok(out)
4567 }
4568
4569 #[doc(hidden)]
4572 pub fn for_each_visible_row_controlled<F>(
4573 &self,
4574 snapshot: Snapshot,
4575 control: &crate::ExecutionControl,
4576 visit: F,
4577 ) -> Result<()>
4578 where
4579 F: FnMut(Row) -> Result<()>,
4580 {
4581 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4582 control.checkpoint()?;
4583 let memtable = self.memtable.visible_versions(snapshot.epoch);
4584 if !memtable.is_empty() {
4585 sources.push(ControlledVisibleSource::memory(memtable));
4586 }
4587 control.checkpoint()?;
4588 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4589 if !mutable.is_empty() {
4590 sources.push(ControlledVisibleSource::memory(mutable));
4591 }
4592 for run in &self.run_refs {
4593 control.checkpoint()?;
4594 let reader = self.open_reader(run.run_id)?;
4595 sources.push(ControlledVisibleSource::run(
4596 reader.into_visible_version_cursor(snapshot.epoch)?,
4597 ));
4598 }
4599 let now_nanos = unix_nanos_now();
4600 merge_controlled_visible_sources(
4601 &mut sources,
4602 control,
4603 |row| self.row_expired_at(row, now_nanos),
4604 visit,
4605 )
4606 }
4607
4608 #[doc(hidden)]
4609 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4610 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4611 let mut fold = |row: Row| {
4612 best.entry(row.row_id.0)
4613 .and_modify(|e| {
4614 if row.committed_epoch > e.0 {
4615 *e = (row.committed_epoch, row.clone());
4616 }
4617 })
4618 .or_insert_with(|| (row.committed_epoch, row));
4619 };
4620 for row in self.memtable.visible_versions(snapshot.epoch) {
4621 fold(row);
4622 }
4623 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4624 fold(row);
4625 }
4626 for rr in &self.run_refs {
4627 let mut reader = self.open_reader(rr.run_id)?;
4628 for row in reader.visible_versions(snapshot.epoch)? {
4629 fold(row);
4630 }
4631 }
4632 let mut out: Vec<Row> = best
4633 .into_values()
4634 .filter_map(|(_, r)| {
4635 if r.deleted || self.row_expired_at(&r, now_nanos) {
4636 None
4637 } else {
4638 Some(r)
4639 }
4640 })
4641 .collect();
4642 out.sort_by_key(|r| r.row_id);
4643 Ok(out)
4644 }
4645
4646 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4653 if self.ttl.is_none()
4654 && self.memtable.is_empty()
4655 && self.mutable_run.is_empty()
4656 && self.run_refs.len() == 1
4657 {
4658 let rr = self.run_refs[0].clone();
4659 let mut reader = self.open_reader(rr.run_id)?;
4660 let idxs = reader.visible_indices(snapshot.epoch)?;
4661 let mut cols = Vec::with_capacity(self.schema.columns.len());
4662 for cdef in &self.schema.columns {
4663 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4664 }
4665 return Ok(cols);
4666 }
4667 let rows = self.visible_rows(snapshot)?;
4669 let mut cols: Vec<(u16, Vec<Value>)> = self
4670 .schema
4671 .columns
4672 .iter()
4673 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4674 .collect();
4675 for r in &rows {
4676 for (cid, vec) in cols.iter_mut() {
4677 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4678 }
4679 }
4680 Ok(cols)
4681 }
4682
4683 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4685 let row_id = self.hot.get(key)?;
4686 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4687 Some(row_id)
4688 } else {
4689 None
4690 }
4691 }
4692
4693 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4698 self.query_at_with_allowed(q, self.snapshot(), None)
4699 }
4700
4701 pub fn query_controlled(
4704 &mut self,
4705 q: &crate::query::Query,
4706 control: &crate::ExecutionControl,
4707 ) -> Result<Vec<Row>> {
4708 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4709 }
4710
4711 pub fn query_at_with_allowed(
4714 &mut self,
4715 q: &crate::query::Query,
4716 snapshot: Snapshot,
4717 allowed: Option<&std::collections::HashSet<RowId>>,
4718 ) -> Result<Vec<Row>> {
4719 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4720 }
4721
4722 #[doc(hidden)]
4723 pub fn query_at_with_allowed_controlled(
4724 &mut self,
4725 q: &crate::query::Query,
4726 snapshot: Snapshot,
4727 allowed: Option<&std::collections::HashSet<RowId>>,
4728 control: &crate::ExecutionControl,
4729 ) -> Result<Vec<Row>> {
4730 self.require_select()?;
4731 self.ensure_indexes_complete_controlled(control, || true)?;
4732 self.validate_native_query(q)?;
4733 self.query_conditions_at(
4734 &q.conditions,
4735 snapshot,
4736 allowed,
4737 q.limit,
4738 q.offset,
4739 None,
4740 unix_nanos_now(),
4741 Some(control),
4742 )
4743 }
4744
4745 #[doc(hidden)]
4746 pub fn query_at_with_allowed_after(
4747 &mut self,
4748 q: &crate::query::Query,
4749 snapshot: Snapshot,
4750 allowed: Option<&std::collections::HashSet<RowId>>,
4751 after_row_id: Option<RowId>,
4752 ) -> Result<Vec<Row>> {
4753 self.query_at_with_allowed_after_at_time(
4754 q,
4755 snapshot,
4756 allowed,
4757 after_row_id,
4758 unix_nanos_now(),
4759 )
4760 }
4761
4762 #[doc(hidden)]
4763 pub fn query_at_with_allowed_after_at_time(
4764 &mut self,
4765 q: &crate::query::Query,
4766 snapshot: Snapshot,
4767 allowed: Option<&std::collections::HashSet<RowId>>,
4768 after_row_id: Option<RowId>,
4769 query_time_nanos: i64,
4770 ) -> Result<Vec<Row>> {
4771 self.require_select()?;
4772 self.ensure_indexes_complete()?;
4773 self.validate_native_query(q)?;
4774 self.query_conditions_at(
4775 &q.conditions,
4776 snapshot,
4777 allowed,
4778 q.limit,
4779 q.offset,
4780 after_row_id,
4781 query_time_nanos,
4782 None,
4783 )
4784 }
4785
4786 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4787 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4788 return Err(MongrelError::InvalidArgument(format!(
4789 "query exceeds {} conditions",
4790 crate::query::MAX_HARD_CONDITIONS
4791 )));
4792 }
4793 if let Some(limit) = q.limit {
4794 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4795 return Err(MongrelError::InvalidArgument(format!(
4796 "query limit must be between 1 and {}",
4797 crate::query::MAX_FINAL_LIMIT
4798 )));
4799 }
4800 }
4801 if q.offset > crate::query::MAX_QUERY_OFFSET {
4802 return Err(MongrelError::InvalidArgument(format!(
4803 "query offset exceeds {}",
4804 crate::query::MAX_QUERY_OFFSET
4805 )));
4806 }
4807 Ok(())
4808 }
4809
4810 #[doc(hidden)]
4813 pub fn query_all_at(
4814 &mut self,
4815 conditions: &[crate::query::Condition],
4816 snapshot: Snapshot,
4817 ) -> Result<Vec<Row>> {
4818 self.require_select()?;
4819 self.ensure_indexes_complete()?;
4820 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4821 return Err(MongrelError::InvalidArgument(format!(
4822 "query exceeds {} conditions",
4823 crate::query::MAX_HARD_CONDITIONS
4824 )));
4825 }
4826 self.query_conditions_at(
4827 conditions,
4828 snapshot,
4829 None,
4830 None,
4831 0,
4832 None,
4833 unix_nanos_now(),
4834 None,
4835 )
4836 }
4837
4838 #[allow(clippy::too_many_arguments)]
4839 fn query_conditions_at(
4840 &self,
4841 conditions: &[crate::query::Condition],
4842 snapshot: Snapshot,
4843 allowed: Option<&std::collections::HashSet<RowId>>,
4844 limit: Option<usize>,
4845 offset: usize,
4846 after_row_id: Option<RowId>,
4847 query_time_nanos: i64,
4848 control: Option<&crate::ExecutionControl>,
4849 ) -> Result<Vec<Row>> {
4850 control
4851 .map(crate::ExecutionControl::checkpoint)
4852 .transpose()?;
4853 crate::trace::QueryTrace::record(|t| {
4854 t.run_count = self.run_refs.len();
4855 t.memtable_rows = self.memtable.len();
4856 t.mutable_run_rows = self.mutable_run.len();
4857 });
4858 if conditions.is_empty() {
4862 crate::trace::QueryTrace::record(|t| {
4863 t.scan_mode = crate::trace::ScanMode::Materialized;
4864 t.row_materialized = true;
4865 });
4866 let mut rows = match control {
4867 Some(control) => self.visible_rows_controlled(snapshot, control)?,
4868 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
4869 };
4870 if let Some(allowed) = allowed {
4871 let mut filtered = Vec::with_capacity(rows.len());
4872 for (index, row) in rows.into_iter().enumerate() {
4873 if index & 255 == 0 {
4874 control
4875 .map(crate::ExecutionControl::checkpoint)
4876 .transpose()?;
4877 }
4878 if allowed.contains(&row.row_id) {
4879 filtered.push(row);
4880 }
4881 }
4882 rows = filtered;
4883 }
4884 if let Some(after_row_id) = after_row_id {
4885 rows.retain(|row| row.row_id > after_row_id);
4886 }
4887 rows.drain(..offset.min(rows.len()));
4888 if let Some(limit) = limit {
4889 rows.truncate(limit);
4890 }
4891 return Ok(rows);
4892 }
4893 crate::trace::QueryTrace::record(|t| {
4894 t.conditions_pushed = conditions.len();
4895 t.scan_mode = crate::trace::ScanMode::Materialized;
4896 t.row_materialized = true;
4897 });
4898 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
4905 ordered.sort_by_key(|c| condition_cost_rank(c));
4906 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
4907 for c in &ordered {
4908 control
4909 .map(crate::ExecutionControl::checkpoint)
4910 .transpose()?;
4911 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
4912 let empty = s.is_empty();
4913 sets.push(s);
4914 if empty {
4915 break;
4916 }
4917 }
4918 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4919 if let Some(allowed) = allowed {
4920 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
4921 }
4922 if let Some(after_row_id) = after_row_id {
4923 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
4924 rids.drain(..first);
4925 }
4926 rids.drain(..offset.min(rids.len()));
4927 if let Some(limit) = limit {
4928 rids.truncate(limit);
4929 }
4930 control
4931 .map(crate::ExecutionControl::checkpoint)
4932 .transpose()?;
4933 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
4934 }
4935
4936 pub fn retrieve(
4938 &mut self,
4939 retriever: &crate::query::Retriever,
4940 ) -> Result<Vec<crate::query::RetrieverHit>> {
4941 self.retrieve_with_allowed(retriever, None)
4942 }
4943
4944 pub fn retrieve_at(
4945 &mut self,
4946 retriever: &crate::query::Retriever,
4947 snapshot: Snapshot,
4948 allowed: Option<&std::collections::HashSet<RowId>>,
4949 ) -> Result<Vec<crate::query::RetrieverHit>> {
4950 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
4951 }
4952
4953 pub fn retrieve_with_allowed(
4956 &mut self,
4957 retriever: &crate::query::Retriever,
4958 allowed: Option<&std::collections::HashSet<RowId>>,
4959 ) -> Result<Vec<crate::query::RetrieverHit>> {
4960 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
4961 }
4962
4963 pub fn retrieve_at_with_allowed(
4964 &mut self,
4965 retriever: &crate::query::Retriever,
4966 snapshot: Snapshot,
4967 allowed: Option<&std::collections::HashSet<RowId>>,
4968 ) -> Result<Vec<crate::query::RetrieverHit>> {
4969 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
4970 }
4971
4972 pub fn retrieve_at_with_allowed_and_context(
4973 &mut self,
4974 retriever: &crate::query::Retriever,
4975 snapshot: Snapshot,
4976 allowed: Option<&std::collections::HashSet<RowId>>,
4977 context: Option<&crate::query::AiExecutionContext>,
4978 ) -> Result<Vec<crate::query::RetrieverHit>> {
4979 self.require_select()?;
4980 self.ensure_indexes_complete()?;
4981 self.validate_retriever(retriever)?;
4982 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
4983 }
4984
4985 pub fn retrieve_at_with_candidate_authorization_and_context(
4986 &mut self,
4987 retriever: &crate::query::Retriever,
4988 snapshot: Snapshot,
4989 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4990 context: Option<&crate::query::AiExecutionContext>,
4991 ) -> Result<Vec<crate::query::RetrieverHit>> {
4992 self.require_select()?;
4993 self.ensure_indexes_complete()?;
4994 self.retrieve_at_with_candidate_authorization_on_generation(
4995 retriever,
4996 snapshot,
4997 authorization,
4998 context,
4999 )
5000 }
5001
5002 #[doc(hidden)]
5003 pub fn retrieve_at_with_candidate_authorization_on_generation(
5004 &self,
5005 retriever: &crate::query::Retriever,
5006 snapshot: Snapshot,
5007 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5008 context: Option<&crate::query::AiExecutionContext>,
5009 ) -> Result<Vec<crate::query::RetrieverHit>> {
5010 self.require_select()?;
5011 self.validate_retriever(retriever)?;
5012 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5013 }
5014
5015 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5016 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5017 let (column_id, k) = match retriever {
5018 Retriever::Ann {
5019 column_id,
5020 query,
5021 k,
5022 } => {
5023 let index = self.ann.get(column_id).ok_or_else(|| {
5024 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5025 })?;
5026 if query.len() != index.dim() {
5027 return Err(MongrelError::InvalidArgument(format!(
5028 "ANN query dimension must be {}, got {}",
5029 index.dim(),
5030 query.len()
5031 )));
5032 }
5033 if query.iter().any(|value| !value.is_finite()) {
5034 return Err(MongrelError::InvalidArgument(
5035 "ANN query values must be finite".into(),
5036 ));
5037 }
5038 (*column_id, *k)
5039 }
5040 Retriever::Sparse {
5041 column_id,
5042 query,
5043 k,
5044 } => {
5045 if !self.sparse.contains_key(column_id) {
5046 return Err(MongrelError::InvalidArgument(format!(
5047 "column {column_id} has no Sparse index"
5048 )));
5049 }
5050 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5051 return Err(MongrelError::InvalidArgument(
5052 "Sparse query must be non-empty with finite weights".into(),
5053 ));
5054 }
5055 if query.len() > MAX_SPARSE_TERMS {
5056 return Err(MongrelError::InvalidArgument(format!(
5057 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5058 )));
5059 }
5060 (*column_id, *k)
5061 }
5062 Retriever::MinHash {
5063 column_id,
5064 members,
5065 k,
5066 } => {
5067 if !self.minhash.contains_key(column_id) {
5068 return Err(MongrelError::InvalidArgument(format!(
5069 "column {column_id} has no MinHash index"
5070 )));
5071 }
5072 if members.is_empty() {
5073 return Err(MongrelError::InvalidArgument(
5074 "MinHash members must not be empty".into(),
5075 ));
5076 }
5077 if members.len() > MAX_SET_MEMBERS {
5078 return Err(MongrelError::InvalidArgument(format!(
5079 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5080 )));
5081 }
5082 let mut total_bytes = 0usize;
5083 for member in members {
5084 let bytes = member.encoded_len();
5085 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5086 return Err(MongrelError::InvalidArgument(format!(
5087 "MinHash member exceeds {} bytes",
5088 crate::query::MAX_SET_MEMBER_BYTES
5089 )));
5090 }
5091 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5092 MongrelError::InvalidArgument("MinHash input size overflow".into())
5093 })?;
5094 }
5095 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5096 return Err(MongrelError::InvalidArgument(format!(
5097 "MinHash input exceeds {} bytes",
5098 crate::query::MAX_SET_INPUT_BYTES
5099 )));
5100 }
5101 (*column_id, *k)
5102 }
5103 };
5104 if k == 0 {
5105 return Err(MongrelError::InvalidArgument(
5106 "retriever k must be > 0".into(),
5107 ));
5108 }
5109 if k > MAX_RETRIEVER_K {
5110 return Err(MongrelError::InvalidArgument(format!(
5111 "retriever k exceeds {MAX_RETRIEVER_K}"
5112 )));
5113 }
5114 debug_assert!(self
5115 .schema
5116 .columns
5117 .iter()
5118 .any(|column| column.id == column_id));
5119 Ok(())
5120 }
5121
5122 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5123 use crate::query::Condition;
5124 match condition {
5125 Condition::Ann {
5126 column_id,
5127 query,
5128 k,
5129 } => self.validate_retriever(&crate::query::Retriever::Ann {
5130 column_id: *column_id,
5131 query: query.clone(),
5132 k: *k,
5133 }),
5134 Condition::SparseMatch {
5135 column_id,
5136 query,
5137 k,
5138 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5139 column_id: *column_id,
5140 query: query.clone(),
5141 k: *k,
5142 }),
5143 Condition::MinHashSimilar {
5144 column_id,
5145 query,
5146 k,
5147 } => {
5148 if !self.minhash.contains_key(column_id) {
5149 return Err(MongrelError::InvalidArgument(format!(
5150 "column {column_id} has no MinHash index"
5151 )));
5152 }
5153 if query.is_empty() || *k == 0 {
5154 return Err(MongrelError::InvalidArgument(
5155 "MinHash query must be non-empty and k must be > 0".into(),
5156 ));
5157 }
5158 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5159 {
5160 return Err(MongrelError::InvalidArgument(format!(
5161 "MinHash query must have <= {} members and k <= {}",
5162 crate::query::MAX_SET_MEMBERS,
5163 crate::query::MAX_RETRIEVER_K
5164 )));
5165 }
5166 Ok(())
5167 }
5168 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5169 Err(MongrelError::InvalidArgument(format!(
5170 "bitmap IN exceeds {} values",
5171 crate::query::MAX_SET_MEMBERS
5172 )))
5173 }
5174 Condition::FmContainsAll { patterns, .. }
5175 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5176 {
5177 Err(MongrelError::InvalidArgument(format!(
5178 "FM query exceeds {} patterns",
5179 crate::query::MAX_HARD_CONDITIONS
5180 )))
5181 }
5182 _ => Ok(()),
5183 }
5184 }
5185
5186 fn retrieve_filtered(
5187 &self,
5188 retriever: &crate::query::Retriever,
5189 snapshot: Snapshot,
5190 hard_filter: Option<&RowIdSet>,
5191 allowed: Option<&std::collections::HashSet<RowId>>,
5192 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5193 context: Option<&crate::query::AiExecutionContext>,
5194 ) -> Result<Vec<crate::query::RetrieverHit>> {
5195 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5196 let started = std::time::Instant::now();
5197 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5198 Retriever::Ann {
5199 column_id,
5200 query,
5201 k,
5202 } => {
5203 let Some(index) = self.ann.get(column_id) else {
5204 return Ok(Vec::new());
5205 };
5206 let cap = ann_candidate_cap(index.len(), context);
5207 if cap == 0 {
5208 return Ok(Vec::new());
5209 }
5210 let mut breadth = (*k).max(1).min(cap);
5211 let mut eligibility = std::collections::HashMap::new();
5212 let mut filtered = loop {
5213 let mut seen = std::collections::HashSet::new();
5214 if let Some(context) = context {
5215 context.checkpoint()?;
5216 }
5217 let raw = index.search_with_context(query, breadth, context)?;
5218 let unchecked: Vec<_> = raw
5219 .iter()
5220 .map(|(row_id, _)| *row_id)
5221 .filter(|row_id| !eligibility.contains_key(row_id))
5222 .filter(|row_id| {
5223 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5224 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5225 })
5226 .collect();
5227 let eligible = self.eligible_and_authorized_candidate_ids(
5228 &unchecked,
5229 *column_id,
5230 snapshot,
5231 candidate_authorization,
5232 context,
5233 )?;
5234 for row_id in unchecked {
5235 eligibility.insert(row_id, eligible.contains(&row_id));
5236 }
5237 let filtered: Vec<_> = raw
5238 .into_iter()
5239 .filter(|(row_id, _)| {
5240 seen.insert(*row_id)
5241 && eligibility.get(row_id).copied().unwrap_or(false)
5242 })
5243 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
5244 .collect();
5245 if filtered.len() >= *k || breadth >= cap {
5246 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5247 crate::trace::QueryTrace::record(|trace| {
5248 trace.ann_candidate_cap_hit = true;
5249 });
5250 }
5251 break filtered;
5252 }
5253 breadth = breadth.saturating_mul(2).min(cap);
5254 };
5255 filtered.truncate(*k);
5256 filtered
5257 }
5258 Retriever::Sparse {
5259 column_id,
5260 query,
5261 k,
5262 } => self
5263 .sparse
5264 .get(column_id)
5265 .map(|index| -> Result<Vec<_>> {
5266 let mut breadth = (*k).max(1);
5267 let mut eligibility = std::collections::HashMap::new();
5268 loop {
5269 if let Some(context) = context {
5270 context.checkpoint()?;
5271 }
5272 let raw = index.search_with_context(query, breadth, context)?;
5273 let unchecked: Vec<_> = raw
5274 .iter()
5275 .map(|(row_id, _)| *row_id)
5276 .filter(|row_id| !eligibility.contains_key(row_id))
5277 .filter(|row_id| {
5278 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5279 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5280 })
5281 .collect();
5282 let eligible = self.eligible_and_authorized_candidate_ids(
5283 &unchecked,
5284 *column_id,
5285 snapshot,
5286 candidate_authorization,
5287 context,
5288 )?;
5289 for row_id in unchecked {
5290 eligibility.insert(row_id, eligible.contains(&row_id));
5291 }
5292 let filtered: Vec<_> = raw
5293 .iter()
5294 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5295 .take(*k)
5296 .map(|(row_id, score)| {
5297 (*row_id, RetrieverScore::SparseDotProduct(*score))
5298 })
5299 .collect();
5300 if filtered.len() >= *k || raw.len() < breadth {
5301 break Ok(filtered);
5302 }
5303 let next = breadth.saturating_mul(2);
5304 if next == breadth {
5305 break Ok(filtered);
5306 }
5307 breadth = next;
5308 }
5309 })
5310 .transpose()?
5311 .unwrap_or_default(),
5312 Retriever::MinHash {
5313 column_id,
5314 members,
5315 k,
5316 } => self
5317 .minhash
5318 .get(column_id)
5319 .map(|index| -> Result<Vec<_>> {
5320 let mut hashes = Vec::with_capacity(members.len());
5321 for member in members {
5322 if let Some(context) = context {
5323 context.consume(crate::query::work_units(
5324 member.encoded_len(),
5325 crate::query::PARSE_WORK_QUANTUM,
5326 ))?;
5327 }
5328 hashes.push(member.hash_v1());
5329 }
5330 let mut breadth = (*k).max(1);
5331 let mut eligibility = std::collections::HashMap::new();
5332 loop {
5333 if let Some(context) = context {
5334 context.checkpoint()?;
5335 }
5336 let raw = index.search_with_context(&hashes, breadth, context)?;
5337 let unchecked: Vec<_> = raw
5338 .iter()
5339 .map(|(row_id, _)| *row_id)
5340 .filter(|row_id| !eligibility.contains_key(row_id))
5341 .filter(|row_id| {
5342 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5343 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5344 })
5345 .collect();
5346 let eligible = self.eligible_and_authorized_candidate_ids(
5347 &unchecked,
5348 *column_id,
5349 snapshot,
5350 candidate_authorization,
5351 context,
5352 )?;
5353 for row_id in unchecked {
5354 eligibility.insert(row_id, eligible.contains(&row_id));
5355 }
5356 let filtered: Vec<_> = raw
5357 .iter()
5358 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5359 .take(*k)
5360 .map(|(row_id, score)| {
5361 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5362 })
5363 .collect();
5364 if filtered.len() >= *k || raw.len() < breadth {
5365 break Ok(filtered);
5366 }
5367 let next = breadth.saturating_mul(2);
5368 if next == breadth {
5369 break Ok(filtered);
5370 }
5371 breadth = next;
5372 }
5373 })
5374 .transpose()?
5375 .unwrap_or_default(),
5376 };
5377 let elapsed = started.elapsed().as_nanos() as u64;
5378 crate::trace::QueryTrace::record(|trace| {
5379 match retriever {
5380 Retriever::Ann { .. } => {
5381 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5382 }
5383 Retriever::Sparse { .. } => {
5384 trace.sparse_candidate_nanos =
5385 trace.sparse_candidate_nanos.saturating_add(elapsed)
5386 }
5387 Retriever::MinHash { .. } => {
5388 trace.minhash_candidate_nanos =
5389 trace.minhash_candidate_nanos.saturating_add(elapsed)
5390 }
5391 }
5392 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5393 });
5394 Ok(scored
5395 .into_iter()
5396 .enumerate()
5397 .map(|(rank, (row_id, score))| RetrieverHit {
5398 row_id,
5399 rank: rank + 1,
5400 score,
5401 })
5402 .collect())
5403 }
5404
5405 fn eligible_candidate_ids(
5406 &self,
5407 candidates: &[RowId],
5408 _column_id: u16,
5409 snapshot: Snapshot,
5410 context: Option<&crate::query::AiExecutionContext>,
5411 ) -> Result<std::collections::HashSet<RowId>> {
5412 if !self.had_deletes
5413 && self.ttl.is_none()
5414 && self.pending_put_cols.is_empty()
5415 && snapshot.epoch == self.snapshot().epoch
5416 {
5417 return Ok(candidates.iter().copied().collect());
5418 }
5419 let mut readers: Vec<_> = self
5420 .run_refs
5421 .iter()
5422 .map(|run| self.open_reader(run.run_id))
5423 .collect::<Result<_>>()?;
5424 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5425 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5426 for &row_id in candidates {
5427 if let Some(context) = context {
5428 context.consume(1)?;
5429 }
5430 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5431 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5432 let overlay = match (mem, mutable) {
5433 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5434 (Some(value), None) | (None, Some(value)) => Some(value),
5435 (None, None) => None,
5436 };
5437 if let Some((_, row)) = overlay {
5438 if !row.deleted && !self.row_expired_at(&row, now) {
5439 eligible.insert(row_id);
5440 }
5441 continue;
5442 }
5443 let mut best: Option<(Epoch, bool, usize)> = None;
5444 for (index, reader) in readers.iter_mut().enumerate() {
5445 if let Some((epoch, deleted)) =
5446 reader.get_version_visibility(row_id, snapshot.epoch)?
5447 {
5448 if best
5449 .as_ref()
5450 .map(|(best_epoch, ..)| epoch > *best_epoch)
5451 .unwrap_or(true)
5452 {
5453 best = Some((epoch, deleted, index));
5454 }
5455 }
5456 }
5457 let Some((_, false, reader_index)) = best else {
5458 continue;
5459 };
5460 if let Some(ttl) = self.ttl {
5461 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5462 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5463 {
5464 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5465 continue;
5466 }
5467 }
5468 }
5469 eligible.insert(row_id);
5470 }
5471 Ok(eligible)
5472 }
5473
5474 fn eligible_and_authorized_candidate_ids(
5475 &self,
5476 candidates: &[RowId],
5477 column_id: u16,
5478 snapshot: Snapshot,
5479 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5480 context: Option<&crate::query::AiExecutionContext>,
5481 ) -> Result<std::collections::HashSet<RowId>> {
5482 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5483 let Some(authorization) = authorization else {
5484 return Ok(eligible);
5485 };
5486 let candidates: Vec<_> = eligible.into_iter().collect();
5487 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5488 }
5489
5490 fn policy_allowed_candidate_ids(
5491 &self,
5492 candidates: &[RowId],
5493 snapshot: Snapshot,
5494 authorization: &crate::security::CandidateAuthorization<'_>,
5495 context: Option<&crate::query::AiExecutionContext>,
5496 ) -> Result<std::collections::HashSet<RowId>> {
5497 let started = std::time::Instant::now();
5498 if candidates.is_empty()
5499 || authorization.principal.is_admin
5500 || !authorization.security.rls_enabled(authorization.table)
5501 {
5502 return Ok(candidates.iter().copied().collect());
5503 }
5504 if let Some(context) = context {
5505 context.checkpoint()?;
5506 }
5507 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5508 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5509 .iter()
5510 .map(|row_id| {
5511 (
5512 *row_id,
5513 Row {
5514 row_id: *row_id,
5515 committed_epoch: snapshot.epoch,
5516 columns: std::collections::HashMap::new(),
5517 deleted: false,
5518 },
5519 )
5520 })
5521 .collect();
5522 let columns = authorization
5523 .security
5524 .select_policy_columns(authorization.table, authorization.principal);
5525 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5526 let mut decoded = 0usize;
5527 for column_id in &columns {
5528 if let Some(context) = context {
5529 context.checkpoint()?;
5530 }
5531 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5532 &row_ids, *column_id, snapshot, query_now, context,
5533 )? {
5534 if let Some(row) = rows.get_mut(&row_id) {
5535 row.columns.insert(*column_id, value);
5536 decoded = decoded.saturating_add(1);
5537 }
5538 }
5539 }
5540 if let Some(context) = context {
5541 context.consume(candidates.len().saturating_add(decoded))?;
5542 }
5543 let allowed = rows
5544 .into_values()
5545 .filter_map(|row| {
5546 authorization
5547 .security
5548 .row_allowed(
5549 authorization.table,
5550 crate::security::PolicyCommand::Select,
5551 &row,
5552 authorization.principal,
5553 false,
5554 )
5555 .then_some(row.row_id)
5556 })
5557 .collect();
5558 crate::trace::QueryTrace::record(|trace| {
5559 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5560 trace.rls_policy_columns_decoded =
5561 trace.rls_policy_columns_decoded.saturating_add(decoded);
5562 trace.authorization_nanos = trace
5563 .authorization_nanos
5564 .saturating_add(started.elapsed().as_nanos() as u64);
5565 });
5566 Ok(allowed)
5567 }
5568
5569 pub fn search(
5571 &mut self,
5572 request: &crate::query::SearchRequest,
5573 ) -> Result<Vec<crate::query::SearchHit>> {
5574 self.search_with_allowed(request, None)
5575 }
5576
5577 pub fn search_at(
5578 &mut self,
5579 request: &crate::query::SearchRequest,
5580 snapshot: Snapshot,
5581 authorized: Option<&std::collections::HashSet<RowId>>,
5582 ) -> Result<Vec<crate::query::SearchHit>> {
5583 self.search_at_with_allowed(request, snapshot, authorized)
5584 }
5585
5586 pub fn search_with_allowed(
5587 &mut self,
5588 request: &crate::query::SearchRequest,
5589 authorized: Option<&std::collections::HashSet<RowId>>,
5590 ) -> Result<Vec<crate::query::SearchHit>> {
5591 self.search_at_with_allowed(request, self.snapshot(), authorized)
5592 }
5593
5594 pub fn search_at_with_allowed(
5595 &mut self,
5596 request: &crate::query::SearchRequest,
5597 snapshot: Snapshot,
5598 authorized: Option<&std::collections::HashSet<RowId>>,
5599 ) -> Result<Vec<crate::query::SearchHit>> {
5600 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5601 }
5602
5603 pub fn search_at_with_allowed_and_context(
5604 &mut self,
5605 request: &crate::query::SearchRequest,
5606 snapshot: Snapshot,
5607 authorized: Option<&std::collections::HashSet<RowId>>,
5608 context: Option<&crate::query::AiExecutionContext>,
5609 ) -> Result<Vec<crate::query::SearchHit>> {
5610 self.ensure_indexes_complete()?;
5611 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5612 }
5613
5614 pub fn search_at_with_candidate_authorization_and_context(
5615 &mut self,
5616 request: &crate::query::SearchRequest,
5617 snapshot: Snapshot,
5618 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5619 context: Option<&crate::query::AiExecutionContext>,
5620 ) -> Result<Vec<crate::query::SearchHit>> {
5621 self.ensure_indexes_complete()?;
5622 self.search_at_with_filters_and_context(
5623 request,
5624 snapshot,
5625 None,
5626 authorization,
5627 context,
5628 None,
5629 )
5630 }
5631
5632 #[doc(hidden)]
5633 pub fn search_at_with_candidate_authorization_on_generation(
5634 &self,
5635 request: &crate::query::SearchRequest,
5636 snapshot: Snapshot,
5637 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5638 context: Option<&crate::query::AiExecutionContext>,
5639 ) -> Result<Vec<crate::query::SearchHit>> {
5640 self.search_at_with_filters_and_context(
5641 request,
5642 snapshot,
5643 None,
5644 authorization,
5645 context,
5646 None,
5647 )
5648 }
5649
5650 #[doc(hidden)]
5651 pub fn search_at_with_candidate_authorization_on_generation_after(
5652 &self,
5653 request: &crate::query::SearchRequest,
5654 snapshot: Snapshot,
5655 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5656 context: Option<&crate::query::AiExecutionContext>,
5657 after: Option<crate::query::SearchAfter>,
5658 ) -> Result<Vec<crate::query::SearchHit>> {
5659 self.search_at_with_filters_and_context(
5660 request,
5661 snapshot,
5662 None,
5663 authorization,
5664 context,
5665 after,
5666 )
5667 }
5668
5669 fn search_at_with_filters_and_context(
5670 &self,
5671 request: &crate::query::SearchRequest,
5672 snapshot: Snapshot,
5673 authorized: Option<&std::collections::HashSet<RowId>>,
5674 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5675 context: Option<&crate::query::AiExecutionContext>,
5676 after: Option<crate::query::SearchAfter>,
5677 ) -> Result<Vec<crate::query::SearchHit>> {
5678 use crate::query::{
5679 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5680 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5681 };
5682 let total_started = std::time::Instant::now();
5683 let rank_offset = after.map_or(0, |after| after.returned_count);
5684 self.require_select()?;
5685 if request.limit == 0 {
5686 return Err(MongrelError::InvalidArgument(
5687 "search limit must be > 0".into(),
5688 ));
5689 }
5690 if request.limit > MAX_FINAL_LIMIT {
5691 return Err(MongrelError::InvalidArgument(format!(
5692 "search limit exceeds {MAX_FINAL_LIMIT}"
5693 )));
5694 }
5695 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5696 return Err(MongrelError::InvalidArgument(
5697 "search-after score must be finite".into(),
5698 ));
5699 }
5700 if request.retrievers.is_empty() {
5701 return Err(MongrelError::InvalidArgument(
5702 "search requires at least one retriever".into(),
5703 ));
5704 }
5705 if request.retrievers.len() > MAX_RETRIEVERS {
5706 return Err(MongrelError::InvalidArgument(format!(
5707 "search exceeds {MAX_RETRIEVERS} retrievers"
5708 )));
5709 }
5710 if request.must.len() > MAX_HARD_CONDITIONS {
5711 return Err(MongrelError::InvalidArgument(format!(
5712 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5713 )));
5714 }
5715 for condition in &request.must {
5716 self.validate_condition(condition)?;
5717 }
5718 if request.must.iter().any(|condition| {
5719 matches!(
5720 condition,
5721 Condition::Ann { .. }
5722 | Condition::SparseMatch { .. }
5723 | Condition::MinHashSimilar { .. }
5724 )
5725 }) {
5726 return Err(MongrelError::InvalidArgument(
5727 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5728 .into(),
5729 ));
5730 }
5731 let mut names = std::collections::HashSet::new();
5732 for named in &request.retrievers {
5733 if named.name.is_empty()
5734 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5735 || !names.insert(named.name.as_str())
5736 {
5737 return Err(MongrelError::InvalidArgument(format!(
5738 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5739 crate::query::MAX_RETRIEVER_NAME_BYTES
5740 )));
5741 }
5742 if !named.weight.is_finite()
5743 || named.weight < 0.0
5744 || named.weight > MAX_RETRIEVER_WEIGHT
5745 {
5746 return Err(MongrelError::InvalidArgument(format!(
5747 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5748 )));
5749 }
5750 self.validate_retriever(&named.retriever)?;
5751 }
5752 let projection = request
5753 .projection
5754 .clone()
5755 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5756 if projection.len() > MAX_PROJECTION_COLUMNS {
5757 return Err(MongrelError::InvalidArgument(format!(
5758 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5759 )));
5760 }
5761 for column_id in &projection {
5762 if !self
5763 .schema
5764 .columns
5765 .iter()
5766 .any(|column| column.id == *column_id)
5767 {
5768 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5769 }
5770 }
5771 if let Some(crate::query::Rerank::ExactVector {
5772 embedding_column,
5773 query,
5774 candidate_limit,
5775 weight,
5776 ..
5777 }) = &request.rerank
5778 {
5779 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5780 {
5781 return Err(MongrelError::InvalidArgument(format!(
5782 "rerank candidate_limit must be between search limit and {}",
5783 crate::query::MAX_RETRIEVER_K
5784 )));
5785 }
5786 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5787 return Err(MongrelError::InvalidArgument(format!(
5788 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5789 )));
5790 }
5791 let column = self
5792 .schema
5793 .columns
5794 .iter()
5795 .find(|column| column.id == *embedding_column)
5796 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5797 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5798 return Err(MongrelError::InvalidArgument(format!(
5799 "rerank column {embedding_column} is not an embedding"
5800 )));
5801 };
5802 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5803 return Err(MongrelError::InvalidArgument(format!(
5804 "rerank query must contain {dim} finite values"
5805 )));
5806 }
5807 }
5808
5809 let hard_filter_started = std::time::Instant::now();
5810 let hard_filter = if request.must.is_empty() {
5811 None
5812 } else {
5813 let mut sets = Vec::with_capacity(request.must.len());
5814 for condition in &request.must {
5815 if let Some(context) = context {
5816 context.checkpoint()?;
5817 }
5818 sets.push(self.resolve_condition(condition, snapshot)?);
5819 }
5820 Some(RowIdSet::intersect_many(sets))
5821 };
5822 crate::trace::QueryTrace::record(|trace| {
5823 trace.hard_filter_nanos = trace
5824 .hard_filter_nanos
5825 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
5826 });
5827 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
5828 return Ok(Vec::new());
5829 }
5830
5831 let constant = match request.fusion {
5832 Fusion::ReciprocalRank { constant } => constant,
5833 };
5834 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
5835 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
5836 let mut fusion_nanos = 0u64;
5837 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
5838 std::collections::HashMap::new();
5839 for named in retrievers {
5840 if named.weight == 0.0 {
5841 continue;
5842 }
5843 if let Some(context) = context {
5844 context.checkpoint()?;
5845 }
5846 let hits = self.retrieve_filtered(
5847 &named.retriever,
5848 snapshot,
5849 hard_filter.as_ref(),
5850 authorized,
5851 candidate_authorization,
5852 context,
5853 )?;
5854 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
5855 let fusion_started = std::time::Instant::now();
5856 for hit in hits {
5857 if let Some(context) = context {
5858 context.consume(1)?;
5859 }
5860 let contribution = named.weight / (constant as f64 + hit.rank as f64);
5861 if !contribution.is_finite() {
5862 return Err(MongrelError::InvalidArgument(
5863 "retriever contribution must be finite".into(),
5864 ));
5865 }
5866 let max_fused_candidates = context.map_or(
5867 crate::query::MAX_FUSED_CANDIDATES,
5868 crate::query::AiExecutionContext::max_fused_candidates,
5869 );
5870 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
5871 return Err(MongrelError::WorkBudgetExceeded);
5872 }
5873 let entry = fused.entry(hit.row_id).or_default();
5874 entry.0 += contribution;
5875 if !entry.0.is_finite() {
5876 return Err(MongrelError::InvalidArgument(
5877 "fused score must be finite".into(),
5878 ));
5879 }
5880 entry.1.push(ComponentScore {
5881 retriever_name: retriever_name.clone(),
5882 rank: hit.rank,
5883 raw_score: hit.score,
5884 contribution,
5885 });
5886 }
5887 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
5888 }
5889 let union_size = fused.len();
5890 let mut ranked: Vec<_> = fused
5891 .into_iter()
5892 .map(|(row_id, (fused_score, components))| {
5893 (row_id, fused_score, components, None, fused_score)
5894 })
5895 .collect();
5896 let order = |(a_row, _, _, _, a_score): &(
5897 RowId,
5898 f64,
5899 Vec<ComponentScore>,
5900 Option<f32>,
5901 f64,
5902 ),
5903 (b_row, _, _, _, b_score): &(
5904 RowId,
5905 f64,
5906 Vec<ComponentScore>,
5907 Option<f32>,
5908 f64,
5909 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
5910 if let Some(crate::query::Rerank::ExactVector {
5911 embedding_column,
5912 query,
5913 metric,
5914 candidate_limit,
5915 weight,
5916 }) = &request.rerank
5917 {
5918 let fused_order = |(a_row, a_score, ..): &(
5919 RowId,
5920 f64,
5921 Vec<ComponentScore>,
5922 Option<f32>,
5923 f64,
5924 ),
5925 (b_row, b_score, ..): &(
5926 RowId,
5927 f64,
5928 Vec<ComponentScore>,
5929 Option<f32>,
5930 f64,
5931 )| {
5932 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
5933 };
5934 let selection_started = std::time::Instant::now();
5935 if let Some(context) = context {
5936 context.consume(ranked.len())?;
5937 }
5938 if ranked.len() > *candidate_limit {
5939 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
5940 ranked.truncate(*candidate_limit);
5941 }
5942 ranked.sort_by(fused_order);
5943 fusion_nanos =
5944 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5945 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5946 if let Some(context) = context {
5947 context.consume(row_ids.len())?;
5948 }
5949 let query_now =
5950 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5951 let gather_started = std::time::Instant::now();
5952 let vectors = self.values_for_rids_batch_at_with_context(
5953 &row_ids,
5954 *embedding_column,
5955 snapshot,
5956 query_now,
5957 context,
5958 )?;
5959 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5960 let vector_work =
5961 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
5962 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
5963 if let Some(context) = context {
5964 context.consume(vector_work)?;
5965 }
5966 query
5967 .iter()
5968 .map(|value| f64::from(*value).powi(2))
5969 .sum::<f64>()
5970 .sqrt()
5971 } else {
5972 0.0
5973 };
5974 let score_started = std::time::Instant::now();
5975 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
5976 for (row_id, value) in vectors {
5977 let Value::Embedding(vector) = value else {
5978 continue;
5979 };
5980 let score = match metric {
5981 crate::query::VectorMetric::DotProduct => {
5982 if let Some(context) = context {
5983 context.consume(vector_work)?;
5984 }
5985 query
5986 .iter()
5987 .zip(&vector)
5988 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5989 .sum::<f64>()
5990 }
5991 crate::query::VectorMetric::Cosine => {
5992 if let Some(context) = context {
5993 context.consume(vector_work.saturating_mul(2))?;
5994 }
5995 let dot = query
5996 .iter()
5997 .zip(&vector)
5998 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5999 .sum::<f64>();
6000 let norm = vector
6001 .iter()
6002 .map(|value| f64::from(*value).powi(2))
6003 .sum::<f64>()
6004 .sqrt();
6005 if query_norm == 0.0 || norm == 0.0 {
6006 0.0
6007 } else {
6008 dot / (query_norm * norm)
6009 }
6010 }
6011 crate::query::VectorMetric::Euclidean => {
6012 if let Some(context) = context {
6013 context.consume(vector_work)?;
6014 }
6015 query
6016 .iter()
6017 .zip(&vector)
6018 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6019 .sum::<f64>()
6020 .sqrt()
6021 }
6022 };
6023 if !score.is_finite() {
6024 return Err(MongrelError::InvalidArgument(
6025 "exact rerank score must be finite".into(),
6026 ));
6027 }
6028 scores.insert(row_id, score as f32);
6029 }
6030 let mut reranked = Vec::with_capacity(ranked.len());
6031 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6032 let Some(score) = scores.get(&row_id).copied() else {
6033 continue;
6034 };
6035 let ordering_score = match metric {
6036 crate::query::VectorMetric::Euclidean => -f64::from(score),
6037 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6038 f64::from(score)
6039 }
6040 };
6041 let final_score = fused_score + *weight * ordering_score;
6042 if !final_score.is_finite() {
6043 return Err(MongrelError::InvalidArgument(
6044 "final rerank score must be finite".into(),
6045 ));
6046 }
6047 reranked.push((row_id, fused_score, components, Some(score), final_score));
6048 }
6049 ranked = reranked;
6050 ranked.sort_by(order);
6051 crate::trace::QueryTrace::record(|trace| {
6052 trace.exact_vector_gather_nanos =
6053 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6054 trace.exact_vector_score_nanos = trace
6055 .exact_vector_score_nanos
6056 .saturating_add(score_started.elapsed().as_nanos() as u64);
6057 });
6058 }
6059 if let Some(after) = after {
6060 ranked.retain(|(row_id, _, _, _, final_score)| {
6061 final_score.total_cmp(&after.final_score).is_lt()
6062 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6063 });
6064 }
6065 let projection_started = std::time::Instant::now();
6066 let sentinel = projection
6067 .first()
6068 .copied()
6069 .or_else(|| self.schema.columns.first().map(|column| column.id));
6070 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6071 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6072 let mut projection_rows = 0usize;
6073 let mut projection_cells = 0usize;
6074 while out.len() < request.limit && !ranked.is_empty() {
6075 if let Some(context) = context {
6076 context.checkpoint()?;
6077 context.consume(ranked.len())?;
6078 }
6079 let needed = request.limit - out.len();
6080 let window_size = ranked
6081 .len()
6082 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6083 let selection_started = std::time::Instant::now();
6084 let mut remainder = if ranked.len() > window_size {
6085 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6086 ranked.split_off(window_size)
6087 } else {
6088 Vec::new()
6089 };
6090 ranked.sort_by(order);
6091 fusion_nanos =
6092 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6093 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6094 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6095 if let Some(context) = context {
6096 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6097 }
6098 projection_rows = projection_rows.saturating_add(row_ids.len());
6099 projection_cells =
6100 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6101 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6102 std::collections::HashMap::new();
6103 if let Some(column_id) = sentinel {
6104 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6105 &row_ids, column_id, snapshot, query_now, context,
6106 )? {
6107 cells.entry(row_id).or_default().insert(column_id, value);
6108 }
6109 }
6110 for &column_id in &projection {
6111 if Some(column_id) == sentinel {
6112 continue;
6113 }
6114 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6115 &row_ids, column_id, snapshot, query_now, context,
6116 )? {
6117 cells.entry(row_id).or_default().insert(column_id, value);
6118 }
6119 }
6120 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6121 ranked.drain(..)
6122 {
6123 let Some(row_cells) = cells.remove(&row_id) else {
6124 continue;
6125 };
6126 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6127 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6128 out.push(SearchHit {
6129 row_id,
6130 cells: projection
6131 .iter()
6132 .filter_map(|column_id| {
6133 row_cells
6134 .get(column_id)
6135 .cloned()
6136 .map(|value| (*column_id, value))
6137 })
6138 .collect(),
6139 components,
6140 fused_score,
6141 exact_rerank_score,
6142 final_score,
6143 final_rank,
6144 });
6145 if out.len() == request.limit {
6146 break;
6147 }
6148 }
6149 ranked.append(&mut remainder);
6150 }
6151 crate::trace::QueryTrace::record(|trace| {
6152 trace.union_size = union_size;
6153 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6154 trace.projection_nanos = trace
6155 .projection_nanos
6156 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6157 trace.total_nanos = trace
6158 .total_nanos
6159 .saturating_add(total_started.elapsed().as_nanos() as u64);
6160 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6161 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6162 if let Some(context) = context {
6163 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6164 }
6165 });
6166 Ok(out)
6167 }
6168
6169 pub fn set_similarity(
6172 &mut self,
6173 request: &crate::query::SetSimilarityRequest,
6174 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6175 self.set_similarity_with_allowed(request, None)
6176 }
6177
6178 pub fn set_similarity_at(
6179 &mut self,
6180 request: &crate::query::SetSimilarityRequest,
6181 snapshot: Snapshot,
6182 allowed: Option<&std::collections::HashSet<RowId>>,
6183 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6184 self.set_similarity_explained_at(request, snapshot, allowed)
6185 .map(|(hits, _)| hits)
6186 }
6187
6188 pub fn ann_rerank(
6190 &mut self,
6191 request: &crate::query::AnnRerankRequest,
6192 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6193 self.ann_rerank_with_allowed(request, None)
6194 }
6195
6196 pub fn ann_rerank_with_allowed(
6197 &mut self,
6198 request: &crate::query::AnnRerankRequest,
6199 allowed: Option<&std::collections::HashSet<RowId>>,
6200 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6201 self.ann_rerank_at(request, self.snapshot(), allowed)
6202 }
6203
6204 pub fn ann_rerank_at(
6205 &mut self,
6206 request: &crate::query::AnnRerankRequest,
6207 snapshot: Snapshot,
6208 allowed: Option<&std::collections::HashSet<RowId>>,
6209 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6210 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6211 }
6212
6213 pub fn ann_rerank_at_with_context(
6214 &mut self,
6215 request: &crate::query::AnnRerankRequest,
6216 snapshot: Snapshot,
6217 allowed: Option<&std::collections::HashSet<RowId>>,
6218 context: Option<&crate::query::AiExecutionContext>,
6219 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6220 self.ensure_indexes_complete()?;
6221 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6222 }
6223
6224 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6225 &mut self,
6226 request: &crate::query::AnnRerankRequest,
6227 snapshot: Snapshot,
6228 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6229 context: Option<&crate::query::AiExecutionContext>,
6230 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6231 self.ensure_indexes_complete()?;
6232 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6233 }
6234
6235 #[doc(hidden)]
6236 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6237 &self,
6238 request: &crate::query::AnnRerankRequest,
6239 snapshot: Snapshot,
6240 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6241 context: Option<&crate::query::AiExecutionContext>,
6242 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6243 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6244 }
6245
6246 fn ann_rerank_at_with_filters_and_context(
6247 &self,
6248 request: &crate::query::AnnRerankRequest,
6249 snapshot: Snapshot,
6250 allowed: Option<&std::collections::HashSet<RowId>>,
6251 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6252 context: Option<&crate::query::AiExecutionContext>,
6253 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6254 use crate::query::{
6255 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6256 };
6257 if request.candidate_k == 0 || request.limit == 0 {
6258 return Err(MongrelError::InvalidArgument(
6259 "candidate_k and limit must be > 0".into(),
6260 ));
6261 }
6262 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6263 return Err(MongrelError::InvalidArgument(format!(
6264 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6265 )));
6266 }
6267 let retriever = Retriever::Ann {
6268 column_id: request.column_id,
6269 query: request.query.clone(),
6270 k: request.candidate_k,
6271 };
6272 self.require_select()?;
6273 self.validate_retriever(&retriever)?;
6274 let hits = self.retrieve_filtered(
6275 &retriever,
6276 snapshot,
6277 None,
6278 allowed,
6279 candidate_authorization,
6280 context,
6281 )?;
6282 let distances: std::collections::HashMap<_, _> = hits
6283 .iter()
6284 .filter_map(|hit| match hit.score {
6285 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
6286 _ => None,
6287 })
6288 .collect();
6289 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6290 if let Some(context) = context {
6291 context.consume(row_ids.len())?;
6292 }
6293 let gather_started = std::time::Instant::now();
6294 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6295 let values = self.values_for_rids_batch_at_with_context(
6296 &row_ids,
6297 request.column_id,
6298 snapshot,
6299 query_now,
6300 context,
6301 )?;
6302 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6303 let score_started = std::time::Instant::now();
6304 let vector_work =
6305 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6306 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6307 if let Some(context) = context {
6308 context.consume(vector_work)?;
6309 }
6310 request
6311 .query
6312 .iter()
6313 .map(|value| f64::from(*value).powi(2))
6314 .sum::<f64>()
6315 .sqrt()
6316 } else {
6317 0.0
6318 };
6319 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6320 for (row_id, value) in values {
6321 let Value::Embedding(vector) = value else {
6322 continue;
6323 };
6324 let exact_score = match request.metric {
6325 VectorMetric::DotProduct => {
6326 if let Some(context) = context {
6327 context.consume(vector_work)?;
6328 }
6329 request
6330 .query
6331 .iter()
6332 .zip(&vector)
6333 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6334 .sum::<f64>()
6335 }
6336 VectorMetric::Cosine => {
6337 if let Some(context) = context {
6338 context.consume(vector_work.saturating_mul(2))?;
6339 }
6340 let dot = request
6341 .query
6342 .iter()
6343 .zip(&vector)
6344 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6345 .sum::<f64>();
6346 let norm = vector
6347 .iter()
6348 .map(|value| f64::from(*value).powi(2))
6349 .sum::<f64>()
6350 .sqrt();
6351 if query_norm == 0.0 || norm == 0.0 {
6352 0.0
6353 } else {
6354 dot / (query_norm * norm)
6355 }
6356 }
6357 VectorMetric::Euclidean => {
6358 if let Some(context) = context {
6359 context.consume(vector_work)?;
6360 }
6361 request
6362 .query
6363 .iter()
6364 .zip(&vector)
6365 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6366 .sum::<f64>()
6367 .sqrt()
6368 }
6369 };
6370 let exact_score = exact_score as f32;
6371 if !exact_score.is_finite() {
6372 return Err(MongrelError::InvalidArgument(
6373 "exact ANN score must be finite".into(),
6374 ));
6375 }
6376 reranked.push(AnnRerankHit {
6377 row_id,
6378 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
6379 exact_score,
6380 });
6381 }
6382 reranked.sort_by(|left, right| {
6383 let score = match request.metric {
6384 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6385 VectorMetric::Cosine | VectorMetric::DotProduct => {
6386 right.exact_score.total_cmp(&left.exact_score)
6387 }
6388 };
6389 score.then_with(|| left.row_id.cmp(&right.row_id))
6390 });
6391 reranked.truncate(request.limit);
6392 crate::trace::QueryTrace::record(|trace| {
6393 trace.exact_vector_gather_nanos =
6394 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6395 trace.exact_vector_score_nanos = trace
6396 .exact_vector_score_nanos
6397 .saturating_add(score_started.elapsed().as_nanos() as u64);
6398 });
6399 Ok(reranked)
6400 }
6401
6402 pub fn set_similarity_with_allowed(
6403 &mut self,
6404 request: &crate::query::SetSimilarityRequest,
6405 allowed: Option<&std::collections::HashSet<RowId>>,
6406 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6407 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6408 .map(|(hits, _)| hits)
6409 }
6410
6411 pub fn set_similarity_explained(
6412 &mut self,
6413 request: &crate::query::SetSimilarityRequest,
6414 ) -> Result<(
6415 Vec<crate::query::SetSimilarityHit>,
6416 crate::query::SetSimilarityTrace,
6417 )> {
6418 self.set_similarity_explained_at(request, self.snapshot(), None)
6419 }
6420
6421 fn set_similarity_explained_at(
6422 &mut self,
6423 request: &crate::query::SetSimilarityRequest,
6424 snapshot: Snapshot,
6425 allowed: Option<&std::collections::HashSet<RowId>>,
6426 ) -> Result<(
6427 Vec<crate::query::SetSimilarityHit>,
6428 crate::query::SetSimilarityTrace,
6429 )> {
6430 self.ensure_indexes_complete()?;
6431 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6432 }
6433
6434 pub fn set_similarity_at_with_context(
6435 &mut self,
6436 request: &crate::query::SetSimilarityRequest,
6437 snapshot: Snapshot,
6438 allowed: Option<&std::collections::HashSet<RowId>>,
6439 context: Option<&crate::query::AiExecutionContext>,
6440 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6441 self.ensure_indexes_complete()?;
6442 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6443 .map(|(hits, _)| hits)
6444 }
6445
6446 pub fn set_similarity_at_with_candidate_authorization_and_context(
6447 &mut self,
6448 request: &crate::query::SetSimilarityRequest,
6449 snapshot: Snapshot,
6450 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6451 context: Option<&crate::query::AiExecutionContext>,
6452 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6453 self.ensure_indexes_complete()?;
6454 self.set_similarity_explained_at_with_context(
6455 request,
6456 snapshot,
6457 None,
6458 authorization,
6459 context,
6460 )
6461 .map(|(hits, _)| hits)
6462 }
6463
6464 #[doc(hidden)]
6465 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6466 &self,
6467 request: &crate::query::SetSimilarityRequest,
6468 snapshot: Snapshot,
6469 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6470 context: Option<&crate::query::AiExecutionContext>,
6471 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6472 self.set_similarity_explained_at_with_context(
6473 request,
6474 snapshot,
6475 None,
6476 authorization,
6477 context,
6478 )
6479 .map(|(hits, _)| hits)
6480 }
6481
6482 fn set_similarity_explained_at_with_context(
6483 &self,
6484 request: &crate::query::SetSimilarityRequest,
6485 snapshot: Snapshot,
6486 allowed: Option<&std::collections::HashSet<RowId>>,
6487 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6488 context: Option<&crate::query::AiExecutionContext>,
6489 ) -> Result<(
6490 Vec<crate::query::SetSimilarityHit>,
6491 crate::query::SetSimilarityTrace,
6492 )> {
6493 use crate::query::{
6494 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6495 MAX_SET_MEMBERS,
6496 };
6497 let mut trace = crate::query::SetSimilarityTrace::default();
6498 if request.members.is_empty() {
6499 return Ok((Vec::new(), trace));
6500 }
6501 if request.candidate_k == 0 || request.limit == 0 {
6502 return Err(MongrelError::InvalidArgument(
6503 "candidate_k and limit must be > 0".into(),
6504 ));
6505 }
6506 if request.candidate_k > MAX_RETRIEVER_K
6507 || request.limit > MAX_FINAL_LIMIT
6508 || request.members.len() > MAX_SET_MEMBERS
6509 {
6510 return Err(MongrelError::InvalidArgument(format!(
6511 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6512 )));
6513 }
6514 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6515 return Err(MongrelError::InvalidArgument(
6516 "min_jaccard must be finite and between 0 and 1".into(),
6517 ));
6518 }
6519 let started = std::time::Instant::now();
6520 let retriever = Retriever::MinHash {
6521 column_id: request.column_id,
6522 members: request.members.clone(),
6523 k: request.candidate_k,
6524 };
6525 self.require_select()?;
6526 self.validate_retriever(&retriever)?;
6527 let hits = self.retrieve_filtered(
6528 &retriever,
6529 snapshot,
6530 None,
6531 allowed,
6532 candidate_authorization,
6533 context,
6534 )?;
6535 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6536 trace.candidate_count = hits.len();
6537 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6538 if let Some(context) = context {
6539 context.consume(row_ids.len())?;
6540 }
6541 let started = std::time::Instant::now();
6542 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6543 let values = self.values_for_rids_batch_at_with_context(
6544 &row_ids,
6545 request.column_id,
6546 snapshot,
6547 query_now,
6548 context,
6549 )?;
6550 trace.gather_us = started.elapsed().as_micros() as u64;
6551 if let Some(context) = context {
6552 context.consume(request.members.len())?;
6553 }
6554 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6555 let estimates: std::collections::HashMap<_, _> = hits
6556 .into_iter()
6557 .filter_map(|hit| match hit.score {
6558 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6559 _ => None,
6560 })
6561 .collect();
6562 let started = std::time::Instant::now();
6563 let mut parsed = Vec::with_capacity(values.len());
6564 for (row_id, value) in values {
6565 let Value::Bytes(bytes) = value else {
6566 continue;
6567 };
6568 if let Some(context) = context {
6569 context.consume(crate::query::work_units(
6570 bytes.len(),
6571 crate::query::PARSE_WORK_QUANTUM,
6572 ))?;
6573 }
6574 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6575 continue;
6576 };
6577 if let Some(context) = context {
6578 context.consume(members.len())?;
6579 }
6580 let stored = members
6581 .into_iter()
6582 .filter_map(|member| match member {
6583 serde_json::Value::String(value) => {
6584 Some(crate::query::SetMember::String(value))
6585 }
6586 serde_json::Value::Number(value) => {
6587 Some(crate::query::SetMember::Number(value))
6588 }
6589 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6590 _ => None,
6591 })
6592 .collect::<std::collections::HashSet<_>>();
6593 parsed.push((row_id, stored));
6594 }
6595 trace.parse_us = started.elapsed().as_micros() as u64;
6596 trace.verified_count = parsed.len();
6597 let started = std::time::Instant::now();
6598 let mut exact = Vec::new();
6599 for (row_id, stored) in parsed {
6600 if let Some(context) = context {
6601 context.consume(query.len().saturating_add(stored.len()))?;
6602 }
6603 let union = query.union(&stored).count();
6604 let score = if union == 0 {
6605 1.0
6606 } else {
6607 query.intersection(&stored).count() as f32 / union as f32
6608 };
6609 if score >= request.min_jaccard {
6610 exact.push(SetSimilarityHit {
6611 row_id,
6612 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6613 exact_jaccard: score,
6614 });
6615 }
6616 }
6617 exact.sort_by(|a, b| {
6618 b.exact_jaccard
6619 .total_cmp(&a.exact_jaccard)
6620 .then_with(|| a.row_id.cmp(&b.row_id))
6621 });
6622 exact.truncate(request.limit);
6623 trace.score_us = started.elapsed().as_micros() as u64;
6624 crate::trace::QueryTrace::record(|query_trace| {
6625 query_trace.exact_set_gather_nanos = query_trace
6626 .exact_set_gather_nanos
6627 .saturating_add(trace.gather_us.saturating_mul(1_000));
6628 query_trace.exact_set_parse_nanos = query_trace
6629 .exact_set_parse_nanos
6630 .saturating_add(trace.parse_us.saturating_mul(1_000));
6631 query_trace.exact_set_score_nanos = query_trace
6632 .exact_set_score_nanos
6633 .saturating_add(trace.score_us.saturating_mul(1_000));
6634 });
6635 Ok((exact, trace))
6636 }
6637
6638 fn values_for_rids_batch_at(
6640 &self,
6641 row_ids: &[u64],
6642 column_id: u16,
6643 snapshot: Snapshot,
6644 now: i64,
6645 ) -> Result<Vec<(RowId, Value)>> {
6646 if self.ttl.is_none()
6647 && self.memtable.is_empty()
6648 && self.mutable_run.is_empty()
6649 && self.run_refs.len() == 1
6650 {
6651 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6652 if row_ids.len().saturating_mul(24) < reader.row_count() {
6657 let mut values = Vec::with_capacity(row_ids.len());
6658 for &raw_row_id in row_ids {
6659 let row_id = RowId(raw_row_id);
6660 if let Some((_, false, Some(value))) =
6661 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6662 {
6663 values.push((row_id, value));
6664 }
6665 }
6666 return Ok(values);
6667 }
6668 let (positions, visible_row_ids) =
6669 reader.visible_positions_with_rids(snapshot.epoch)?;
6670 let requested: Vec<(RowId, usize)> = row_ids
6671 .iter()
6672 .filter_map(|raw| {
6673 visible_row_ids
6674 .binary_search(&(*raw as i64))
6675 .ok()
6676 .map(|index| (RowId(*raw), positions[index]))
6677 })
6678 .collect();
6679 let values = reader.gather_column(
6680 column_id,
6681 &requested
6682 .iter()
6683 .map(|(_, position)| *position)
6684 .collect::<Vec<_>>(),
6685 )?;
6686 return Ok(requested
6687 .into_iter()
6688 .zip(values)
6689 .map(|((row_id, _), value)| (row_id, value))
6690 .collect());
6691 }
6692 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6693 }
6694
6695 fn values_for_rids_batch_at_with_context(
6696 &self,
6697 row_ids: &[u64],
6698 column_id: u16,
6699 snapshot: Snapshot,
6700 now: i64,
6701 context: Option<&crate::query::AiExecutionContext>,
6702 ) -> Result<Vec<(RowId, Value)>> {
6703 let Some(context) = context else {
6704 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6705 };
6706 let mut values = Vec::with_capacity(row_ids.len());
6707 for chunk in row_ids.chunks(256) {
6708 context.checkpoint()?;
6709 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6710 }
6711 Ok(values)
6712 }
6713
6714 fn values_for_rids_at(
6716 &self,
6717 row_ids: &[u64],
6718 column_id: u16,
6719 snapshot: Snapshot,
6720 now: i64,
6721 ) -> Result<Vec<(RowId, Value)>> {
6722 let mut readers: Vec<_> = self
6723 .run_refs
6724 .iter()
6725 .map(|run| self.open_reader(run.run_id))
6726 .collect::<Result<_>>()?;
6727 let mut out = Vec::with_capacity(row_ids.len());
6728 for &raw_row_id in row_ids {
6729 let row_id = RowId(raw_row_id);
6730 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6731 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6732 let overlay = match (mem, mutable) {
6733 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6734 (a_epoch, a)
6735 } else {
6736 (b_epoch, b)
6737 }),
6738 (Some(value), None) | (None, Some(value)) => Some(value),
6739 (None, None) => None,
6740 };
6741 if let Some((_, row)) = overlay {
6742 if !row.deleted && !self.row_expired_at(&row, now) {
6743 if let Some(value) = row.columns.get(&column_id) {
6744 out.push((row_id, value.clone()));
6745 }
6746 }
6747 continue;
6748 }
6749
6750 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6751 for (index, reader) in readers.iter_mut().enumerate() {
6752 if let Some((epoch, deleted, value)) =
6753 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6754 {
6755 if best
6756 .as_ref()
6757 .map(|(best_epoch, ..)| epoch > *best_epoch)
6758 .unwrap_or(true)
6759 {
6760 best = Some((epoch, deleted, value, index));
6761 }
6762 }
6763 }
6764 let Some((_, false, Some(value), reader_index)) = best else {
6765 continue;
6766 };
6767 if let Some(ttl) = self.ttl {
6768 if ttl.column_id != column_id {
6769 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6770 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6771 {
6772 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6773 continue;
6774 }
6775 }
6776 } else if let Value::Int64(timestamp) = value {
6777 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6778 continue;
6779 }
6780 }
6781 }
6782 out.push((row_id, value));
6783 }
6784 Ok(out)
6785 }
6786
6787 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6792 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6793 }
6794
6795 pub fn rows_for_rids_with_context(
6796 &self,
6797 rids: &[u64],
6798 snapshot: Snapshot,
6799 context: &crate::query::AiExecutionContext,
6800 ) -> Result<Vec<Row>> {
6801 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6802 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
6803 }
6804
6805 fn rows_for_rids_at_time(
6806 &self,
6807 rids: &[u64],
6808 snapshot: Snapshot,
6809 ttl_now: i64,
6810 control: Option<&crate::ExecutionControl>,
6811 ) -> Result<Vec<Row>> {
6812 use std::collections::HashMap;
6813 let mut rows = Vec::with_capacity(rids.len());
6814 let tier_size = self.memtable.len() + self.mutable_run.len();
6831 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
6832 if rids.len().saturating_mul(24) < tier_size {
6833 for &rid in rids {
6834 if overlay.len() & 255 == 0 {
6835 control
6836 .map(crate::ExecutionControl::checkpoint)
6837 .transpose()?;
6838 }
6839 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
6840 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
6841 let newest = match (mem, mrun) {
6842 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
6843 (Some((_, mr)), None) => Some(mr),
6844 (None, Some((_, rr))) => Some(rr),
6845 (None, None) => None,
6846 };
6847 if let Some(row) = newest {
6848 overlay.insert(rid, row);
6849 }
6850 }
6851 } else {
6852 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
6853 overlay
6854 .entry(row.row_id.0)
6855 .and_modify(|e| {
6856 if row.committed_epoch > e.committed_epoch {
6857 *e = row.clone();
6858 }
6859 })
6860 .or_insert(row);
6861 };
6862 for (index, row) in self
6863 .memtable
6864 .visible_versions(snapshot.epoch)
6865 .into_iter()
6866 .enumerate()
6867 {
6868 if index & 255 == 0 {
6869 control
6870 .map(crate::ExecutionControl::checkpoint)
6871 .transpose()?;
6872 }
6873 fold_newest(row, &mut overlay);
6874 }
6875 for (index, row) in self
6876 .mutable_run
6877 .visible_versions(snapshot.epoch)
6878 .into_iter()
6879 .enumerate()
6880 {
6881 if index & 255 == 0 {
6882 control
6883 .map(crate::ExecutionControl::checkpoint)
6884 .transpose()?;
6885 }
6886 fold_newest(row, &mut overlay);
6887 }
6888 }
6889 if self.run_refs.len() == 1 {
6890 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6891 if rids.len().saturating_mul(24) < reader.row_count() {
6899 for (index, &rid) in rids.iter().enumerate() {
6900 if index & 255 == 0 {
6901 control
6902 .map(crate::ExecutionControl::checkpoint)
6903 .transpose()?;
6904 }
6905 if let Some(r) = overlay.get(&rid) {
6906 if !r.deleted {
6907 rows.push(r.clone());
6908 }
6909 continue;
6910 }
6911 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
6912 if !row.deleted {
6913 rows.push(row);
6914 }
6915 }
6916 }
6917 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6918 return Ok(rows);
6919 }
6920 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6929 enum Src {
6932 Overlay,
6933 Run,
6934 }
6935 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
6936 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
6937 for (index, rid) in rids.iter().enumerate() {
6938 if index & 255 == 0 {
6939 control
6940 .map(crate::ExecutionControl::checkpoint)
6941 .transpose()?;
6942 }
6943 if overlay.contains_key(rid) {
6944 plan.push(Src::Overlay);
6945 continue;
6946 }
6947 match vis_rids.binary_search(&(*rid as i64)) {
6948 Ok(i) => {
6949 plan.push(Src::Run);
6950 fetch.push(positions[i]);
6951 }
6952 Err(_) => { }
6953 }
6954 }
6955 let fetched = reader.materialize_batch(&fetch)?;
6956 let mut fetched_iter = fetched.into_iter();
6957 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
6958 if index & 255 == 0 {
6959 control
6960 .map(crate::ExecutionControl::checkpoint)
6961 .transpose()?;
6962 }
6963 match src {
6964 Src::Overlay => {
6965 if let Some(r) = overlay.get(rid) {
6966 if !r.deleted {
6967 rows.push(r.clone());
6968 }
6969 }
6970 }
6971 Src::Run => {
6972 if let Some(row) = fetched_iter.next() {
6973 if !row.deleted {
6974 rows.push(row);
6975 }
6976 }
6977 }
6978 }
6979 }
6980 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6981 return Ok(rows);
6982 }
6983 let mut readers: Vec<_> = self
6987 .run_refs
6988 .iter()
6989 .map(|rr| self.open_reader(rr.run_id))
6990 .collect::<Result<Vec<_>>>()?;
6991 for (index, rid) in rids.iter().enumerate() {
6992 if index & 255 == 0 {
6993 control
6994 .map(crate::ExecutionControl::checkpoint)
6995 .transpose()?;
6996 }
6997 if let Some(r) = overlay.get(rid) {
6998 if !r.deleted {
6999 rows.push(r.clone());
7000 }
7001 continue;
7002 }
7003 let mut best: Option<(Epoch, Row)> = None;
7004 for reader in readers.iter_mut() {
7005 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7006 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7007 best = Some((epoch, row));
7008 }
7009 }
7010 }
7011 if let Some((_, r)) = best {
7012 if !r.deleted {
7013 rows.push(r);
7014 }
7015 }
7016 }
7017 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7018 Ok(rows)
7019 }
7020
7021 pub fn indexes_complete(&self) -> bool {
7031 self.indexes_complete
7032 }
7033
7034 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7036 self.index_build_policy
7037 }
7038
7039 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7043 self.index_build_policy = policy;
7044 }
7045
7046 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7051 if !self.indexes_complete {
7055 return None;
7056 }
7057 let b = self.bitmap.get(&column_id)?;
7058 let result: Vec<Vec<u8>> = b
7059 .keys()
7060 .into_iter()
7061 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7062 .collect();
7063 Some(result)
7064 }
7065
7066 pub fn fk_join_row_ids(
7067 &self,
7068 fk_column_id: u16,
7069 pk_values: &[Vec<u8>],
7070 fk_conditions: &[crate::query::Condition],
7071 snapshot: Snapshot,
7072 ) -> Result<Vec<u64>> {
7073 let Some(b) = self.bitmap.get(&fk_column_id) else {
7074 return Ok(Vec::new());
7075 };
7076 let mut join_set = {
7077 let mut acc = roaring::RoaringBitmap::new();
7078 for v in pk_values {
7079 acc |= b.get(v);
7080 }
7081 RowIdSet::from_roaring(acc)
7082 };
7083 if !fk_conditions.is_empty() {
7084 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7085 sets.push(join_set);
7086 for c in fk_conditions {
7087 sets.push(self.resolve_condition(c, snapshot)?);
7088 }
7089 join_set = RowIdSet::intersect_many(sets);
7090 }
7091 Ok(join_set.into_sorted_vec())
7092 }
7093
7094 pub fn fk_join_count(
7100 &self,
7101 fk_column_id: u16,
7102 pk_values: &[Vec<u8>],
7103 fk_conditions: &[crate::query::Condition],
7104 snapshot: Snapshot,
7105 ) -> Result<u64> {
7106 let Some(b) = self.bitmap.get(&fk_column_id) else {
7107 return Ok(0);
7108 };
7109 let mut acc = roaring::RoaringBitmap::new();
7110 for v in pk_values {
7111 acc |= b.get(v);
7112 }
7113 if fk_conditions.is_empty() {
7114 return Ok(acc.len());
7115 }
7116 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7117 sets.push(RowIdSet::from_roaring(acc));
7118 for c in fk_conditions {
7119 sets.push(self.resolve_condition(c, snapshot)?);
7120 }
7121 Ok(RowIdSet::intersect_many(sets).len() as u64)
7122 }
7123
7124 fn resolve_condition(
7129 &self,
7130 c: &crate::query::Condition,
7131 snapshot: Snapshot,
7132 ) -> Result<RowIdSet> {
7133 self.resolve_condition_with_allowed(c, snapshot, None)
7134 }
7135
7136 fn resolve_condition_with_allowed(
7137 &self,
7138 c: &crate::query::Condition,
7139 snapshot: Snapshot,
7140 allowed: Option<&std::collections::HashSet<RowId>>,
7141 ) -> Result<RowIdSet> {
7142 use crate::query::Condition;
7143 self.validate_condition(c)?;
7144 Ok(match c {
7145 Condition::Pk(key) => {
7146 let lookup = self
7147 .schema
7148 .primary_key()
7149 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7150 .unwrap_or_else(|| key.clone());
7151 self.hot
7152 .get(&lookup)
7153 .map(|r| RowIdSet::one(r.0))
7154 .unwrap_or_else(RowIdSet::empty)
7155 }
7156 Condition::BitmapEq { column_id, value } => {
7157 let lookup = self.index_lookup_key_bytes(*column_id, value);
7158 self.bitmap
7159 .get(column_id)
7160 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7161 .unwrap_or_else(RowIdSet::empty)
7162 }
7163 Condition::BitmapIn { column_id, values } => {
7164 let bm = self.bitmap.get(column_id);
7165 let mut acc = roaring::RoaringBitmap::new();
7166 if let Some(b) = bm {
7167 for v in values {
7168 let lookup = self.index_lookup_key_bytes(*column_id, v);
7169 acc |= b.get(&lookup);
7170 }
7171 }
7172 RowIdSet::from_roaring(acc)
7173 }
7174 Condition::BytesPrefix { column_id, prefix } => {
7175 if let Some(b) = self.bitmap.get(column_id) {
7180 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7181 let mut acc = roaring::RoaringBitmap::new();
7182 for key in b.keys() {
7183 if key.starts_with(&lookup_prefix) {
7184 acc |= b.get(&key);
7185 }
7186 }
7187 RowIdSet::from_roaring(acc)
7188 } else {
7189 RowIdSet::empty()
7190 }
7191 }
7192 Condition::FmContains { column_id, pattern } => self
7193 .fm
7194 .get(column_id)
7195 .map(|f| {
7196 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7197 })
7198 .unwrap_or_else(RowIdSet::empty),
7199 Condition::FmContainsAll {
7200 column_id,
7201 patterns,
7202 } => {
7203 if let Some(f) = self.fm.get(column_id) {
7206 let sets: Vec<RowIdSet> = patterns
7207 .iter()
7208 .map(|pat| {
7209 RowIdSet::from_unsorted(
7210 f.locate(pat).into_iter().map(|r| r.0).collect(),
7211 )
7212 })
7213 .collect();
7214 RowIdSet::intersect_many(sets)
7215 } else {
7216 RowIdSet::empty()
7217 }
7218 }
7219 Condition::Ann {
7220 column_id,
7221 query,
7222 k,
7223 } => RowIdSet::from_unsorted(
7224 self.retrieve_filtered(
7225 &crate::query::Retriever::Ann {
7226 column_id: *column_id,
7227 query: query.clone(),
7228 k: *k,
7229 },
7230 snapshot,
7231 None,
7232 allowed,
7233 None,
7234 None,
7235 )?
7236 .into_iter()
7237 .map(|hit| hit.row_id.0)
7238 .collect(),
7239 ),
7240 Condition::SparseMatch {
7241 column_id,
7242 query,
7243 k,
7244 } => RowIdSet::from_unsorted(
7245 self.retrieve_filtered(
7246 &crate::query::Retriever::Sparse {
7247 column_id: *column_id,
7248 query: query.clone(),
7249 k: *k,
7250 },
7251 snapshot,
7252 None,
7253 allowed,
7254 None,
7255 None,
7256 )?
7257 .into_iter()
7258 .map(|hit| hit.row_id.0)
7259 .collect(),
7260 ),
7261 Condition::MinHashSimilar {
7262 column_id,
7263 query,
7264 k,
7265 } => match self.minhash.get(column_id) {
7266 Some(index) => {
7267 let candidates = index.candidate_row_ids(query);
7268 let eligible =
7269 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7270 RowIdSet::from_unsorted(
7271 index
7272 .search_filtered(query, *k, |row_id| {
7273 eligible.contains(&row_id)
7274 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7275 })
7276 .into_iter()
7277 .map(|(row_id, _)| row_id.0)
7278 .collect(),
7279 )
7280 }
7281 None => RowIdSet::empty(),
7282 },
7283 Condition::Range { column_id, lo, hi } => {
7284 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7293 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7294 } else if self.run_refs.len() == 1 {
7295 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7296 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7297 } else {
7298 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7299 };
7300 set.remove_many(self.overlay_rid_set(snapshot));
7301 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7302 set
7303 }
7304 Condition::RangeF64 {
7305 column_id,
7306 lo,
7307 lo_inclusive,
7308 hi,
7309 hi_inclusive,
7310 } => {
7311 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7314 RowIdSet::from_unsorted(
7315 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7316 .into_iter()
7317 .collect(),
7318 )
7319 } else if self.run_refs.len() == 1 {
7320 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7321 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7322 } else {
7323 return self.range_scan_f64(
7324 *column_id,
7325 *lo,
7326 *lo_inclusive,
7327 *hi,
7328 *hi_inclusive,
7329 snapshot,
7330 );
7331 };
7332 set.remove_many(self.overlay_rid_set(snapshot));
7333 self.range_scan_overlay_f64(
7334 &mut set,
7335 *column_id,
7336 *lo,
7337 *lo_inclusive,
7338 *hi,
7339 *hi_inclusive,
7340 snapshot,
7341 );
7342 set
7343 }
7344 Condition::IsNull { column_id } => {
7345 let mut set = if self.run_refs.len() == 1 {
7346 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7347 r.null_row_id_set(*column_id, true)?
7348 } else {
7349 return self.null_scan(*column_id, true, snapshot);
7350 };
7351 set.remove_many(self.overlay_rid_set(snapshot));
7352 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7353 set
7354 }
7355 Condition::IsNotNull { column_id } => {
7356 let mut set = if self.run_refs.len() == 1 {
7357 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7358 r.null_row_id_set(*column_id, false)?
7359 } else {
7360 return self.null_scan(*column_id, false, snapshot);
7361 };
7362 set.remove_many(self.overlay_rid_set(snapshot));
7363 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7364 set
7365 }
7366 })
7367 }
7368
7369 fn range_scan_i64(
7377 &self,
7378 column_id: u16,
7379 lo: i64,
7380 hi: i64,
7381 snapshot: Snapshot,
7382 ) -> Result<RowIdSet> {
7383 let mut row_ids = Vec::new();
7384 let overlay_rids = self.overlay_rid_set(snapshot);
7385 for rr in &self.run_refs {
7386 let mut reader = self.open_reader(rr.run_id)?;
7387 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7388 for rid in matched {
7389 if !overlay_rids.contains(&rid) {
7390 row_ids.push(rid);
7391 }
7392 }
7393 }
7394 let mut s = RowIdSet::from_unsorted(row_ids);
7395 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7396 Ok(s)
7397 }
7398
7399 fn range_scan_f64(
7402 &self,
7403 column_id: u16,
7404 lo: f64,
7405 lo_inclusive: bool,
7406 hi: f64,
7407 hi_inclusive: bool,
7408 snapshot: Snapshot,
7409 ) -> Result<RowIdSet> {
7410 let mut row_ids = Vec::new();
7411 let overlay_rids = self.overlay_rid_set(snapshot);
7412 for rr in &self.run_refs {
7413 let mut reader = self.open_reader(rr.run_id)?;
7414 let matched = reader.range_row_ids_visible_f64(
7415 column_id,
7416 lo,
7417 lo_inclusive,
7418 hi,
7419 hi_inclusive,
7420 snapshot.epoch,
7421 )?;
7422 for rid in matched {
7423 if !overlay_rids.contains(&rid) {
7424 row_ids.push(rid);
7425 }
7426 }
7427 }
7428 let mut s = RowIdSet::from_unsorted(row_ids);
7429 self.range_scan_overlay_f64(
7430 &mut s,
7431 column_id,
7432 lo,
7433 lo_inclusive,
7434 hi,
7435 hi_inclusive,
7436 snapshot,
7437 );
7438 Ok(s)
7439 }
7440
7441 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7443 let mut s = HashSet::new();
7444 for row in self.memtable.visible_versions(snapshot.epoch) {
7445 s.insert(row.row_id.0);
7446 }
7447 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7448 s.insert(row.row_id.0);
7449 }
7450 s
7451 }
7452
7453 fn range_scan_overlay_i64(
7454 &self,
7455 s: &mut RowIdSet,
7456 column_id: u16,
7457 lo: i64,
7458 hi: i64,
7459 snapshot: Snapshot,
7460 ) {
7461 let mut newest: HashMap<u64, &Row> = HashMap::new();
7466 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7467 let memtable = self.memtable.visible_versions(snapshot.epoch);
7468 for r in &mutable {
7469 newest.entry(r.row_id.0).or_insert(r);
7470 }
7471 for r in &memtable {
7472 newest.insert(r.row_id.0, r);
7473 }
7474 for row in newest.values() {
7475 if !row.deleted {
7476 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7477 if *v >= lo && *v <= hi {
7478 s.insert(row.row_id.0);
7479 }
7480 }
7481 }
7482 }
7483 }
7484
7485 #[allow(clippy::too_many_arguments)]
7486 fn range_scan_overlay_f64(
7487 &self,
7488 s: &mut RowIdSet,
7489 column_id: u16,
7490 lo: f64,
7491 lo_inclusive: bool,
7492 hi: f64,
7493 hi_inclusive: bool,
7494 snapshot: Snapshot,
7495 ) {
7496 let mut newest: HashMap<u64, &Row> = HashMap::new();
7499 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7500 let memtable = self.memtable.visible_versions(snapshot.epoch);
7501 for r in &mutable {
7502 newest.entry(r.row_id.0).or_insert(r);
7503 }
7504 for r in &memtable {
7505 newest.insert(r.row_id.0, r);
7506 }
7507 for row in newest.values() {
7508 if !row.deleted {
7509 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7510 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7511 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7512 if ok_lo && ok_hi {
7513 s.insert(row.row_id.0);
7514 }
7515 }
7516 }
7517 }
7518 }
7519
7520 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7523 let mut row_ids = Vec::new();
7524 let overlay_rids = self.overlay_rid_set(snapshot);
7525 for rr in &self.run_refs {
7526 let mut reader = self.open_reader(rr.run_id)?;
7527 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7528 for rid in matched {
7529 if !overlay_rids.contains(&rid) {
7530 row_ids.push(rid);
7531 }
7532 }
7533 }
7534 let mut s = RowIdSet::from_unsorted(row_ids);
7535 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7536 Ok(s)
7537 }
7538
7539 fn null_scan_overlay(
7543 &self,
7544 s: &mut RowIdSet,
7545 column_id: u16,
7546 want_nulls: bool,
7547 snapshot: Snapshot,
7548 ) {
7549 let mut newest: HashMap<u64, &Row> = HashMap::new();
7550 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7551 let memtable = self.memtable.visible_versions(snapshot.epoch);
7552 for r in &mutable {
7553 newest.entry(r.row_id.0).or_insert(r);
7554 }
7555 for r in &memtable {
7556 newest.insert(r.row_id.0, r);
7557 }
7558 for row in newest.values() {
7559 if row.deleted {
7560 continue;
7561 }
7562 let is_null = !row.columns.contains_key(&column_id)
7563 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7564 if is_null == want_nulls {
7565 s.insert(row.row_id.0);
7566 }
7567 }
7568 }
7569
7570 pub fn snapshot(&self) -> Snapshot {
7571 Snapshot::at(self.epoch.visible())
7572 }
7573
7574 pub fn data_generation(&self) -> u64 {
7576 self.data_generation
7577 }
7578
7579 pub(crate) fn bump_data_generation(&mut self) {
7580 self.data_generation = self.data_generation.wrapping_add(1);
7581 }
7582
7583 pub fn table_id(&self) -> u64 {
7585 self.table_id
7586 }
7587
7588 fn seal_generations(&mut self) {
7594 self.memtable.seal();
7595 self.mutable_run.seal();
7596 self.hot.seal();
7597 for index in self.bitmap.values_mut() {
7598 index.seal();
7599 }
7600 for index in self.ann.values_mut() {
7601 index.seal();
7602 }
7603 for index in self.fm.values_mut() {
7604 index.seal();
7605 }
7606 for index in self.sparse.values_mut() {
7607 index.seal();
7608 }
7609 for index in self.minhash.values_mut() {
7610 index.seal();
7611 }
7612 self.pk_by_row.seal();
7613 }
7614
7615 fn capture_read_generation(&self) -> ReadGeneration {
7620 let visible_through = self.current_epoch();
7621 ReadGeneration {
7622 schema: Arc::new(self.schema.clone()),
7623 base_runs: Arc::new(self.run_refs.clone()),
7624 deltas: TableDeltas {
7625 memtable: self.memtable.clone(),
7626 mutable_run: self.mutable_run.clone(),
7627 hot: self.hot.clone(),
7628 pk_by_row: self.pk_by_row.clone(),
7629 },
7630 indexes: Arc::new(IndexGeneration::capture(
7631 &self.bitmap,
7632 &self.learned_range,
7633 &self.fm,
7634 &self.ann,
7635 &self.sparse,
7636 &self.minhash,
7637 visible_through,
7638 )),
7639 visible_through,
7640 }
7641 }
7642
7643 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
7648 self.ensure_indexes_complete()?;
7649 self.seal_generations();
7650 let view = Arc::new(self.capture_read_generation());
7651 self.published.store(Arc::clone(&view));
7652 Ok(view)
7653 }
7654
7655 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
7661 self.published.load_full()
7662 }
7663
7664 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
7666 &self.pins
7667 }
7668
7669 pub fn version_gc_floor(&self) -> Epoch {
7674 self.min_active_snapshot()
7675 .unwrap_or_else(|| self.current_epoch())
7676 }
7677
7678 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
7685 let mut report = self.pins.report();
7686 let transaction_floor = [
7687 self.pinned.keys().next().copied(),
7688 self.snapshots.min_pinned(),
7689 ]
7690 .into_iter()
7691 .flatten()
7692 .min();
7693 if let Some(epoch) = transaction_floor {
7694 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
7695 }
7696 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
7697 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
7698 }
7699 report
7700 }
7701
7702 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7703 self.publish_read_generation()?;
7704 let mut generation = self.clone();
7705 generation.read_only = true;
7706 generation.wal = WalSink::ReadOnly;
7707 generation.pending_delete_rids.clear();
7708 generation.pending_put_cols.clear();
7709 generation.pending_rows.clear();
7710 generation.pending_rows_auto_inc.clear();
7711 generation.pending_dels.clear();
7712 generation.pending_truncate = None;
7713 generation.agg_cache = Arc::new(HashMap::new());
7714 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
7717 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
7720 crate::retention::PinSource::ReadGeneration,
7721 self.current_epoch(),
7722 )));
7723 Ok(generation)
7724 }
7725
7726 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7727 (std::mem::size_of::<Self>() as u64)
7728 .saturating_add(self.memtable.approx_bytes())
7729 .saturating_add(self.mutable_run.approx_bytes())
7730 .saturating_add(self.live_count.saturating_mul(64))
7731 }
7732
7733 pub fn pin_snapshot(&mut self) -> Snapshot {
7736 let e = self.epoch.visible();
7737 *self.pinned.entry(e).or_insert(0) += 1;
7738 Snapshot::at(e)
7739 }
7740
7741 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7743 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7744 *count -= 1;
7745 if *count == 0 {
7746 self.pinned.remove(&snap.epoch);
7747 }
7748 }
7749 }
7750
7751 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7763 let local = self.pinned.keys().next().copied();
7764 let global = self.snapshots.min_pinned();
7765 let history = self.snapshots.history_floor(self.current_epoch());
7766 let pinned = self.pins.oldest_pinned();
7767 [local, global, history, pinned].into_iter().flatten().min()
7768 }
7769
7770 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7774 self.ensure_writable()?;
7775 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7776 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7777 }
7778
7779 pub fn clear_ttl(&mut self) -> Result<()> {
7780 self.ensure_writable()?;
7781 self.apply_ttl_policy_at(None, self.current_epoch())
7782 }
7783
7784 pub fn ttl(&self) -> Option<TtlPolicy> {
7785 self.ttl
7786 }
7787
7788 pub(crate) fn prepare_ttl_policy(
7789 &self,
7790 column_name: &str,
7791 duration_nanos: u64,
7792 ) -> Result<TtlPolicy> {
7793 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7794 return Err(MongrelError::InvalidArgument(
7795 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7796 ));
7797 }
7798 let column = self
7799 .schema
7800 .columns
7801 .iter()
7802 .find(|column| column.name == column_name)
7803 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7804 if column.ty != TypeId::TimestampNanos {
7805 return Err(MongrelError::Schema(format!(
7806 "TTL column {column_name} must be TimestampNanos, is {:?}",
7807 column.ty
7808 )));
7809 }
7810 Ok(TtlPolicy {
7811 column_id: column.id,
7812 duration_nanos,
7813 })
7814 }
7815
7816 pub(crate) fn apply_ttl_policy_at(
7817 &mut self,
7818 policy: Option<TtlPolicy>,
7819 epoch: Epoch,
7820 ) -> Result<()> {
7821 if let Some(policy) = policy {
7822 let column = self
7823 .schema
7824 .columns
7825 .iter()
7826 .find(|column| column.id == policy.column_id)
7827 .ok_or_else(|| {
7828 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
7829 })?;
7830 if column.ty != TypeId::TimestampNanos
7831 || policy.duration_nanos == 0
7832 || policy.duration_nanos > i64::MAX as u64
7833 {
7834 return Err(MongrelError::Schema("invalid TTL policy".into()));
7835 }
7836 }
7837 self.ttl = policy;
7838 self.agg_cache = Arc::new(HashMap::new());
7839 self.clear_result_cache();
7840 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7841 self.persist_manifest(epoch)
7842 }
7843
7844 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
7845 let Some(policy) = self.ttl else {
7846 return false;
7847 };
7848 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
7849 return false;
7850 };
7851 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
7852 }
7853
7854 pub fn current_epoch(&self) -> Epoch {
7855 self.epoch.visible()
7856 }
7857
7858 pub fn memtable_len(&self) -> usize {
7859 self.memtable.len()
7860 }
7861
7862 pub fn count(&self) -> u64 {
7865 if self.ttl.is_none()
7866 && self.pending_put_cols.is_empty()
7867 && self.pending_delete_rids.is_empty()
7868 && self.pending_rows.is_empty()
7869 && self.pending_dels.is_empty()
7870 && self.pending_truncate.is_none()
7871 {
7872 self.live_count
7873 } else {
7874 self.visible_rows(self.snapshot())
7875 .map(|rows| rows.len() as u64)
7876 .unwrap_or(self.live_count)
7877 }
7878 }
7879
7880 pub fn count_conditions(
7884 &mut self,
7885 conditions: &[crate::query::Condition],
7886 snapshot: Snapshot,
7887 ) -> Result<Option<u64>> {
7888 use crate::query::Condition;
7889 if self.ttl.is_some() {
7890 if conditions.is_empty() {
7891 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
7892 }
7893 let mut sets = Vec::with_capacity(conditions.len());
7894 for condition in conditions {
7895 sets.push(self.resolve_condition(condition, snapshot)?);
7896 }
7897 let survivors = RowIdSet::intersect_many(sets);
7898 let rows = self.visible_rows(snapshot)?;
7899 return Ok(Some(
7900 rows.into_iter()
7901 .filter(|row| survivors.contains(row.row_id.0))
7902 .count() as u64,
7903 ));
7904 }
7905 if conditions.is_empty() {
7906 return Ok(Some(self.count()));
7907 }
7908 let served = |c: &Condition| {
7909 matches!(
7910 c,
7911 Condition::Pk(_)
7912 | Condition::BitmapEq { .. }
7913 | Condition::BitmapIn { .. }
7914 | Condition::BytesPrefix { .. }
7915 | Condition::FmContains { .. }
7916 | Condition::FmContainsAll { .. }
7917 | Condition::Ann { .. }
7918 | Condition::Range { .. }
7919 | Condition::RangeF64 { .. }
7920 | Condition::SparseMatch { .. }
7921 | Condition::MinHashSimilar { .. }
7922 | Condition::IsNull { .. }
7923 | Condition::IsNotNull { .. }
7924 )
7925 };
7926 if !conditions.iter().all(served) {
7927 return Ok(None);
7928 }
7929 self.ensure_indexes_complete()?;
7930 if !self.pending_put_cols.is_empty()
7931 || !self.pending_delete_rids.is_empty()
7932 || !self.pending_rows.is_empty()
7933 || !self.pending_dels.is_empty()
7934 || self.pending_truncate.is_some()
7935 {
7936 let mut sets = Vec::with_capacity(conditions.len());
7937 for condition in conditions {
7938 sets.push(self.resolve_condition(condition, snapshot)?);
7939 }
7940 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7941 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
7942 }
7943 let mut sets = Vec::with_capacity(conditions.len());
7944 for condition in conditions {
7945 sets.push(self.resolve_condition(condition, snapshot)?);
7946 }
7947 let mut rids = RowIdSet::intersect_many(sets);
7948 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7958 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
7959 }
7960 let count = rids.len() as u64;
7961 crate::trace::QueryTrace::record(|t| {
7962 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
7963 t.survivor_count = Some(count as usize);
7964 t.conditions_pushed = conditions.len();
7965 });
7966 Ok(Some(count))
7967 }
7968
7969 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
7974 let mut out = Vec::new();
7975 for row in self.memtable.visible_versions(snapshot.epoch) {
7976 if row.deleted {
7977 out.push(row.row_id.0);
7978 }
7979 }
7980 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7981 if row.deleted {
7982 out.push(row.row_id.0);
7983 }
7984 }
7985 out
7986 }
7987
7988 pub fn bulk_load_columns(
7997 &mut self,
7998 user_columns: Vec<(u16, columnar::NativeColumn)>,
7999 ) -> Result<Epoch> {
8000 self.bulk_load_columns_with(user_columns, 3, false, true)
8001 }
8002
8003 pub fn bulk_load_fast(
8010 &mut self,
8011 user_columns: Vec<(u16, columnar::NativeColumn)>,
8012 ) -> Result<Epoch> {
8013 self.bulk_load_columns_with(user_columns, -1, true, false)
8014 }
8015
8016 fn bulk_load_columns_with(
8017 &mut self,
8018 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8019 zstd_level: i32,
8020 force_plain: bool,
8021 lz4: bool,
8022 ) -> Result<Epoch> {
8023 self.ensure_writable()?;
8024 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8025 if n == 0 {
8026 return Ok(self.current_epoch());
8027 }
8028 let epoch = self.commit_new_epoch()?;
8029 let live_before = self.live_count;
8030 self.spill_mutable_run(epoch)?;
8032 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8033 && self.indexes_complete
8034 && self.run_refs.is_empty()
8035 && self.memtable.is_empty()
8036 && self.mutable_run.is_empty();
8037 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8040 self.validate_columns_not_null(&user_columns, n)?;
8041 let winner_idx = self
8042 .bulk_pk_winner_indices(&user_columns, n)
8043 .filter(|idx| idx.len() != n);
8044 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8045 match winner_idx.as_deref() {
8046 Some(idx) => {
8047 let compacted = user_columns
8048 .iter()
8049 .map(|(id, c)| (*id, c.gather(idx)))
8050 .collect();
8051 (compacted, idx.len())
8052 }
8053 None => (user_columns, n),
8054 };
8055 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8056 let first = self.allocator.alloc_range(write_n as u64)?.0;
8057 for rid in first..first + write_n as u64 {
8058 self.reservoir.offer(rid);
8059 }
8060 let run_id = self.alloc_run_id()?;
8061 let path = self.run_path(run_id);
8062 let mut writer =
8063 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8064 if force_plain {
8065 writer = writer.with_plain();
8066 } else if lz4 {
8067 writer = writer.with_lz4();
8070 } else {
8071 writer = writer.with_zstd_level(zstd_level);
8072 }
8073 if let Some(kek) = &self.kek {
8074 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8075 }
8076 let header = match self.create_run_file(run_id)? {
8077 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8078 None => writer.write_native(&path, &write_columns, write_n, first)?,
8079 };
8080 self.run_refs.push(RunRef {
8081 run_id: run_id as u128,
8082 level: 0,
8083 epoch_created: epoch.0,
8084 row_count: header.row_count,
8085 });
8086 self.live_count = self.live_count.saturating_add(write_n as u64);
8087 if eager_index_build {
8088 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8089 self.index_columns_bulk(&write_columns, &row_ids);
8090 self.indexes_complete = true;
8091 self.build_learned_ranges()?;
8092 } else {
8093 self.indexes_complete = false;
8097 }
8098 self.mark_flushed(epoch)?;
8099 self.persist_manifest(epoch)?;
8100 if eager_index_build {
8101 self.checkpoint_indexes(epoch);
8102 }
8103 self.clear_result_cache();
8104 self.data_generation = self.data_generation.wrapping_add(1);
8105 Ok(epoch)
8106 }
8107
8108 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8126 let n = row_ids.len();
8127 if n == 0 {
8128 return;
8129 }
8130 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8131 columns.iter().map(|(id, c)| (*id, c)).collect();
8132 let ty_of: std::collections::HashMap<u16, TypeId> = self
8133 .schema
8134 .columns
8135 .iter()
8136 .map(|c| (c.id, c.ty.clone()))
8137 .collect();
8138 let pk_id = self.schema.primary_key().map(|c| c.id);
8139
8140 for (i, &rid) in row_ids.iter().enumerate() {
8141 let row_id = RowId(rid);
8142 if let Some(pid) = pk_id {
8143 if let Some(col) = by_id.get(&pid) {
8144 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8145 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8146 self.insert_hot_pk(key, row_id);
8147 }
8148 }
8149 }
8150 for idef in &self.schema.indexes {
8151 let Some(col) = by_id.get(&idef.column_id) else {
8152 continue;
8153 };
8154 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8155 match idef.kind {
8156 IndexKind::Bitmap => {
8157 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8158 if let Some(key) =
8159 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8160 {
8161 b.insert(key, row_id);
8162 }
8163 }
8164 }
8165 IndexKind::FmIndex => {
8166 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8167 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8168 f.insert(bytes.to_vec(), row_id);
8169 }
8170 }
8171 }
8172 IndexKind::Sparse => {
8173 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8174 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8175 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8176 s.insert(&terms, row_id);
8177 }
8178 }
8179 }
8180 }
8181 IndexKind::MinHash => {
8182 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8183 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8184 let tokens = crate::index::token_hashes_from_bytes(bytes);
8185 mh.insert(&tokens, row_id);
8186 }
8187 }
8188 }
8189 _ => {}
8190 }
8191 }
8192 }
8193 }
8194
8195 pub fn visible_columns_native(
8200 &self,
8201 snapshot: Snapshot,
8202 projection: Option<&[u16]>,
8203 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8204 self.visible_columns_native_inner(snapshot, projection, None)
8205 }
8206
8207 pub fn visible_columns_native_with_control(
8208 &self,
8209 snapshot: Snapshot,
8210 projection: Option<&[u16]>,
8211 control: &crate::ExecutionControl,
8212 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8213 self.visible_columns_native_inner(snapshot, projection, Some(control))
8214 }
8215
8216 fn visible_columns_native_inner(
8217 &self,
8218 snapshot: Snapshot,
8219 projection: Option<&[u16]>,
8220 control: Option<&crate::ExecutionControl>,
8221 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8222 execution_checkpoint(control, 0)?;
8223 let wanted: Vec<u16> = match projection {
8224 Some(p) => p.to_vec(),
8225 None => self.schema.columns.iter().map(|c| c.id).collect(),
8226 };
8227 if self.ttl.is_none()
8228 && self.memtable.is_empty()
8229 && self.mutable_run.is_empty()
8230 && self.run_refs.len() == 1
8231 {
8232 let rr = self.run_refs[0].clone();
8233 let mut reader = self.open_reader(rr.run_id)?;
8234 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8235 execution_checkpoint(control, 0)?;
8236 let all_visible = idxs.len() == reader.row_count();
8237 if reader.has_mmap() && control.is_none() {
8243 use rayon::prelude::*;
8244 let valid: Vec<u16> = wanted
8247 .iter()
8248 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8249 .copied()
8250 .collect();
8251 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8253 .par_iter()
8254 .filter_map(|cid| {
8255 reader
8256 .column_native_shared(*cid)
8257 .ok()
8258 .map(|col| (*cid, col))
8259 })
8260 .collect();
8261 let cols = decoded
8262 .into_iter()
8263 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8264 .collect();
8265 return Ok(cols);
8266 }
8267 let mut cols = Vec::with_capacity(wanted.len());
8268 for (index, cid) in wanted.iter().enumerate() {
8269 execution_checkpoint(control, index)?;
8270 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8271 Some(c) => c,
8272 None => continue,
8273 };
8274 let col = reader.column_native(cdef.id)?;
8275 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8276 }
8277 return Ok(cols);
8278 }
8279 let vcols = self.visible_columns(snapshot)?;
8280 execution_checkpoint(control, 0)?;
8281 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8282 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8283 .into_iter()
8284 .filter(|(id, _)| want_set.contains(id))
8285 .map(|(id, vals)| {
8286 let ty = self
8287 .schema
8288 .columns
8289 .iter()
8290 .find(|c| c.id == id)
8291 .map(|c| c.ty.clone())
8292 .unwrap_or(TypeId::Bytes);
8293 (id, columnar::values_to_native(ty, &vals))
8294 })
8295 .collect();
8296 Ok(out)
8297 }
8298
8299 pub fn run_count(&self) -> usize {
8300 self.run_refs.len()
8301 }
8302
8303 pub fn memtable_is_empty(&self) -> bool {
8305 self.memtable.is_empty()
8306 }
8307
8308 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8312 self.page_cache.stats()
8313 }
8314
8315 pub fn reset_page_cache_stats(&self) {
8317 self.page_cache.reset_stats();
8318 }
8319
8320 pub fn run_ids(&self) -> Vec<u128> {
8323 self.run_refs.iter().map(|r| r.run_id).collect()
8324 }
8325
8326 pub fn single_run_is_clean(&self) -> bool {
8330 if self.ttl.is_some() || self.run_refs.len() != 1 {
8331 return false;
8332 }
8333 self.open_reader(self.run_refs[0].run_id)
8334 .map(|r| r.is_clean())
8335 .unwrap_or(false)
8336 }
8337
8338 fn resolve_footprint(
8345 &self,
8346 conditions: &[crate::query::Condition],
8347 snapshot: Snapshot,
8348 ) -> roaring::RoaringBitmap {
8349 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8350 return roaring::RoaringBitmap::new();
8351 }
8352 if self.run_refs.is_empty() {
8353 return roaring::RoaringBitmap::new();
8354 }
8355 if self.run_refs.len() == 1 {
8357 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8358 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8359 return rids.to_roaring_lossy();
8360 }
8361 }
8362 }
8363 roaring::RoaringBitmap::new()
8364 }
8365
8366 pub fn query_columns_native_cached(
8377 &mut self,
8378 conditions: &[crate::query::Condition],
8379 projection: Option<&[u16]>,
8380 snapshot: Snapshot,
8381 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8382 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8383 }
8384
8385 pub fn query_columns_native_cached_with_control(
8386 &mut self,
8387 conditions: &[crate::query::Condition],
8388 projection: Option<&[u16]>,
8389 snapshot: Snapshot,
8390 control: &crate::ExecutionControl,
8391 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8392 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8393 }
8394
8395 fn query_columns_native_cached_inner(
8396 &mut self,
8397 conditions: &[crate::query::Condition],
8398 projection: Option<&[u16]>,
8399 snapshot: Snapshot,
8400 control: Option<&crate::ExecutionControl>,
8401 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8402 execution_checkpoint(control, 0)?;
8403 if self.ttl.is_some() {
8406 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8407 }
8408 if conditions.is_empty() {
8409 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8410 }
8411 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8415 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8416 crate::trace::QueryTrace::record(|t| {
8417 t.result_cache_hit = true;
8418 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8419 });
8420 return Ok(Some((*hit).clone()));
8421 }
8422 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8423 execution_checkpoint(control, 0)?;
8424 if let Some(cols) = &res {
8425 let footprint = self.resolve_footprint(conditions, snapshot);
8426 let condition_cols = crate::query::condition_columns(conditions);
8427 execution_checkpoint(control, 0)?;
8428 self.result_cache.lock().insert(
8429 key,
8430 CachedEntry {
8431 data: CachedData::Columns(Arc::new(cols.clone())),
8432 footprint,
8433 condition_cols,
8434 },
8435 );
8436 }
8437 Ok(res)
8438 }
8439
8440 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8445 if self.ttl.is_some() {
8446 return self.query(q);
8447 }
8448 if q.conditions.is_empty() {
8449 return self.query(q);
8450 }
8451 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8452 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8453 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8454 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8455 crate::trace::QueryTrace::record(|t| {
8456 t.result_cache_hit = true;
8457 t.scan_mode = crate::trace::ScanMode::Materialized;
8458 });
8459 return Ok((*hit).clone());
8460 }
8461 let rows = self.query(q)?;
8462 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8463 let condition_cols = crate::query::condition_columns(&q.conditions);
8464 self.result_cache.lock().insert(
8465 key,
8466 CachedEntry {
8467 data: CachedData::Rows(Arc::new(rows.clone())),
8468 footprint,
8469 condition_cols,
8470 },
8471 );
8472 Ok(rows)
8473 }
8474
8475 #[allow(clippy::type_complexity)]
8490 pub fn query_columns_native_traced(
8491 &mut self,
8492 conditions: &[crate::query::Condition],
8493 projection: Option<&[u16]>,
8494 snapshot: Snapshot,
8495 ) -> Result<(
8496 Option<Vec<(u16, columnar::NativeColumn)>>,
8497 crate::trace::QueryTrace,
8498 )> {
8499 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8500 self.query_columns_native(conditions, projection, snapshot)
8501 });
8502 Ok((result?, trace))
8503 }
8504
8505 #[allow(clippy::type_complexity)]
8508 pub fn query_columns_native_cached_traced(
8509 &mut self,
8510 conditions: &[crate::query::Condition],
8511 projection: Option<&[u16]>,
8512 snapshot: Snapshot,
8513 ) -> Result<(
8514 Option<Vec<(u16, columnar::NativeColumn)>>,
8515 crate::trace::QueryTrace,
8516 )> {
8517 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8518 self.query_columns_native_cached(conditions, projection, snapshot)
8519 });
8520 Ok((result?, trace))
8521 }
8522
8523 pub fn native_page_cursor_traced(
8525 &self,
8526 snapshot: Snapshot,
8527 projection: Vec<(u16, TypeId)>,
8528 conditions: &[crate::query::Condition],
8529 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8530 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8531 self.native_page_cursor(snapshot, projection, conditions)
8532 });
8533 Ok((result?, trace))
8534 }
8535
8536 pub fn native_multi_run_cursor_traced(
8538 &self,
8539 snapshot: Snapshot,
8540 projection: Vec<(u16, TypeId)>,
8541 conditions: &[crate::query::Condition],
8542 ) -> Result<(
8543 Option<crate::cursor::MultiRunCursor>,
8544 crate::trace::QueryTrace,
8545 )> {
8546 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8547 self.native_multi_run_cursor(snapshot, projection, conditions)
8548 });
8549 Ok((result?, trace))
8550 }
8551
8552 pub fn count_conditions_traced(
8554 &mut self,
8555 conditions: &[crate::query::Condition],
8556 snapshot: Snapshot,
8557 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8558 let (result, trace) =
8559 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8560 Ok((result?, trace))
8561 }
8562
8563 pub fn query_traced(
8565 &mut self,
8566 q: &crate::query::Query,
8567 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8568 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8569 Ok((result?, trace))
8570 }
8571
8572 pub fn query_columns_native(
8577 &mut self,
8578 conditions: &[crate::query::Condition],
8579 projection: Option<&[u16]>,
8580 snapshot: Snapshot,
8581 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8582 self.query_columns_native_inner(conditions, projection, snapshot, None)
8583 }
8584
8585 pub fn query_columns_native_with_control(
8586 &mut self,
8587 conditions: &[crate::query::Condition],
8588 projection: Option<&[u16]>,
8589 snapshot: Snapshot,
8590 control: &crate::ExecutionControl,
8591 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8592 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8593 }
8594
8595 fn query_columns_native_inner(
8596 &mut self,
8597 conditions: &[crate::query::Condition],
8598 projection: Option<&[u16]>,
8599 snapshot: Snapshot,
8600 control: Option<&crate::ExecutionControl>,
8601 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8602 use crate::query::Condition;
8603 execution_checkpoint(control, 0)?;
8604 if self.ttl.is_some() {
8607 return Ok(None);
8608 }
8609 if conditions.is_empty() {
8610 return Ok(None);
8611 }
8612 self.ensure_indexes_complete()?;
8613
8614 let served = |c: &Condition| {
8619 matches!(
8620 c,
8621 Condition::Pk(_)
8622 | Condition::BitmapEq { .. }
8623 | Condition::BitmapIn { .. }
8624 | Condition::BytesPrefix { .. }
8625 | Condition::FmContains { .. }
8626 | Condition::FmContainsAll { .. }
8627 | Condition::Ann { .. }
8628 | Condition::Range { .. }
8629 | Condition::RangeF64 { .. }
8630 | Condition::SparseMatch { .. }
8631 | Condition::MinHashSimilar { .. }
8632 | Condition::IsNull { .. }
8633 | Condition::IsNotNull { .. }
8634 )
8635 };
8636 if !conditions.iter().all(served) {
8637 return Ok(None);
8638 }
8639 let fast_path =
8640 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8641 crate::trace::QueryTrace::record(|t| {
8642 t.run_count = self.run_refs.len();
8643 t.memtable_rows = self.memtable.len();
8644 t.mutable_run_rows = self.mutable_run.len();
8645 t.conditions_pushed = conditions.len();
8646 t.learned_range_used = conditions.iter().any(|c| match c {
8647 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8648 self.learned_range.contains_key(column_id)
8649 }
8650 _ => false,
8651 });
8652 });
8653 let col_ids: Vec<u16> = projection
8655 .map(|p| p.to_vec())
8656 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8657 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8658 .iter()
8659 .map(|&cid| {
8660 let ty = self
8661 .schema
8662 .columns
8663 .iter()
8664 .find(|c| c.id == cid)
8665 .map(|c| c.ty.clone())
8666 .unwrap_or(TypeId::Bytes);
8667 (cid, ty)
8668 })
8669 .collect();
8670
8671 if fast_path {
8677 let needs_column = conditions.iter().any(|c| match c {
8680 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8681 Condition::RangeF64 { column_id, .. } => {
8682 !self.learned_range.contains_key(column_id)
8683 }
8684 _ => false,
8685 });
8686 let mut reader_opt: Option<RunReader> = if needs_column {
8687 Some(self.open_reader(self.run_refs[0].run_id)?)
8688 } else {
8689 None
8690 };
8691 let mut sets: Vec<RowIdSet> = Vec::new();
8692 for (index, c) in conditions.iter().enumerate() {
8693 execution_checkpoint(control, index)?;
8694 let s = match c {
8695 Condition::Range { column_id, lo, hi }
8696 if !self.learned_range.contains_key(column_id) =>
8697 {
8698 if reader_opt.is_none() {
8699 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8700 }
8701 reader_opt
8702 .as_mut()
8703 .expect("reader opened for range")
8704 .range_row_id_set_i64(*column_id, *lo, *hi)?
8705 }
8706 Condition::RangeF64 {
8707 column_id,
8708 lo,
8709 lo_inclusive,
8710 hi,
8711 hi_inclusive,
8712 } if !self.learned_range.contains_key(column_id) => {
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_f64(
8720 *column_id,
8721 *lo,
8722 *lo_inclusive,
8723 *hi,
8724 *hi_inclusive,
8725 )?
8726 }
8727 _ => self.resolve_condition(c, snapshot)?,
8728 };
8729 sets.push(s);
8730 }
8731 let candidates = RowIdSet::intersect_many(sets);
8732 crate::trace::QueryTrace::record(|t| {
8733 t.survivor_count = Some(candidates.len());
8734 });
8735 if candidates.is_empty() {
8736 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8737 .iter()
8738 .map(|&id| {
8739 (
8740 id,
8741 columnar::null_native(
8742 proj_pairs
8743 .iter()
8744 .find(|(c, _)| c == &id)
8745 .map(|(_, t)| t.clone())
8746 .unwrap_or(TypeId::Bytes),
8747 0,
8748 ),
8749 )
8750 })
8751 .collect();
8752 return Ok(Some(cols));
8753 }
8754 let mut reader = match reader_opt.take() {
8755 Some(r) => r,
8756 None => self.open_reader(self.run_refs[0].run_id)?,
8757 };
8758 let candidate_ids = candidates.into_sorted_vec();
8759 let (positions, fast_rid) = if let Some(positions) =
8760 reader.positions_for_row_ids_fast(&candidate_ids)
8761 {
8762 (positions, true)
8763 } else {
8764 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8765 match col {
8766 columnar::NativeColumn::Int64 { data, .. } => {
8767 let mut p = Vec::with_capacity(candidate_ids.len());
8768 for (index, rid) in candidate_ids.iter().enumerate() {
8769 execution_checkpoint(control, index)?;
8770 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8771 p.push(position);
8772 }
8773 }
8774 p.sort_unstable();
8775 (p, false)
8776 }
8777 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8778 }
8779 };
8780 crate::trace::QueryTrace::record(|t| {
8781 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8782 t.fast_row_id_map = fast_rid;
8783 });
8784 let mut cols = Vec::with_capacity(col_ids.len());
8785 for (index, cid) in col_ids.iter().enumerate() {
8786 execution_checkpoint(control, index)?;
8787 let col = reader.column_native(*cid)?;
8788 cols.push((*cid, col.gather(&positions)));
8789 }
8790 return Ok(Some(cols));
8791 }
8792
8793 if !self.run_refs.is_empty() {
8806 use crate::cursor::{
8807 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8808 };
8809 let remaining: usize;
8810 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
8811 let c = self
8812 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
8813 .expect("single-run cursor should build when run_refs.len() == 1");
8814 remaining = c.remaining_rows();
8815 Box::new(c)
8816 } else {
8817 let c = self
8818 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
8819 .expect("multi-run cursor should build when run_refs.len() >= 1");
8820 remaining = c.remaining_rows();
8821 Box::new(c)
8822 };
8823 crate::trace::QueryTrace::record(|t| {
8824 if t.survivor_count.is_none() {
8825 t.survivor_count = Some(remaining);
8826 }
8827 });
8828 let cols = match control {
8829 Some(control) => {
8830 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
8831 }
8832 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
8833 };
8834 return Ok(Some(cols));
8835 }
8836
8837 crate::trace::QueryTrace::record(|t| {
8842 t.scan_mode = crate::trace::ScanMode::Materialized;
8843 t.row_materialized = true;
8844 });
8845 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8846 for (index, c) in conditions.iter().enumerate() {
8847 execution_checkpoint(control, index)?;
8848 sets.push(self.resolve_condition(c, snapshot)?);
8849 }
8850 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8851 let rows = self.rows_for_rids(&rids, snapshot)?;
8852 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
8853 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
8854 execution_checkpoint(control, index)?;
8855 let vals: Vec<Value> = rows
8856 .iter()
8857 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8858 .collect();
8859 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
8860 }
8861 Ok(Some(cols))
8862 }
8863
8864 pub fn native_page_cursor(
8879 &self,
8880 snapshot: Snapshot,
8881 projection: Vec<(u16, TypeId)>,
8882 conditions: &[crate::query::Condition],
8883 ) -> Result<Option<NativePageCursor>> {
8884 use crate::cursor::build_page_plans;
8885 if self.ttl.is_some() {
8886 return Ok(None);
8887 }
8888 if !conditions.is_empty() && !self.indexes_complete {
8891 return Ok(None);
8892 }
8893 if self.run_refs.len() != 1 {
8894 return Ok(None);
8895 }
8896 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8897 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8898
8899 let overlay_rids: HashSet<u64> = {
8902 let mut s = HashSet::new();
8903 for row in self.memtable.visible_versions(snapshot.epoch) {
8904 s.insert(row.row_id.0);
8905 }
8906 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8907 s.insert(row.row_id.0);
8908 }
8909 s
8910 };
8911
8912 let survivors = if conditions.is_empty() {
8916 None
8917 } else {
8918 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
8919 };
8920
8921 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
8928 survivors.clone()
8929 } else if let Some(s) = &survivors {
8930 let mut run_set = s.clone();
8931 run_set.remove_many(overlay_rids.iter().copied());
8932 Some(run_set)
8933 } else {
8934 Some(RowIdSet::from_unsorted(
8935 rids.iter()
8936 .map(|&r| r as u64)
8937 .filter(|r| !overlay_rids.contains(r))
8938 .collect(),
8939 ))
8940 };
8941
8942 let overlay_rows = if overlay_rids.is_empty() {
8943 Vec::new()
8944 } else {
8945 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8946 self.overlay_visible_rows(snapshot, bound)
8947 };
8948
8949 let plans = if positions.is_empty() {
8951 Vec::new()
8952 } else {
8953 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
8954 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
8955 };
8956
8957 let overlay = if overlay_rows.is_empty() {
8959 None
8960 } else {
8961 let filtered =
8962 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8963 if filtered.is_empty() {
8964 None
8965 } else {
8966 Some(self.materialize_overlay(&filtered, &projection))
8967 }
8968 };
8969
8970 let overlay_row_count = overlay
8971 .as_ref()
8972 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8973 .unwrap_or(0);
8974 crate::trace::QueryTrace::record(|t| {
8975 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
8976 t.run_count = self.run_refs.len();
8977 t.memtable_rows = self.memtable.len();
8978 t.mutable_run_rows = self.mutable_run.len();
8979 t.overlay_rows = overlay_row_count;
8980 t.conditions_pushed = conditions.len();
8981 t.pages_decoded = plans
8982 .iter()
8983 .map(|p| p.positions.len())
8984 .sum::<usize>()
8985 .min(1);
8986 });
8987
8988 Ok(Some(NativePageCursor::new_with_overlay(
8989 reader, projection, plans, overlay,
8990 )))
8991 }
8992 #[allow(clippy::type_complexity)]
9002 pub fn native_multi_run_cursor(
9003 &self,
9004 snapshot: Snapshot,
9005 projection: Vec<(u16, TypeId)>,
9006 conditions: &[crate::query::Condition],
9007 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9008 use crate::cursor::{MultiRunCursor, RunStream};
9009 use crate::sorted_run::SYS_ROW_ID;
9010 use std::collections::{BinaryHeap, HashMap, HashSet};
9011 if self.ttl.is_some() {
9012 return Ok(None);
9013 }
9014 if !conditions.is_empty() && !self.indexes_complete {
9017 return Ok(None);
9018 }
9019 if self.run_refs.is_empty() {
9020 return Ok(None);
9021 }
9022
9023 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9025 Vec::with_capacity(self.run_refs.len());
9026 for rr in &self.run_refs {
9027 let mut reader = self.open_reader(rr.run_id)?;
9028 let (rids, eps, del) = reader.system_columns_native()?;
9029 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9030 run_meta.push((reader, rids, eps, del, page_rows));
9031 }
9032
9033 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9037 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9038 for i in 0..rids.len() {
9039 let rid = rids[i] as u64;
9040 let e = eps[i] as u64;
9041 if e > snapshot.epoch.0 {
9042 continue;
9043 }
9044 let is_del = del[i] != 0;
9045 best.entry(rid)
9046 .and_modify(|cur| {
9047 if e > cur.0 {
9048 *cur = (e, run_idx, i, is_del);
9049 }
9050 })
9051 .or_insert((e, run_idx, i, is_del));
9052 }
9053 }
9054
9055 let overlay_rids: HashSet<u64> = {
9057 let mut s = HashSet::new();
9058 for row in self.memtable.visible_versions(snapshot.epoch) {
9059 s.insert(row.row_id.0);
9060 }
9061 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9062 s.insert(row.row_id.0);
9063 }
9064 s
9065 };
9066
9067 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9069 None
9070 } else {
9071 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9072 for c in conditions {
9073 sets.push(self.resolve_condition(c, snapshot)?);
9074 }
9075 Some(RowIdSet::intersect_many(sets))
9076 };
9077
9078 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9082 for (rid, (_, run_idx, pos, deleted)) in &best {
9083 if *deleted {
9084 continue;
9085 }
9086 if overlay_rids.contains(rid) {
9087 continue;
9088 }
9089 if let Some(s) = &survivors {
9090 if !s.contains(*rid) {
9091 continue;
9092 }
9093 }
9094 per_run[*run_idx].push((*rid, *pos));
9095 }
9096 for v in per_run.iter_mut() {
9097 v.sort_unstable_by_key(|&(rid, _)| rid);
9098 }
9099
9100 let mut streams = Vec::with_capacity(run_meta.len());
9102 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9103 let mut total = 0usize;
9104 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9105 let mut starts = Vec::with_capacity(page_rows.len());
9106 let mut acc = 0usize;
9107 for &r in &page_rows {
9108 starts.push(acc);
9109 acc += r;
9110 }
9111 let mut survivors_vec: Vec<(u64, usize, usize)> =
9112 Vec::with_capacity(per_run[run_idx].len());
9113 for &(rid, pos) in &per_run[run_idx] {
9114 let page_seq = match starts.partition_point(|&s| s <= pos) {
9115 0 => continue,
9116 p => p - 1,
9117 };
9118 let within = pos - starts[page_seq];
9119 survivors_vec.push((rid, page_seq, within));
9120 }
9121 total += survivors_vec.len();
9122 if let Some(&(rid, _, _)) = survivors_vec.first() {
9123 heap.push(std::cmp::Reverse((rid, run_idx)));
9124 }
9125 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9126 }
9127
9128 let overlay_rows = if overlay_rids.is_empty() {
9130 Vec::new()
9131 } else {
9132 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9133 self.overlay_visible_rows(snapshot, bound)
9134 };
9135 let overlay = if overlay_rows.is_empty() {
9136 None
9137 } else {
9138 let filtered =
9139 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9140 if filtered.is_empty() {
9141 None
9142 } else {
9143 Some(self.materialize_overlay(&filtered, &projection))
9144 }
9145 };
9146
9147 let overlay_row_count = overlay
9148 .as_ref()
9149 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9150 .unwrap_or(0);
9151 crate::trace::QueryTrace::record(|t| {
9152 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9153 t.run_count = self.run_refs.len();
9154 t.memtable_rows = self.memtable.len();
9155 t.mutable_run_rows = self.mutable_run.len();
9156 t.overlay_rows = overlay_row_count;
9157 t.conditions_pushed = conditions.len();
9158 t.survivor_count = Some(total);
9159 });
9160
9161 Ok(Some(MultiRunCursor::new(
9162 streams, projection, heap, total, overlay,
9163 )))
9164 }
9165
9166 fn overlay_materialization_bound<'a>(
9178 conditions: &[crate::query::Condition],
9179 survivors: &'a Option<RowIdSet>,
9180 ) -> Option<&'a RowIdSet> {
9181 use crate::query::Condition;
9182 let has_range = conditions
9183 .iter()
9184 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9185 if has_range {
9186 None
9187 } else {
9188 survivors.as_ref()
9189 }
9190 }
9191
9192 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9204 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9205 let mut fold = |row: Row| {
9206 if let Some(b) = bound {
9207 if !b.contains(row.row_id.0) {
9208 return;
9209 }
9210 }
9211 best.entry(row.row_id.0)
9212 .and_modify(|(be, br)| {
9213 if row.committed_epoch > *be {
9214 *be = row.committed_epoch;
9215 *br = row.clone();
9216 }
9217 })
9218 .or_insert_with(|| (row.committed_epoch, row));
9219 };
9220 for row in self.memtable.visible_versions(snapshot.epoch) {
9221 fold(row);
9222 }
9223 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9224 fold(row);
9225 }
9226 let mut out: Vec<Row> = best
9227 .into_values()
9228 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9229 .collect();
9230 out.sort_by_key(|r| r.row_id);
9231 out
9232 }
9233
9234 fn filter_overlay_rows(
9242 &self,
9243 rows: Vec<Row>,
9244 conditions: &[crate::query::Condition],
9245 survivors: Option<&RowIdSet>,
9246 snapshot: Snapshot,
9247 ) -> Result<Vec<Row>> {
9248 if conditions.is_empty() {
9249 return Ok(rows);
9250 }
9251 use crate::query::Condition;
9252 let all_index_served = !conditions
9256 .iter()
9257 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9258 if all_index_served {
9259 return Ok(rows
9260 .into_iter()
9261 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9262 .collect());
9263 }
9264 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9267 for c in conditions {
9268 let s = match c {
9269 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9270 _ => self.resolve_condition(c, snapshot)?,
9271 };
9272 per_cond_sets.push(s);
9273 }
9274 Ok(rows
9275 .into_iter()
9276 .filter(|row| {
9277 conditions.iter().enumerate().all(|(i, c)| match c {
9278 Condition::Range { column_id, lo, hi } => {
9279 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9280 }
9281 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9282 match row.columns.get(column_id) {
9283 Some(Value::Float64(v)) => {
9284 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9285 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9286 lo_ok && hi_ok
9287 }
9288 _ => false,
9289 }
9290 }
9291 _ => per_cond_sets[i].contains(row.row_id.0),
9292 })
9293 })
9294 .collect())
9295 }
9296
9297 fn materialize_overlay(
9300 &self,
9301 rows: &[Row],
9302 projection: &[(u16, TypeId)],
9303 ) -> Vec<columnar::NativeColumn> {
9304 if projection.is_empty() {
9305 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9306 }
9307 let mut cols = Vec::with_capacity(projection.len());
9308 for (cid, ty) in projection {
9309 let vals: Vec<Value> = rows
9310 .iter()
9311 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9312 .collect();
9313 cols.push(columnar::values_to_native(ty.clone(), &vals));
9314 }
9315 cols
9316 }
9317
9318 fn resolve_survivor_rids(
9323 &self,
9324 conditions: &[crate::query::Condition],
9325 reader: &mut RunReader,
9326 snapshot: Snapshot,
9327 ) -> Result<RowIdSet> {
9328 use crate::query::Condition;
9329 let mut sets: Vec<RowIdSet> = Vec::new();
9330 for c in conditions {
9331 self.validate_condition(c)?;
9332 let s: RowIdSet = match c {
9333 Condition::Pk(key) => {
9334 let lookup = self
9335 .schema
9336 .primary_key()
9337 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9338 .unwrap_or_else(|| key.clone());
9339 self.hot
9340 .get(&lookup)
9341 .map(|r| RowIdSet::one(r.0))
9342 .unwrap_or_else(RowIdSet::empty)
9343 }
9344 Condition::BitmapEq { column_id, value } => {
9345 let lookup = self.index_lookup_key_bytes(*column_id, value);
9346 self.bitmap
9347 .get(column_id)
9348 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9349 .unwrap_or_else(RowIdSet::empty)
9350 }
9351 Condition::BitmapIn { column_id, values } => {
9352 let bm = self.bitmap.get(column_id);
9353 let mut acc = roaring::RoaringBitmap::new();
9354 if let Some(b) = bm {
9355 for v in values {
9356 let lookup = self.index_lookup_key_bytes(*column_id, v);
9357 acc |= b.get(&lookup);
9358 }
9359 }
9360 RowIdSet::from_roaring(acc)
9361 }
9362 Condition::BytesPrefix { column_id, prefix } => {
9363 if let Some(b) = self.bitmap.get(column_id) {
9364 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9365 let mut acc = roaring::RoaringBitmap::new();
9366 for key in b.keys() {
9367 if key.starts_with(&lookup_prefix) {
9368 acc |= b.get(&key);
9369 }
9370 }
9371 RowIdSet::from_roaring(acc)
9372 } else {
9373 RowIdSet::empty()
9374 }
9375 }
9376 Condition::FmContains { column_id, pattern } => self
9377 .fm
9378 .get(column_id)
9379 .map(|f| {
9380 RowIdSet::from_unsorted(
9381 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9382 )
9383 })
9384 .unwrap_or_else(RowIdSet::empty),
9385 Condition::FmContainsAll {
9386 column_id,
9387 patterns,
9388 } => {
9389 if let Some(f) = self.fm.get(column_id) {
9390 let sets: Vec<RowIdSet> = patterns
9391 .iter()
9392 .map(|pat| {
9393 RowIdSet::from_unsorted(
9394 f.locate(pat).into_iter().map(|r| r.0).collect(),
9395 )
9396 })
9397 .collect();
9398 RowIdSet::intersect_many(sets)
9399 } else {
9400 RowIdSet::empty()
9401 }
9402 }
9403 Condition::Ann {
9404 column_id,
9405 query,
9406 k,
9407 } => RowIdSet::from_unsorted(
9408 self.retrieve_filtered(
9409 &crate::query::Retriever::Ann {
9410 column_id: *column_id,
9411 query: query.clone(),
9412 k: *k,
9413 },
9414 snapshot,
9415 None,
9416 None,
9417 None,
9418 None,
9419 )?
9420 .into_iter()
9421 .map(|hit| hit.row_id.0)
9422 .collect(),
9423 ),
9424 Condition::SparseMatch {
9425 column_id,
9426 query,
9427 k,
9428 } => RowIdSet::from_unsorted(
9429 self.retrieve_filtered(
9430 &crate::query::Retriever::Sparse {
9431 column_id: *column_id,
9432 query: query.clone(),
9433 k: *k,
9434 },
9435 snapshot,
9436 None,
9437 None,
9438 None,
9439 None,
9440 )?
9441 .into_iter()
9442 .map(|hit| hit.row_id.0)
9443 .collect(),
9444 ),
9445 Condition::MinHashSimilar {
9446 column_id,
9447 query,
9448 k,
9449 } => match self.minhash.get(column_id) {
9450 Some(index) => {
9451 let candidates = index.candidate_row_ids(query);
9452 let eligible =
9453 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9454 RowIdSet::from_unsorted(
9455 index
9456 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9457 .into_iter()
9458 .map(|(row_id, _)| row_id.0)
9459 .collect(),
9460 )
9461 }
9462 None => RowIdSet::empty(),
9463 },
9464 Condition::Range { column_id, lo, hi } => {
9465 if let Some(li) = self.learned_range.get(column_id) {
9466 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9467 } else {
9468 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9469 }
9470 }
9471 Condition::RangeF64 {
9472 column_id,
9473 lo,
9474 lo_inclusive,
9475 hi,
9476 hi_inclusive,
9477 } => {
9478 if let Some(li) = self.learned_range.get(column_id) {
9479 RowIdSet::from_unsorted(
9480 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9481 .into_iter()
9482 .collect(),
9483 )
9484 } else {
9485 reader.range_row_id_set_f64(
9486 *column_id,
9487 *lo,
9488 *lo_inclusive,
9489 *hi,
9490 *hi_inclusive,
9491 )?
9492 }
9493 }
9494 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9495 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9496 };
9497 sets.push(s);
9498 }
9499 Ok(RowIdSet::intersect_many(sets))
9500 }
9501
9502 pub fn scan_cursor(
9523 &self,
9524 snapshot: Snapshot,
9525 projection: Vec<(u16, TypeId)>,
9526 conditions: &[crate::query::Condition],
9527 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9528 if self.ttl.is_some() {
9529 return Ok(None);
9530 }
9531 if !conditions.is_empty() && !self.indexes_complete {
9537 return Ok(None);
9538 }
9539 if self.run_refs.len() == 1 {
9540 Ok(self
9541 .native_page_cursor(snapshot, projection, conditions)?
9542 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9543 } else {
9544 Ok(self
9545 .native_multi_run_cursor(snapshot, projection, conditions)?
9546 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9547 }
9548 }
9549
9550 pub fn aggregate_native(
9564 &self,
9565 snapshot: Snapshot,
9566 column: Option<u16>,
9567 conditions: &[crate::query::Condition],
9568 agg: NativeAgg,
9569 ) -> Result<Option<NativeAggResult>> {
9570 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9571 }
9572
9573 pub fn aggregate_native_with_control(
9574 &self,
9575 snapshot: Snapshot,
9576 column: Option<u16>,
9577 conditions: &[crate::query::Condition],
9578 agg: NativeAgg,
9579 control: &crate::ExecutionControl,
9580 ) -> Result<Option<NativeAggResult>> {
9581 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9582 }
9583
9584 fn aggregate_native_inner(
9585 &self,
9586 snapshot: Snapshot,
9587 column: Option<u16>,
9588 conditions: &[crate::query::Condition],
9589 agg: NativeAgg,
9590 control: Option<&crate::ExecutionControl>,
9591 ) -> Result<Option<NativeAggResult>> {
9592 execution_checkpoint(control, 0)?;
9593 if self.ttl.is_some() {
9594 return Ok(None);
9595 }
9596 if self.run_refs.len() == 1 && conditions.is_empty() {
9598 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9599 return Ok(Some(res));
9600 }
9601 }
9602 if matches!(agg, NativeAgg::Count) && column.is_none() {
9606 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
9607 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
9608 }
9609 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9610 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
9611 }
9612 let cid = match column {
9615 Some(c) => c,
9616 None => return Ok(None),
9617 };
9618 let ty = self.column_type(cid);
9619 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
9620 execution_checkpoint(control, 0)?;
9621 return match ty {
9622 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9623 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9624 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9625 }
9626 TypeId::Float64 => {
9627 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9628 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9629 }
9630 _ => Ok(None),
9631 };
9632 }
9633 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9635 execution_checkpoint(control, 0)?;
9636 match ty {
9637 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9638 let mut count = 0u64;
9639 let mut sum = 0i128;
9640 let mut mn = i64::MAX;
9641 let mut mx = i64::MIN;
9642 for row in &rows {
9643 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
9644 count += 1;
9645 sum += i128::from(*v);
9646 mn = mn.min(*v);
9647 mx = mx.max(*v);
9648 }
9649 }
9650 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9651 }
9652 TypeId::Float64 => {
9653 let mut count = 0u64;
9654 let mut sum = 0.0f64;
9655 let mut mn = f64::INFINITY;
9656 let mut mx = f64::NEG_INFINITY;
9657 for row in &rows {
9658 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
9659 count += 1;
9660 sum += *v;
9661 mn = mn.min(*v);
9662 mx = mx.max(*v);
9663 }
9664 }
9665 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9666 }
9667 _ => Ok(None),
9668 }
9669 }
9670
9671 fn visible_rows_filtered(
9673 &self,
9674 snapshot: Snapshot,
9675 conditions: &[crate::query::Condition],
9676 control: Option<&crate::ExecutionControl>,
9677 ) -> Result<Vec<Row>> {
9678 let rows = if let Some(control) = control {
9679 self.visible_rows_controlled(snapshot, control)?
9680 } else {
9681 self.visible_rows(snapshot)?
9682 };
9683 if conditions.is_empty() {
9684 return Ok(rows);
9685 }
9686 Ok(rows
9687 .into_iter()
9688 .filter(|row| {
9689 conditions
9690 .iter()
9691 .all(|cond| condition_matches_row(cond, row, &self.schema))
9692 })
9693 .collect())
9694 }
9695
9696 fn aggregate_from_stats(
9704 &self,
9705 snapshot: Snapshot,
9706 column: Option<u16>,
9707 agg: NativeAgg,
9708 ) -> Result<Option<NativeAggResult>> {
9709 let cid = match (agg, column) {
9710 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9711 _ => return Ok(None), };
9713 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9714 return Ok(None);
9715 };
9716 let Some(cs) = stats.get(&cid) else {
9717 return Ok(None);
9718 };
9719 match agg {
9720 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9722 self.live_count.saturating_sub(cs.null_count),
9723 ))),
9724 NativeAgg::Min | NativeAgg::Max => {
9725 let bound = if agg == NativeAgg::Min {
9726 &cs.min
9727 } else {
9728 &cs.max
9729 };
9730 match bound {
9731 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9732 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9733 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9738 None => Ok(None),
9739 }
9740 }
9741 _ => Ok(None),
9742 }
9743 }
9744
9745 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9754 if self.ttl.is_some() {
9755 return Ok(None);
9756 }
9757 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9758 return Ok(None);
9759 }
9760 self.ensure_indexes_complete()?;
9763 let reader = self.open_reader(self.run_refs[0].run_id)?;
9764 if self.live_count != reader.row_count() as u64 {
9765 return Ok(None);
9766 }
9767 let Some(bm) = self.bitmap.get(&column_id) else {
9768 return Ok(None); };
9770 let mut distinct = bm.value_count() as u64;
9771 if !bm.get(&Value::Null.encode_key()).is_empty() {
9774 distinct = distinct.saturating_sub(1);
9775 }
9776 Ok(Some(distinct))
9777 }
9778
9779 pub fn aggregate_incremental(
9791 &mut self,
9792 cache_key: u64,
9793 conditions: &[crate::query::Condition],
9794 column: Option<u16>,
9795 agg: NativeAgg,
9796 ) -> Result<IncrementalAggResult> {
9797 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9798 }
9799
9800 pub fn aggregate_incremental_with_control(
9801 &mut self,
9802 cache_key: u64,
9803 conditions: &[crate::query::Condition],
9804 column: Option<u16>,
9805 agg: NativeAgg,
9806 control: &crate::ExecutionControl,
9807 ) -> Result<IncrementalAggResult> {
9808 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9809 }
9810
9811 fn aggregate_incremental_inner(
9812 &mut self,
9813 cache_key: u64,
9814 conditions: &[crate::query::Condition],
9815 column: Option<u16>,
9816 agg: NativeAgg,
9817 control: Option<&crate::ExecutionControl>,
9818 ) -> Result<IncrementalAggResult> {
9819 execution_checkpoint(control, 0)?;
9820 let snap = self.snapshot();
9821 let cur_wm = self.allocator.current().0;
9822 let cur_epoch = snap.epoch.0;
9823 let incremental_ok = self.ttl.is_none()
9830 && !self.had_deletes
9831 && self.memtable.is_empty()
9832 && self.mutable_run.is_empty();
9833
9834 if incremental_ok {
9837 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
9838 if cached.epoch == cur_epoch {
9839 return Ok(IncrementalAggResult {
9840 state: cached.state,
9841 incremental: true,
9842 delta_rows: 0,
9843 });
9844 }
9845 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
9846 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
9847 let mut delta_rids = Vec::with_capacity(delta_len);
9848 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
9849 execution_checkpoint(control, index)?;
9850 delta_rids.push(row_id);
9851 }
9852 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
9853 execution_checkpoint(control, 0)?;
9854 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9855 let delta_state = agg_state_from_rows(
9856 &delta_rows,
9857 conditions,
9858 &index_sets,
9859 column,
9860 agg,
9861 &self.schema,
9862 control,
9863 )?;
9864 let merged = cached.state.merge(delta_state);
9865 let delta_n = delta_rids.len() as u64;
9866 Arc::make_mut(&mut self.agg_cache).insert(
9867 cache_key,
9868 CachedAgg {
9869 state: merged.clone(),
9870 watermark: cur_wm,
9871 epoch: cur_epoch,
9872 },
9873 );
9874 return Ok(IncrementalAggResult {
9875 state: merged,
9876 incremental: true,
9877 delta_rows: delta_n,
9878 });
9879 }
9880 }
9881 }
9882
9883 let cursor_ok =
9888 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9889 let state = if cursor_ok && agg != NativeAgg::Avg {
9890 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
9891 Some(result) => {
9892 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
9893 }
9894 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
9895 }
9896 } else {
9897 self.agg_state_full_scan(conditions, column, agg, snap, control)?
9898 };
9899 if incremental_ok {
9901 Arc::make_mut(&mut self.agg_cache).insert(
9902 cache_key,
9903 CachedAgg {
9904 state: state.clone(),
9905 watermark: cur_wm,
9906 epoch: cur_epoch,
9907 },
9908 );
9909 }
9910 Ok(IncrementalAggResult {
9911 state,
9912 incremental: false,
9913 delta_rows: 0,
9914 })
9915 }
9916
9917 fn agg_state_full_scan(
9920 &self,
9921 conditions: &[crate::query::Condition],
9922 column: Option<u16>,
9923 agg: NativeAgg,
9924 snap: Snapshot,
9925 control: Option<&crate::ExecutionControl>,
9926 ) -> Result<AggState> {
9927 execution_checkpoint(control, 0)?;
9928 let rows = self.visible_rows(snap)?;
9929 execution_checkpoint(control, 0)?;
9930 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9931 agg_state_from_rows(
9932 &rows,
9933 conditions,
9934 &index_sets,
9935 column,
9936 agg,
9937 &self.schema,
9938 control,
9939 )
9940 }
9941
9942 fn resolve_index_conditions(
9945 &self,
9946 conditions: &[crate::query::Condition],
9947 snapshot: Snapshot,
9948 ) -> Result<Vec<RowIdSet>> {
9949 use crate::query::Condition;
9950 let mut sets = Vec::new();
9951 for c in conditions {
9952 if matches!(
9953 c,
9954 Condition::Ann { .. }
9955 | Condition::SparseMatch { .. }
9956 | Condition::MinHashSimilar { .. }
9957 ) {
9958 sets.push(self.resolve_condition(c, snapshot)?);
9959 }
9960 }
9961 Ok(sets)
9962 }
9963
9964 fn column_type(&self, cid: u16) -> TypeId {
9965 self.schema
9966 .columns
9967 .iter()
9968 .find(|c| c.id == cid)
9969 .map(|c| c.ty.clone())
9970 .unwrap_or(TypeId::Bytes)
9971 }
9972
9973 pub fn approx_aggregate(
9982 &mut self,
9983 conditions: &[crate::query::Condition],
9984 column: Option<u16>,
9985 agg: ApproxAgg,
9986 z: f64,
9987 ) -> Result<Option<ApproxResult>> {
9988 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
9989 }
9990
9991 pub fn approx_aggregate_with_candidate_authorization(
9994 &mut self,
9995 conditions: &[crate::query::Condition],
9996 column: Option<u16>,
9997 agg: ApproxAgg,
9998 z: f64,
9999 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10000 ) -> Result<Option<ApproxResult>> {
10001 use crate::query::Condition;
10002 self.ensure_reservoir_complete()?;
10003 let snapshot = self.snapshot();
10004 let n_pop = self.count();
10005 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10006 if sample_rids.is_empty() {
10007 return Ok(None);
10008 }
10009 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10011 let s = live_sample.len();
10012 if s == 0 {
10013 return Ok(None);
10014 }
10015 let authorized = authorization
10016 .map(|authorization| {
10017 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10018 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10019 })
10020 .transpose()?;
10021
10022 let mut index_sets: Vec<RowIdSet> = Vec::new();
10025 for c in conditions {
10026 if matches!(
10027 c,
10028 Condition::Ann { .. }
10029 | Condition::SparseMatch { .. }
10030 | Condition::MinHashSimilar { .. }
10031 ) {
10032 index_sets.push(self.resolve_condition(c, snapshot)?);
10033 }
10034 }
10035
10036 let cid = match (agg, column) {
10038 (ApproxAgg::Count, _) => None,
10039 (_, Some(c)) => Some(c),
10040 _ => return Ok(None),
10041 };
10042 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10043 for r in &live_sample {
10044 if authorized
10045 .as_ref()
10046 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10047 {
10048 continue;
10049 }
10050 if !conditions
10052 .iter()
10053 .all(|c| condition_matches_row(c, r, &self.schema))
10054 {
10055 continue;
10056 }
10057 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10059 continue;
10060 }
10061 if let Some(cid) = cid {
10062 let mut cells = r
10063 .columns
10064 .get(&cid)
10065 .cloned()
10066 .map(|value| vec![(cid, value)])
10067 .unwrap_or_default();
10068 if let Some(authorization) = authorization {
10069 authorization.security.apply_masks_to_cells(
10070 authorization.table,
10071 &mut cells,
10072 authorization.principal,
10073 );
10074 }
10075 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10076 passing_vals.push(v);
10077 } } else {
10079 passing_vals.push(0.0); }
10081 }
10082 let m = passing_vals.len();
10083
10084 let (point, half) = match agg {
10085 ApproxAgg::Count => {
10086 let p = m as f64 / s as f64;
10088 let point = n_pop as f64 * p;
10089 let var = if s > 1 {
10090 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10091 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10092 } else {
10093 0.0
10094 };
10095 (point, z * var.sqrt())
10096 }
10097 ApproxAgg::Sum => {
10098 let y: Vec<f64> = live_sample
10100 .iter()
10101 .map(|r| {
10102 let passes_row = authorized
10103 .as_ref()
10104 .is_none_or(|authorized| authorized.contains(&r.row_id))
10105 && conditions
10106 .iter()
10107 .all(|c| condition_matches_row(c, r, &self.schema))
10108 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10109 if passes_row {
10110 cid.and_then(|cid| {
10111 let mut cells = r
10112 .columns
10113 .get(&cid)
10114 .cloned()
10115 .map(|value| vec![(cid, value)])
10116 .unwrap_or_default();
10117 if let Some(authorization) = authorization {
10118 authorization.security.apply_masks_to_cells(
10119 authorization.table,
10120 &mut cells,
10121 authorization.principal,
10122 );
10123 }
10124 as_f64(cells.first().map(|(_, value)| value))
10125 })
10126 .unwrap_or(0.0)
10127 } else {
10128 0.0
10129 }
10130 })
10131 .collect();
10132 let mean_y = y.iter().sum::<f64>() / s as f64;
10133 let point = n_pop as f64 * mean_y;
10134 let var = if s > 1 {
10135 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10136 let var_y = ss / (s - 1) as f64;
10137 n_pop as f64 * n_pop as f64 * var_y / s as f64
10138 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10139 } else {
10140 0.0
10141 };
10142 (point, z * var.sqrt())
10143 }
10144 ApproxAgg::Avg => {
10145 if m == 0 {
10146 return Ok(Some(ApproxResult {
10147 point: 0.0,
10148 ci_low: 0.0,
10149 ci_high: 0.0,
10150 n_population: n_pop,
10151 n_sample_live: s,
10152 n_passing: 0,
10153 }));
10154 }
10155 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10156 let half = if m > 1 {
10157 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10158 let sd = (ss / (m - 1) as f64).sqrt();
10159 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10160 z * sd / (m as f64).sqrt() * fpc.sqrt()
10161 } else {
10162 0.0
10163 };
10164 (mean, half)
10165 }
10166 };
10167
10168 Ok(Some(ApproxResult {
10169 point,
10170 ci_low: point - half,
10171 ci_high: point + half,
10172 n_population: n_pop,
10173 n_sample_live: s,
10174 n_passing: m,
10175 }))
10176 }
10177
10178 pub fn exact_column_stats(
10186 &self,
10187 _snapshot: Snapshot,
10188 projection: &[u16],
10189 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10190 if self.ttl.is_some()
10191 || !(self.memtable.is_empty()
10192 && self.mutable_run.is_empty()
10193 && self.run_refs.len() == 1)
10194 {
10195 return Ok(None);
10196 }
10197 let reader = self.open_reader(self.run_refs[0].run_id)?;
10198 if self.live_count != reader.row_count() as u64 {
10199 return Ok(None);
10200 }
10201 let mut out = HashMap::new();
10202 for &cid in projection {
10203 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10204 Some(c) => c,
10205 None => continue,
10206 };
10207 let Some(stats) = reader.column_page_stats(cid) else {
10209 out.insert(
10210 cid,
10211 ColumnStat {
10212 min: None,
10213 max: None,
10214 null_count: self.live_count,
10215 },
10216 );
10217 continue;
10218 };
10219 let stat = match cdef.ty {
10220 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10221 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10222 min: mn.map(Value::Int64),
10223 max: mx.map(Value::Int64),
10224 null_count: n,
10225 })
10226 }
10227 TypeId::Float64 => {
10228 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10229 min: mn.map(Value::Float64),
10230 max: mx.map(Value::Float64),
10231 null_count: n,
10232 })
10233 }
10234 _ => None,
10235 };
10236 if let Some(s) = stat {
10237 out.insert(cid, s);
10238 }
10239 }
10240 Ok(Some(out))
10241 }
10242
10243 pub fn dir(&self) -> &Path {
10244 &self.dir
10245 }
10246
10247 pub fn schema(&self) -> &Schema {
10248 &self.schema
10249 }
10250
10251 pub(crate) fn set_catalog_name(&mut self, name: String) {
10252 self.name = name;
10253 }
10254
10255 pub(crate) fn prepare_alter_column(
10256 &mut self,
10257 column_name: &str,
10258 change: &AlterColumn,
10259 ) -> Result<(ColumnDef, Option<Schema>)> {
10260 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10261 return Err(MongrelError::InvalidArgument(
10262 "ALTER COLUMN requires committing staged writes first".into(),
10263 ));
10264 }
10265 let old = self
10266 .schema
10267 .columns
10268 .iter()
10269 .find(|c| c.name == column_name)
10270 .cloned()
10271 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10272 let mut next = old.clone();
10273
10274 if let Some(name) = &change.name {
10275 let trimmed = name.trim();
10276 if trimmed.is_empty() {
10277 return Err(MongrelError::InvalidArgument(
10278 "ALTER COLUMN name must not be empty".into(),
10279 ));
10280 }
10281 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10282 return Err(MongrelError::Schema(format!(
10283 "column {trimmed} already exists"
10284 )));
10285 }
10286 next.name = trimmed.to_string();
10287 }
10288
10289 if let Some(ty) = &change.ty {
10290 next.ty = ty.clone();
10291 }
10292 if let Some(flags) = change.flags {
10293 validate_alter_column_flags(old.flags, flags)?;
10294 next.flags = flags;
10295 }
10296
10297 if let Some(default_change) = &change.default_value {
10298 next.default_value = default_change.clone();
10299 }
10300 if let Some(source_change) = &change.embedding_source {
10301 next.embedding_source = source_change.clone();
10302 }
10303
10304 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10305 if old.flags.contains(ColumnFlags::NULLABLE)
10306 && !next.flags.contains(ColumnFlags::NULLABLE)
10307 && self.column_has_nulls(old.id)?
10308 {
10309 return Err(MongrelError::InvalidArgument(format!(
10310 "column '{}' contains NULL values",
10311 old.name
10312 )));
10313 }
10314 if next == old {
10315 return Ok((next, None));
10316 }
10317 let mut schema = self.schema.clone();
10318 let index = schema
10319 .columns
10320 .iter()
10321 .position(|column| column.id == next.id)
10322 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10323 schema.columns[index] = next.clone();
10324 schema.schema_id = schema
10325 .schema_id
10326 .checked_add(1)
10327 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10328 schema.validate_auto_increment()?;
10329 schema.validate_defaults()?;
10330 Ok((next, Some(schema)))
10331 }
10332
10333 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10334 self.schema = schema;
10335 self.auto_inc = resolve_auto_inc(&self.schema);
10336 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10337 self.clear_result_cache();
10338 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10339 }
10340
10341 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10342 checkpoint_current_schema(self)
10343 }
10344
10345 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10346 self.ensure_writable()?;
10347 let previous_schema = self.schema.clone();
10348 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10349 if let Some(schema) = schema {
10350 self.apply_altered_schema_prepared(schema);
10351 self.checkpoint_standalone_schema_change(previous_schema)?;
10352 }
10353 Ok(column)
10354 }
10355
10356 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10357 if self.live_count == 0 {
10358 return Ok(false);
10359 }
10360 let snap = self.snapshot();
10361 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10362 Ok(columns
10363 .first()
10364 .map(|(_, col)| col.null_count(col.len()) != 0)
10365 .unwrap_or(true))
10366 }
10367
10368 fn has_stored_versions(&self) -> bool {
10369 !self.memtable.is_empty()
10370 || !self.mutable_run.is_empty()
10371 || self.run_refs.iter().any(|r| r.row_count > 0)
10372 || !self.retiring.is_empty()
10373 }
10374
10375 pub fn add_column(
10380 &mut self,
10381 name: &str,
10382 ty: TypeId,
10383 flags: ColumnFlags,
10384 default_value: Option<crate::schema::DefaultExpr>,
10385 ) -> Result<u16> {
10386 self.add_column_with_id(name, ty, flags, default_value, None)
10387 }
10388
10389 pub fn add_column_with_id(
10390 &mut self,
10391 name: &str,
10392 ty: TypeId,
10393 flags: ColumnFlags,
10394 default_value: Option<crate::schema::DefaultExpr>,
10395 requested_id: Option<u16>,
10396 ) -> Result<u16> {
10397 self.ensure_writable()?;
10398 if self.schema.columns.iter().any(|c| c.name == name) {
10399 return Err(MongrelError::Schema(format!(
10400 "column {name} already exists"
10401 )));
10402 }
10403 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10404 if self.schema.columns.iter().any(|c| c.id == id) {
10405 return Err(MongrelError::Schema(format!(
10406 "column id {id} already exists"
10407 )));
10408 }
10409 id
10410 } else {
10411 self.schema
10412 .columns
10413 .iter()
10414 .map(|c| c.id)
10415 .max()
10416 .unwrap_or(0)
10417 .checked_add(1)
10418 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10419 };
10420 let previous_schema = self.schema.clone();
10421 let mut next_schema = previous_schema.clone();
10422 next_schema.columns.push(ColumnDef {
10423 id,
10424 name: name.to_string(),
10425 ty,
10426 flags,
10427 default_value,
10428 embedding_source: None,
10429 });
10430 next_schema.schema_id = next_schema
10431 .schema_id
10432 .checked_add(1)
10433 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10434 next_schema.validate_auto_increment()?;
10435 next_schema.validate_defaults()?;
10436 self.apply_altered_schema_prepared(next_schema);
10437 self.checkpoint_standalone_schema_change(previous_schema)?;
10438 Ok(id)
10439 }
10440
10441 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10450 self.ensure_writable()?;
10451 let cid = self
10452 .schema
10453 .columns
10454 .iter()
10455 .find(|c| c.name == column_name)
10456 .map(|c| c.id)
10457 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10458 let ty = self
10459 .schema
10460 .columns
10461 .iter()
10462 .find(|c| c.id == cid)
10463 .map(|c| c.ty.clone())
10464 .unwrap_or(TypeId::Int64);
10465 if !matches!(
10466 ty,
10467 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10468 ) {
10469 return Err(MongrelError::Schema(format!(
10470 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10471 )));
10472 }
10473 if self
10474 .schema
10475 .indexes
10476 .iter()
10477 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10478 {
10479 return Ok(()); }
10481 let previous_schema = self.schema.clone();
10482 let previous_learned_range = Arc::clone(&self.learned_range);
10483 let mut next_schema = previous_schema.clone();
10484 next_schema.indexes.push(IndexDef {
10485 name: format!("{}_learned_range", column_name),
10486 column_id: cid,
10487 kind: IndexKind::LearnedRange,
10488 predicate: None,
10489 options: Default::default(),
10490 });
10491 next_schema.schema_id = next_schema
10492 .schema_id
10493 .checked_add(1)
10494 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10495 self.apply_altered_schema_prepared(next_schema);
10496 if let Err(error) = self.build_learned_ranges() {
10497 self.apply_altered_schema_prepared(previous_schema);
10498 self.learned_range = previous_learned_range;
10499 return Err(error);
10500 }
10501 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10502 if !matches!(
10503 &error,
10504 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10505 ) {
10506 self.learned_range = previous_learned_range;
10507 }
10508 return Err(error);
10509 }
10510 Ok(())
10511 }
10512
10513 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10514 let mut schema_published = false;
10515 let schema_result = match self._root_guard.as_deref() {
10516 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10517 schema_published = true;
10518 }),
10519 None => write_schema_with_after(&self.dir, &self.schema, || {
10520 schema_published = true;
10521 }),
10522 };
10523 if schema_result.is_err() && !schema_published {
10524 self.apply_altered_schema_prepared(previous_schema);
10525 return schema_result;
10526 }
10527
10528 let manifest_result = self.persist_manifest(self.current_epoch());
10529 match (schema_result, manifest_result) {
10530 (_, Ok(())) => Ok(()),
10531 (Ok(()), Err(error)) => {
10532 self.poison_after_maintenance_publish_failure();
10533 Err(MongrelError::DurableCommit {
10534 epoch: self.current_epoch().0,
10535 message: format!(
10536 "schema is durable but matching manifest publication failed: {error}"
10537 ),
10538 })
10539 }
10540 (Err(schema_error), Err(manifest_error)) => {
10541 self.poison_after_maintenance_publish_failure();
10542 Err(MongrelError::CommitOutcomeUnknown {
10543 epoch: self.current_epoch().0,
10544 message: format!(
10545 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10546 ),
10547 })
10548 }
10549 }
10550 }
10551
10552 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10555 self.sync_byte_threshold = threshold;
10556 if let WalSink::Private(w) = &mut self.wal {
10557 w.set_sync_byte_threshold(threshold);
10558 }
10559 }
10560
10561 pub fn page_cache_flush(&self) {
10565 self.page_cache.flush_to_disk();
10566 }
10567
10568 pub fn page_cache_len(&self) -> usize {
10570 self.page_cache.len()
10571 }
10572
10573 pub fn decoded_cache_len(&self) -> usize {
10576 self.decoded_cache.len()
10577 }
10578
10579 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10582 self.memtable.drain_sorted()
10583 }
10584
10585 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10586 self.runs_dir().join(format!("r-{run_id}.sr"))
10587 }
10588
10589 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10590 match self.runs_root.as_deref() {
10591 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10592 None => Ok(None),
10593 }
10594 }
10595
10596 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10597 match self.runs_root.as_deref() {
10598 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10599 None => Ok(None),
10600 }
10601 }
10602
10603 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10604 match self.runs_root.as_deref() {
10605 Some(root) => match root.remove_file(name) {
10606 Ok(()) => Ok(()),
10607 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10608 Err(error) => Err(error.into()),
10609 },
10610 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10611 Ok(()) => Ok(()),
10612 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10613 Err(error) => Err(error.into()),
10614 },
10615 }
10616 }
10617
10618 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10619 match self.runs_root.as_deref() {
10620 Some(root) => root
10621 .rename_file_new(source, destination)
10622 .map_err(Into::into),
10623 None => crate::durable_file::rename(
10624 &self.runs_dir().join(source),
10625 &self.runs_dir().join(destination),
10626 )
10627 .map_err(Into::into),
10628 }
10629 }
10630
10631 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10632 self.run_refs.iter().map(|run| run.run_id)
10633 }
10634
10635 pub(crate) fn table_dir(&self) -> &Path {
10636 &self.dir
10637 }
10638
10639 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10640 &self.schema
10641 }
10642
10643 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10644 let id = self.next_run_id;
10645 self.next_run_id = self
10646 .next_run_id
10647 .checked_add(1)
10648 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10649 Ok(id)
10650 }
10651
10652 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10653 self.run_refs.push(run_ref);
10654 }
10655
10656 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10666 self.retiring.push(crate::manifest::RetiredRun {
10667 run_id,
10668 retire_epoch,
10669 });
10670 }
10671
10672 pub(crate) fn reap_retiring(
10676 &mut self,
10677 min_active: Epoch,
10678 backup_pinned: &std::collections::HashSet<u128>,
10679 ) -> Result<usize> {
10680 if self.retiring.is_empty() {
10681 return Ok(0);
10682 }
10683 let mut reaped = 0;
10684 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
10685 for r in std::mem::take(&mut self.retiring) {
10691 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
10692 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
10693 reaped += 1;
10694 } else {
10695 kept.push(r);
10696 }
10697 }
10698 self.retiring = kept;
10699 if reaped > 0 {
10700 self.persist_manifest(self.current_epoch())?;
10701 }
10702 Ok(reaped)
10703 }
10704
10705 pub(crate) fn has_reapable_retiring(
10706 &self,
10707 min_active: Epoch,
10708 backup_pinned: &std::collections::HashSet<u128>,
10709 ) -> bool {
10710 self.retiring
10711 .iter()
10712 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
10713 }
10714
10715 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
10716 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
10717 return false;
10718 }
10719 self.live_count = self.live_count.saturating_add(run_ref.row_count);
10720 self.run_refs.push(run_ref);
10721 self.indexes_complete = false;
10722 true
10723 }
10724
10725 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
10726 self.kek.as_ref()
10727 }
10728
10729 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
10730 let mut reader = match self.runs_root.as_deref() {
10731 Some(root) => RunReader::open_file_with_cache(
10732 root.open_regular(format!("r-{run_id}.sr"))?,
10733 self.schema.clone(),
10734 self.kek.clone(),
10735 Some(self.page_cache.clone()),
10736 Some(self.decoded_cache.clone()),
10737 self.table_id,
10738 Some(&self.verified_runs),
10739 None,
10740 )?,
10741 None => RunReader::open_with_cache(
10742 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
10743 self.schema.clone(),
10744 self.kek.clone(),
10745 Some(self.page_cache.clone()),
10746 Some(self.decoded_cache.clone()),
10747 self.table_id,
10748 Some(&self.verified_runs),
10749 )?,
10750 };
10751 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
10755 reader.set_uniform_epoch(Epoch(rr.epoch_created));
10756 }
10757 Ok(reader)
10758 }
10759
10760 pub(crate) fn run_refs(&self) -> &[RunRef] {
10761 &self.run_refs
10762 }
10763
10764 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10765 self.retiring.iter().map(|run| run.run_id)
10766 }
10767
10768 pub(crate) fn runs_dir(&self) -> PathBuf {
10769 self.runs_root
10770 .as_deref()
10771 .and_then(|root| root.io_path().ok())
10772 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
10773 }
10774
10775 pub(crate) fn wal_dir(&self) -> PathBuf {
10776 self.dir.join(WAL_DIR)
10777 }
10778
10779 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
10780 self.run_refs = refs;
10781 }
10782
10783 pub(crate) fn compaction_zstd_level(&self) -> i32 {
10784 self.compaction_zstd_level
10785 }
10786
10787 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
10788 self.kek.clone()
10789 }
10790
10791 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10795 self.kek.as_ref().map(|k| k.derive_idx_key())
10796 }
10797
10798 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10802 self.kek.as_ref().map(|k| *k.derive_meta_key())
10803 }
10804
10805 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
10808 self.column_keys
10809 .iter()
10810 .map(|(&id, &(_, scheme))| (id, scheme))
10811 .collect()
10812 }
10813
10814 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
10819 self.tokenize_value_enc(column_id, v)
10820 }
10821
10822 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
10823 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10824 let (key, scheme) = self.column_keys.get(&column_id)?;
10825 let token: Vec<u8> = match (*scheme, v) {
10826 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
10827 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
10828 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
10829 _ => hmac_token(key, &v.encode_key()).to_vec(),
10830 };
10831 Some(Value::Bytes(token))
10832 }
10833
10834 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
10836 self.index_lookup_key_bytes(column_id, &v.encode_key())
10837 }
10838
10839 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
10842 {
10843 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
10844 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
10845 if *scheme == SCHEME_HMAC_EQ {
10846 return hmac_token(key, encoded).to_vec();
10847 }
10848 }
10849 }
10850 let _ = column_id;
10851 encoded.to_vec()
10852 }
10853}
10854
10855fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
10856 let columnar::NativeColumn::Int64 { data, validity } = col else {
10857 return false;
10858 };
10859 if data.len() < n || !columnar::all_non_null(validity, n) {
10860 return false;
10861 }
10862 data.iter()
10863 .take(n)
10864 .zip(data.iter().skip(1))
10865 .all(|(a, b)| a < b)
10866}
10867
10868#[derive(Debug, Clone)]
10872pub struct ColumnStat {
10873 pub min: Option<Value>,
10874 pub max: Option<Value>,
10875 pub null_count: u64,
10876}
10877
10878#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10880pub enum NativeAgg {
10881 Count,
10882 Sum,
10883 Min,
10884 Max,
10885 Avg,
10886}
10887
10888#[derive(Debug, Clone, PartialEq)]
10890pub enum NativeAggResult {
10891 Count(u64),
10892 Int(i64),
10893 Float(f64),
10894 Null,
10896}
10897
10898#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10900pub enum ApproxAgg {
10901 Count,
10902 Sum,
10903 Avg,
10904}
10905
10906#[derive(Debug, Clone)]
10910pub struct ApproxResult {
10911 pub point: f64,
10913 pub ci_low: f64,
10915 pub ci_high: f64,
10917 pub n_population: u64,
10919 pub n_sample_live: usize,
10921 pub n_passing: usize,
10923}
10924
10925#[derive(Debug, Clone, PartialEq)]
10930pub enum AggState {
10931 Count(u64),
10933 SumI {
10935 sum: i128,
10936 count: u64,
10937 },
10938 SumF {
10940 sum: f64,
10941 count: u64,
10942 },
10943 AvgI {
10945 sum: i128,
10946 count: u64,
10947 },
10948 AvgF {
10950 sum: f64,
10951 count: u64,
10952 },
10953 MinI(i64),
10955 MaxI(i64),
10956 MinF(f64),
10958 MaxF(f64),
10959 Empty,
10961}
10962
10963impl AggState {
10964 pub fn merge(self, other: AggState) -> AggState {
10966 use AggState::*;
10967 match (self, other) {
10968 (Empty, x) | (x, Empty) => x,
10969 (Count(a), Count(b)) => Count(a + b),
10970 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
10971 sum: sa + sb,
10972 count: ca + cb,
10973 },
10974 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
10975 sum: sa + sb,
10976 count: ca + cb,
10977 },
10978 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
10979 sum: sa + sb,
10980 count: ca + cb,
10981 },
10982 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
10983 sum: sa + sb,
10984 count: ca + cb,
10985 },
10986 (MinI(a), MinI(b)) => MinI(a.min(b)),
10987 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
10988 (MinF(a), MinF(b)) => MinF(a.min(b)),
10989 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
10990 _ => Empty, }
10992 }
10993
10994 pub fn point(&self) -> Option<f64> {
10996 match self {
10997 AggState::Count(n) => Some(*n as f64),
10998 AggState::SumI { sum, .. } => Some(*sum as f64),
10999 AggState::SumF { sum, .. } => Some(*sum),
11000 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11001 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11002 AggState::MinI(n) => Some(*n as f64),
11003 AggState::MaxI(n) => Some(*n as f64),
11004 AggState::MinF(n) => Some(*n),
11005 AggState::MaxF(n) => Some(*n),
11006 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11007 }
11008 }
11009
11010 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11014 let is_float = matches!(ty, Some(TypeId::Float64));
11015 match (agg, result) {
11016 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11017 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11018 sum: x as i128,
11019 count: 1, },
11021 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11022 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11023 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11024 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11025 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11026 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11027 (NativeAgg::Count, _) => AggState::Empty,
11028 (_, NativeAggResult::Null) => AggState::Empty,
11029 _ => {
11030 let _ = is_float;
11031 AggState::Empty
11032 }
11033 }
11034 }
11035}
11036
11037#[derive(Debug, Clone)]
11040pub struct CachedAgg {
11041 pub state: AggState,
11042 pub watermark: u64,
11043 pub epoch: u64,
11044}
11045
11046#[derive(Debug, Clone)]
11048pub struct IncrementalAggResult {
11049 pub state: AggState,
11051 pub incremental: bool,
11054 pub delta_rows: u64,
11056}
11057
11058fn agg_state_from_rows(
11062 rows: &[Row],
11063 conditions: &[crate::query::Condition],
11064 index_sets: &[RowIdSet],
11065 column: Option<u16>,
11066 agg: NativeAgg,
11067 schema: &Schema,
11068 control: Option<&crate::ExecutionControl>,
11069) -> Result<AggState> {
11070 let mut count: u64 = 0;
11071 let mut sum_i: i128 = 0;
11072 let mut sum_f: f64 = 0.0;
11073 let mut mn_i: i64 = i64::MAX;
11074 let mut mx_i: i64 = i64::MIN;
11075 let mut mn_f: f64 = f64::INFINITY;
11076 let mut mx_f: f64 = f64::NEG_INFINITY;
11077 let mut saw_int = false;
11078 let mut saw_float = false;
11079 for (index, r) in rows.iter().enumerate() {
11080 execution_checkpoint(control, index)?;
11081 if !conditions
11082 .iter()
11083 .all(|c| condition_matches_row(c, r, schema))
11084 {
11085 continue;
11086 }
11087 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11088 continue;
11089 }
11090 match agg {
11091 NativeAgg::Count => match column {
11092 None => count += 1,
11094 Some(cid) => match r.columns.get(&cid) {
11097 None | Some(Value::Null) => {}
11098 Some(_) => count += 1,
11099 },
11100 },
11101 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11102 Some(Value::Int64(n)) => {
11103 count += 1;
11104 sum_i += *n as i128;
11105 mn_i = mn_i.min(*n);
11106 mx_i = mx_i.max(*n);
11107 saw_int = true;
11108 }
11109 Some(Value::Float64(f)) => {
11110 count += 1;
11111 sum_f += f;
11112 mn_f = mn_f.min(*f);
11113 mx_f = mx_f.max(*f);
11114 saw_float = true;
11115 }
11116 _ => {}
11117 },
11118 }
11119 }
11120 Ok(match agg {
11121 NativeAgg::Count => {
11122 if count == 0 {
11123 AggState::Empty
11124 } else {
11125 AggState::Count(count)
11126 }
11127 }
11128 NativeAgg::Sum => {
11129 if count == 0 {
11130 AggState::Empty
11131 } else if saw_int {
11132 AggState::SumI { sum: sum_i, count }
11133 } else {
11134 AggState::SumF { sum: sum_f, count }
11135 }
11136 }
11137 NativeAgg::Avg => {
11138 if count == 0 {
11139 AggState::Empty
11140 } else if saw_int {
11141 AggState::AvgI { sum: sum_i, count }
11142 } else {
11143 AggState::AvgF { sum: sum_f, count }
11144 }
11145 }
11146 NativeAgg::Min => {
11147 if !saw_int && !saw_float {
11148 AggState::Empty
11149 } else if saw_int {
11150 AggState::MinI(mn_i)
11151 } else {
11152 AggState::MinF(mn_f)
11153 }
11154 }
11155 NativeAgg::Max => {
11156 if !saw_int && !saw_float {
11157 AggState::Empty
11158 } else if saw_int {
11159 AggState::MaxI(mx_i)
11160 } else {
11161 AggState::MaxF(mx_f)
11162 }
11163 }
11164 })
11165}
11166
11167fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11171 use crate::query::Condition;
11172 match c {
11173 Condition::Pk(key) => match schema.primary_key() {
11174 Some(pk) => row
11175 .columns
11176 .get(&pk.id)
11177 .map(|v| v.encode_key() == *key)
11178 .unwrap_or(false),
11179 None => false,
11180 },
11181 Condition::BitmapEq { column_id, value } => row
11182 .columns
11183 .get(column_id)
11184 .map(|v| v.encode_key() == *value)
11185 .unwrap_or(false),
11186 Condition::BitmapIn { column_id, values } => {
11187 let key = row.columns.get(column_id).map(|v| v.encode_key());
11188 match key {
11189 Some(k) => values.contains(&k),
11190 None => false,
11191 }
11192 }
11193 Condition::BytesPrefix { column_id, prefix } => row
11194 .columns
11195 .get(column_id)
11196 .map(|v| v.encode_key().starts_with(prefix))
11197 .unwrap_or(false),
11198 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11199 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11200 _ => false,
11201 },
11202 Condition::RangeF64 {
11203 column_id,
11204 lo,
11205 lo_inclusive,
11206 hi,
11207 hi_inclusive,
11208 } => match row.columns.get(column_id) {
11209 Some(Value::Float64(n)) => {
11210 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11211 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11212 lo_ok && hi_ok
11213 }
11214 _ => false,
11215 },
11216 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11217 Some(Value::Bytes(b)) => {
11218 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11219 }
11220 _ => false,
11221 },
11222 Condition::FmContainsAll {
11223 column_id,
11224 patterns,
11225 } => match row.columns.get(column_id) {
11226 Some(Value::Bytes(b)) => patterns
11227 .iter()
11228 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11229 _ => false,
11230 },
11231 Condition::Ann { .. }
11232 | Condition::SparseMatch { .. }
11233 | Condition::MinHashSimilar { .. } => true,
11234 Condition::IsNull { column_id } => {
11235 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11236 }
11237 Condition::IsNotNull { column_id } => {
11238 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11239 }
11240 }
11241}
11242
11243fn as_f64(v: Option<&Value>) -> Option<f64> {
11245 match v {
11246 Some(Value::Int64(n)) => Some(*n as f64),
11247 Some(Value::Float64(f)) => Some(*f),
11248 _ => None,
11249 }
11250}
11251
11252fn accumulate_int(
11256 cursor: &mut dyn crate::cursor::Cursor,
11257 control: Option<&crate::ExecutionControl>,
11258) -> Result<(u64, i128, i64, i64)> {
11259 let mut count: u64 = 0;
11260 let mut sum: i128 = 0;
11261 let mut mn: i64 = i64::MAX;
11262 let mut mx: i64 = i64::MIN;
11263 while let Some(cols) = cursor.next_batch()? {
11264 execution_checkpoint(control, 0)?;
11265 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11266 if crate::columnar::all_non_null(validity, data.len()) {
11267 count += data.len() as u64;
11269 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11270 execution_checkpoint(control, chunk_index * 1024)?;
11271 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11272 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11273 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11274 }
11275 } else {
11276 for (i, &v) in data.iter().enumerate() {
11277 execution_checkpoint(control, i)?;
11278 if crate::columnar::validity_bit(validity, i) {
11279 count += 1;
11280 sum += v as i128;
11281 mn = mn.min(v);
11282 mx = mx.max(v);
11283 }
11284 }
11285 }
11286 }
11287 }
11288 Ok((count, sum, mn, mx))
11289}
11290
11291fn accumulate_float(
11293 cursor: &mut dyn crate::cursor::Cursor,
11294 control: Option<&crate::ExecutionControl>,
11295) -> Result<(u64, f64, f64, f64)> {
11296 let mut count: u64 = 0;
11297 let mut sum: f64 = 0.0;
11298 let mut mn: f64 = f64::INFINITY;
11299 let mut mx: f64 = f64::NEG_INFINITY;
11300 while let Some(cols) = cursor.next_batch()? {
11301 execution_checkpoint(control, 0)?;
11302 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11303 if crate::columnar::all_non_null(validity, data.len()) {
11304 count += data.len() as u64;
11305 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11306 execution_checkpoint(control, chunk_index * 1024)?;
11307 sum += chunk.iter().sum::<f64>();
11308 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11309 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11310 }
11311 } else {
11312 for (i, &v) in data.iter().enumerate() {
11313 execution_checkpoint(control, i)?;
11314 if crate::columnar::validity_bit(validity, i) {
11315 count += 1;
11316 sum += v;
11317 mn = mn.min(v);
11318 mx = mx.max(v);
11319 }
11320 }
11321 }
11322 }
11323 }
11324 Ok((count, sum, mn, mx))
11325}
11326
11327#[inline]
11328fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11329 if index.is_multiple_of(256) {
11330 control
11331 .map(crate::ExecutionControl::checkpoint)
11332 .transpose()?;
11333 }
11334 Ok(())
11335}
11336
11337fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11338 if count == 0 && !matches!(agg, NativeAgg::Count) {
11339 return NativeAggResult::Null;
11340 }
11341 match agg {
11342 NativeAgg::Count => NativeAggResult::Count(count),
11343 NativeAgg::Sum => match sum.try_into() {
11346 Ok(v) => NativeAggResult::Int(v),
11347 Err(_) => NativeAggResult::Null,
11348 },
11349 NativeAgg::Min => NativeAggResult::Int(mn),
11350 NativeAgg::Max => NativeAggResult::Int(mx),
11351 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11352 }
11353}
11354
11355fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11356 if count == 0 && !matches!(agg, NativeAgg::Count) {
11357 return NativeAggResult::Null;
11358 }
11359 match agg {
11360 NativeAgg::Count => NativeAggResult::Count(count),
11361 NativeAgg::Sum => NativeAggResult::Float(sum),
11362 NativeAgg::Min => NativeAggResult::Float(mn),
11363 NativeAgg::Max => NativeAggResult::Float(mx),
11364 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11365 }
11366}
11367
11368fn agg_int(
11371 stats: &[crate::page::PageStat],
11372 decode: fn(Option<&[u8]>) -> Option<i64>,
11373) -> Option<(Option<i64>, Option<i64>, u64)> {
11374 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11375 let mut any = false;
11376 for s in stats {
11377 if let Some(v) = decode(s.min.as_deref()) {
11378 mn = mn.min(v);
11379 any = true;
11380 }
11381 if let Some(v) = decode(s.max.as_deref()) {
11382 mx = mx.max(v);
11383 any = true;
11384 }
11385 nulls += s.null_count;
11386 }
11387 any.then_some((Some(mn), Some(mx), nulls))
11388}
11389
11390fn agg_float(
11392 stats: &[crate::page::PageStat],
11393 decode: fn(Option<&[u8]>) -> Option<f64>,
11394) -> Option<(Option<f64>, Option<f64>, u64)> {
11395 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11396 let mut any = false;
11397 for s in stats {
11398 if let Some(v) = decode(s.min.as_deref()) {
11399 mn = mn.min(v);
11400 any = true;
11401 }
11402 if let Some(v) = decode(s.max.as_deref()) {
11403 mx = mx.max(v);
11404 any = true;
11405 }
11406 nulls += s.null_count;
11407 }
11408 any.then_some((Some(mn), Some(mx), nulls))
11409}
11410
11411type SecondaryIndexes = (
11413 HashMap<u16, BitmapIndex>,
11414 HashMap<u16, AnnIndex>,
11415 HashMap<u16, FmIndex>,
11416 HashMap<u16, SparseIndex>,
11417 HashMap<u16, MinHashIndex>,
11418);
11419
11420fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11421 let mut bitmap = HashMap::new();
11422 let mut ann = HashMap::new();
11423 let mut fm = HashMap::new();
11424 let mut sparse = HashMap::new();
11425 let mut minhash = HashMap::new();
11426 for idef in &schema.indexes {
11427 match idef.kind {
11428 IndexKind::Bitmap => {
11429 bitmap.insert(idef.column_id, BitmapIndex::new());
11430 }
11431 IndexKind::Ann => {
11432 let dim = schema
11433 .columns
11434 .iter()
11435 .find(|c| c.id == idef.column_id)
11436 .and_then(|c| match c.ty {
11437 TypeId::Embedding { dim } => Some(dim as usize),
11438 _ => None,
11439 })
11440 .unwrap_or(0);
11441 let options = idef.options.ann.clone().unwrap_or_default();
11442 ann.insert(
11443 idef.column_id,
11444 AnnIndex::with_options(
11445 dim,
11446 options.m,
11447 options.ef_construction,
11448 options.ef_search,
11449 ),
11450 );
11451 }
11452 IndexKind::FmIndex => {
11453 fm.insert(idef.column_id, FmIndex::new());
11454 }
11455 IndexKind::Sparse => {
11456 sparse.insert(idef.column_id, SparseIndex::new());
11457 }
11458 IndexKind::MinHash => {
11459 let options = idef.options.minhash.clone().unwrap_or_default();
11460 minhash.insert(
11461 idef.column_id,
11462 MinHashIndex::with_options(options.permutations, options.bands),
11463 );
11464 }
11465 _ => {}
11466 }
11467 }
11468 (bitmap, ann, fm, sparse, minhash)
11469}
11470
11471const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11472 | ColumnFlags::AUTO_INCREMENT
11473 | ColumnFlags::ENCRYPTED
11474 | ColumnFlags::ENCRYPTED_INDEXABLE
11475 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11476
11477fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11478 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11479 return Err(MongrelError::Schema(
11480 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11481 ));
11482 }
11483 Ok(())
11484}
11485
11486fn validate_alter_column_type(
11487 schema: &Schema,
11488 old: &ColumnDef,
11489 next: &ColumnDef,
11490 has_stored_versions: bool,
11491) -> Result<()> {
11492 if old.ty == next.ty {
11493 return Ok(());
11494 }
11495 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11496 return Err(MongrelError::Schema(format!(
11497 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11498 old.name
11499 )));
11500 }
11501 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11502 return Ok(());
11503 }
11504 Err(MongrelError::Schema(format!(
11505 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11506 old.ty, next.ty
11507 )))
11508}
11509
11510fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11511 matches!(
11512 (old, new),
11513 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11514 )
11515}
11516
11517fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11523 let mut prev: Option<i64> = None;
11524 for r in rows {
11525 match r.columns.get(&pk_id) {
11526 Some(Value::Int64(v)) => {
11527 if prev.is_some_and(|p| p >= *v) {
11528 return false;
11529 }
11530 prev = Some(*v);
11531 }
11532 _ => return false,
11533 }
11534 }
11535 true
11536}
11537
11538#[allow(clippy::too_many_arguments)]
11539fn index_into(
11540 schema: &Schema,
11541 row: &Row,
11542 hot: &mut HotIndex,
11543 bitmap: &mut HashMap<u16, BitmapIndex>,
11544 ann: &mut HashMap<u16, AnnIndex>,
11545 fm: &mut HashMap<u16, FmIndex>,
11546 sparse: &mut HashMap<u16, SparseIndex>,
11547 minhash: &mut HashMap<u16, MinHashIndex>,
11548) {
11549 for idef in &schema.indexes {
11550 let Some(val) = row.columns.get(&idef.column_id) else {
11551 continue;
11552 };
11553 match idef.kind {
11554 IndexKind::Bitmap => {
11555 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11556 b.insert(val.encode_key(), row.row_id);
11557 }
11558 }
11559 IndexKind::Ann => {
11560 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11561 a.insert_validated(v, row.row_id);
11562 }
11563 }
11564 IndexKind::FmIndex => {
11565 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11566 f.insert(b.clone(), row.row_id);
11567 }
11568 }
11569 IndexKind::Sparse => {
11570 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11571 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11574 s.insert(&terms, row.row_id);
11575 }
11576 }
11577 }
11578 IndexKind::MinHash => {
11579 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11580 let tokens = crate::index::token_hashes_from_bytes(b);
11583 mh.insert(&tokens, row.row_id);
11584 }
11585 }
11586 _ => {}
11587 }
11588 }
11589 if let Some(pk_col) = schema.primary_key() {
11590 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11591 hot.insert(pk_val.encode_key(), row.row_id);
11592 }
11593 }
11594}
11595
11596#[allow(clippy::too_many_arguments)]
11599fn index_into_single(
11600 idef: &IndexDef,
11601 _schema: &Schema,
11602 row: &Row,
11603 _hot: &mut HotIndex,
11604 bitmap: &mut HashMap<u16, BitmapIndex>,
11605 ann: &mut HashMap<u16, AnnIndex>,
11606 fm: &mut HashMap<u16, FmIndex>,
11607 sparse: &mut HashMap<u16, SparseIndex>,
11608 minhash: &mut HashMap<u16, MinHashIndex>,
11609) {
11610 let Some(val) = row.columns.get(&idef.column_id) else {
11611 return;
11612 };
11613 match idef.kind {
11614 IndexKind::Bitmap => {
11615 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11616 b.insert(val.encode_key(), row.row_id);
11617 }
11618 }
11619 IndexKind::Ann => {
11620 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11621 a.insert_validated(v, row.row_id);
11622 }
11623 }
11624 IndexKind::FmIndex => {
11625 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11626 f.insert(b.clone(), row.row_id);
11627 }
11628 }
11629 IndexKind::Sparse => {
11630 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11631 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11632 s.insert(&terms, row.row_id);
11633 }
11634 }
11635 }
11636 IndexKind::MinHash => {
11637 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11638 let tokens = crate::index::token_hashes_from_bytes(b);
11639 mh.insert(&tokens, row.row_id);
11640 }
11641 }
11642 _ => {}
11643 }
11644}
11645
11646fn eval_partial_predicate(
11652 pred: &str,
11653 columns_map: &HashMap<u16, &Value>,
11654 name_to_id: &HashMap<&str, u16>,
11655) -> bool {
11656 let lower = pred.trim().to_ascii_lowercase();
11657 if let Some(rest) = lower.strip_suffix(" is not null") {
11659 let col_name = rest.trim();
11660 if let Some(col_id) = name_to_id.get(col_name) {
11661 return columns_map
11662 .get(col_id)
11663 .is_some_and(|v| !matches!(v, Value::Null));
11664 }
11665 }
11666 if let Some(rest) = lower.strip_suffix(" is null") {
11668 let col_name = rest.trim();
11669 if let Some(col_id) = name_to_id.get(col_name) {
11670 return columns_map
11671 .get(col_id)
11672 .is_none_or(|v| matches!(v, Value::Null));
11673 }
11674 }
11675 true
11678}
11679
11680#[allow(dead_code)]
11686fn bulk_index_key(
11687 column_keys: &HashMap<u16, ([u8; 32], u8)>,
11688 column_id: u16,
11689 ty: TypeId,
11690 col: &columnar::NativeColumn,
11691 i: usize,
11692) -> Option<Vec<u8>> {
11693 let encoded = columnar::encode_key_native(ty, col, i)?;
11694 {
11695 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11696 if let Some((key, scheme)) = column_keys.get(&column_id) {
11697 return Some(match (*scheme, col) {
11698 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
11699 (_, columnar::NativeColumn::Int64 { data, .. }) => {
11700 ope_token_i64(key, data[i]).to_vec()
11701 }
11702 (_, columnar::NativeColumn::Float64 { data, .. }) => {
11703 ope_token_f64(key, data[i]).to_vec()
11704 }
11705 _ => hmac_token(key, &encoded).to_vec(),
11706 });
11707 }
11708 }
11709 Some(encoded)
11710}
11711
11712pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
11713 write_schema_with_after(dir, schema, || {})
11714}
11715
11716pub(crate) fn write_schema_durable(
11717 root: &crate::durable_file::DurableRoot,
11718 schema: &Schema,
11719) -> Result<()> {
11720 write_schema_durable_with_after(root, schema, || {})
11721}
11722
11723fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
11724where
11725 F: FnOnce(),
11726{
11727 let json = serde_json::to_string_pretty(schema)
11728 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
11729 crate::durable_file::write_atomic_with_after(
11730 &dir.join(SCHEMA_FILENAME),
11731 json.as_bytes(),
11732 after_publish,
11733 )?;
11734 Ok(())
11735}
11736
11737fn write_schema_durable_with_after<F>(
11738 root: &crate::durable_file::DurableRoot,
11739 schema: &Schema,
11740 after_publish: F,
11741) -> Result<()>
11742where
11743 F: FnOnce(),
11744{
11745 let json = serde_json::to_string_pretty(schema)
11746 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
11747 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
11748 Ok(())
11749}
11750
11751fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
11752 let mut schema_published = false;
11753 let schema_result = match table._root_guard.as_deref() {
11754 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
11755 schema_published = true;
11756 }),
11757 None => write_schema_with_after(&table.dir, &table.schema, || {
11758 schema_published = true;
11759 }),
11760 };
11761 if schema_result.is_err() && !schema_published {
11762 return schema_result;
11763 }
11764 match table.persist_manifest(table.current_epoch()) {
11765 Ok(()) => Ok(()),
11766 Err(manifest_error) => Err(match schema_result {
11767 Ok(()) => manifest_error,
11768 Err(schema_error) => MongrelError::Other(format!(
11769 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11770 )),
11771 }),
11772 }
11773}
11774
11775fn read_schema(dir: &Path) -> Result<Schema> {
11776 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
11777 read_schema_file(file)
11778}
11779
11780fn read_schema_file(file: std::fs::File) -> Result<Schema> {
11781 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
11782 use std::io::Read;
11783
11784 let length = file.metadata()?.len();
11785 if length > MAX_SCHEMA_BYTES {
11786 return Err(MongrelError::ResourceLimitExceeded {
11787 resource: "schema bytes",
11788 requested: usize::try_from(length).unwrap_or(usize::MAX),
11789 limit: MAX_SCHEMA_BYTES as usize,
11790 });
11791 }
11792 let mut bytes = Vec::with_capacity(length as usize);
11793 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
11794 if bytes.len() as u64 != length {
11795 return Err(MongrelError::Schema(
11796 "schema length changed while reading".into(),
11797 ));
11798 }
11799 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
11800}
11801
11802fn preflight_standalone_open(
11803 dir: &Path,
11804 runs_root: Option<&crate::durable_file::DurableRoot>,
11805 idx_root: Option<&crate::durable_file::DurableRoot>,
11806 manifest: &Manifest,
11807 schema: &Schema,
11808 records: &[crate::wal::Record],
11809 kek: Option<Arc<Kek>>,
11810) -> Result<()> {
11811 crate::wal::validate_shared_transaction_framing(records)?;
11812 if manifest.schema_id > schema.schema_id
11813 || manifest.flushed_epoch > manifest.current_epoch
11814 || manifest.global_idx_epoch > manifest.current_epoch
11815 || manifest.next_row_id == u64::MAX
11816 || manifest.auto_inc_next < 0
11817 || manifest.auto_inc_next == i64::MAX
11818 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
11819 {
11820 return Err(MongrelError::InvalidArgument(
11821 "manifest counters or schema identity are invalid".into(),
11822 ));
11823 }
11824 let mut run_ids = HashSet::new();
11825 let mut maximum_row_id = None::<u64>;
11826 for run in &manifest.runs {
11827 if run.run_id >= u64::MAX as u128
11828 || !run_ids.insert(run.run_id)
11829 || run.epoch_created > manifest.current_epoch
11830 {
11831 return Err(MongrelError::InvalidArgument(
11832 "manifest contains an invalid or duplicate active run".into(),
11833 ));
11834 }
11835 let mut reader = match runs_root {
11836 Some(root) => RunReader::open_file(
11837 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
11838 schema.clone(),
11839 kek.clone(),
11840 )?,
11841 None => RunReader::open(
11842 dir.join(RUNS_DIR)
11843 .join(format!("r-{}.sr", run.run_id as u64)),
11844 schema.clone(),
11845 kek.clone(),
11846 )?,
11847 };
11848 let header = reader.header();
11849 if header.run_id != run.run_id
11850 || header.level != run.level
11851 || header.row_count != run.row_count
11852 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
11853 || header.is_uniform_epoch() && header.epoch_created != 0
11854 || header.schema_id > schema.schema_id
11855 {
11856 return Err(MongrelError::InvalidArgument(format!(
11857 "run {} differs from its manifest",
11858 run.run_id
11859 )));
11860 }
11861 if header.row_count != 0 {
11862 maximum_row_id = Some(
11863 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
11864 );
11865 }
11866 reader.validate_all_pages()?;
11867 }
11868 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
11869 return Err(MongrelError::InvalidArgument(
11870 "manifest next_row_id does not advance beyond persisted rows".into(),
11871 ));
11872 }
11873 for run in &manifest.retiring {
11874 if run.run_id >= u64::MAX as u128
11875 || run.retire_epoch > manifest.current_epoch
11876 || !run_ids.insert(run.run_id)
11877 {
11878 return Err(MongrelError::InvalidArgument(
11879 "manifest contains an invalid or duplicate retired run".into(),
11880 ));
11881 }
11882 }
11883 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
11884 match idx_root {
11885 Some(root) => {
11886 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
11887 }
11888 None => {
11889 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
11890 }
11891 }
11892
11893 let committed = records
11894 .iter()
11895 .filter_map(|record| match record.op {
11896 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
11897 _ => None,
11898 })
11899 .collect::<HashMap<_, _>>();
11900 for record in records {
11901 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
11902 continue;
11903 };
11904 match &record.op {
11905 Op::Put { table_id, rows } => {
11906 if *table_id != manifest.table_id {
11907 return Err(MongrelError::CorruptWal {
11908 offset: record.seq.0,
11909 reason: format!(
11910 "private WAL record references table {table_id}, expected {}",
11911 manifest.table_id
11912 ),
11913 });
11914 }
11915 let rows: Vec<Row> =
11916 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
11917 offset: record.seq.0,
11918 reason: format!("committed Put payload could not be decoded: {error}"),
11919 })?;
11920 for row in rows {
11921 if row.deleted || row.row_id.0 == u64::MAX {
11922 return Err(MongrelError::CorruptWal {
11923 offset: record.seq.0,
11924 reason: "committed Put contains an invalid row identity".into(),
11925 });
11926 }
11927 let cells = row.columns.into_iter().collect::<Vec<_>>();
11928 schema
11929 .validate_values(&cells)
11930 .map_err(|error| MongrelError::CorruptWal {
11931 offset: record.seq.0,
11932 reason: format!("committed Put violates table schema: {error}"),
11933 })?;
11934 if schema.auto_increment_column().is_some_and(|column| {
11935 matches!(
11936 cells.iter().find(|(id, _)| *id == column.id),
11937 Some((_, Value::Int64(value))) if *value == i64::MAX
11938 )
11939 }) {
11940 return Err(MongrelError::CorruptWal {
11941 offset: record.seq.0,
11942 reason: "committed Put exhausts AUTO_INCREMENT".into(),
11943 });
11944 }
11945 }
11946 }
11947 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
11948 if *table_id != manifest.table_id =>
11949 {
11950 return Err(MongrelError::CorruptWal {
11951 offset: record.seq.0,
11952 reason: format!(
11953 "private WAL record references table {table_id}, expected {}",
11954 manifest.table_id
11955 ),
11956 });
11957 }
11958 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
11959 return Err(MongrelError::CorruptWal {
11960 offset: record.seq.0,
11961 reason: "private WAL contains shared spilled-run metadata".into(),
11962 });
11963 }
11964 _ => {}
11965 }
11966 }
11967 Ok(())
11968}
11969
11970fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
11971 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
11972}
11973
11974fn wal_segment_number(path: &Path) -> Option<u64> {
11975 path.file_stem()
11976 .and_then(|stem| stem.to_str())
11977 .and_then(|stem| stem.strip_prefix("seg-"))
11978 .and_then(|number| number.parse().ok())
11979}
11980
11981fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
11982 let n = list_wal_numbers(wal_dir)?;
11983 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
11984}
11985
11986fn next_wal_number(wal_dir: &Path) -> Result<u32> {
11987 list_wal_numbers(wal_dir)?
11988 .map(|maximum| {
11989 maximum
11990 .checked_add(1)
11991 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
11992 })
11993 .unwrap_or(Ok(0))
11994}
11995
11996fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
11997 let mut max_n = None;
11998 let entries = match std::fs::read_dir(wal_dir) {
11999 Ok(entries) => entries,
12000 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12001 Err(error) => return Err(error.into()),
12002 };
12003 for entry in entries {
12004 let entry = entry?;
12005 let fname = entry.file_name();
12006 let Some(s) = fname.to_str() else {
12007 continue;
12008 };
12009 let Some(stripped) = s.strip_prefix("seg-") else {
12010 continue;
12011 };
12012 let Some(number) = stripped.strip_suffix(".wal") else {
12013 return Err(MongrelError::CorruptWal {
12014 offset: 0,
12015 reason: format!("malformed WAL segment name {s:?}"),
12016 });
12017 };
12018 let n = number
12019 .parse::<u32>()
12020 .map_err(|_| MongrelError::CorruptWal {
12021 offset: 0,
12022 reason: format!("malformed WAL segment name {s:?}"),
12023 })?;
12024 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12025 return Err(MongrelError::CorruptWal {
12026 offset: n as u64,
12027 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12028 });
12029 }
12030 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12031 }
12032 Ok(max_n)
12033}