1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, EpochGuard, MaintenanceReceipt, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, IndexGeneration, MinHashIndex,
18 SparseIndex,
19};
20use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
21use crate::memtable::{Memtable, Row, Value};
22use crate::mutable_run::MutableRun;
23use crate::row_id_set::RowIdSet;
24use crate::rowid::{RowId, RowIdAllocator};
25use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
26use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
27use crate::txn::{GroupCommit, OwnedRow};
28use crate::wal::{Op, SharedWal, Wal};
29use crate::{MongrelError, Result};
30use arc_swap::ArcSwap;
31use std::cmp::Reverse;
32use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
33use std::path::{Path, PathBuf};
34use std::sync::atomic::AtomicBool;
35use std::sync::Arc;
36use zeroize::Zeroizing;
37
38pub const WAL_DIR: &str = "_wal";
39pub const RUNS_DIR: &str = "_runs";
40pub const CACHE_DIR: &str = "_cache";
41pub const META_DIR: &str = "_meta";
42pub const RCACHE_DIR: &str = "_rcache";
43pub const KEYS_FILENAME: &str = "keys";
44pub const SCHEMA_FILENAME: &str = "schema.json";
45
46fn derive_next_run_id(
47 dir: &Path,
48 runs_root: Option<&crate::durable_file::DurableRoot>,
49 active: &[RunRef],
50 retiring: &[crate::manifest::RetiredRun],
51) -> Result<u64> {
52 let mut maximum = 0_u64;
53 for run_id in active
54 .iter()
55 .map(|run| run.run_id)
56 .chain(retiring.iter().map(|run| run.run_id))
57 {
58 let run_id = u64::try_from(run_id)
59 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
60 maximum = maximum.max(run_id);
61 }
62 let names = match runs_root {
63 Some(root) => root.list_regular_files(".")?,
64 None => std::fs::read_dir(dir.join(RUNS_DIR))?
65 .map(|entry| entry.map(|entry| entry.file_name()))
66 .collect::<std::io::Result<Vec<_>>>()?,
67 };
68 for name in names {
69 let Some(name) = name.to_str() else {
70 continue;
71 };
72 let Some(digits) = name
73 .strip_prefix("r-")
74 .and_then(|name| name.strip_suffix(".sr"))
75 else {
76 continue;
77 };
78 let Ok(run_id) = digits.parse::<u64>() else {
79 continue;
80 };
81 if name == format!("r-{run_id}.sr") {
82 maximum = maximum.max(run_id);
83 }
84 }
85 maximum
86 .checked_add(1)
87 .map(|next| next.max(1))
88 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
89}
90
91enum ControlledVisibleCandidate {
92 Memory(Row),
93 Run(RunVisibleVersion),
94}
95
96impl ControlledVisibleCandidate {
97 fn row_id(&self) -> RowId {
98 match self {
99 Self::Memory(row) => row.row_id,
100 Self::Run(version) => version.row_id,
101 }
102 }
103
104 fn committed_epoch(&self) -> Epoch {
105 match self {
106 Self::Memory(row) => row.committed_epoch,
107 Self::Run(version) => version.committed_epoch,
108 }
109 }
110
111 fn deleted(&self) -> bool {
112 match self {
113 Self::Memory(row) => row.deleted,
114 Self::Run(version) => version.deleted,
115 }
116 }
117}
118
119enum ControlledVisibleCursor {
120 Memory(std::vec::IntoIter<Row>),
121 Run(Box<RunVisibleVersionCursor>),
122 #[cfg(test)]
123 Synthetic {
124 next: u64,
125 end: u64,
126 },
127}
128
129struct ControlledVisibleSource {
130 cursor: ControlledVisibleCursor,
131 current: Option<ControlledVisibleCandidate>,
132}
133
134impl ControlledVisibleSource {
135 fn memory(rows: Vec<Row>) -> Self {
136 Self {
137 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
138 current: None,
139 }
140 }
141
142 fn run(cursor: RunVisibleVersionCursor) -> Self {
143 Self {
144 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
145 current: None,
146 }
147 }
148
149 #[cfg(test)]
150 fn synthetic(end: u64) -> Self {
151 Self {
152 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
153 current: None,
154 }
155 }
156
157 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
158 self.current = match &mut self.cursor {
159 ControlledVisibleCursor::Memory(rows) => {
160 rows.next().map(ControlledVisibleCandidate::Memory)
161 }
162 ControlledVisibleCursor::Run(cursor) => cursor
163 .next_visible_version(control)?
164 .map(ControlledVisibleCandidate::Run),
165 #[cfg(test)]
166 ControlledVisibleCursor::Synthetic { next, end } => {
167 if *next > *end {
168 None
169 } else {
170 let row = Row::new(RowId(*next), Epoch(1));
171 *next += 1;
172 Some(ControlledVisibleCandidate::Memory(row))
173 }
174 }
175 };
176 Ok(())
177 }
178
179 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate> {
180 let current = self.current.take().ok_or_else(|| {
181 MongrelError::Other("controlled visible source was not primed".into())
182 })?;
183 self.advance(control)?;
184 Ok(current)
185 }
186
187 fn materialize(
188 &mut self,
189 candidate: ControlledVisibleCandidate,
190 control: &crate::ExecutionControl,
191 ) -> Result<Row> {
192 match candidate {
193 ControlledVisibleCandidate::Memory(row) => Ok(row),
194 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
195 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
196 _ => Err(MongrelError::Other(
197 "run candidate escaped its controlled cursor".into(),
198 )),
199 },
200 }
201 }
202}
203
204fn merge_controlled_visible_sources(
205 sources: &mut [ControlledVisibleSource],
206 control: &crate::ExecutionControl,
207 mut expired: impl FnMut(&Row) -> bool,
208 mut visit: impl FnMut(Row) -> Result<()>,
209) -> Result<()> {
210 let mut heap = BinaryHeap::new();
211 for (source_index, source) in sources.iter_mut().enumerate() {
212 source.advance(control)?;
213 if let Some(candidate) = &source.current {
214 heap.push(Reverse((candidate.row_id(), source_index)));
215 }
216 }
217 let mut merged = 0_usize;
218 while let Some(Reverse((row_id, source_index))) = heap.pop() {
219 if merged.is_multiple_of(256) {
220 control.checkpoint()?;
221 }
222 merged += 1;
223 let mut best_source = source_index;
224 let mut best = sources[source_index].pop(control)?;
225 if let Some(next) = &sources[source_index].current {
226 heap.push(Reverse((next.row_id(), source_index)));
227 }
228 while heap
229 .peek()
230 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
231 {
232 let Some(Reverse((_, source_index))) = heap.pop() else {
233 break;
234 };
235 let candidate = sources[source_index].pop(control)?;
236 if candidate.committed_epoch() > best.committed_epoch() {
237 best = candidate;
238 best_source = source_index;
239 }
240 if let Some(next) = &sources[source_index].current {
241 heap.push(Reverse((next.row_id(), source_index)));
242 }
243 }
244 if best.deleted() {
245 continue;
246 }
247 let row = sources[best_source].materialize(best, control)?;
248 if !expired(&row) {
249 visit(row)?;
250 }
251 }
252 control.checkpoint()
253}
254
255#[cfg(test)]
256mod controlled_visible_cursor_tests {
257 use super::*;
258
259 #[test]
260 fn streams_more_than_one_million_rows_without_a_source_cap() {
261 let control = crate::ExecutionControl::new(None);
262 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
263 let mut count = 0_u64;
264 let mut last = 0_u64;
265 merge_controlled_visible_sources(
266 &mut sources,
267 &control,
268 |_| false,
269 |row| {
270 count += 1;
271 assert!(row.row_id.0 > last);
272 last = row.row_id.0;
273 Ok(())
274 },
275 )
276 .unwrap();
277 assert_eq!(count, 1_000_001);
278 assert_eq!(last, 1_000_001);
279 }
280
281 #[test]
282 fn merge_orders_rows_and_honors_newest_tombstones() {
283 let control = crate::ExecutionControl::new(None);
284 let older = vec![
285 Row::new(RowId(1), Epoch(1)),
286 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
287 Row::new(RowId(4), Epoch(1)),
288 ];
289 let mut deleted = Row::new(RowId(1), Epoch(2));
290 deleted.deleted = true;
291 let newer = vec![
292 deleted,
293 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
294 Row::new(RowId(3), Epoch(2)),
295 ];
296 let mut sources = vec![
297 ControlledVisibleSource::memory(older),
298 ControlledVisibleSource::memory(newer),
299 ];
300 let mut rows = Vec::new();
301 merge_controlled_visible_sources(
302 &mut sources,
303 &control,
304 |_| false,
305 |row| {
306 rows.push(row);
307 Ok(())
308 },
309 )
310 .unwrap();
311 assert_eq!(
312 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
313 vec![2, 3, 4]
314 );
315 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
316 }
317}
318
319fn iso_now_bytes() -> Vec<u8> {
322 let secs = std::time::SystemTime::now()
323 .duration_since(std::time::UNIX_EPOCH)
324 .map(|d| d.as_secs() as i64)
325 .unwrap_or(0);
326 let days = secs.div_euclid(86_400);
327 let rem = secs.rem_euclid(86_400);
328 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
329 let (year, month, day) = civil_from_days(days);
330 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
331}
332
333pub(crate) fn unix_nanos_now() -> i64 {
334 std::time::SystemTime::now()
335 .duration_since(std::time::UNIX_EPOCH)
336 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
337 .unwrap_or(0)
338}
339
340fn ann_candidate_cap(
341 index_len: usize,
342 context: Option<&crate::query::AiExecutionContext>,
343) -> usize {
344 index_len
345 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
346 .min(context.map_or(
347 crate::query::MAX_RAW_INDEX_CANDIDATES,
348 crate::query::AiExecutionContext::max_fused_candidates,
349 ))
350}
351
352#[cfg(test)]
353mod ann_candidate_cap_tests {
354 use super::*;
355
356 #[test]
357 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
358 assert_eq!(
359 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
360 crate::query::MAX_RAW_INDEX_CANDIDATES,
361 );
362 let context = crate::query::AiExecutionContext::with_limits(
363 std::time::Duration::from_secs(1),
364 usize::MAX,
365 17,
366 );
367 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
368 }
369}
370
371fn civil_from_days(z: i64) -> (i64, u32, u32) {
372 let z = z + 719_468;
373 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
374 let doe = z - era * 146_097;
375 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
376 let y = yoe + era * 400;
377 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
378 let mp = (5 * doy + 2) / 153;
379 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
380 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
381 (if m <= 2 { y + 1 } else { y }, m, d)
382}
383
384pub fn clustered_row_id(primary_key: &Value) -> RowId {
391 let mut hash: u64 = 0xcbf29ce484222325;
392 for byte in primary_key.encode_key() {
393 hash ^= u64::from(byte);
394 hash = hash.wrapping_mul(0x100000001b3);
395 }
396 RowId(hash.max(1))
397}
398
399const DEFAULT_SYNC_BYTE_THRESHOLD: u64 = 0; pub(crate) const PAGE_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; pub(crate) const DECODED_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; const DEFAULT_MUTABLE_RUN_SPILL_BYTES: u64 = 8 * 1024 * 1024;
406
407#[derive(Clone, Copy, Debug)]
422struct AutoIncState {
423 column_id: u16,
424 next: i64,
425 seeded: bool,
426}
427
428pub(crate) struct RecoveryMetadataPlan {
429 live_count: u64,
430 auto_inc: Option<AutoIncState>,
431 changed: bool,
432}
433
434type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
435
436fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
439 schema.auto_increment_column().map(|c| AutoIncState {
440 column_id: c.id,
441 next: 0,
442 seeded: false,
443 })
444}
445
446#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
459pub enum IndexBuildPolicy {
460 #[default]
462 Deferred,
463 Eager,
465}
466
467#[derive(Clone)]
468struct ReversePkSegment {
469 values: HashMap<RowId, Vec<u8>>,
470 removed: HashSet<RowId>,
471}
472
473#[derive(Clone)]
474struct ReversePkMap {
475 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
476 active: ReversePkSegment,
477}
478
479impl ReversePkMap {
480 fn new() -> Self {
481 Self {
482 frozen: Arc::new(Vec::new()),
483 active: ReversePkSegment {
484 values: HashMap::new(),
485 removed: HashSet::new(),
486 },
487 }
488 }
489
490 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
491 let mut map = Self::new();
492 map.active.values.extend(entries);
493 map
494 }
495
496 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
497 self.active.removed.remove(&row_id);
498 self.active.values.insert(row_id, key);
499 }
500
501 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
502 if let Some(key) = self.active.values.get(row_id) {
503 return Some(key);
504 }
505 if self.active.removed.contains(row_id) {
506 return None;
507 }
508 for segment in self.frozen.iter().rev() {
509 if let Some(key) = segment.values.get(row_id) {
510 return Some(key);
511 }
512 if segment.removed.contains(row_id) {
513 return None;
514 }
515 }
516 None
517 }
518
519 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
520 let previous = self.get(row_id).cloned();
521 self.active.values.remove(row_id);
522 self.active.removed.insert(*row_id);
523 previous
524 }
525
526 fn clear(&mut self) {
527 *self = Self::new();
528 }
529
530 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
531 let mut entries = HashMap::new();
532 for segment in self
533 .frozen
534 .iter()
535 .map(Arc::as_ref)
536 .chain(std::iter::once(&self.active))
537 {
538 for row_id in &segment.removed {
539 entries.remove(row_id);
540 }
541 entries.extend(
542 segment
543 .values
544 .iter()
545 .map(|(row_id, key)| (*row_id, key.clone())),
546 );
547 }
548 entries
549 }
550
551 fn seal(&mut self) {
552 if self.active.values.is_empty() && self.active.removed.is_empty() {
553 return;
554 }
555 let active = std::mem::replace(
556 &mut self.active,
557 ReversePkSegment {
558 values: HashMap::new(),
559 removed: HashSet::new(),
560 },
561 );
562 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
563 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
564 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
565 values: self.entries(),
566 removed: HashSet::new(),
567 })]);
568 }
569 }
570}
571
572#[derive(Clone)]
585pub struct ReadGeneration {
586 schema: Arc<Schema>,
587 base_runs: Arc<Vec<RunRef>>,
588 deltas: TableDeltas,
589 indexes: Arc<IndexGeneration>,
590 visible_through: Epoch,
591}
592
593pub(crate) enum SecondaryIndexArtifact {
597 Bitmap(u16, BitmapIndex),
598 LearnedRange(u16, ColumnLearnedRange),
599 Fm(u16, Box<FmIndex>),
600 Ann(u16, AnnIndex),
601 Sparse(u16, SparseIndex),
602 MinHash(u16, MinHashIndex),
603}
604
605#[derive(Clone)]
611pub struct TableDeltas {
612 memtable: Memtable,
613 mutable_run: MutableRun,
614 hot: HotIndex,
615 pk_by_row: ReversePkMap,
616}
617
618impl ReadGeneration {
619 fn empty(schema: &Schema) -> Self {
622 Self {
623 schema: Arc::new(schema.clone()),
624 base_runs: Arc::new(Vec::new()),
625 deltas: TableDeltas {
626 memtable: Memtable::new(),
627 mutable_run: MutableRun::new(),
628 hot: HotIndex::new(),
629 pk_by_row: ReversePkMap::new(),
630 },
631 indexes: Arc::new(IndexGeneration::default()),
632 visible_through: Epoch(0),
633 }
634 }
635
636 pub fn schema(&self) -> &Arc<Schema> {
638 &self.schema
639 }
640
641 pub fn base_runs(&self) -> &[RunRef] {
643 &self.base_runs
644 }
645
646 pub fn indexes(&self) -> &Arc<IndexGeneration> {
648 &self.indexes
649 }
650
651 pub fn visible_through(&self) -> Epoch {
653 self.visible_through
654 }
655
656 pub fn deltas(&self) -> &TableDeltas {
660 &self.deltas
661 }
662}
663
664impl TableDeltas {
665 pub fn approx_bytes(&self) -> u64 {
668 self.memtable
669 .approx_bytes()
670 .saturating_add(self.mutable_run.approx_bytes())
671 }
672
673 pub fn memtable_len(&self) -> usize {
675 self.memtable.len()
676 }
677
678 pub fn mutable_run_len(&self) -> usize {
680 self.mutable_run.len()
681 }
682
683 pub fn hot_len(&self) -> usize {
685 self.hot.len()
686 }
687
688 pub fn reverse_pk_len(&self) -> usize {
690 self.pk_by_row.entries().len()
691 }
692}
693
694#[derive(Clone)]
696pub struct Table {
697 dir: PathBuf,
698 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
699 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
700 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
701 table_id: u64,
702 name: String,
706 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
711 read_only: bool,
714 durable_commit_failed: bool,
718 wal: WalSink,
719 memtable: Memtable,
720 mutable_run: MutableRun,
725 mutable_run_spill_bytes: u64,
727 compaction_zstd_level: i32,
730 allocator: RowIdAllocator,
731 epoch: Arc<EpochAuthority>,
732 data_generation: u64,
735 schema: Schema,
736 hot: HotIndex,
737 kek: Option<Arc<Kek>>,
740 column_keys: HashMap<u16, ([u8; 32], u8)>,
744 run_refs: Vec<RunRef>,
745 retiring: Vec<crate::manifest::RetiredRun>,
748 next_run_id: u64,
749 sync_byte_threshold: u64,
750 current_txn_id: u64,
755 pending_private_mutations: bool,
759 bitmap: HashMap<u16, BitmapIndex>,
760 ann: HashMap<u16, AnnIndex>,
761 fm: HashMap<u16, FmIndex>,
762 sparse: HashMap<u16, SparseIndex>,
763 minhash: HashMap<u16, MinHashIndex>,
764 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
767 pk_by_row: ReversePkMap,
769 pinned: BTreeMap<Epoch, usize>,
772 pub(crate) live_count: u64,
775 reservoir: crate::reservoir::Reservoir,
778 reservoir_complete: bool,
786 had_deletes: bool,
790 agg_cache: Arc<HashMap<u64, CachedAgg>>,
794 global_idx_epoch: u64,
798 indexes_complete: bool,
803 index_build_policy: IndexBuildPolicy,
805 pk_by_row_complete: bool,
812 flushed_epoch: u64,
815 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
818 snapshots: Arc<crate::retention::SnapshotRegistry>,
821 commit_lock: Arc<parking_lot::Mutex<()>>,
823 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
826 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
836 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
845 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
847 pending_delete_rids: roaring::RoaringBitmap,
850 pending_put_cols: std::collections::HashSet<u16>,
853 pending_rows: Vec<Row>,
859 pending_rows_auto_inc: Vec<bool>,
860 pending_dels: Vec<RowId>,
863 pending_truncate: Option<Epoch>,
867 auto_inc: Option<AutoIncState>,
870 ttl: Option<TtlPolicy>,
873 pins: Arc<crate::retention::PinRegistry>,
880 published: Arc<ArcSwap<ReadGeneration>>,
885 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
890}
891
892const _: () = {
899 const fn assert_sync<T: ?Sized + Sync>() {}
900 assert_sync::<Table>();
901};
902
903enum CachedData {
909 Rows(Arc<Vec<Row>>),
910 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
911}
912
913impl CachedData {
914 fn approx_bytes(&self) -> u64 {
915 match self {
916 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
917 CachedData::Columns(c) => c
918 .iter()
919 .map(|(_, c)| c.approx_bytes())
920 .sum::<u64>()
921 .saturating_add(c.len() as u64 * 16),
922 }
923 }
924}
925
926struct CachedEntry {
930 data: CachedData,
931 footprint: roaring::RoaringBitmap,
932 condition_cols: Vec<u16>,
933}
934
935struct ResultCache {
946 entries: std::collections::HashMap<u64, CachedEntry>,
947 order: std::collections::VecDeque<u64>,
948 bytes: u64,
949 max_bytes: u64,
950 dir: Option<std::path::PathBuf>,
951 #[allow(dead_code)]
952 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
953}
954
955#[derive(serde::Serialize, serde::Deserialize)]
957struct SerializedEntry {
958 condition_cols: Vec<u16>,
959 footprint_bits: Vec<u32>,
960 data: SerializedData,
961}
962
963#[derive(serde::Serialize, serde::Deserialize)]
964enum SerializedData {
965 Rows(Vec<Row>),
966 Columns(Vec<(u16, columnar::NativeColumn)>),
967}
968
969impl SerializedEntry {
970 fn from_entry(entry: &CachedEntry) -> Self {
971 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
972 let data = match &entry.data {
973 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
974 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
975 };
976 Self {
977 condition_cols: entry.condition_cols.clone(),
978 footprint_bits,
979 data,
980 }
981 }
982
983 fn into_entry(self) -> Option<CachedEntry> {
984 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
985 let data = match self.data {
986 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
987 SerializedData::Columns(c) => {
988 if !c.iter().all(|(_, col)| col.validate()) {
991 return None;
992 }
993 CachedData::Columns(Arc::new(c))
994 }
995 };
996 Some(CachedEntry {
997 data,
998 footprint,
999 condition_cols: self.condition_cols,
1000 })
1001 }
1002}
1003
1004impl ResultCache {
1005 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1006
1007 fn new() -> Self {
1008 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1009 }
1010
1011 fn with_max_bytes(max_bytes: u64) -> Self {
1012 Self {
1013 entries: std::collections::HashMap::new(),
1014 order: std::collections::VecDeque::new(),
1015 bytes: 0,
1016 max_bytes,
1017 dir: None,
1018 cache_dek: None,
1019 }
1020 }
1021
1022 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1023 let _ = std::fs::create_dir_all(&dir);
1024 self.dir = Some(dir);
1025 self
1026 }
1027
1028 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1029 self.cache_dek = dek;
1030 self
1031 }
1032
1033 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1034 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1035 }
1036
1037 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1041 let Some(path) = self.disk_path(key) else {
1042 return;
1043 };
1044 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
1045 Ok(s) => s,
1046 Err(_) => return,
1047 };
1048 let on_disk = if let Some(dek) = &self.cache_dek {
1050 match self.encrypt_cache(&serialized, dek) {
1051 Some(b) => b,
1052 None => return,
1053 }
1054 } else {
1055 serialized
1056 };
1057 let tmp = path.with_extension("tmp");
1058 use std::io::Write;
1059 let write = || -> std::io::Result<()> {
1060 let mut f = std::fs::File::create(&tmp)?;
1061 f.write_all(&on_disk)?;
1062 f.flush()?;
1063 Ok(())
1064 };
1065 if write().is_err() {
1066 let _ = std::fs::remove_file(&tmp);
1067 return;
1068 }
1069 let _ = std::fs::rename(&tmp, &path);
1070 }
1071
1072 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1074 let path = self.disk_path(key)?;
1075 let bytes = std::fs::read(&path).ok()?;
1076 let plaintext = if let Some(dek) = &self.cache_dek {
1077 self.decrypt_cache(&bytes, dek)?
1078 } else {
1079 bytes
1080 };
1081 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1082 serialized.into_entry()
1083 }
1084
1085 fn remove_from_disk(&self, key: u64) {
1087 if let Some(path) = self.disk_path(key) {
1088 let _ = std::fs::remove_file(&path);
1089 }
1090 }
1091
1092 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1094 use crate::encryption::Cipher;
1095 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1096 let mut nonce = [0u8; 12];
1097 crate::encryption::fill_random(&mut nonce).ok()?;
1098 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1099 let mut out = Vec::with_capacity(12 + ct.len());
1100 out.extend_from_slice(&nonce);
1101 out.extend_from_slice(&ct);
1102 Some(out)
1103 }
1104
1105 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1107 use crate::encryption::Cipher;
1108 if bytes.len() < 28 {
1109 return None;
1110 }
1111 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1112 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1113 let ct = &bytes[12..];
1114 cipher.decrypt_page(&nonce, ct).ok()
1115 }
1116
1117 fn load_persistent(&mut self) {
1120 let Some(dir) = self.dir.as_ref().cloned() else {
1121 return;
1122 };
1123 let entries = match std::fs::read_dir(&dir) {
1124 Ok(e) => e,
1125 Err(_) => return,
1126 };
1127 for entry in entries.flatten() {
1128 let path = entry.path();
1129 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1131 let _ = std::fs::remove_file(&path);
1132 continue;
1133 }
1134 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1135 continue;
1136 }
1137 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1138 Some(s) => s,
1139 None => continue,
1140 };
1141 let key = match u64::from_str_radix(stem, 16) {
1142 Ok(k) => k,
1143 Err(_) => continue,
1144 };
1145 let bytes = match std::fs::read(&path) {
1146 Ok(b) => b,
1147 Err(_) => continue,
1148 };
1149 let plaintext = if let Some(dek) = &self.cache_dek {
1151 match self.decrypt_cache(&bytes, dek) {
1152 Some(p) => p,
1153 None => {
1154 let _ = std::fs::remove_file(&path);
1155 continue;
1156 }
1157 }
1158 } else {
1159 bytes
1160 };
1161 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1162 Ok(serialized) => {
1163 if let Some(entry) = serialized.into_entry() {
1164 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1165 self.entries.insert(key, entry);
1166 self.order.push_back(key);
1167 } else {
1168 let _ = std::fs::remove_file(&path);
1169 }
1170 }
1171 Err(_) => {
1172 let _ = std::fs::remove_file(&path);
1173 }
1174 }
1175 }
1176 self.evict();
1177 }
1178
1179 fn set_max_bytes(&mut self, max_bytes: u64) {
1180 self.max_bytes = max_bytes;
1181 self.evict();
1182 }
1183
1184 fn touch(&mut self, key: u64) {
1186 self.order.retain(|k| *k != key);
1187 self.order.push_back(key);
1188 }
1189
1190 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1191 let res = self.entries.get(&key).and_then(|e| match &e.data {
1192 CachedData::Rows(r) => Some(r.clone()),
1193 CachedData::Columns(_) => None,
1194 });
1195 if res.is_some() {
1196 self.touch(key);
1197 return res;
1198 }
1199 if let Some(entry) = self.load_from_disk(key) {
1201 let res = match &entry.data {
1202 CachedData::Rows(r) => Some(r.clone()),
1203 CachedData::Columns(_) => None,
1204 };
1205 if res.is_some() {
1206 let approx = entry.data.approx_bytes();
1207 self.bytes = self.bytes.saturating_add(approx);
1208 self.entries.insert(key, entry);
1209 self.order.push_back(key);
1210 self.evict();
1211 return res;
1212 }
1213 }
1214 None
1215 }
1216
1217 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1218 let res = self.entries.get(&key).and_then(|e| match &e.data {
1219 CachedData::Columns(c) => Some(c.clone()),
1220 CachedData::Rows(_) => None,
1221 });
1222 if res.is_some() {
1223 self.touch(key);
1224 return res;
1225 }
1226 if let Some(entry) = self.load_from_disk(key) {
1228 let res = match &entry.data {
1229 CachedData::Columns(c) => Some(c.clone()),
1230 CachedData::Rows(_) => None,
1231 };
1232 if res.is_some() {
1233 let approx = entry.data.approx_bytes();
1234 self.bytes = self.bytes.saturating_add(approx);
1235 self.entries.insert(key, entry);
1236 self.order.push_back(key);
1237 self.evict();
1238 return res;
1239 }
1240 }
1241 None
1242 }
1243
1244 fn insert(&mut self, key: u64, entry: CachedEntry) {
1245 let approx = entry.data.approx_bytes();
1246 if self.entries.remove(&key).is_some() {
1247 self.order.retain(|k| *k != key);
1248 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1249 }
1250 self.store_to_disk(key, &entry);
1252 self.bytes = self.bytes.saturating_add(approx);
1253 self.entries.insert(key, entry);
1254 self.order.push_back(key);
1255 self.evict();
1256 }
1257
1258 fn invalidate(
1267 &mut self,
1268 delete_rids: &roaring::RoaringBitmap,
1269 put_cols: &std::collections::HashSet<u16>,
1270 ) {
1271 if self.entries.is_empty() {
1272 return;
1273 }
1274 let has_deletes = !delete_rids.is_empty();
1275 let to_remove: std::collections::HashSet<u64> = self
1276 .entries
1277 .iter()
1278 .filter(|(_, e)| {
1279 let delete_hit = if e.footprint.is_empty() {
1280 has_deletes
1281 } else {
1282 e.footprint.intersection_len(delete_rids) > 0
1283 };
1284 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1285 delete_hit || col_hit
1286 })
1287 .map(|(&k, _)| k)
1288 .collect();
1289 for key in &to_remove {
1290 if let Some(e) = self.entries.remove(key) {
1291 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1292 }
1293 self.remove_from_disk(*key);
1294 }
1295 if !to_remove.is_empty() {
1296 self.order.retain(|k| !to_remove.contains(k));
1297 }
1298 }
1299
1300 fn clear(&mut self) {
1301 if let Some(dir) = &self.dir {
1303 if let Ok(entries) = std::fs::read_dir(dir) {
1304 for entry in entries.flatten() {
1305 let path = entry.path();
1306 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1307 let _ = std::fs::remove_file(&path);
1308 }
1309 }
1310 }
1311 }
1312 self.entries.clear();
1313 self.order.clear();
1314 self.bytes = 0;
1315 }
1316
1317 fn evict(&mut self) {
1318 while self.bytes > self.max_bytes {
1319 let Some(k) = self.order.pop_front() else {
1320 break;
1321 };
1322 if let Some(e) = self.entries.remove(&k) {
1323 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1324 self.remove_from_disk(k);
1328 }
1329 }
1330 }
1331}
1332
1333type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1340
1341fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1342 let _ = kek;
1343 {
1344 if let Some(k) = kek {
1345 return (
1346 Some(k.derive_table_wal_key(_table_id)),
1347 Some(k.derive_cache_key()),
1348 );
1349 }
1350 }
1351 (None, None)
1352}
1353
1354fn read_table_encryption_salt_root(
1355 root: &crate::durable_file::DurableRoot,
1356) -> Result<[u8; crate::encryption::SALT_LEN]> {
1357 use std::io::Read;
1358
1359 let mut file = root
1360 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1361 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1362 let length = file.metadata()?.len();
1363 if length != crate::encryption::SALT_LEN as u64 {
1364 return Err(MongrelError::InvalidArgument(format!(
1365 "salt file is {length} bytes, expected {}",
1366 crate::encryption::SALT_LEN
1367 )));
1368 }
1369 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1370 file.read_exact(&mut salt)?;
1371 Ok(salt)
1372}
1373
1374fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1376 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1377}
1378
1379fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1380 let Some(kek) = kek else {
1381 return HashMap::new();
1382 };
1383 {
1384 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1385 schema
1386 .columns
1387 .iter()
1388 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1389 .map(|c| {
1390 let scheme = if schema
1391 .indexes
1392 .iter()
1393 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1394 {
1395 SCHEME_OPE_RANGE
1396 } else {
1397 SCHEME_HMAC_EQ
1398 };
1399 let key: [u8; 32] = *kek.derive_column_key(c.id);
1400 (c.id, (key, scheme))
1401 })
1402 .collect()
1403 }
1404}
1405
1406pub(crate) struct SharedCtx {
1411 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1412 pub epoch: Arc<EpochAuthority>,
1413 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1414 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1415 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1416 pub kek: Option<Arc<Kek>>,
1417 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1423 pub shared: Option<SharedWalCtx>,
1427 pub table_name: Option<String>,
1430 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1433 pub read_only: bool,
1435}
1436
1437#[derive(Clone)]
1443pub(crate) struct SharedWalCtx {
1444 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1445 pub group: Arc<GroupCommit>,
1446 pub poisoned: Arc<AtomicBool>,
1447 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1448 pub change_wake: tokio::sync::broadcast::Sender<()>,
1449 pub lifecycle: Arc<crate::core::LifecycleController>,
1452}
1453
1454enum WalSink {
1457 Private(Wal),
1458 Shared(SharedWalCtx),
1459 ReadOnly,
1460}
1461
1462impl Clone for WalSink {
1463 fn clone(&self) -> Self {
1464 match self {
1465 Self::Shared(shared) => Self::Shared(shared.clone()),
1466 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1467 }
1468 }
1469}
1470
1471impl SharedCtx {
1472 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1476 let n_shards = if cache_dir.is_some() {
1480 1
1481 } else {
1482 crate::cache::CACHE_SHARDS
1483 };
1484 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1485 let page_cache = if let Some(d) = cache_dir {
1486 Arc::new(crate::cache::Sharded::new(1, || {
1487 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1488 }))
1489 } else {
1490 Arc::new(crate::cache::Sharded::new(n_shards, || {
1491 crate::cache::PageCache::new(per_shard)
1492 }))
1493 };
1494 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1495 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1496 crate::cache::CACHE_SHARDS,
1497 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1498 ));
1499 Self {
1500 root_guard: None,
1501 epoch: Arc::new(EpochAuthority::new(0)),
1502 page_cache,
1503 decoded_cache,
1504 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1505 kek,
1506 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1507 shared: None,
1508 table_name: None,
1509 auth: None,
1510 read_only: false,
1511 }
1512 }
1513}
1514
1515fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1519 use crate::query::Condition;
1520 match c {
1521 Condition::Pk(_)
1523 | Condition::BitmapEq { .. }
1524 | Condition::BitmapIn { .. }
1525 | Condition::BytesPrefix { .. }
1526 | Condition::IsNull { .. }
1527 | Condition::IsNotNull { .. } => 0,
1528 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1530 1
1531 }
1532 Condition::FmContains { .. }
1534 | Condition::FmContainsAll { .. }
1535 | Condition::Ann { .. }
1536 | Condition::SparseMatch { .. } => 2,
1537 }
1538}
1539
1540impl Table {
1541 pub(crate) fn build_secondary_index_artifact<F>(
1547 &self,
1548 definition: &IndexDef,
1549 rows: &[Row],
1550 mut checkpoint: F,
1551 ) -> Result<SecondaryIndexArtifact>
1552 where
1553 F: FnMut(usize, usize) -> Result<()>,
1554 {
1555 let mut build_schema = self.schema.clone();
1556 build_schema.indexes = vec![definition.clone()];
1557 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
1558 let name_to_id: HashMap<&str, u16> = self
1559 .schema
1560 .columns
1561 .iter()
1562 .map(|column| (column.name.as_str(), column.id))
1563 .collect();
1564 let mut hot = HotIndex::new();
1565 let mut accepted = Vec::with_capacity(rows.len());
1566
1567 for (offset, row) in rows.iter().enumerate() {
1568 checkpoint(offset, rows.len())?;
1569 if row.deleted {
1570 continue;
1571 }
1572 if let Some(predicate) = &definition.predicate {
1573 let columns: HashMap<u16, &Value> =
1574 row.columns.iter().map(|(id, value)| (*id, value)).collect();
1575 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
1576 continue;
1577 }
1578 }
1579 accepted.push(row);
1580 if definition.kind == IndexKind::LearnedRange {
1581 continue;
1582 }
1583 let effective = if self.column_keys.is_empty() {
1584 row.clone()
1585 } else {
1586 self.tokenized_for_indexes(row)
1587 };
1588 if definition.kind == IndexKind::Ann {
1589 if let (Some(index), Some(vector)) = (
1590 ann.get_mut(&definition.column_id),
1591 effective
1592 .columns
1593 .get(&definition.column_id)
1594 .and_then(Value::as_embedding),
1595 ) {
1596 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
1597 checkpoint(offset, rows.len())
1598 })?;
1599 }
1600 } else {
1601 index_into_single(
1602 definition,
1603 &build_schema,
1604 &effective,
1605 &mut hot,
1606 &mut bitmap,
1607 &mut ann,
1608 &mut fm,
1609 &mut sparse,
1610 &mut minhash,
1611 );
1612 }
1613 }
1614 checkpoint(rows.len(), rows.len())?;
1615
1616 let column_id = definition.column_id;
1617 match definition.kind {
1618 IndexKind::Bitmap => bitmap
1619 .remove(&column_id)
1620 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
1621 IndexKind::Ann => ann
1622 .remove(&column_id)
1623 .map(|index| SecondaryIndexArtifact::Ann(column_id, index)),
1624 IndexKind::FmIndex => fm
1625 .remove(&column_id)
1626 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
1627 IndexKind::Sparse => sparse
1628 .remove(&column_id)
1629 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
1630 IndexKind::MinHash => minhash
1631 .remove(&column_id)
1632 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
1633 IndexKind::LearnedRange => {
1634 let epsilon = definition
1635 .options
1636 .learned_range
1637 .as_ref()
1638 .map(|options| options.epsilon)
1639 .unwrap_or(16);
1640 let column = self
1641 .schema
1642 .columns
1643 .iter()
1644 .find(|column| column.id == column_id)
1645 .ok_or_else(|| {
1646 MongrelError::Schema(format!(
1647 "index {} references unknown column {column_id}",
1648 definition.name
1649 ))
1650 })?;
1651 let index = match column.ty {
1652 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1653 let pairs: Vec<_> = accepted
1654 .iter()
1655 .filter_map(|row| match row.columns.get(&column_id) {
1656 Some(Value::Int64(value)) if !row.deleted => {
1657 Some((*value, row.row_id.0))
1658 }
1659 _ => None,
1660 })
1661 .collect();
1662 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
1663 }
1664 TypeId::Float32 | TypeId::Float64 => {
1665 let pairs: Vec<_> = accepted
1666 .iter()
1667 .filter_map(|row| match row.columns.get(&column_id) {
1668 Some(Value::Float64(value)) if !row.deleted => {
1669 Some((*value, row.row_id.0))
1670 }
1671 _ => None,
1672 })
1673 .collect();
1674 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
1675 }
1676 ref ty => {
1677 return Err(MongrelError::Schema(format!(
1678 "LearnedRange index {} does not support {ty:?}",
1679 definition.name
1680 )));
1681 }
1682 };
1683 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
1684 }
1685 }
1686 .ok_or_else(|| {
1687 MongrelError::Other(format!(
1688 "failed to construct hidden index {}",
1689 definition.name
1690 ))
1691 })
1692 }
1693
1694 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1695 let dir = dir.as_ref().to_path_buf();
1696 crate::durable_file::create_directory_all(&dir)?;
1697 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1698 let pinned = root.io_path()?;
1699 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1700 ctx.root_guard = Some(root);
1701 Self::create_in(&pinned, schema, table_id, ctx)
1702 }
1703
1704 pub fn create_encrypted(
1715 dir: impl AsRef<Path>,
1716 schema: Schema,
1717 table_id: u64,
1718 passphrase: &str,
1719 ) -> Result<Self> {
1720 let dir = dir.as_ref().to_path_buf();
1721 crate::durable_file::create_directory_all(&dir)?;
1722 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1723 root.create_directory_all(META_DIR)?;
1724 let salt = crate::encryption::random_salt()?;
1725 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1726 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1727 let pinned = root.io_path()?;
1728 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1729 ctx.root_guard = Some(root);
1730 Self::create_in(&pinned, schema, table_id, ctx)
1731 }
1732
1733 pub fn create_with_key(
1738 dir: impl AsRef<Path>,
1739 schema: Schema,
1740 table_id: u64,
1741 key: &[u8],
1742 ) -> Result<Self> {
1743 let dir = dir.as_ref().to_path_buf();
1744 crate::durable_file::create_directory_all(&dir)?;
1745 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1746 root.create_directory_all(META_DIR)?;
1747 let salt = crate::encryption::random_salt()?;
1748 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1749 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1750 let pinned = root.io_path()?;
1751 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1752 ctx.root_guard = Some(root);
1753 Self::create_in(&pinned, schema, table_id, ctx)
1754 }
1755
1756 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1758 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1759 let salt = read_table_encryption_salt_root(&root)?;
1760 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1761 let pinned = root.io_path()?;
1762 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1763 ctx.root_guard = Some(root);
1764 Self::open_in(&pinned, ctx)
1765 }
1766
1767 pub(crate) fn create_in(
1768 dir: impl AsRef<Path>,
1769 schema: Schema,
1770 table_id: u64,
1771 ctx: SharedCtx,
1772 ) -> Result<Self> {
1773 schema.validate_auto_increment()?;
1774 schema.validate_defaults()?;
1775 schema.validate_ai()?;
1776 for index in &schema.indexes {
1777 index.validate_options()?;
1778 }
1779 let dir = dir.as_ref().to_path_buf();
1780 let runs_root = match ctx.root_guard.as_ref() {
1781 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1782 None => {
1783 crate::durable_file::create_directory_all(&dir)?;
1784 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1785 None
1786 }
1787 };
1788 match ctx.root_guard.as_deref() {
1789 Some(root) => write_schema_durable(root, &schema)?,
1790 None => write_schema(&dir, &schema)?,
1791 }
1792 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1793 let (wal, current_txn_id) = match ctx.shared.clone() {
1796 Some(s) => (WalSink::Shared(s), 0),
1797 None => {
1798 let pinned_wal_root = match ctx.root_guard.as_deref() {
1799 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1800 None => None,
1801 };
1802 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1803 root.io_path()?
1804 } else {
1805 let wal_dir = dir.join(WAL_DIR);
1806 crate::durable_file::create_directory_all(&wal_dir)?;
1807 wal_dir
1808 };
1809 let mut w = if let Some(ref dk) = wal_dek {
1810 Wal::create_with_cipher(
1811 wal_dir.join("seg-000000.wal"),
1812 Epoch(0),
1813 Some(make_cipher(dk)),
1814 0,
1815 )?
1816 } else {
1817 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1818 };
1819 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1820 (WalSink::Private(w), 1)
1821 }
1822 };
1823 let mut manifest = Manifest::new(table_id, schema.schema_id);
1824 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1829 match ctx.root_guard.as_deref() {
1830 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1831 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1832 }
1833 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1834 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1835 let auto_inc = resolve_auto_inc(&schema);
1836 let rcache_dir = dir.join(RCACHE_DIR);
1837 let initial_view = ReadGeneration::empty(&schema);
1838 Ok(Self {
1839 dir,
1840 _root_guard: ctx.root_guard,
1841 runs_root,
1842 idx_root: None,
1843 table_id,
1844 name: ctx.table_name.unwrap_or_default(),
1845 auth: ctx.auth,
1846 read_only: ctx.read_only,
1847 durable_commit_failed: false,
1848 wal,
1849 memtable: Memtable::new(),
1850 mutable_run: MutableRun::new(),
1851 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1852 compaction_zstd_level: 3,
1853 allocator: RowIdAllocator::new(0),
1854 epoch: ctx.epoch,
1855 data_generation: 0,
1856 schema,
1857 hot: HotIndex::new(),
1858 kek: ctx.kek,
1859 column_keys,
1860 run_refs: Vec::new(),
1861 retiring: Vec::new(),
1862 next_run_id: 1,
1863 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1864 current_txn_id,
1865 pending_private_mutations: false,
1866 bitmap,
1867 ann,
1868 fm,
1869 sparse,
1870 minhash,
1871 learned_range: Arc::new(HashMap::new()),
1872 pk_by_row: ReversePkMap::new(),
1873 pinned: BTreeMap::new(),
1874 live_count: 0,
1875 reservoir: crate::reservoir::Reservoir::default(),
1876 reservoir_complete: true,
1877 had_deletes: false,
1878 agg_cache: Arc::new(HashMap::new()),
1879 global_idx_epoch: 0,
1880 indexes_complete: true,
1881 index_build_policy: IndexBuildPolicy::default(),
1882 pk_by_row_complete: false,
1883 flushed_epoch: 0,
1884 page_cache: ctx.page_cache,
1885 decoded_cache: ctx.decoded_cache,
1886 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1887 snapshots: ctx.snapshots,
1888 commit_lock: ctx.commit_lock,
1889 result_cache: Arc::new(parking_lot::Mutex::new(
1890 ResultCache::new()
1891 .with_dir(rcache_dir)
1892 .with_cache_dek(cache_dek.clone()),
1893 )),
1894 pending_delete_rids: roaring::RoaringBitmap::new(),
1895 pending_put_cols: std::collections::HashSet::new(),
1896 pending_rows: Vec::new(),
1897 pending_rows_auto_inc: Vec::new(),
1898 pending_dels: Vec::new(),
1899 pending_truncate: None,
1900 wal_dek,
1901 auto_inc,
1902 ttl: None,
1903 pins: Arc::new(crate::retention::PinRegistry::new()),
1904 published: Arc::new(ArcSwap::from_pointee(initial_view)),
1905 read_generation_pin: None,
1906 })
1907 }
1908
1909 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1913 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1914 let pinned = root.io_path()?;
1915 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1916 ctx.root_guard = Some(root);
1917 Self::open_in(&pinned, ctx)
1918 }
1919
1920 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1923 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1924 let salt = read_table_encryption_salt_root(&root)?;
1925 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1926 let pinned = root.io_path()?;
1927 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1928 ctx.root_guard = Some(root);
1929 let t = Self::open_in(&pinned, ctx)?;
1930 Ok(t)
1931 }
1932
1933 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1934 let dir = dir.as_ref().to_path_buf();
1935 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1936 let mut manifest = match ctx.root_guard.as_ref() {
1937 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1938 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1939 };
1940 let schema: Schema = match ctx.root_guard.as_ref() {
1941 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1942 None => read_schema(&dir)?,
1943 };
1944 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1950 let runs_root = match ctx.root_guard.as_ref() {
1951 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1952 None => None,
1953 };
1954 let idx_root = match ctx.root_guard.as_ref() {
1955 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1956 Ok(root) => Some(Arc::new(root)),
1957 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1958 Err(error) => return Err(error.into()),
1959 },
1960 None => None,
1961 };
1962 schema.validate_auto_increment()?;
1963 schema.validate_defaults()?;
1964 schema.validate_ai()?;
1965 for index in &schema.indexes {
1966 index.validate_options()?;
1967 }
1968 let replay_epoch = Epoch(manifest.current_epoch);
1969 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1970 let private_replayed = if ctx.shared.is_none() {
1971 match latest_wal_segment(&dir.join(WAL_DIR))? {
1972 Some(path) => {
1973 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1974 crate::wal::replay_with_cipher(path, cipher)?
1975 }
1976 None => Vec::new(),
1977 }
1978 } else {
1979 Vec::new()
1980 };
1981 if ctx.shared.is_none() {
1982 preflight_standalone_open(
1983 &dir,
1984 runs_root.as_deref(),
1985 idx_root.as_deref(),
1986 &manifest,
1987 &schema,
1988 &private_replayed,
1989 ctx.kek.clone(),
1990 )?;
1991 }
1992 let next_run_id = derive_next_run_id(
1993 &dir,
1994 runs_root.as_deref(),
1995 &manifest.runs,
1996 &manifest.retiring,
1997 )?;
1998 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2002 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2003 None => {
2004 let replayed = private_replayed;
2005 let wal_dir = dir.join(WAL_DIR);
2011 crate::durable_file::create_directory_all(&wal_dir)?;
2012 let segment = next_wal_segment(&wal_dir)?;
2013 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2014 let temporary = wal_dir.join(format!(
2015 ".recovery-{}-{}-{segment_no:06}.tmp",
2016 std::process::id(),
2017 std::time::SystemTime::now()
2018 .duration_since(std::time::UNIX_EPOCH)
2019 .unwrap_or_default()
2020 .as_nanos()
2021 ));
2022 let mut w = Wal::create_with_cipher(
2023 &temporary,
2024 replay_epoch,
2025 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2026 segment_no,
2027 )?;
2028 for record in &replayed {
2029 w.append_txn(record.txn_id, record.op.clone())?;
2030 }
2031 let mut w = w.publish_as(segment)?;
2032 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2033 let next_txn_id = replayed
2034 .iter()
2035 .map(|record| record.txn_id)
2036 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2037 .max()
2038 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2039 .unwrap_or(1);
2040 (WalSink::Private(w), replayed, next_txn_id)
2041 }
2042 };
2043
2044 let mut memtable = Memtable::new();
2045 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2046 let persisted_epoch = manifest.current_epoch;
2047 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2054 s.next = manifest.auto_inc_next;
2055 s.seeded = manifest.auto_inc_next > 0;
2056 s
2057 });
2058
2059 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2066 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2067 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2068 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2069 std::collections::BTreeMap::new();
2070 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2071 let mut recovered_epoch = manifest.current_epoch;
2072 let mut recovered_manifest_dirty = schema_manifest_repair;
2073 let mut saw_delete = false;
2074 for record in replayed {
2075 let txn_id = record.txn_id;
2076 match record.op {
2077 Op::Put { rows, .. } => {
2078 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2079 for row in &rows {
2080 allocator.advance_to(row.row_id)?;
2081 if let Some(ai) = auto_inc.as_mut() {
2082 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2083 let next = n.checked_add(1).ok_or_else(|| {
2084 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2085 })?;
2086 if next > ai.next {
2087 ai.next = next;
2088 }
2089 }
2090 }
2091 }
2092 staged_puts.entry(txn_id).or_default().extend(rows);
2093 }
2094 Op::Delete { row_ids, .. } => {
2095 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2096 }
2097 Op::TxnCommit { epoch, .. } => {
2098 let commit_epoch = Epoch(epoch);
2099 recovered_epoch = recovered_epoch.max(epoch);
2100 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2101 memtable = Memtable::new();
2102 replayed_puts.clear();
2103 replayed_deletes.clear();
2104 manifest.runs.clear();
2105 manifest.retiring.clear();
2106 manifest.live_count = 0;
2107 manifest.global_idx_epoch = 0;
2108 manifest.current_epoch = manifest.current_epoch.max(epoch);
2109 recovered_manifest_dirty = true;
2110 saw_delete = true;
2111 }
2112 if let Some(puts) = staged_puts.remove(&txn_id) {
2113 if commit_epoch.0 > manifest.flushed_epoch {
2114 for row in &puts {
2115 memtable.upsert(row.clone());
2116 }
2117 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2118 }
2119 }
2120 if let Some(dels) = staged_deletes.remove(&txn_id) {
2121 saw_delete = true;
2122 if commit_epoch.0 > manifest.flushed_epoch {
2123 for rid in dels {
2124 memtable.tombstone(rid, commit_epoch);
2125 replayed_deletes.push((rid, commit_epoch));
2126 }
2127 }
2128 }
2129 }
2130 Op::TxnAbort => {
2131 staged_puts.remove(&txn_id);
2132 staged_deletes.remove(&txn_id);
2133 staged_truncates.remove(&txn_id);
2134 }
2135 Op::TruncateTable { .. } => {
2136 staged_puts.remove(&txn_id);
2137 staged_deletes.remove(&txn_id);
2138 staged_truncates.insert(txn_id);
2139 }
2140 Op::ExternalTableState { .. }
2141 | Op::Flush { .. }
2142 | Op::Ddl(_)
2143 | Op::BeforeImage { .. }
2144 | Op::CommitTimestamp { .. }
2145 | Op::SpilledRows { .. } => {}
2146 }
2147 }
2148
2149 let rcache_dir = dir.join(RCACHE_DIR);
2150 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2151 let initial_view = ReadGeneration::empty(&schema);
2152 let mut db = Self {
2153 dir,
2154 _root_guard: ctx.root_guard,
2155 runs_root,
2156 idx_root,
2157 table_id: manifest.table_id,
2158 name: ctx.table_name.unwrap_or_default(),
2159 auth: ctx.auth,
2160 read_only: ctx.read_only,
2161 durable_commit_failed: false,
2162 wal,
2163 memtable,
2164 mutable_run: MutableRun::new(),
2165 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2166 compaction_zstd_level: 3,
2167 allocator,
2168 epoch: ctx.epoch,
2169 data_generation: persisted_epoch,
2170 schema,
2171 hot: HotIndex::new(),
2172 kek: ctx.kek,
2173 column_keys,
2174 run_refs: manifest.runs.clone(),
2175 retiring: manifest.retiring.clone(),
2176 next_run_id,
2177 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2178 current_txn_id,
2179 pending_private_mutations: false,
2180 bitmap: HashMap::new(),
2181 ann: HashMap::new(),
2182 fm: HashMap::new(),
2183 sparse: HashMap::new(),
2184 minhash: HashMap::new(),
2185 learned_range: Arc::new(HashMap::new()),
2186 pk_by_row: ReversePkMap::new(),
2187 pinned: BTreeMap::new(),
2188 live_count: manifest.live_count,
2189 reservoir: crate::reservoir::Reservoir::default(),
2190 reservoir_complete: false,
2191 had_deletes: saw_delete
2192 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2193 != manifest.live_count,
2194 agg_cache: Arc::new(HashMap::new()),
2195 global_idx_epoch: manifest.global_idx_epoch,
2196 indexes_complete: true,
2197 index_build_policy: IndexBuildPolicy::default(),
2198 pk_by_row_complete: false,
2199 flushed_epoch: manifest.flushed_epoch,
2200 page_cache: ctx.page_cache,
2201 decoded_cache: ctx.decoded_cache,
2202 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2203 snapshots: ctx.snapshots,
2204 commit_lock: ctx.commit_lock,
2205 result_cache: Arc::new(parking_lot::Mutex::new(
2206 ResultCache::new()
2207 .with_dir(rcache_dir)
2208 .with_cache_dek(cache_dek.clone()),
2209 )),
2210 pending_delete_rids: roaring::RoaringBitmap::new(),
2211 pending_put_cols: std::collections::HashSet::new(),
2212 pending_rows: Vec::new(),
2213 pending_rows_auto_inc: Vec::new(),
2214 pending_dels: Vec::new(),
2215 pending_truncate: None,
2216 wal_dek,
2217 auto_inc,
2218 ttl: manifest.ttl,
2219 pins: Arc::new(crate::retention::PinRegistry::new()),
2220 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2221 read_generation_pin: None,
2222 };
2223
2224 db.epoch.advance_recovered(Epoch(recovered_epoch));
2227
2228 let checkpoint = match db.idx_root.as_deref() {
2233 Some(root) => {
2234 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2235 }
2236 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2237 };
2238 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2239 c.epoch_built == manifest.global_idx_epoch
2240 && manifest.global_idx_epoch > 0
2241 && manifest
2242 .runs
2243 .iter()
2244 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2245 });
2246 if let Some(loaded) = checkpoint {
2247 if checkpoint_valid {
2248 db.hot = loaded.hot;
2249 db.bitmap = loaded.bitmap;
2250 db.ann = loaded.ann;
2251 db.fm = loaded.fm;
2252 db.sparse = loaded.sparse;
2253 db.minhash = loaded.minhash;
2254 db.learned_range = Arc::new(loaded.learned_range);
2255 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
2260 for (cid, idx) in bitmap0 {
2261 db.bitmap.entry(cid).or_insert(idx);
2262 }
2263 for (cid, idx) in ann0 {
2264 db.ann.entry(cid).or_insert(idx);
2265 }
2266 for (cid, idx) in fm0 {
2267 db.fm.entry(cid).or_insert(idx);
2268 }
2269 for (cid, idx) in sparse0 {
2270 db.sparse.entry(cid).or_insert(idx);
2271 }
2272 for (cid, idx) in minhash0 {
2273 db.minhash.entry(cid).or_insert(idx);
2274 }
2275 }
2278 }
2279 if !checkpoint_valid {
2280 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
2281 db.bitmap = bitmap;
2282 db.ann = ann;
2283 db.fm = fm;
2284 db.sparse = sparse;
2285 db.minhash = minhash;
2286 db.rebuild_indexes_from_runs()?;
2287 db.build_learned_ranges()?;
2288 }
2289
2290 for (epoch, group) in replayed_puts {
2295 let (losers, winner_pks) = db.partition_pk_winners(&group);
2296 for (key, &row_id) in &winner_pks {
2297 if let Some(old_rid) = db.hot.get(key) {
2298 if old_rid != row_id {
2299 db.tombstone_row(old_rid, epoch, false);
2300 }
2301 }
2302 }
2303 for &loser_rid in &losers {
2304 db.tombstone_row(loser_rid, epoch, false);
2305 }
2306 for (key, row_id) in winner_pks {
2307 db.insert_hot_pk(key, row_id);
2308 }
2309 if db.schema.primary_key().is_none() {
2310 for r in &group {
2311 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2312 }
2313 }
2314 for r in &group {
2315 if !losers.contains(&r.row_id) {
2316 db.index_row(r);
2317 }
2318 }
2319 }
2320 for (rid, epoch) in &replayed_deletes {
2324 db.remove_hot_for_row(*rid, *epoch);
2325 }
2326
2327 if recovered_manifest_dirty {
2328 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2329 db.live_count = rows.len() as u64;
2330 db.persist_manifest(Epoch(recovered_epoch))?;
2331 }
2332
2333 db.result_cache.lock().load_persistent();
2340 Ok(db)
2341 }
2342
2343 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2349 if self.reservoir_complete {
2350 return Ok(());
2351 }
2352 self.rebuild_reservoir()?;
2353 self.reservoir_complete = true;
2354 Ok(())
2355 }
2356
2357 fn rebuild_reservoir(&mut self) -> Result<()> {
2360 let snap = self.snapshot();
2361 let rows = self.visible_rows(snap)?;
2362 self.reservoir.reset();
2363 for r in rows {
2364 self.reservoir.offer(r.row_id.0);
2365 }
2366 Ok(())
2367 }
2368
2369 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2370 self.rebuild_indexes_from_runs_inner(None)
2371 }
2372
2373 fn rebuild_indexes_from_runs_inner(
2374 &mut self,
2375 control: Option<&crate::ExecutionControl>,
2376 ) -> Result<()> {
2377 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
2380 crate::retention::PinSource::OnlineIndexBuild,
2381 self.current_epoch(),
2382 );
2383 self.hot = HotIndex::new();
2384 self.pk_by_row.clear();
2385 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2386 self.bitmap = bitmap;
2387 self.ann = ann;
2388 self.fm = fm;
2389 self.sparse = sparse;
2390 self.minhash = minhash;
2391 let snapshot = Epoch(u64::MAX);
2392 let ttl_now = unix_nanos_now();
2393 let mut scanned = 0_usize;
2394 for rr in self.run_refs.clone() {
2395 if let Some(control) = control {
2396 control.checkpoint()?;
2397 }
2398 let mut reader = self.open_reader(rr.run_id)?;
2399 for row in reader.visible_rows(snapshot)? {
2400 if scanned.is_multiple_of(256) {
2401 if let Some(control) = control {
2402 control.checkpoint()?;
2403 }
2404 }
2405 scanned += 1;
2406 if self.row_expired_at(&row, ttl_now) {
2407 continue;
2408 }
2409 let tok_row = self.tokenized_for_indexes(&row);
2410 index_into(
2411 &self.schema,
2412 &tok_row,
2413 &mut self.hot,
2414 &mut self.bitmap,
2415 &mut self.ann,
2416 &mut self.fm,
2417 &mut self.sparse,
2418 &mut self.minhash,
2419 );
2420 }
2421 }
2422 for row in self.mutable_run.visible_versions(snapshot) {
2423 if scanned.is_multiple_of(256) {
2424 if let Some(control) = control {
2425 control.checkpoint()?;
2426 }
2427 }
2428 scanned += 1;
2429 if row.deleted {
2430 self.remove_hot_for_row(row.row_id, snapshot);
2431 } else if !self.row_expired_at(&row, ttl_now) {
2432 self.index_row(&row);
2433 }
2434 }
2435 for row in self.memtable.visible_versions(snapshot) {
2436 if scanned.is_multiple_of(256) {
2437 if let Some(control) = control {
2438 control.checkpoint()?;
2439 }
2440 }
2441 scanned += 1;
2442 if row.deleted {
2443 self.remove_hot_for_row(row.row_id, snapshot);
2444 } else if !self.row_expired_at(&row, ttl_now) {
2445 self.index_row(&row);
2446 }
2447 }
2448 self.refresh_pk_by_row_from_hot();
2449 Ok(())
2450 }
2451
2452 fn refresh_pk_by_row_from_hot(&mut self) {
2453 self.pk_by_row_complete = true;
2454 if self.schema.primary_key().is_none() {
2455 self.pk_by_row.clear();
2456 return;
2457 }
2458 self.pk_by_row = ReversePkMap::from_entries(
2464 self.hot
2465 .entries()
2466 .into_iter()
2467 .map(|(key, row_id)| (row_id, key)),
2468 );
2469 }
2470
2471 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2472 if self.schema.primary_key().is_some() {
2473 self.pk_by_row.insert(row_id, key.clone());
2474 }
2475 self.hot.insert(key, row_id);
2476 }
2477
2478 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2482 self.build_learned_ranges_inner(None)
2483 }
2484
2485 fn build_learned_ranges_inner(
2486 &mut self,
2487 control: Option<&crate::ExecutionControl>,
2488 ) -> Result<()> {
2489 self.learned_range = Arc::new(HashMap::new());
2490 if self.run_refs.len() != 1 {
2491 return Ok(());
2492 }
2493 let cols: Vec<(u16, usize)> = self
2494 .schema
2495 .indexes
2496 .iter()
2497 .filter(|i| i.kind == IndexKind::LearnedRange)
2498 .map(|i| {
2499 (
2500 i.column_id,
2501 i.options
2502 .learned_range
2503 .as_ref()
2504 .map(|options| options.epsilon)
2505 .unwrap_or(16),
2506 )
2507 })
2508 .collect();
2509 if cols.is_empty() {
2510 return Ok(());
2511 }
2512 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2513 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2514 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2515 _ => return Ok(()),
2516 };
2517 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2518 if column_index % 256 == 0 {
2519 if let Some(control) = control {
2520 control.checkpoint()?;
2521 }
2522 }
2523 let ty = self
2524 .schema
2525 .columns
2526 .iter()
2527 .find(|c| c.id == cid)
2528 .map(|c| c.ty.clone())
2529 .unwrap_or(TypeId::Int64);
2530 match ty {
2531 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2532 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2533 let pairs: Vec<(i64, u64)> = data
2534 .iter()
2535 .zip(row_ids.iter())
2536 .map(|(v, r)| (*v, *r))
2537 .collect();
2538 Arc::make_mut(&mut self.learned_range).insert(
2539 cid,
2540 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2541 );
2542 }
2543 }
2544 TypeId::Float64 => {
2545 if let columnar::NativeColumn::Float64 { data, .. } =
2546 reader.column_native(cid)?
2547 {
2548 let pairs: Vec<(f64, u64)> = data
2549 .iter()
2550 .zip(row_ids.iter())
2551 .map(|(v, r)| (*v, *r))
2552 .collect();
2553 Arc::make_mut(&mut self.learned_range).insert(
2554 cid,
2555 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2556 );
2557 }
2558 }
2559 _ => {}
2560 }
2561 }
2562 Ok(())
2563 }
2564
2565 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2572 if self.indexes_complete {
2573 crate::trace::QueryTrace::record(|t| {
2574 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2575 });
2576 return Ok(());
2577 }
2578 crate::trace::QueryTrace::record(|t| {
2579 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2580 });
2581 self.rebuild_indexes_from_runs()?;
2582 self.build_learned_ranges()?;
2583 self.indexes_complete = true;
2584 let epoch = self.current_epoch();
2585 self.checkpoint_indexes(epoch);
2586 Ok(())
2587 }
2588
2589 #[doc(hidden)]
2592 pub fn ensure_indexes_complete_controlled<F>(
2593 &mut self,
2594 control: &crate::ExecutionControl,
2595 before_publish: F,
2596 ) -> Result<bool>
2597 where
2598 F: FnOnce() -> bool,
2599 {
2600 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2601 .map(|(changed, _)| changed)
2602 }
2603
2604 #[doc(hidden)]
2607 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2608 &mut self,
2609 control: &crate::ExecutionControl,
2610 before_publish: F,
2611 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2612 where
2613 F: FnOnce() -> bool,
2614 {
2615 if self.indexes_complete {
2616 crate::trace::QueryTrace::record(|trace| {
2617 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2618 });
2619 return Ok((false, None));
2620 }
2621 crate::trace::QueryTrace::record(|trace| {
2622 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2623 });
2624 control.checkpoint()?;
2625 let maintenance_epoch = self.current_epoch();
2626 self.rebuild_indexes_from_runs_inner(Some(control))?;
2627 self.build_learned_ranges_inner(Some(control))?;
2628 control.checkpoint()?;
2629 if !before_publish() {
2630 return Err(MongrelError::Cancelled);
2631 }
2632 self.indexes_complete = true;
2633 self.checkpoint_indexes(maintenance_epoch);
2634 Ok((
2635 true,
2636 Some(MaintenanceReceipt {
2637 epoch: maintenance_epoch,
2638 }),
2639 ))
2640 }
2641
2642 fn pending_epoch(&self) -> Epoch {
2643 Epoch(self.epoch.visible().0 + 1)
2644 }
2645
2646 fn is_shared(&self) -> bool {
2649 matches!(self.wal, WalSink::Shared(_))
2650 }
2651
2652 fn ensure_txn_id(&mut self) -> Result<u64> {
2656 if self.current_txn_id == 0 {
2657 let id = match &self.wal {
2658 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2659 WalSink::Private(_) => {
2660 return Err(MongrelError::Full(
2661 "standalone transaction id namespace exhausted".into(),
2662 ))
2663 }
2664 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2665 };
2666 self.current_txn_id = id;
2667 }
2668 Ok(self.current_txn_id)
2669 }
2670
2671 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2674 self.ensure_writable()?;
2675 let txn_id = self.ensure_txn_id()?;
2676 let table_id = self.table_id;
2677 match &mut self.wal {
2678 WalSink::Private(w) => {
2679 w.append_txn(txn_id, op)?;
2680 self.pending_private_mutations = true;
2681 }
2682 WalSink::Shared(s) => {
2683 s.wal.lock().append(txn_id, table_id, op)?;
2684 }
2685 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2686 }
2687 Ok(())
2688 }
2689
2690 fn ensure_writable(&self) -> Result<()> {
2691 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2692 return Err(MongrelError::ReadOnlyReplica);
2693 }
2694 if self.durable_commit_failed {
2695 return Err(MongrelError::Other(
2696 "table poisoned by post-commit failure; reopen required".into(),
2697 ));
2698 }
2699 Ok(())
2700 }
2701
2702 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2713 match &self.auth {
2714 Some(checker) => checker.check(&self.name, perm),
2715 None => Ok(()),
2716 }
2717 }
2718 pub fn require_select(&self) -> Result<()> {
2723 self.require(crate::auth_state::RequiredPermission::Select)
2724 }
2725 fn require_insert(&self) -> Result<()> {
2726 self.require(crate::auth_state::RequiredPermission::Insert)
2727 }
2728 #[allow(dead_code)]
2732 fn require_update(&self) -> Result<()> {
2733 self.require(crate::auth_state::RequiredPermission::Update)
2734 }
2735 fn require_delete(&self) -> Result<()> {
2736 self.require(crate::auth_state::RequiredPermission::Delete)
2737 }
2738
2739 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2742 self.require_insert()?;
2743 Ok(self.put_returning(columns)?.0)
2744 }
2745
2746 pub fn put_returning(
2751 &mut self,
2752 mut columns: Vec<(u16, Value)>,
2753 ) -> Result<(RowId, Option<i64>)> {
2754 self.require_insert()?;
2755 let assigned = self.fill_auto_inc(&mut columns)?;
2756 self.apply_defaults(&mut columns)?;
2757 self.schema.validate_values(&columns)?;
2758 let row_id = if self.schema.clustered {
2763 self.derive_clustered_row_id(&columns)?
2764 } else {
2765 self.allocator.alloc()?
2766 };
2767 let epoch = self.pending_epoch();
2768 let mut row = Row::new(row_id, epoch);
2769 for (col_id, val) in columns {
2770 row.columns.insert(col_id, val);
2771 }
2772 self.commit_rows(vec![row], assigned.is_some())?;
2773 Ok((row_id, assigned))
2774 }
2775
2776 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2779 self.require_insert()?;
2780 Ok(self
2781 .put_batch_returning(batch)?
2782 .into_iter()
2783 .map(|(r, _)| r)
2784 .collect())
2785 }
2786
2787 pub fn put_batch_returning(
2790 &mut self,
2791 batch: Vec<Vec<(u16, Value)>>,
2792 ) -> Result<Vec<(RowId, Option<i64>)>> {
2793 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2794 for mut cols in batch {
2795 let assigned = self.fill_auto_inc(&mut cols)?;
2796 self.apply_defaults(&mut cols)?;
2797 filled.push((cols, assigned));
2798 }
2799 for (cols, _) in &filled {
2800 self.schema.validate_values(cols)?;
2801 }
2802 let epoch = self.pending_epoch();
2803 let mut rows = Vec::with_capacity(filled.len());
2804 let mut ids = Vec::with_capacity(filled.len());
2805 let first_row_id = if self.schema.clustered {
2806 None
2807 } else {
2808 let count = u64::try_from(filled.len())
2809 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2810 Some(self.allocator.alloc_range(count)?.0)
2811 };
2812 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2813 let row_id = match first_row_id {
2814 Some(first) => RowId(first + row_index as u64),
2815 None => self.derive_clustered_row_id(&cols)?,
2816 };
2817 let mut row = Row::new(row_id, epoch);
2818 for (c, v) in cols {
2819 row.columns.insert(c, v);
2820 }
2821 ids.push((row_id, assigned));
2822 rows.push(row);
2823 }
2824 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2825 self.commit_rows(rows, all_auto_generated)?;
2826 Ok(ids)
2827 }
2828
2829 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2835 self.ensure_writable()?;
2836 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2837 return Ok(None);
2838 };
2839 let pos = columns.iter().position(|(c, _)| *c == cid);
2840 let assigned = match pos {
2841 Some(i) => match &columns[i].1 {
2842 Value::Null => {
2843 let next = self.alloc_auto_inc_value()?;
2844 columns[i].1 = Value::Int64(next);
2845 Some(next)
2846 }
2847 Value::Int64(n) => {
2848 self.advance_auto_inc_past(*n)?;
2849 None
2850 }
2851 other => {
2852 return Err(MongrelError::InvalidArgument(format!(
2853 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2854 other
2855 )))
2856 }
2857 },
2858 None => {
2859 let next = self.alloc_auto_inc_value()?;
2860 columns.push((cid, Value::Int64(next)));
2861 Some(next)
2862 }
2863 };
2864 Ok(assigned)
2865 }
2866
2867 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2873 for col in &self.schema.columns {
2874 let Some(expr) = &col.default_value else {
2875 continue;
2876 };
2877 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2879 continue;
2880 }
2881 let pos = columns.iter().position(|(c, _)| *c == col.id);
2882 let needs_default = match pos {
2883 None => true,
2884 Some(i) => matches!(columns[i].1, Value::Null),
2885 };
2886 if !needs_default {
2887 continue;
2888 }
2889 let v = match expr {
2890 crate::schema::DefaultExpr::Static(v) => v.clone(),
2891 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2892 crate::schema::DefaultExpr::Uuid => {
2893 let mut buf = [0u8; 16];
2894 getrandom::getrandom(&mut buf)
2895 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2896 Value::Uuid(buf)
2897 }
2898 };
2899 match pos {
2900 None => columns.push((col.id, v)),
2901 Some(i) => columns[i].1 = v,
2902 }
2903 }
2904 Ok(())
2905 }
2906
2907 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2909 self.ensure_auto_inc_seeded()?;
2910 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2912 let v = ai.next;
2913 ai.next = ai
2914 .next
2915 .checked_add(1)
2916 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2917 Ok(v)
2918 }
2919
2920 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2923 self.ensure_auto_inc_seeded()?;
2924 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2925 let floor = used
2926 .checked_add(1)
2927 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2928 .max(1);
2929 if ai.next < floor {
2930 ai.next = floor;
2931 }
2932 Ok(())
2933 }
2934
2935 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2940 let needs_seed = match self.auto_inc {
2941 Some(ai) => !ai.seeded,
2942 None => return Ok(()),
2943 };
2944 if !needs_seed {
2945 return Ok(());
2946 }
2947 if self.seed_empty_auto_inc() {
2948 return Ok(());
2949 }
2950 let cid = self
2951 .auto_inc
2952 .as_ref()
2953 .expect("auto-inc column present")
2954 .column_id;
2955 let max = self.scan_max_int64(cid)?;
2956 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2957 let floor = max
2958 .checked_add(1)
2959 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2960 .max(1);
2961 if ai.next < floor {
2962 ai.next = floor;
2963 }
2964 ai.seeded = true;
2965 Ok(())
2966 }
2967
2968 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2969 if n == 0 || self.auto_inc.is_none() {
2970 return Ok(None);
2971 }
2972 self.ensure_auto_inc_seeded()?;
2973 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2974 let start = ai.next;
2975 let count = i64::try_from(n)
2976 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2977 ai.next = ai
2978 .next
2979 .checked_add(count)
2980 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2981 Ok(Some(start))
2982 }
2983
2984 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2988 let mut max: i64 = 0;
2989 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2990 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2991 if *n > max {
2992 max = *n;
2993 }
2994 }
2995 }
2996 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2997 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2998 if *n > max {
2999 max = *n;
3000 }
3001 }
3002 }
3003 for rr in self.run_refs.clone() {
3004 let reader = self.open_reader(rr.run_id)?;
3005 if let Some(stats) = reader.column_page_stats(column_id) {
3006 for s in stats {
3007 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3008 if n > max {
3009 max = n;
3010 }
3011 }
3012 }
3013 } else if reader.has_column(column_id) {
3014 if let columnar::NativeColumn::Int64 { data, validity } =
3015 reader.column_native_shared(column_id)?
3016 {
3017 for (i, n) in data.iter().enumerate() {
3018 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3019 {
3020 max = *n;
3021 }
3022 }
3023 }
3024 }
3025 }
3026 Ok(max)
3027 }
3028
3029 fn seed_empty_auto_inc(&mut self) -> bool {
3030 let Some(ai) = self.auto_inc.as_mut() else {
3031 return false;
3032 };
3033 if ai.seeded || self.live_count != 0 {
3034 return false;
3035 }
3036 if ai.next < 1 {
3037 ai.next = 1;
3038 }
3039 ai.seeded = true;
3040 true
3041 }
3042
3043 fn advance_auto_inc_from_native_columns(
3044 &mut self,
3045 columns: &[(u16, columnar::NativeColumn)],
3046 n: usize,
3047 live_before: u64,
3048 ) -> Result<()> {
3049 let Some(ai) = self.auto_inc.as_mut() else {
3050 return Ok(());
3051 };
3052 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3053 return Ok(());
3054 };
3055 let columnar::NativeColumn::Int64 { data, validity } = col else {
3056 return Err(MongrelError::InvalidArgument(format!(
3057 "AUTO_INCREMENT column {} must be Int64",
3058 ai.column_id
3059 )));
3060 };
3061 let max = if native_int64_strictly_increasing(col, n) {
3062 data.get(n.saturating_sub(1)).copied()
3063 } else {
3064 data.iter()
3065 .take(n)
3066 .enumerate()
3067 .filter_map(|(i, v)| {
3068 if validity.is_empty() || columnar::validity_bit(validity, i) {
3069 Some(*v)
3070 } else {
3071 None
3072 }
3073 })
3074 .max()
3075 };
3076 if let Some(max) = max {
3077 let floor = max
3078 .checked_add(1)
3079 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3080 .max(1);
3081 if ai.next < floor {
3082 ai.next = floor;
3083 }
3084 if ai.seeded || live_before == 0 {
3085 ai.seeded = true;
3086 }
3087 }
3088 Ok(())
3089 }
3090
3091 fn fill_auto_inc_native_columns(
3092 &mut self,
3093 columns: &mut Vec<(u16, columnar::NativeColumn)>,
3094 n: usize,
3095 ) -> Result<()> {
3096 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3097 return Ok(());
3098 };
3099 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
3100 if let Some(start) = self.alloc_auto_inc_range(n)? {
3101 columns.push((
3102 cid,
3103 columnar::NativeColumn::Int64 {
3104 data: (start..start.saturating_add(n as i64)).collect(),
3105 validity: vec![0xFF; n.div_ceil(8)],
3106 },
3107 ));
3108 }
3109 return Ok(());
3110 };
3111
3112 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
3113 return Err(MongrelError::InvalidArgument(format!(
3114 "AUTO_INCREMENT column {cid} must be Int64"
3115 )));
3116 };
3117 if data.len() < n {
3118 return Err(MongrelError::InvalidArgument(format!(
3119 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
3120 data.len()
3121 )));
3122 }
3123 if columnar::all_non_null(validity, n) {
3124 return Ok(());
3125 }
3126 if validity.iter().all(|b| *b == 0) {
3127 if let Some(start) = self.alloc_auto_inc_range(n)? {
3128 for (i, slot) in data.iter_mut().take(n).enumerate() {
3129 *slot = start.saturating_add(i as i64);
3130 }
3131 *validity = vec![0xFF; n.div_ceil(8)];
3132 }
3133 return Ok(());
3134 }
3135
3136 let new_validity = vec![0xFF; data.len().div_ceil(8)];
3137 for (i, slot) in data.iter_mut().enumerate().take(n) {
3138 if columnar::validity_bit(validity, i) {
3139 self.advance_auto_inc_past(*slot)?;
3140 } else {
3141 *slot = self.alloc_auto_inc_value()?;
3142 }
3143 }
3144 *validity = new_validity;
3145 Ok(())
3146 }
3147
3148 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
3162 self.ensure_writable()?;
3163 if self.auto_inc.is_none() {
3164 return Ok(None);
3165 }
3166 Ok(Some(self.alloc_auto_inc_value()?))
3167 }
3168
3169 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
3175 let payload = bincode::serialize(&rows)?;
3176 self.wal_append_data(Op::Put {
3177 table_id: self.table_id,
3178 rows: payload,
3179 })?;
3180 if self.is_shared() {
3181 self.pending_rows_auto_inc
3182 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
3183 self.pending_rows.extend(rows);
3184 } else {
3185 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
3186 }
3187 Ok(())
3188 }
3189
3190 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
3193 self.ensure_indexes_complete()
3194 }
3195
3196 pub(crate) fn prepare_durable_publish_controlled(
3197 &mut self,
3198 control: &crate::ExecutionControl,
3199 ) -> Result<()> {
3200 self.ensure_indexes_complete_controlled(control, || true)?;
3201 Ok(())
3202 }
3203
3204 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
3205 self.apply_put_rows_inner_prepared(rows, true);
3206 }
3207
3208 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
3209 if check_existing_pk {
3210 self.ensure_indexes_complete()?;
3211 }
3212 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
3213 Ok(())
3214 }
3215
3216 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
3220 if rows.len() == 1 {
3224 let row = rows.into_iter().next().expect("len checked");
3225 self.apply_put_row_single(row, check_existing_pk);
3226 return;
3227 }
3228 let pk_id = self.schema.primary_key().map(|c| c.id);
3245 let probe = match pk_id {
3246 Some(pid) => {
3247 check_existing_pk
3248 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
3249 }
3250 None => false,
3251 };
3252 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
3255 for r in rows {
3256 for &cid in r.columns.keys() {
3257 self.pending_put_cols.insert(cid);
3258 }
3259 match pk_id {
3260 Some(pid) if probe || maintain_pk_by_row => {
3261 if let Some(pk_val) = r.columns.get(&pid) {
3262 let key = self.index_lookup_key(pid, pk_val);
3263 if probe {
3264 if let Some(old_rid) = self.hot.get(&key) {
3265 if old_rid != r.row_id {
3266 self.tombstone_row(old_rid, r.committed_epoch, true);
3267 }
3268 }
3269 }
3270 if maintain_pk_by_row {
3271 self.pk_by_row.insert(r.row_id, key);
3272 }
3273 }
3274 }
3275 Some(_) => {}
3276 None => {
3277 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3278 }
3279 }
3280 self.index_row(&r);
3281 self.reservoir.offer(r.row_id.0);
3282 self.memtable.upsert(r);
3283 self.live_count = self.live_count.saturating_add(1);
3286 }
3287 self.data_generation = self.data_generation.wrapping_add(1);
3288 }
3289
3290 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
3294 for &cid in row.columns.keys() {
3295 self.pending_put_cols.insert(cid);
3296 }
3297 let epoch = row.committed_epoch;
3298 if let Some(pk_col) = self.schema.primary_key() {
3299 let pk_id = pk_col.id;
3300 if let Some(pk_val) = row.columns.get(&pk_id) {
3301 let maintain_pk_by_row = self.pk_by_row_complete;
3305 if check_existing_pk || maintain_pk_by_row {
3306 let key = self.index_lookup_key(pk_id, pk_val);
3307 if check_existing_pk {
3308 if let Some(old_rid) = self.hot.get(&key) {
3309 if old_rid != row.row_id {
3310 self.tombstone_row(old_rid, epoch, true);
3311 }
3312 }
3313 }
3314 if maintain_pk_by_row {
3315 self.pk_by_row.insert(row.row_id, key);
3316 }
3317 }
3318 }
3319 } else {
3320 self.hot
3321 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3322 }
3323 self.index_row(&row);
3324 self.reservoir.offer(row.row_id.0);
3325 self.memtable.upsert(row);
3326 self.live_count = self.live_count.saturating_add(1);
3327 self.data_generation = self.data_generation.wrapping_add(1);
3328 }
3329
3330 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3333 self.allocator.alloc()
3334 }
3335
3336 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3342 let pk = self.schema.primary_key().ok_or_else(|| {
3343 MongrelError::Schema("clustered table requires a single-column primary key".into())
3344 })?;
3345 let pk_val = columns
3346 .iter()
3347 .find(|(id, _)| *id == pk.id)
3348 .map(|(_, v)| v)
3349 .ok_or_else(|| {
3350 MongrelError::Schema(format!(
3351 "clustered table missing primary key column {} ({})",
3352 pk.id, pk.name
3353 ))
3354 })?;
3355 Ok(clustered_row_id(pk_val))
3356 }
3357
3358 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3366 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3367 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3368 }
3369 let n = rows.len();
3370 for r in rows {
3371 for &cid in r.columns.keys() {
3372 self.pending_put_cols.insert(cid);
3373 }
3374 }
3375 let (losers, winner_pks) = self.partition_pk_winners(rows);
3376 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3377 for (key, &row_id) in &winner_pks {
3379 if let Some(old_rid) = self.hot.get(key) {
3380 if old_rid != row_id {
3381 self.tombstone_row(old_rid, epoch, true);
3382 }
3383 }
3384 }
3385 for &loser_rid in &losers {
3388 self.tombstone_row(loser_rid, epoch, false);
3389 }
3390 for (key, row_id) in winner_pks {
3392 self.insert_hot_pk(key, row_id);
3393 }
3394 if self.schema.primary_key().is_none() {
3395 for r in rows {
3396 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3397 }
3398 }
3399 for r in rows {
3400 self.allocator.advance_to(r.row_id)?;
3401 if !losers.contains(&r.row_id) {
3402 self.index_row(r);
3403 }
3404 }
3405 for r in rows {
3406 if !losers.contains(&r.row_id) {
3407 self.reservoir.offer(r.row_id.0);
3408 }
3409 }
3410 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3411 self.data_generation = self.data_generation.wrapping_add(1);
3412 Ok(())
3413 }
3414
3415 pub(crate) fn recover_apply(
3420 &mut self,
3421 rows: Vec<Row>,
3422 deletes: Vec<(RowId, Epoch)>,
3423 ) -> Result<()> {
3424 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3428 std::collections::BTreeMap::new();
3429 for row in rows {
3430 if row.row_id.0 >= u64::MAX - 1 {
3431 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3432 }
3433 self.allocator.advance_to(row.row_id)?;
3434 if let Some(ai) = self.auto_inc.as_mut() {
3439 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3440 let next = n.checked_add(1).ok_or_else(|| {
3441 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3442 })?;
3443 if next > ai.next {
3444 ai.next = next;
3445 }
3446 }
3447 }
3448 by_epoch.entry(row.committed_epoch).or_default().push(row);
3449 }
3450 for (epoch, group) in by_epoch {
3451 let (losers, winner_pks) = self.partition_pk_winners(&group);
3452 for (key, &row_id) in &winner_pks {
3454 if let Some(old_rid) = self.hot.get(key) {
3455 if old_rid != row_id {
3456 self.tombstone_row(old_rid, epoch, false);
3457 }
3458 }
3459 }
3460 for (key, row_id) in winner_pks {
3461 self.insert_hot_pk(key, row_id);
3462 }
3463 if self.schema.primary_key().is_none() {
3464 for r in &group {
3465 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3466 }
3467 }
3468 for r in &group {
3469 if !losers.contains(&r.row_id) {
3470 self.memtable.upsert(r.clone());
3471 self.index_row(r);
3472 }
3473 }
3474 }
3475 for (rid, epoch) in deletes {
3476 self.memtable.tombstone(rid, epoch);
3477 self.remove_hot_for_row(rid, epoch);
3478 }
3479 self.reservoir_complete = false;
3482 Ok(())
3483 }
3484
3485 pub(crate) fn flushed_epoch(&self) -> u64 {
3487 self.flushed_epoch
3488 }
3489
3490 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3491 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3492 }
3493
3494 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3496 self.schema.validate_values(cells)
3497 }
3498
3499 fn validate_columns_not_null(
3503 &self,
3504 columns: &[(u16, columnar::NativeColumn)],
3505 n: usize,
3506 ) -> Result<()> {
3507 let by_id: HashMap<u16, &columnar::NativeColumn> =
3508 columns.iter().map(|(id, c)| (*id, c)).collect();
3509 for col in &self.schema.columns {
3510 if !col.flags.contains(ColumnFlags::NULLABLE) {
3511 match by_id.get(&col.id) {
3512 None => {
3513 return Err(MongrelError::InvalidArgument(format!(
3514 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3515 col.name, col.id
3516 )));
3517 }
3518 Some(c) => {
3519 if c.null_count(n) != 0 {
3520 return Err(MongrelError::InvalidArgument(format!(
3521 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3522 col.name, col.id
3523 )));
3524 }
3525 }
3526 }
3527 }
3528 if let TypeId::Enum { variants } = &col.ty {
3529 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3530 if by_id.contains_key(&col.id) {
3531 return Err(MongrelError::InvalidArgument(format!(
3532 "column '{}' ({}) enum requires a bytes column",
3533 col.name, col.id
3534 )));
3535 }
3536 continue;
3537 };
3538 for index in 0..n {
3539 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3540 continue;
3541 };
3542 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3543 return Err(MongrelError::InvalidArgument(format!(
3544 "column '{}' ({}) enum value {:?} is not one of {:?}",
3545 col.name,
3546 col.id,
3547 String::from_utf8_lossy(value),
3548 variants
3549 )));
3550 }
3551 }
3552 }
3553 }
3554 Ok(())
3555 }
3556
3557 fn bulk_pk_winner_indices(
3562 &self,
3563 columns: &[(u16, columnar::NativeColumn)],
3564 n: usize,
3565 ) -> Option<Vec<usize>> {
3566 let pk_col = self.schema.primary_key()?;
3567 let pk_id = pk_col.id;
3568 let pk_ty = pk_col.ty.clone();
3569 let by_id: HashMap<u16, &columnar::NativeColumn> =
3570 columns.iter().map(|(id, c)| (*id, c)).collect();
3571 let pk_native = by_id.get(&pk_id)?;
3572 if native_int64_strictly_increasing(pk_native, n) {
3573 return None;
3574 }
3575 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3577 let mut null_pk_rows: Vec<usize> = Vec::new();
3578 for i in 0..n {
3579 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3580 Some(key) => {
3581 last.insert(key, i);
3582 }
3583 None => null_pk_rows.push(i),
3584 }
3585 }
3586 let mut winners: HashSet<usize> = last.values().copied().collect();
3587 for i in null_pk_rows {
3588 winners.insert(i);
3589 }
3590 Some((0..n).filter(|i| winners.contains(i)).collect())
3591 }
3592
3593 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3595 self.require_delete()?;
3596 let epoch = self.pending_epoch();
3597 self.wal_append_data(Op::Delete {
3598 table_id: self.table_id,
3599 row_ids: vec![row_id],
3600 })?;
3601 if self.is_shared() {
3602 self.pending_dels.push(row_id);
3603 } else {
3604 self.apply_delete(row_id, epoch);
3605 }
3606 Ok(())
3607 }
3608
3609 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3610 let pre = self.get(row_id, self.snapshot());
3611 self.delete(row_id)?;
3612 Ok(pre.map(|row| {
3613 let mut columns: Vec<_> = row.columns.into_iter().collect();
3614 columns.sort_by_key(|(id, _)| *id);
3615 OwnedRow { columns }
3616 }))
3617 }
3618
3619 pub fn truncate(&mut self) -> Result<()> {
3621 self.require_delete()?;
3622 let epoch = self.pending_epoch();
3623 self.wal_append_data(Op::TruncateTable {
3624 table_id: self.table_id,
3625 })?;
3626 self.pending_rows.clear();
3627 self.pending_rows_auto_inc.clear();
3628 self.pending_dels.clear();
3629 self.pending_truncate = Some(epoch);
3630 Ok(())
3631 }
3632
3633 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3635 self.run_refs.clear();
3640 self.retiring.clear();
3641 self.memtable = Memtable::new();
3642 self.mutable_run = MutableRun::new();
3643 self.hot = HotIndex::new();
3644 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3645 self.bitmap = bitmap;
3646 self.ann = ann;
3647 self.fm = fm;
3648 self.sparse = sparse;
3649 self.minhash = minhash;
3650 self.learned_range = Arc::new(HashMap::new());
3651 self.pk_by_row.clear();
3652 self.pk_by_row_complete = false;
3653 self.live_count = 0;
3654 self.reservoir = crate::reservoir::Reservoir::default();
3655 self.reservoir_complete = true;
3656 self.had_deletes = true;
3657 self.agg_cache = Arc::new(HashMap::new());
3658 self.global_idx_epoch = 0;
3659 self.indexes_complete = true;
3660 self.pending_delete_rids.clear();
3661 self.pending_put_cols.clear();
3662 self.pending_rows.clear();
3663 self.pending_rows_auto_inc.clear();
3664 self.pending_dels.clear();
3665 self.clear_result_cache();
3666 self.invalidate_index_checkpoint();
3667 self.data_generation = self.data_generation.wrapping_add(1);
3668 }
3669
3670 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3673 self.remove_hot_for_row(row_id, epoch);
3674 self.tombstone_row(row_id, epoch, true);
3675 self.data_generation = self.data_generation.wrapping_add(1);
3676 }
3677
3678 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3682 let tombstone = Row {
3683 row_id,
3684 committed_epoch: epoch,
3685 columns: std::collections::HashMap::new(),
3686 deleted: true,
3687 };
3688 self.memtable.upsert(tombstone);
3689 self.pk_by_row.remove(&row_id);
3690 if adjust_live_count {
3691 self.live_count = self.live_count.saturating_sub(1);
3692 }
3693 self.pending_delete_rids.insert(row_id.0 as u32);
3695 self.had_deletes = true;
3698 self.agg_cache = Arc::new(HashMap::new());
3699 }
3700
3701 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3705 let Some(pk_col) = self.schema.primary_key() else {
3706 return;
3707 };
3708 if self.pk_by_row_complete {
3711 if let Some(key) = self.pk_by_row.remove(&row_id) {
3712 if self.hot.get(&key) == Some(row_id) {
3713 self.hot.remove(&key);
3714 }
3715 }
3716 return;
3717 }
3718 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3737 if self.indexes_complete {
3738 let pk_val = self
3739 .memtable
3740 .get_version(row_id, lookup_epoch)
3741 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3742 .or_else(|| {
3743 self.mutable_run
3744 .get_version(row_id, lookup_epoch)
3745 .filter(|(_, r)| !r.deleted)
3746 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3747 })
3748 .or_else(|| {
3749 self.run_refs.iter().find_map(|rr| {
3750 let mut reader = self.open_reader(rr.run_id).ok()?;
3751 let (_, deleted, val) = reader
3752 .get_version_column(row_id, lookup_epoch, pk_col.id)
3753 .ok()??;
3754 if deleted {
3755 return None;
3756 }
3757 val
3758 })
3759 });
3760 if let Some(pk_val) = pk_val {
3761 let key = self.index_lookup_key(pk_col.id, &pk_val);
3762 if self.hot.get(&key) == Some(row_id) {
3763 self.hot.remove(&key);
3764 }
3765 return;
3766 }
3767 }
3768 self.refresh_pk_by_row_from_hot();
3773 if let Some(key) = self.pk_by_row.remove(&row_id) {
3774 if self.hot.get(&key) == Some(row_id) {
3775 self.hot.remove(&key);
3776 }
3777 }
3778 }
3779
3780 fn partition_pk_winners(
3785 &self,
3786 rows: &[Row],
3787 ) -> (
3788 std::collections::HashSet<RowId>,
3789 std::collections::HashMap<Vec<u8>, RowId>,
3790 ) {
3791 let mut losers = std::collections::HashSet::new();
3792 let Some(pk_col) = self.schema.primary_key() else {
3793 return (losers, std::collections::HashMap::new());
3794 };
3795 let pk_id = pk_col.id;
3796 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3797 std::collections::HashMap::new();
3798 for r in rows {
3799 let Some(pk_val) = r.columns.get(&pk_id) else {
3800 continue;
3801 };
3802 let key = self.index_lookup_key(pk_id, pk_val);
3803 if let Some(&old_rid) = winners.get(&key) {
3804 losers.insert(old_rid);
3805 }
3806 winners.insert(key, r.row_id);
3807 }
3808 (losers, winners)
3809 }
3810
3811 fn index_row(&mut self, row: &Row) {
3812 if row.deleted {
3813 return;
3814 }
3815 let any_predicate = self
3823 .schema
3824 .indexes
3825 .iter()
3826 .any(|idx| idx.predicate.is_some());
3827 if any_predicate {
3828 let columns_map: HashMap<u16, &Value> =
3829 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3830 let name_to_id: HashMap<&str, u16> = self
3831 .schema
3832 .columns
3833 .iter()
3834 .map(|c| (c.name.as_str(), c.id))
3835 .collect();
3836 for idx in &self.schema.indexes {
3837 if let Some(pred) = &idx.predicate {
3838 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3839 continue; }
3841 }
3842 index_into_single(
3844 idx,
3845 &self.schema,
3846 row,
3847 &mut self.hot,
3848 &mut self.bitmap,
3849 &mut self.ann,
3850 &mut self.fm,
3851 &mut self.sparse,
3852 &mut self.minhash,
3853 );
3854 }
3855 return;
3856 }
3857 if self.column_keys.is_empty() {
3861 index_into(
3862 &self.schema,
3863 row,
3864 &mut self.hot,
3865 &mut self.bitmap,
3866 &mut self.ann,
3867 &mut self.fm,
3868 &mut self.sparse,
3869 &mut self.minhash,
3870 );
3871 return;
3872 }
3873 let effective_row = self.tokenized_for_indexes(row);
3874 index_into(
3875 &self.schema,
3876 &effective_row,
3877 &mut self.hot,
3878 &mut self.bitmap,
3879 &mut self.ann,
3880 &mut self.fm,
3881 &mut self.sparse,
3882 &mut self.minhash,
3883 );
3884 }
3885
3886 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3892 if self.column_keys.is_empty() {
3893 return row.clone();
3894 }
3895 {
3896 use crate::encryption::SCHEME_HMAC_EQ;
3897 let mut tok = row.clone();
3898 for (&cid, &(_, scheme)) in &self.column_keys {
3899 if scheme != SCHEME_HMAC_EQ {
3900 continue;
3901 }
3902 if let Some(v) = tok.columns.get(&cid).cloned() {
3903 if let Some(t) = self.tokenize_value(cid, &v) {
3904 tok.columns.insert(cid, t);
3905 }
3906 }
3907 }
3908 tok
3909 }
3910 }
3911
3912 pub fn commit(&mut self) -> Result<Epoch> {
3917 self.commit_inner(None)
3918 }
3919
3920 #[doc(hidden)]
3923 pub fn commit_controlled<F>(
3924 &mut self,
3925 control: &crate::ExecutionControl,
3926 mut before_commit: F,
3927 ) -> Result<Epoch>
3928 where
3929 F: FnMut() -> Result<()>,
3930 {
3931 self.commit_inner(Some((control, &mut before_commit)))
3932 }
3933
3934 fn commit_inner(
3935 &mut self,
3936 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3937 ) -> Result<Epoch> {
3938 self.ensure_writable()?;
3939 if !self.has_pending_mutations() {
3940 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3941 return Err(MongrelError::Full(
3942 "standalone transaction id namespace exhausted".into(),
3943 ));
3944 }
3945 return Ok(self.epoch.visible());
3946 }
3947 self.commit_new_epoch_inner(controlled)
3948 }
3949
3950 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3954 self.commit_new_epoch_inner(None)
3955 }
3956
3957 fn commit_new_epoch_inner(
3958 &mut self,
3959 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3960 ) -> Result<Epoch> {
3961 self.ensure_writable()?;
3962 if self.is_shared() {
3963 self.commit_shared(controlled)
3964 } else {
3965 self.commit_private(controlled)
3966 }
3967 }
3968
3969 fn commit_private(
3971 &mut self,
3972 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3973 ) -> Result<Epoch> {
3974 let commit_lock = Arc::clone(&self.commit_lock);
3978 let _g = commit_lock.lock();
3979 let txn_id = self.ensure_txn_id()?;
3982 if let Some((control, before_commit)) = controlled {
3983 control.checkpoint()?;
3984 before_commit()?;
3985 }
3986 let new_epoch = self.epoch.bump_assigned();
3987 let epoch_authority = Arc::clone(&self.epoch);
3988 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3989 let wal_result = match &mut self.wal {
3993 WalSink::Private(w) => w
3994 .append_txn(
3995 txn_id,
3996 Op::TxnCommit {
3997 epoch: new_epoch.0,
3998 added_runs: Vec::new(),
3999 },
4000 )
4001 .and_then(|_| w.sync()),
4002 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
4003 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
4004 };
4005 if let Err(error) = wal_result {
4006 self.durable_commit_failed = true;
4007 return Err(MongrelError::CommitOutcomeUnknown {
4008 epoch: new_epoch.0,
4009 message: error.to_string(),
4010 });
4011 }
4012 if let Some(epoch) = self.pending_truncate.take() {
4015 self.apply_truncate(epoch);
4016 }
4017 self.invalidate_pending_cache();
4018 let publish_result = self.persist_manifest(new_epoch);
4019 self.epoch.publish_in_order(new_epoch);
4023 epoch_guard.disarm();
4024 if let Err(error) = publish_result {
4025 self.durable_commit_failed = true;
4026 return Err(MongrelError::DurableCommit {
4027 epoch: new_epoch.0,
4028 message: error.to_string(),
4029 });
4030 }
4031 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
4032 self.pending_private_mutations = false;
4033 self.data_generation = self.data_generation.wrapping_add(1);
4034 Ok(new_epoch)
4035 }
4036
4037 fn commit_shared(
4043 &mut self,
4044 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4045 ) -> Result<Epoch> {
4046 use std::sync::atomic::Ordering;
4047 let s = match &self.wal {
4048 WalSink::Shared(s) => s.clone(),
4049 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
4050 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4051 };
4052 if s.poisoned.load(Ordering::Relaxed) {
4053 return Err(MongrelError::Other(
4054 "database poisoned by fsync error".into(),
4055 ));
4056 }
4057 let commit_lock = Arc::clone(&self.commit_lock);
4064 let _g = commit_lock.lock();
4065 if !self.pending_rows.is_empty() {
4066 match controlled.as_ref() {
4067 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
4068 None => self.prepare_durable_publish()?,
4069 }
4070 }
4071 let txn_id = self.ensure_txn_id()?;
4074 let mut wal = s.wal.lock();
4075 if let Some((control, before_commit)) = controlled {
4076 control.checkpoint()?;
4077 before_commit()?;
4078 }
4079 let new_epoch = self.epoch.bump_assigned();
4080 let epoch_authority = Arc::clone(&self.epoch);
4081 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
4082 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
4083 Ok(commit_seq) => commit_seq,
4084 Err(error) => {
4085 s.poisoned.store(true, Ordering::Relaxed);
4086 s.lifecycle.poison();
4087 return Err(MongrelError::CommitOutcomeUnknown {
4088 epoch: new_epoch.0,
4089 message: error.to_string(),
4090 });
4091 }
4092 };
4093 drop(wal);
4094 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
4095 s.poisoned.store(true, Ordering::Relaxed);
4096 s.lifecycle.poison();
4097 return Err(MongrelError::CommitOutcomeUnknown {
4098 epoch: new_epoch.0,
4099 message: error.to_string(),
4100 });
4101 }
4102
4103 if self.pending_truncate.take().is_some() {
4106 self.apply_truncate(new_epoch);
4107 }
4108 let mut rows = std::mem::take(&mut self.pending_rows);
4109 if !rows.is_empty() {
4110 for r in &mut rows {
4111 r.committed_epoch = new_epoch;
4112 }
4113 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
4114 let all_auto_generated =
4115 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
4116 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
4117 } else {
4118 self.pending_rows_auto_inc.clear();
4119 }
4120 let dels = std::mem::take(&mut self.pending_dels);
4121 for rid in dels {
4122 self.apply_delete(rid, new_epoch);
4123 }
4124
4125 self.invalidate_pending_cache();
4126 let publish_result = self.persist_manifest(new_epoch);
4127 self.epoch.publish_in_order(new_epoch);
4128 epoch_guard.disarm();
4129 let _ = s.change_wake.send(());
4130 if let Err(error) = publish_result {
4131 self.durable_commit_failed = true;
4132 s.poisoned.store(true, Ordering::Relaxed);
4133 s.lifecycle.poison();
4134 return Err(MongrelError::DurableCommit {
4135 epoch: new_epoch.0,
4136 message: error.to_string(),
4137 });
4138 }
4139 self.current_txn_id = 0;
4141 self.data_generation = self.data_generation.wrapping_add(1);
4142 Ok(new_epoch)
4143 }
4144
4145 pub fn flush(&mut self) -> Result<Epoch> {
4153 self.flush_with_outcome().map(|(epoch, _)| epoch)
4154 }
4155
4156 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
4158 self.flush_with_outcome_inner(None)
4159 }
4160
4161 #[doc(hidden)]
4164 pub fn flush_with_outcome_controlled<F>(
4165 &mut self,
4166 control: &crate::ExecutionControl,
4167 mut before_commit: F,
4168 ) -> Result<(Epoch, bool)>
4169 where
4170 F: FnMut() -> Result<()>,
4171 {
4172 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
4173 }
4174
4175 fn flush_with_outcome_inner(
4176 &mut self,
4177 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
4178 ) -> Result<(Epoch, bool)> {
4179 match controlled.as_ref() {
4180 Some((control, _)) => {
4181 self.ensure_indexes_complete_controlled(control, || true)?;
4182 }
4183 None => self.ensure_indexes_complete()?,
4184 }
4185 let committed = self.has_pending_mutations();
4186 let epoch = self.commit_inner(controlled)?;
4187 let finish: Result<(Epoch, bool)> = (|| {
4188 let rows = self.memtable.drain_sorted();
4189 if !rows.is_empty() {
4190 self.mutable_run.insert_many(rows);
4191 }
4192 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
4193 self.spill_mutable_run(epoch)?;
4194 self.mark_flushed(epoch)?;
4198 self.persist_manifest(epoch)?;
4199 self.build_learned_ranges()?;
4200 self.checkpoint_indexes(epoch);
4203 }
4204 Ok((epoch, committed))
4207 })();
4208 let outcome = match finish {
4209 Err(error) if committed => Err(MongrelError::DurableCommit {
4210 epoch: epoch.0,
4211 message: error.to_string(),
4212 }),
4213 result => result,
4214 };
4215 if outcome.is_ok() {
4216 let _ = self.publish_read_generation();
4222 }
4223 outcome
4224 }
4225
4226 fn has_pending_mutations(&self) -> bool {
4227 self.pending_private_mutations
4228 || !self.pending_rows.is_empty()
4229 || !self.pending_dels.is_empty()
4230 || self.pending_truncate.is_some()
4231 }
4232
4233 pub fn has_pending_writes(&self) -> bool {
4234 self.has_pending_mutations()
4235 }
4236
4237 pub fn force_flush(&mut self) -> Result<Epoch> {
4246 let saved = self.mutable_run_spill_bytes;
4247 self.mutable_run_spill_bytes = 1;
4248 let result = self.flush();
4249 self.mutable_run_spill_bytes = saved;
4250 result
4251 }
4252
4253 pub fn close(&mut self) -> Result<()> {
4260 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
4261 self.force_flush()?;
4262 }
4263 Ok(())
4264 }
4265
4266 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
4273 let op = Op::Flush {
4274 table_id: self.table_id,
4275 flushed_epoch: epoch.0,
4276 };
4277 match &mut self.wal {
4278 WalSink::Private(w) => {
4279 w.append_system(op)?;
4280 w.sync()?;
4281 }
4282 WalSink::Shared(s) => {
4283 s.wal.lock().append_system(op)?;
4288 }
4289 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4290 }
4291 self.flushed_epoch = epoch.0;
4292 if matches!(self.wal, WalSink::Private(_)) {
4293 self.rotate_wal(epoch)?;
4294 }
4295 Ok(())
4296 }
4297
4298 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
4302 if self.mutable_run.is_empty() {
4303 return Ok(());
4304 }
4305 let run_id = self.alloc_run_id()?;
4306 let rows = self.mutable_run.drain_sorted();
4307 if rows.is_empty() {
4308 return Ok(());
4309 }
4310 let path = self.run_path(run_id);
4311 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4312 if let Some(kek) = &self.kek {
4313 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4314 }
4315 let header = match self.create_run_file(run_id)? {
4316 Some(file) => writer.write_file(file, &rows)?,
4317 None => writer.write(&path, &rows)?,
4318 };
4319 self.run_refs.push(RunRef {
4320 run_id: run_id as u128,
4321 level: 0,
4322 epoch_created: epoch.0,
4323 row_count: header.row_count,
4324 });
4325 Ok(())
4326 }
4327
4328 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4332 self.mutable_run_spill_bytes = bytes.max(1);
4333 }
4334
4335 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4339 self.compaction_zstd_level = level;
4340 }
4341
4342 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4346 self.result_cache.lock().set_max_bytes(max_bytes);
4347 }
4348
4349 pub(crate) fn clear_result_cache(&mut self) {
4353 self.result_cache.lock().clear();
4354 }
4355
4356 pub fn mutable_run_len(&self) -> usize {
4358 self.mutable_run.len()
4359 }
4360
4361 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4364 self.mutable_run.drain_sorted()
4365 }
4366
4367 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4369 let mut snapshot = self.mutable_run.clone();
4370 snapshot.drain_sorted()
4371 }
4372
4373 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4378 self.ensure_writable()?;
4379 let n = batch.len();
4380 if n == 0 {
4381 return Ok(self.current_epoch());
4382 }
4383 for row in &batch {
4384 self.schema.validate_values(row)?;
4385 }
4386 let epoch = self.commit_new_epoch()?;
4387 let live_before = self.live_count;
4388 self.spill_mutable_run(epoch)?;
4392 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4393 && self.indexes_complete
4394 && self.run_refs.is_empty()
4395 && self.memtable.is_empty()
4396 && self.mutable_run.is_empty();
4397 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4403 use rayon::prelude::*;
4404 self.schema
4405 .columns
4406 .par_iter()
4407 .map(|cdef| {
4408 (
4409 cdef.id,
4410 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4411 )
4412 })
4413 .collect::<Vec<_>>()
4414 };
4415 drop(batch);
4416 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4421 self.validate_columns_not_null(&user_columns, n)?;
4422 let winner_idx = self
4423 .bulk_pk_winner_indices(&user_columns, n)
4424 .filter(|idx| idx.len() != n);
4425 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4426 match winner_idx.as_deref() {
4427 Some(idx) => {
4428 let compacted = user_columns
4429 .iter()
4430 .map(|(id, c)| (*id, c.gather(idx)))
4431 .collect();
4432 (compacted, idx.len())
4433 }
4434 None => (std::mem::take(&mut user_columns), n),
4435 };
4436 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4437 let first = self.allocator.alloc_range(write_n as u64)?.0;
4438 for rid in first..first + write_n as u64 {
4439 self.reservoir.offer(rid);
4440 }
4441 let run_id = self.alloc_run_id()?;
4442 let path = self.run_path(run_id);
4443 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4444 .clean(true)
4445 .with_lz4()
4446 .with_native_endian();
4447 if let Some(kek) = &self.kek {
4448 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4449 }
4450 let header = match self.create_run_file(run_id)? {
4451 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4452 None => writer.write_native(&path, &write_columns, write_n, first)?,
4453 };
4454 self.run_refs.push(RunRef {
4455 run_id: run_id as u128,
4456 level: 0,
4457 epoch_created: epoch.0,
4458 row_count: header.row_count,
4459 });
4460 self.live_count = self.live_count.saturating_add(write_n as u64);
4461 if eager_index_build {
4462 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4463 self.index_columns_bulk(&write_columns, &row_ids);
4464 self.indexes_complete = true;
4465 self.build_learned_ranges()?;
4466 } else {
4467 self.indexes_complete = false;
4468 }
4469 self.mark_flushed(epoch)?;
4470 self.persist_manifest(epoch)?;
4471 if eager_index_build {
4472 self.checkpoint_indexes(epoch);
4473 }
4474 self.clear_result_cache();
4475 Ok(epoch)
4476 }
4477
4478 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4481 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4482 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4483 let segment_no = segment
4486 .file_stem()
4487 .and_then(|s| s.to_str())
4488 .and_then(|s| s.strip_prefix("seg-"))
4489 .and_then(|s| s.parse::<u64>().ok())
4490 .unwrap_or(0);
4491 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4492 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4493 wal.sync()?;
4494 self.wal = WalSink::Private(wal);
4495 Ok(())
4496 }
4497
4498 pub(crate) fn invalidate_pending_cache(&mut self) {
4503 self.result_cache
4504 .lock()
4505 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4506 self.pending_delete_rids.clear();
4507 self.pending_put_cols.clear();
4508 }
4509
4510 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4511 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4512 m.current_epoch = epoch.0;
4513 m.next_row_id = self.allocator.current().0;
4514 m.runs = self.run_refs.clone();
4515 m.live_count = self.live_count;
4516 m.global_idx_epoch = self.global_idx_epoch;
4517 m.flushed_epoch = self.flushed_epoch;
4518 m.retiring = self.retiring.clone();
4519 m.auto_inc_next = match self.auto_inc {
4523 Some(ai) if ai.seeded => ai.next,
4524 _ => 0,
4525 };
4526 m.ttl = self.ttl;
4527 let meta_dek = self.manifest_meta_dek();
4528 match self._root_guard.as_deref() {
4529 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4530 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4531 }
4532 Ok(())
4533 }
4534
4535 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4536 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4540 let live_count = u64::try_from(rows.len())
4541 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4542 let auto_inc = match self.auto_inc {
4543 Some(mut state) => {
4544 let maximum = self.scan_max_int64(state.column_id)?;
4545 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4546 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4547 })?;
4548 state.next = state.next.max(after_maximum).max(1);
4549 state.seeded = true;
4550 Some(state)
4551 }
4552 None => None,
4553 };
4554 Ok(RecoveryMetadataPlan {
4555 live_count,
4556 auto_inc,
4557 changed: live_count != self.live_count
4558 || auto_inc.is_some_and(|planned| {
4559 self.auto_inc.is_none_or(|current| {
4560 current.next != planned.next || current.seeded != planned.seeded
4561 })
4562 }),
4563 })
4564 }
4565
4566 pub(crate) fn apply_recovered_metadata(
4567 &mut self,
4568 plan: RecoveryMetadataPlan,
4569 epoch: Epoch,
4570 ) -> Result<()> {
4571 if !plan.changed {
4572 return Ok(());
4573 }
4574 self.live_count = plan.live_count;
4575 self.auto_inc = plan.auto_inc;
4576 self.persist_manifest(epoch)
4577 }
4578
4579 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4585 if !self.indexes_complete {
4588 return;
4589 }
4590 if crate::catalog::inject_hook("index.publish.before").is_err() {
4593 return;
4594 }
4595 if self.idx_root.is_none() {
4596 if let Some(root) = self._root_guard.as_ref() {
4597 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4598 return;
4599 };
4600 self.idx_root = Some(Arc::new(idx_root));
4601 }
4602 }
4603 let snap = global_idx::IndexSnapshot {
4604 hot: &self.hot,
4605 bitmap: &self.bitmap,
4606 ann: &self.ann,
4607 fm: &self.fm,
4608 sparse: &self.sparse,
4609 minhash: &self.minhash,
4610 learned_range: &self.learned_range,
4611 };
4612 let idx_dek = self.idx_dek();
4614 let written = match self.idx_root.as_deref() {
4615 Some(root) => global_idx::write_atomic_root(
4616 root,
4617 self.table_id,
4618 epoch.0,
4619 snap,
4620 idx_dek.as_deref(),
4621 ),
4622 None => global_idx::write_atomic(
4623 &self.dir,
4624 self.table_id,
4625 epoch.0,
4626 snap,
4627 idx_dek.as_deref(),
4628 ),
4629 };
4630 if written.is_ok() {
4631 self.global_idx_epoch = epoch.0;
4632 let _ = self.persist_manifest(epoch);
4633 let _ = crate::catalog::inject_hook("index.publish.after");
4635 }
4636 }
4637
4638 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4641 self.global_idx_epoch = 0;
4642 if let Some(root) = self.idx_root.as_deref() {
4643 let _ = root.remove_file(global_idx::IDX_FILENAME);
4644 } else {
4645 global_idx::remove(&self.dir);
4646 }
4647 let _ = self.persist_manifest(self.epoch.visible());
4648 }
4649
4650 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4655 self.indexes_complete = false;
4656 self.global_idx_epoch = 0;
4657 if let Some(root) = self.idx_root.as_deref() {
4658 let _ = root.remove_file(global_idx::IDX_FILENAME);
4659 } else {
4660 global_idx::remove(&self.dir);
4661 }
4662 }
4663
4664 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4665 self.indexes_complete = true;
4666 }
4667
4668 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4674 self.durable_commit_failed = true;
4675 if let WalSink::Shared(shared) = &self.wal {
4676 shared
4677 .poisoned
4678 .store(true, std::sync::atomic::Ordering::Relaxed);
4679 }
4680 }
4681
4682 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4686 self.durable_commit_failed = true;
4687 }
4688
4689 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4692 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4693 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4694 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4695 best = Some((epoch, row));
4696 }
4697 }
4698 for rr in &self.run_refs {
4699 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4700 continue;
4701 };
4702 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4703 continue;
4704 };
4705 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4706 best = Some((epoch, row));
4707 }
4708 }
4709 let now_nanos = unix_nanos_now();
4710 match best {
4711 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4712 Some((_, r)) => Some(r),
4713 None => None,
4714 }
4715 }
4716
4717 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4721 self.visible_rows_at_time(snapshot, unix_nanos_now())
4722 }
4723
4724 #[doc(hidden)]
4727 pub fn visible_rows_controlled(
4728 &self,
4729 snapshot: Snapshot,
4730 control: &crate::ExecutionControl,
4731 ) -> Result<Vec<Row>> {
4732 let mut out = Vec::new();
4733 self.for_each_visible_row_controlled(snapshot, control, |row| {
4734 out.push(row);
4735 Ok(())
4736 })?;
4737 Ok(out)
4738 }
4739
4740 #[doc(hidden)]
4743 pub fn for_each_visible_row_controlled<F>(
4744 &self,
4745 snapshot: Snapshot,
4746 control: &crate::ExecutionControl,
4747 visit: F,
4748 ) -> Result<()>
4749 where
4750 F: FnMut(Row) -> Result<()>,
4751 {
4752 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4753 control.checkpoint()?;
4754 let memtable = self.memtable.visible_versions(snapshot.epoch);
4755 if !memtable.is_empty() {
4756 sources.push(ControlledVisibleSource::memory(memtable));
4757 }
4758 control.checkpoint()?;
4759 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4760 if !mutable.is_empty() {
4761 sources.push(ControlledVisibleSource::memory(mutable));
4762 }
4763 for run in &self.run_refs {
4764 control.checkpoint()?;
4765 let reader = self.open_reader(run.run_id)?;
4766 sources.push(ControlledVisibleSource::run(
4767 reader.into_visible_version_cursor(snapshot.epoch)?,
4768 ));
4769 }
4770 let now_nanos = unix_nanos_now();
4771 merge_controlled_visible_sources(
4772 &mut sources,
4773 control,
4774 |row| self.row_expired_at(row, now_nanos),
4775 visit,
4776 )
4777 }
4778
4779 #[doc(hidden)]
4780 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4781 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4782 let mut fold = |row: Row| {
4783 best.entry(row.row_id.0)
4784 .and_modify(|e| {
4785 if row.committed_epoch > e.0 {
4786 *e = (row.committed_epoch, row.clone());
4787 }
4788 })
4789 .or_insert_with(|| (row.committed_epoch, row));
4790 };
4791 for row in self.memtable.visible_versions(snapshot.epoch) {
4792 fold(row);
4793 }
4794 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4795 fold(row);
4796 }
4797 for rr in &self.run_refs {
4798 let mut reader = self.open_reader(rr.run_id)?;
4799 for row in reader.visible_versions(snapshot.epoch)? {
4800 fold(row);
4801 }
4802 }
4803 let mut out: Vec<Row> = best
4804 .into_values()
4805 .filter_map(|(_, r)| {
4806 if r.deleted || self.row_expired_at(&r, now_nanos) {
4807 None
4808 } else {
4809 Some(r)
4810 }
4811 })
4812 .collect();
4813 out.sort_by_key(|r| r.row_id);
4814 Ok(out)
4815 }
4816
4817 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4824 if self.ttl.is_none()
4825 && self.memtable.is_empty()
4826 && self.mutable_run.is_empty()
4827 && self.run_refs.len() == 1
4828 {
4829 let rr = self.run_refs[0].clone();
4830 let mut reader = self.open_reader(rr.run_id)?;
4831 let idxs = reader.visible_indices(snapshot.epoch)?;
4832 let mut cols = Vec::with_capacity(self.schema.columns.len());
4833 for cdef in &self.schema.columns {
4834 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4835 }
4836 return Ok(cols);
4837 }
4838 let rows = self.visible_rows(snapshot)?;
4840 let mut cols: Vec<(u16, Vec<Value>)> = self
4841 .schema
4842 .columns
4843 .iter()
4844 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4845 .collect();
4846 for r in &rows {
4847 for (cid, vec) in cols.iter_mut() {
4848 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4849 }
4850 }
4851 Ok(cols)
4852 }
4853
4854 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4856 let row_id = self.hot.get(key)?;
4857 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4858 Some(row_id)
4859 } else {
4860 None
4861 }
4862 }
4863
4864 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4869 self.query_at_with_allowed(q, self.snapshot(), None)
4870 }
4871
4872 pub fn query_controlled(
4875 &mut self,
4876 q: &crate::query::Query,
4877 control: &crate::ExecutionControl,
4878 ) -> Result<Vec<Row>> {
4879 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
4880 }
4881
4882 pub fn query_at_with_allowed(
4885 &mut self,
4886 q: &crate::query::Query,
4887 snapshot: Snapshot,
4888 allowed: Option<&std::collections::HashSet<RowId>>,
4889 ) -> Result<Vec<Row>> {
4890 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4891 }
4892
4893 #[doc(hidden)]
4894 pub fn query_at_with_allowed_controlled(
4895 &mut self,
4896 q: &crate::query::Query,
4897 snapshot: Snapshot,
4898 allowed: Option<&std::collections::HashSet<RowId>>,
4899 control: &crate::ExecutionControl,
4900 ) -> Result<Vec<Row>> {
4901 self.require_select()?;
4902 self.ensure_indexes_complete_controlled(control, || true)?;
4903 self.validate_native_query(q)?;
4904 self.query_conditions_at(
4905 &q.conditions,
4906 snapshot,
4907 allowed,
4908 q.limit,
4909 q.offset,
4910 None,
4911 unix_nanos_now(),
4912 Some(control),
4913 )
4914 }
4915
4916 #[doc(hidden)]
4917 pub fn query_at_with_allowed_after(
4918 &mut self,
4919 q: &crate::query::Query,
4920 snapshot: Snapshot,
4921 allowed: Option<&std::collections::HashSet<RowId>>,
4922 after_row_id: Option<RowId>,
4923 ) -> Result<Vec<Row>> {
4924 self.query_at_with_allowed_after_at_time(
4925 q,
4926 snapshot,
4927 allowed,
4928 after_row_id,
4929 unix_nanos_now(),
4930 )
4931 }
4932
4933 #[doc(hidden)]
4934 pub fn query_at_with_allowed_after_at_time(
4935 &mut self,
4936 q: &crate::query::Query,
4937 snapshot: Snapshot,
4938 allowed: Option<&std::collections::HashSet<RowId>>,
4939 after_row_id: Option<RowId>,
4940 query_time_nanos: i64,
4941 ) -> Result<Vec<Row>> {
4942 self.require_select()?;
4943 self.ensure_indexes_complete()?;
4944 self.validate_native_query(q)?;
4945 self.query_conditions_at(
4946 &q.conditions,
4947 snapshot,
4948 allowed,
4949 q.limit,
4950 q.offset,
4951 after_row_id,
4952 query_time_nanos,
4953 None,
4954 )
4955 }
4956
4957 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
4958 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4959 return Err(MongrelError::InvalidArgument(format!(
4960 "query exceeds {} conditions",
4961 crate::query::MAX_HARD_CONDITIONS
4962 )));
4963 }
4964 if let Some(limit) = q.limit {
4965 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4966 return Err(MongrelError::InvalidArgument(format!(
4967 "query limit must be between 1 and {}",
4968 crate::query::MAX_FINAL_LIMIT
4969 )));
4970 }
4971 }
4972 if q.offset > crate::query::MAX_QUERY_OFFSET {
4973 return Err(MongrelError::InvalidArgument(format!(
4974 "query offset exceeds {}",
4975 crate::query::MAX_QUERY_OFFSET
4976 )));
4977 }
4978 Ok(())
4979 }
4980
4981 #[doc(hidden)]
4984 pub fn query_all_at(
4985 &mut self,
4986 conditions: &[crate::query::Condition],
4987 snapshot: Snapshot,
4988 ) -> Result<Vec<Row>> {
4989 self.require_select()?;
4990 self.ensure_indexes_complete()?;
4991 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4992 return Err(MongrelError::InvalidArgument(format!(
4993 "query exceeds {} conditions",
4994 crate::query::MAX_HARD_CONDITIONS
4995 )));
4996 }
4997 self.query_conditions_at(
4998 conditions,
4999 snapshot,
5000 None,
5001 None,
5002 0,
5003 None,
5004 unix_nanos_now(),
5005 None,
5006 )
5007 }
5008
5009 #[allow(clippy::too_many_arguments)]
5010 fn query_conditions_at(
5011 &self,
5012 conditions: &[crate::query::Condition],
5013 snapshot: Snapshot,
5014 allowed: Option<&std::collections::HashSet<RowId>>,
5015 limit: Option<usize>,
5016 offset: usize,
5017 after_row_id: Option<RowId>,
5018 query_time_nanos: i64,
5019 control: Option<&crate::ExecutionControl>,
5020 ) -> Result<Vec<Row>> {
5021 control
5022 .map(crate::ExecutionControl::checkpoint)
5023 .transpose()?;
5024 crate::trace::QueryTrace::record(|t| {
5025 t.run_count = self.run_refs.len();
5026 t.memtable_rows = self.memtable.len();
5027 t.mutable_run_rows = self.mutable_run.len();
5028 });
5029 if conditions.is_empty() {
5033 crate::trace::QueryTrace::record(|t| {
5034 t.scan_mode = crate::trace::ScanMode::Materialized;
5035 t.row_materialized = true;
5036 });
5037 let mut rows = match control {
5038 Some(control) => self.visible_rows_controlled(snapshot, control)?,
5039 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
5040 };
5041 if let Some(allowed) = allowed {
5042 let mut filtered = Vec::with_capacity(rows.len());
5043 for (index, row) in rows.into_iter().enumerate() {
5044 if index & 255 == 0 {
5045 control
5046 .map(crate::ExecutionControl::checkpoint)
5047 .transpose()?;
5048 }
5049 if allowed.contains(&row.row_id) {
5050 filtered.push(row);
5051 }
5052 }
5053 rows = filtered;
5054 }
5055 if let Some(after_row_id) = after_row_id {
5056 rows.retain(|row| row.row_id > after_row_id);
5057 }
5058 rows.drain(..offset.min(rows.len()));
5059 if let Some(limit) = limit {
5060 rows.truncate(limit);
5061 }
5062 return Ok(rows);
5063 }
5064 crate::trace::QueryTrace::record(|t| {
5065 t.conditions_pushed = conditions.len();
5066 t.scan_mode = crate::trace::ScanMode::Materialized;
5067 t.row_materialized = true;
5068 });
5069 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
5076 ordered.sort_by_key(|c| condition_cost_rank(c));
5077 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
5078 for c in &ordered {
5079 control
5080 .map(crate::ExecutionControl::checkpoint)
5081 .transpose()?;
5082 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
5083 let empty = s.is_empty();
5084 sets.push(s);
5085 if empty {
5086 break;
5087 }
5088 }
5089 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
5090 if let Some(allowed) = allowed {
5091 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
5092 }
5093 if let Some(after_row_id) = after_row_id {
5094 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
5095 rids.drain(..first);
5096 }
5097 rids.drain(..offset.min(rids.len()));
5098 if let Some(limit) = limit {
5099 rids.truncate(limit);
5100 }
5101 control
5102 .map(crate::ExecutionControl::checkpoint)
5103 .transpose()?;
5104 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
5105 }
5106
5107 pub fn retrieve(
5109 &mut self,
5110 retriever: &crate::query::Retriever,
5111 ) -> Result<Vec<crate::query::RetrieverHit>> {
5112 self.retrieve_with_allowed(retriever, None)
5113 }
5114
5115 pub fn retrieve_at(
5116 &mut self,
5117 retriever: &crate::query::Retriever,
5118 snapshot: Snapshot,
5119 allowed: Option<&std::collections::HashSet<RowId>>,
5120 ) -> Result<Vec<crate::query::RetrieverHit>> {
5121 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
5122 }
5123
5124 pub fn retrieve_with_allowed(
5127 &mut self,
5128 retriever: &crate::query::Retriever,
5129 allowed: Option<&std::collections::HashSet<RowId>>,
5130 ) -> Result<Vec<crate::query::RetrieverHit>> {
5131 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
5132 }
5133
5134 pub fn retrieve_at_with_allowed(
5135 &mut self,
5136 retriever: &crate::query::Retriever,
5137 snapshot: Snapshot,
5138 allowed: Option<&std::collections::HashSet<RowId>>,
5139 ) -> Result<Vec<crate::query::RetrieverHit>> {
5140 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
5141 }
5142
5143 pub fn retrieve_at_with_allowed_and_context(
5144 &mut self,
5145 retriever: &crate::query::Retriever,
5146 snapshot: Snapshot,
5147 allowed: Option<&std::collections::HashSet<RowId>>,
5148 context: Option<&crate::query::AiExecutionContext>,
5149 ) -> Result<Vec<crate::query::RetrieverHit>> {
5150 self.require_select()?;
5151 self.ensure_indexes_complete()?;
5152 self.validate_retriever(retriever)?;
5153 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
5154 }
5155
5156 pub fn retrieve_at_with_candidate_authorization_and_context(
5157 &mut self,
5158 retriever: &crate::query::Retriever,
5159 snapshot: Snapshot,
5160 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5161 context: Option<&crate::query::AiExecutionContext>,
5162 ) -> Result<Vec<crate::query::RetrieverHit>> {
5163 self.require_select()?;
5164 self.ensure_indexes_complete()?;
5165 self.retrieve_at_with_candidate_authorization_on_generation(
5166 retriever,
5167 snapshot,
5168 authorization,
5169 context,
5170 )
5171 }
5172
5173 #[doc(hidden)]
5174 pub fn retrieve_at_with_candidate_authorization_on_generation(
5175 &self,
5176 retriever: &crate::query::Retriever,
5177 snapshot: Snapshot,
5178 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5179 context: Option<&crate::query::AiExecutionContext>,
5180 ) -> Result<Vec<crate::query::RetrieverHit>> {
5181 self.require_select()?;
5182 self.validate_retriever(retriever)?;
5183 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
5184 }
5185
5186 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
5187 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
5188 let (column_id, k) = match retriever {
5189 Retriever::Ann {
5190 column_id,
5191 query,
5192 k,
5193 } => {
5194 let index = self.ann.get(column_id).ok_or_else(|| {
5195 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
5196 })?;
5197 if query.len() != index.dim() {
5198 return Err(MongrelError::InvalidArgument(format!(
5199 "ANN query dimension must be {}, got {}",
5200 index.dim(),
5201 query.len()
5202 )));
5203 }
5204 if query.iter().any(|value| !value.is_finite()) {
5205 return Err(MongrelError::InvalidArgument(
5206 "ANN query values must be finite".into(),
5207 ));
5208 }
5209 (*column_id, *k)
5210 }
5211 Retriever::Sparse {
5212 column_id,
5213 query,
5214 k,
5215 } => {
5216 if !self.sparse.contains_key(column_id) {
5217 return Err(MongrelError::InvalidArgument(format!(
5218 "column {column_id} has no Sparse index"
5219 )));
5220 }
5221 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
5222 return Err(MongrelError::InvalidArgument(
5223 "Sparse query must be non-empty with finite weights".into(),
5224 ));
5225 }
5226 if query.len() > MAX_SPARSE_TERMS {
5227 return Err(MongrelError::InvalidArgument(format!(
5228 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
5229 )));
5230 }
5231 (*column_id, *k)
5232 }
5233 Retriever::MinHash {
5234 column_id,
5235 members,
5236 k,
5237 } => {
5238 if !self.minhash.contains_key(column_id) {
5239 return Err(MongrelError::InvalidArgument(format!(
5240 "column {column_id} has no MinHash index"
5241 )));
5242 }
5243 if members.is_empty() {
5244 return Err(MongrelError::InvalidArgument(
5245 "MinHash members must not be empty".into(),
5246 ));
5247 }
5248 if members.len() > MAX_SET_MEMBERS {
5249 return Err(MongrelError::InvalidArgument(format!(
5250 "MinHash query exceeds {MAX_SET_MEMBERS} members"
5251 )));
5252 }
5253 let mut total_bytes = 0usize;
5254 for member in members {
5255 let bytes = member.encoded_len();
5256 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
5257 return Err(MongrelError::InvalidArgument(format!(
5258 "MinHash member exceeds {} bytes",
5259 crate::query::MAX_SET_MEMBER_BYTES
5260 )));
5261 }
5262 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
5263 MongrelError::InvalidArgument("MinHash input size overflow".into())
5264 })?;
5265 }
5266 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
5267 return Err(MongrelError::InvalidArgument(format!(
5268 "MinHash input exceeds {} bytes",
5269 crate::query::MAX_SET_INPUT_BYTES
5270 )));
5271 }
5272 (*column_id, *k)
5273 }
5274 };
5275 if k == 0 {
5276 return Err(MongrelError::InvalidArgument(
5277 "retriever k must be > 0".into(),
5278 ));
5279 }
5280 if k > MAX_RETRIEVER_K {
5281 return Err(MongrelError::InvalidArgument(format!(
5282 "retriever k exceeds {MAX_RETRIEVER_K}"
5283 )));
5284 }
5285 debug_assert!(self
5286 .schema
5287 .columns
5288 .iter()
5289 .any(|column| column.id == column_id));
5290 Ok(())
5291 }
5292
5293 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
5294 use crate::query::Condition;
5295 match condition {
5296 Condition::Ann {
5297 column_id,
5298 query,
5299 k,
5300 } => self.validate_retriever(&crate::query::Retriever::Ann {
5301 column_id: *column_id,
5302 query: query.clone(),
5303 k: *k,
5304 }),
5305 Condition::SparseMatch {
5306 column_id,
5307 query,
5308 k,
5309 } => self.validate_retriever(&crate::query::Retriever::Sparse {
5310 column_id: *column_id,
5311 query: query.clone(),
5312 k: *k,
5313 }),
5314 Condition::MinHashSimilar {
5315 column_id,
5316 query,
5317 k,
5318 } => {
5319 if !self.minhash.contains_key(column_id) {
5320 return Err(MongrelError::InvalidArgument(format!(
5321 "column {column_id} has no MinHash index"
5322 )));
5323 }
5324 if query.is_empty() || *k == 0 {
5325 return Err(MongrelError::InvalidArgument(
5326 "MinHash query must be non-empty and k must be > 0".into(),
5327 ));
5328 }
5329 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
5330 {
5331 return Err(MongrelError::InvalidArgument(format!(
5332 "MinHash query must have <= {} members and k <= {}",
5333 crate::query::MAX_SET_MEMBERS,
5334 crate::query::MAX_RETRIEVER_K
5335 )));
5336 }
5337 Ok(())
5338 }
5339 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
5340 Err(MongrelError::InvalidArgument(format!(
5341 "bitmap IN exceeds {} values",
5342 crate::query::MAX_SET_MEMBERS
5343 )))
5344 }
5345 Condition::FmContainsAll { patterns, .. }
5346 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
5347 {
5348 Err(MongrelError::InvalidArgument(format!(
5349 "FM query exceeds {} patterns",
5350 crate::query::MAX_HARD_CONDITIONS
5351 )))
5352 }
5353 _ => Ok(()),
5354 }
5355 }
5356
5357 fn retrieve_filtered(
5358 &self,
5359 retriever: &crate::query::Retriever,
5360 snapshot: Snapshot,
5361 hard_filter: Option<&RowIdSet>,
5362 allowed: Option<&std::collections::HashSet<RowId>>,
5363 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5364 context: Option<&crate::query::AiExecutionContext>,
5365 ) -> Result<Vec<crate::query::RetrieverHit>> {
5366 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
5367 let started = std::time::Instant::now();
5368 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
5369 Retriever::Ann {
5370 column_id,
5371 query,
5372 k,
5373 } => {
5374 let Some(index) = self.ann.get(column_id) else {
5375 return Ok(Vec::new());
5376 };
5377 let cap = ann_candidate_cap(index.len(), context);
5378 if cap == 0 {
5379 return Ok(Vec::new());
5380 }
5381 let mut breadth = (*k).max(1).min(cap);
5382 let mut eligibility = std::collections::HashMap::new();
5383 let mut filtered = loop {
5384 let mut seen = std::collections::HashSet::new();
5385 if let Some(context) = context {
5386 context.checkpoint()?;
5387 }
5388 let raw = index.search_with_context(query, breadth, context)?;
5389 let unchecked: Vec<_> = raw
5390 .iter()
5391 .map(|(row_id, _)| *row_id)
5392 .filter(|row_id| !eligibility.contains_key(row_id))
5393 .filter(|row_id| {
5394 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5395 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5396 })
5397 .collect();
5398 let eligible = self.eligible_and_authorized_candidate_ids(
5399 &unchecked,
5400 *column_id,
5401 snapshot,
5402 candidate_authorization,
5403 context,
5404 )?;
5405 for row_id in unchecked {
5406 eligibility.insert(row_id, eligible.contains(&row_id));
5407 }
5408 let filtered: Vec<_> = raw
5409 .into_iter()
5410 .filter(|(row_id, _)| {
5411 seen.insert(*row_id)
5412 && eligibility.get(row_id).copied().unwrap_or(false)
5413 })
5414 .map(|(row_id, score)| {
5415 let score = match score {
5416 crate::index::AnnDistance::Hamming(d) => {
5417 RetrieverScore::AnnHammingDistance(d)
5418 }
5419 crate::index::AnnDistance::Cosine(d) => {
5420 RetrieverScore::AnnCosineDistance(d)
5421 }
5422 };
5423 (row_id, score)
5424 })
5425 .collect();
5426 if filtered.len() >= *k || breadth >= cap {
5427 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5428 crate::trace::QueryTrace::record(|trace| {
5429 trace.ann_candidate_cap_hit = true;
5430 });
5431 }
5432 break filtered;
5433 }
5434 breadth = breadth.saturating_mul(2).min(cap);
5435 };
5436 filtered.truncate(*k);
5437 filtered
5438 }
5439 Retriever::Sparse {
5440 column_id,
5441 query,
5442 k,
5443 } => self
5444 .sparse
5445 .get(column_id)
5446 .map(|index| -> Result<Vec<_>> {
5447 let mut breadth = (*k).max(1);
5448 let mut eligibility = std::collections::HashMap::new();
5449 loop {
5450 if let Some(context) = context {
5451 context.checkpoint()?;
5452 }
5453 let raw = index.search_with_context(query, breadth, context)?;
5454 let unchecked: Vec<_> = raw
5455 .iter()
5456 .map(|(row_id, _)| *row_id)
5457 .filter(|row_id| !eligibility.contains_key(row_id))
5458 .filter(|row_id| {
5459 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5460 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5461 })
5462 .collect();
5463 let eligible = self.eligible_and_authorized_candidate_ids(
5464 &unchecked,
5465 *column_id,
5466 snapshot,
5467 candidate_authorization,
5468 context,
5469 )?;
5470 for row_id in unchecked {
5471 eligibility.insert(row_id, eligible.contains(&row_id));
5472 }
5473 let filtered: Vec<_> = raw
5474 .iter()
5475 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5476 .take(*k)
5477 .map(|(row_id, score)| {
5478 (*row_id, RetrieverScore::SparseDotProduct(*score))
5479 })
5480 .collect();
5481 if filtered.len() >= *k || raw.len() < breadth {
5482 break Ok(filtered);
5483 }
5484 let next = breadth.saturating_mul(2);
5485 if next == breadth {
5486 break Ok(filtered);
5487 }
5488 breadth = next;
5489 }
5490 })
5491 .transpose()?
5492 .unwrap_or_default(),
5493 Retriever::MinHash {
5494 column_id,
5495 members,
5496 k,
5497 } => self
5498 .minhash
5499 .get(column_id)
5500 .map(|index| -> Result<Vec<_>> {
5501 let mut hashes = Vec::with_capacity(members.len());
5502 for member in members {
5503 if let Some(context) = context {
5504 context.consume(crate::query::work_units(
5505 member.encoded_len(),
5506 crate::query::PARSE_WORK_QUANTUM,
5507 ))?;
5508 }
5509 hashes.push(member.hash_v1());
5510 }
5511 let mut breadth = (*k).max(1);
5512 let mut eligibility = std::collections::HashMap::new();
5513 loop {
5514 if let Some(context) = context {
5515 context.checkpoint()?;
5516 }
5517 let raw = index.search_with_context(&hashes, breadth, context)?;
5518 let unchecked: Vec<_> = raw
5519 .iter()
5520 .map(|(row_id, _)| *row_id)
5521 .filter(|row_id| !eligibility.contains_key(row_id))
5522 .filter(|row_id| {
5523 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5524 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5525 })
5526 .collect();
5527 let eligible = self.eligible_and_authorized_candidate_ids(
5528 &unchecked,
5529 *column_id,
5530 snapshot,
5531 candidate_authorization,
5532 context,
5533 )?;
5534 for row_id in unchecked {
5535 eligibility.insert(row_id, eligible.contains(&row_id));
5536 }
5537 let filtered: Vec<_> = raw
5538 .iter()
5539 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5540 .take(*k)
5541 .map(|(row_id, score)| {
5542 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5543 })
5544 .collect();
5545 if filtered.len() >= *k || raw.len() < breadth {
5546 break Ok(filtered);
5547 }
5548 let next = breadth.saturating_mul(2);
5549 if next == breadth {
5550 break Ok(filtered);
5551 }
5552 breadth = next;
5553 }
5554 })
5555 .transpose()?
5556 .unwrap_or_default(),
5557 };
5558 let elapsed = started.elapsed().as_nanos() as u64;
5559 crate::trace::QueryTrace::record(|trace| {
5560 match retriever {
5561 Retriever::Ann { .. } => {
5562 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5563 }
5564 Retriever::Sparse { .. } => {
5565 trace.sparse_candidate_nanos =
5566 trace.sparse_candidate_nanos.saturating_add(elapsed)
5567 }
5568 Retriever::MinHash { .. } => {
5569 trace.minhash_candidate_nanos =
5570 trace.minhash_candidate_nanos.saturating_add(elapsed)
5571 }
5572 }
5573 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5574 });
5575 Ok(scored
5576 .into_iter()
5577 .enumerate()
5578 .map(|(rank, (row_id, score))| RetrieverHit {
5579 row_id,
5580 rank: rank + 1,
5581 score,
5582 })
5583 .collect())
5584 }
5585
5586 fn eligible_candidate_ids(
5587 &self,
5588 candidates: &[RowId],
5589 _column_id: u16,
5590 snapshot: Snapshot,
5591 context: Option<&crate::query::AiExecutionContext>,
5592 ) -> Result<std::collections::HashSet<RowId>> {
5593 if !self.had_deletes
5594 && self.ttl.is_none()
5595 && self.pending_put_cols.is_empty()
5596 && snapshot.epoch == self.snapshot().epoch
5597 {
5598 return Ok(candidates.iter().copied().collect());
5599 }
5600 let mut readers: Vec<_> = self
5601 .run_refs
5602 .iter()
5603 .map(|run| self.open_reader(run.run_id))
5604 .collect::<Result<_>>()?;
5605 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5606 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5607 for &row_id in candidates {
5608 if let Some(context) = context {
5609 context.consume(1)?;
5610 }
5611 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5612 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5613 let overlay = match (mem, mutable) {
5614 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5615 (Some(value), None) | (None, Some(value)) => Some(value),
5616 (None, None) => None,
5617 };
5618 if let Some((_, row)) = overlay {
5619 if !row.deleted && !self.row_expired_at(&row, now) {
5620 eligible.insert(row_id);
5621 }
5622 continue;
5623 }
5624 let mut best: Option<(Epoch, bool, usize)> = None;
5625 for (index, reader) in readers.iter_mut().enumerate() {
5626 if let Some((epoch, deleted)) =
5627 reader.get_version_visibility(row_id, snapshot.epoch)?
5628 {
5629 if best
5630 .as_ref()
5631 .map(|(best_epoch, ..)| epoch > *best_epoch)
5632 .unwrap_or(true)
5633 {
5634 best = Some((epoch, deleted, index));
5635 }
5636 }
5637 }
5638 let Some((_, false, reader_index)) = best else {
5639 continue;
5640 };
5641 if let Some(ttl) = self.ttl {
5642 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5643 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5644 {
5645 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5646 continue;
5647 }
5648 }
5649 }
5650 eligible.insert(row_id);
5651 }
5652 Ok(eligible)
5653 }
5654
5655 fn eligible_and_authorized_candidate_ids(
5656 &self,
5657 candidates: &[RowId],
5658 column_id: u16,
5659 snapshot: Snapshot,
5660 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5661 context: Option<&crate::query::AiExecutionContext>,
5662 ) -> Result<std::collections::HashSet<RowId>> {
5663 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5664 let Some(authorization) = authorization else {
5665 return Ok(eligible);
5666 };
5667 let candidates: Vec<_> = eligible.into_iter().collect();
5668 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5669 }
5670
5671 fn policy_allowed_candidate_ids(
5672 &self,
5673 candidates: &[RowId],
5674 snapshot: Snapshot,
5675 authorization: &crate::security::CandidateAuthorization<'_>,
5676 context: Option<&crate::query::AiExecutionContext>,
5677 ) -> Result<std::collections::HashSet<RowId>> {
5678 let started = std::time::Instant::now();
5679 if candidates.is_empty()
5680 || authorization.principal.is_admin
5681 || !authorization.security.rls_enabled(authorization.table)
5682 {
5683 return Ok(candidates.iter().copied().collect());
5684 }
5685 if let Some(context) = context {
5686 context.checkpoint()?;
5687 }
5688 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5689 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5690 .iter()
5691 .map(|row_id| {
5692 (
5693 *row_id,
5694 Row {
5695 row_id: *row_id,
5696 committed_epoch: snapshot.epoch,
5697 columns: std::collections::HashMap::new(),
5698 deleted: false,
5699 },
5700 )
5701 })
5702 .collect();
5703 let columns = authorization
5704 .security
5705 .select_policy_columns(authorization.table, authorization.principal);
5706 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5707 let mut decoded = 0usize;
5708 for column_id in &columns {
5709 if let Some(context) = context {
5710 context.checkpoint()?;
5711 }
5712 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5713 &row_ids, *column_id, snapshot, query_now, context,
5714 )? {
5715 if let Some(row) = rows.get_mut(&row_id) {
5716 row.columns.insert(*column_id, value);
5717 decoded = decoded.saturating_add(1);
5718 }
5719 }
5720 }
5721 if let Some(context) = context {
5722 context.consume(candidates.len().saturating_add(decoded))?;
5723 }
5724 let allowed = rows
5725 .into_values()
5726 .filter_map(|row| {
5727 authorization
5728 .security
5729 .row_allowed(
5730 authorization.table,
5731 crate::security::PolicyCommand::Select,
5732 &row,
5733 authorization.principal,
5734 false,
5735 )
5736 .then_some(row.row_id)
5737 })
5738 .collect();
5739 crate::trace::QueryTrace::record(|trace| {
5740 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5741 trace.rls_policy_columns_decoded =
5742 trace.rls_policy_columns_decoded.saturating_add(decoded);
5743 trace.authorization_nanos = trace
5744 .authorization_nanos
5745 .saturating_add(started.elapsed().as_nanos() as u64);
5746 });
5747 Ok(allowed)
5748 }
5749
5750 pub fn search(
5752 &mut self,
5753 request: &crate::query::SearchRequest,
5754 ) -> Result<Vec<crate::query::SearchHit>> {
5755 self.search_with_allowed(request, None)
5756 }
5757
5758 pub fn search_at(
5759 &mut self,
5760 request: &crate::query::SearchRequest,
5761 snapshot: Snapshot,
5762 authorized: Option<&std::collections::HashSet<RowId>>,
5763 ) -> Result<Vec<crate::query::SearchHit>> {
5764 self.search_at_with_allowed(request, snapshot, authorized)
5765 }
5766
5767 pub fn search_with_allowed(
5768 &mut self,
5769 request: &crate::query::SearchRequest,
5770 authorized: Option<&std::collections::HashSet<RowId>>,
5771 ) -> Result<Vec<crate::query::SearchHit>> {
5772 self.search_at_with_allowed(request, self.snapshot(), authorized)
5773 }
5774
5775 pub fn search_at_with_allowed(
5776 &mut self,
5777 request: &crate::query::SearchRequest,
5778 snapshot: Snapshot,
5779 authorized: Option<&std::collections::HashSet<RowId>>,
5780 ) -> Result<Vec<crate::query::SearchHit>> {
5781 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5782 }
5783
5784 pub fn search_at_with_allowed_and_context(
5785 &mut self,
5786 request: &crate::query::SearchRequest,
5787 snapshot: Snapshot,
5788 authorized: Option<&std::collections::HashSet<RowId>>,
5789 context: Option<&crate::query::AiExecutionContext>,
5790 ) -> Result<Vec<crate::query::SearchHit>> {
5791 self.ensure_indexes_complete()?;
5792 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5793 }
5794
5795 pub fn search_at_with_candidate_authorization_and_context(
5796 &mut self,
5797 request: &crate::query::SearchRequest,
5798 snapshot: Snapshot,
5799 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5800 context: Option<&crate::query::AiExecutionContext>,
5801 ) -> Result<Vec<crate::query::SearchHit>> {
5802 self.ensure_indexes_complete()?;
5803 self.search_at_with_filters_and_context(
5804 request,
5805 snapshot,
5806 None,
5807 authorization,
5808 context,
5809 None,
5810 )
5811 }
5812
5813 #[doc(hidden)]
5814 pub fn search_at_with_candidate_authorization_on_generation(
5815 &self,
5816 request: &crate::query::SearchRequest,
5817 snapshot: Snapshot,
5818 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5819 context: Option<&crate::query::AiExecutionContext>,
5820 ) -> Result<Vec<crate::query::SearchHit>> {
5821 self.search_at_with_filters_and_context(
5822 request,
5823 snapshot,
5824 None,
5825 authorization,
5826 context,
5827 None,
5828 )
5829 }
5830
5831 #[doc(hidden)]
5832 pub fn search_at_with_candidate_authorization_on_generation_after(
5833 &self,
5834 request: &crate::query::SearchRequest,
5835 snapshot: Snapshot,
5836 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5837 context: Option<&crate::query::AiExecutionContext>,
5838 after: Option<crate::query::SearchAfter>,
5839 ) -> Result<Vec<crate::query::SearchHit>> {
5840 self.search_at_with_filters_and_context(
5841 request,
5842 snapshot,
5843 None,
5844 authorization,
5845 context,
5846 after,
5847 )
5848 }
5849
5850 fn search_at_with_filters_and_context(
5851 &self,
5852 request: &crate::query::SearchRequest,
5853 snapshot: Snapshot,
5854 authorized: Option<&std::collections::HashSet<RowId>>,
5855 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5856 context: Option<&crate::query::AiExecutionContext>,
5857 after: Option<crate::query::SearchAfter>,
5858 ) -> Result<Vec<crate::query::SearchHit>> {
5859 use crate::query::{
5860 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5861 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5862 };
5863 let total_started = std::time::Instant::now();
5864 let rank_offset = after.map_or(0, |after| after.returned_count);
5865 self.require_select()?;
5866 if request.limit == 0 {
5867 return Err(MongrelError::InvalidArgument(
5868 "search limit must be > 0".into(),
5869 ));
5870 }
5871 if request.limit > MAX_FINAL_LIMIT {
5872 return Err(MongrelError::InvalidArgument(format!(
5873 "search limit exceeds {MAX_FINAL_LIMIT}"
5874 )));
5875 }
5876 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5877 return Err(MongrelError::InvalidArgument(
5878 "search-after score must be finite".into(),
5879 ));
5880 }
5881 if request.retrievers.is_empty() {
5882 return Err(MongrelError::InvalidArgument(
5883 "search requires at least one retriever".into(),
5884 ));
5885 }
5886 if request.retrievers.len() > MAX_RETRIEVERS {
5887 return Err(MongrelError::InvalidArgument(format!(
5888 "search exceeds {MAX_RETRIEVERS} retrievers"
5889 )));
5890 }
5891 if request.must.len() > MAX_HARD_CONDITIONS {
5892 return Err(MongrelError::InvalidArgument(format!(
5893 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5894 )));
5895 }
5896 for condition in &request.must {
5897 self.validate_condition(condition)?;
5898 }
5899 if request.must.iter().any(|condition| {
5900 matches!(
5901 condition,
5902 Condition::Ann { .. }
5903 | Condition::SparseMatch { .. }
5904 | Condition::MinHashSimilar { .. }
5905 )
5906 }) {
5907 return Err(MongrelError::InvalidArgument(
5908 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5909 .into(),
5910 ));
5911 }
5912 let mut names = std::collections::HashSet::new();
5913 for named in &request.retrievers {
5914 if named.name.is_empty()
5915 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5916 || !names.insert(named.name.as_str())
5917 {
5918 return Err(MongrelError::InvalidArgument(format!(
5919 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5920 crate::query::MAX_RETRIEVER_NAME_BYTES
5921 )));
5922 }
5923 if !named.weight.is_finite()
5924 || named.weight < 0.0
5925 || named.weight > MAX_RETRIEVER_WEIGHT
5926 {
5927 return Err(MongrelError::InvalidArgument(format!(
5928 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5929 )));
5930 }
5931 self.validate_retriever(&named.retriever)?;
5932 }
5933 let projection = request
5934 .projection
5935 .clone()
5936 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5937 if projection.len() > MAX_PROJECTION_COLUMNS {
5938 return Err(MongrelError::InvalidArgument(format!(
5939 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5940 )));
5941 }
5942 for column_id in &projection {
5943 if !self
5944 .schema
5945 .columns
5946 .iter()
5947 .any(|column| column.id == *column_id)
5948 {
5949 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5950 }
5951 }
5952 if let Some(crate::query::Rerank::ExactVector {
5953 embedding_column,
5954 query,
5955 candidate_limit,
5956 weight,
5957 ..
5958 }) = &request.rerank
5959 {
5960 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5961 {
5962 return Err(MongrelError::InvalidArgument(format!(
5963 "rerank candidate_limit must be between search limit and {}",
5964 crate::query::MAX_RETRIEVER_K
5965 )));
5966 }
5967 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5968 return Err(MongrelError::InvalidArgument(format!(
5969 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5970 )));
5971 }
5972 let column = self
5973 .schema
5974 .columns
5975 .iter()
5976 .find(|column| column.id == *embedding_column)
5977 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5978 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5979 return Err(MongrelError::InvalidArgument(format!(
5980 "rerank column {embedding_column} is not an embedding"
5981 )));
5982 };
5983 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5984 return Err(MongrelError::InvalidArgument(format!(
5985 "rerank query must contain {dim} finite values"
5986 )));
5987 }
5988 }
5989
5990 let hard_filter_started = std::time::Instant::now();
5991 let hard_filter = if request.must.is_empty() {
5992 None
5993 } else {
5994 let mut sets = Vec::with_capacity(request.must.len());
5995 for condition in &request.must {
5996 if let Some(context) = context {
5997 context.checkpoint()?;
5998 }
5999 sets.push(self.resolve_condition(condition, snapshot)?);
6000 }
6001 Some(RowIdSet::intersect_many(sets))
6002 };
6003 crate::trace::QueryTrace::record(|trace| {
6004 trace.hard_filter_nanos = trace
6005 .hard_filter_nanos
6006 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
6007 });
6008 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
6009 return Ok(Vec::new());
6010 }
6011
6012 let constant = match request.fusion {
6013 Fusion::ReciprocalRank { constant } => constant,
6014 };
6015 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
6016 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
6017 let mut fusion_nanos = 0u64;
6018 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
6019 std::collections::HashMap::new();
6020 for named in retrievers {
6021 if named.weight == 0.0 {
6022 continue;
6023 }
6024 if let Some(context) = context {
6025 context.checkpoint()?;
6026 }
6027 let hits = self.retrieve_filtered(
6028 &named.retriever,
6029 snapshot,
6030 hard_filter.as_ref(),
6031 authorized,
6032 candidate_authorization,
6033 context,
6034 )?;
6035 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
6036 let fusion_started = std::time::Instant::now();
6037 for hit in hits {
6038 if let Some(context) = context {
6039 context.consume(1)?;
6040 }
6041 let contribution = named.weight / (constant as f64 + hit.rank as f64);
6042 if !contribution.is_finite() {
6043 return Err(MongrelError::InvalidArgument(
6044 "retriever contribution must be finite".into(),
6045 ));
6046 }
6047 let max_fused_candidates = context.map_or(
6048 crate::query::MAX_FUSED_CANDIDATES,
6049 crate::query::AiExecutionContext::max_fused_candidates,
6050 );
6051 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
6052 return Err(MongrelError::WorkBudgetExceeded);
6053 }
6054 let entry = fused.entry(hit.row_id).or_default();
6055 entry.0 += contribution;
6056 if !entry.0.is_finite() {
6057 return Err(MongrelError::InvalidArgument(
6058 "fused score must be finite".into(),
6059 ));
6060 }
6061 entry.1.push(ComponentScore {
6062 retriever_name: retriever_name.clone(),
6063 rank: hit.rank,
6064 raw_score: hit.score,
6065 contribution,
6066 });
6067 }
6068 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
6069 }
6070 let union_size = fused.len();
6071 let mut ranked: Vec<_> = fused
6072 .into_iter()
6073 .map(|(row_id, (fused_score, components))| {
6074 (row_id, fused_score, components, None, fused_score)
6075 })
6076 .collect();
6077 let order = |(a_row, _, _, _, a_score): &(
6078 RowId,
6079 f64,
6080 Vec<ComponentScore>,
6081 Option<f32>,
6082 f64,
6083 ),
6084 (b_row, _, _, _, b_score): &(
6085 RowId,
6086 f64,
6087 Vec<ComponentScore>,
6088 Option<f32>,
6089 f64,
6090 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
6091 if let Some(crate::query::Rerank::ExactVector {
6092 embedding_column,
6093 query,
6094 metric,
6095 candidate_limit,
6096 weight,
6097 }) = &request.rerank
6098 {
6099 let fused_order = |(a_row, a_score, ..): &(
6100 RowId,
6101 f64,
6102 Vec<ComponentScore>,
6103 Option<f32>,
6104 f64,
6105 ),
6106 (b_row, b_score, ..): &(
6107 RowId,
6108 f64,
6109 Vec<ComponentScore>,
6110 Option<f32>,
6111 f64,
6112 )| {
6113 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
6114 };
6115 let selection_started = std::time::Instant::now();
6116 if let Some(context) = context {
6117 context.consume(ranked.len())?;
6118 }
6119 if ranked.len() > *candidate_limit {
6120 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
6121 ranked.truncate(*candidate_limit);
6122 }
6123 ranked.sort_by(fused_order);
6124 fusion_nanos =
6125 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6126 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6127 if let Some(context) = context {
6128 context.consume(row_ids.len())?;
6129 }
6130 let query_now =
6131 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6132 let gather_started = std::time::Instant::now();
6133 let vectors = self.values_for_rids_batch_at_with_context(
6134 &row_ids,
6135 *embedding_column,
6136 snapshot,
6137 query_now,
6138 context,
6139 )?;
6140 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6141 let vector_work =
6142 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
6143 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
6144 if let Some(context) = context {
6145 context.consume(vector_work)?;
6146 }
6147 query
6148 .iter()
6149 .map(|value| f64::from(*value).powi(2))
6150 .sum::<f64>()
6151 .sqrt()
6152 } else {
6153 0.0
6154 };
6155 let score_started = std::time::Instant::now();
6156 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
6157 for (row_id, value) in vectors {
6158 let Some(vector) = value.as_embedding() else {
6159 continue;
6160 };
6161 let score = match metric {
6162 crate::query::VectorMetric::DotProduct => {
6163 if let Some(context) = context {
6164 context.consume(vector_work)?;
6165 }
6166 query
6167 .iter()
6168 .zip(vector)
6169 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6170 .sum::<f64>()
6171 }
6172 crate::query::VectorMetric::Cosine => {
6173 if let Some(context) = context {
6174 context.consume(vector_work.saturating_mul(2))?;
6175 }
6176 let dot = query
6177 .iter()
6178 .zip(vector)
6179 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6180 .sum::<f64>();
6181 let norm = vector
6182 .iter()
6183 .map(|value| f64::from(*value).powi(2))
6184 .sum::<f64>()
6185 .sqrt();
6186 if query_norm == 0.0 || norm == 0.0 {
6187 0.0
6188 } else {
6189 dot / (query_norm * norm)
6190 }
6191 }
6192 crate::query::VectorMetric::Euclidean => {
6193 if let Some(context) = context {
6194 context.consume(vector_work)?;
6195 }
6196 query
6197 .iter()
6198 .zip(vector)
6199 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6200 .sum::<f64>()
6201 .sqrt()
6202 }
6203 };
6204 if !score.is_finite() {
6205 return Err(MongrelError::InvalidArgument(
6206 "exact rerank score must be finite".into(),
6207 ));
6208 }
6209 scores.insert(row_id, score as f32);
6210 }
6211 let mut reranked = Vec::with_capacity(ranked.len());
6212 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
6213 let Some(score) = scores.get(&row_id).copied() else {
6214 continue;
6215 };
6216 let ordering_score = match metric {
6217 crate::query::VectorMetric::Euclidean => -f64::from(score),
6218 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
6219 f64::from(score)
6220 }
6221 };
6222 let final_score = fused_score + *weight * ordering_score;
6223 if !final_score.is_finite() {
6224 return Err(MongrelError::InvalidArgument(
6225 "final rerank score must be finite".into(),
6226 ));
6227 }
6228 reranked.push((row_id, fused_score, components, Some(score), final_score));
6229 }
6230 ranked = reranked;
6231 ranked.sort_by(order);
6232 crate::trace::QueryTrace::record(|trace| {
6233 trace.exact_vector_gather_nanos =
6234 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6235 trace.exact_vector_score_nanos = trace
6236 .exact_vector_score_nanos
6237 .saturating_add(score_started.elapsed().as_nanos() as u64);
6238 });
6239 }
6240 if let Some(after) = after {
6241 ranked.retain(|(row_id, _, _, _, final_score)| {
6242 final_score.total_cmp(&after.final_score).is_lt()
6243 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
6244 });
6245 }
6246 let projection_started = std::time::Instant::now();
6247 let sentinel = projection
6248 .first()
6249 .copied()
6250 .or_else(|| self.schema.columns.first().map(|column| column.id));
6251 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6252 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
6253 let mut projection_rows = 0usize;
6254 let mut projection_cells = 0usize;
6255 while out.len() < request.limit && !ranked.is_empty() {
6256 if let Some(context) = context {
6257 context.checkpoint()?;
6258 context.consume(ranked.len())?;
6259 }
6260 let needed = request.limit - out.len();
6261 let window_size = ranked
6262 .len()
6263 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
6264 let selection_started = std::time::Instant::now();
6265 let mut remainder = if ranked.len() > window_size {
6266 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
6267 ranked.split_off(window_size)
6268 } else {
6269 Vec::new()
6270 };
6271 ranked.sort_by(order);
6272 fusion_nanos =
6273 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
6274 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
6275 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
6276 if let Some(context) = context {
6277 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
6278 }
6279 projection_rows = projection_rows.saturating_add(row_ids.len());
6280 projection_cells =
6281 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
6282 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
6283 std::collections::HashMap::new();
6284 if let Some(column_id) = sentinel {
6285 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6286 &row_ids, column_id, snapshot, query_now, context,
6287 )? {
6288 cells.entry(row_id).or_default().insert(column_id, value);
6289 }
6290 }
6291 for &column_id in &projection {
6292 if Some(column_id) == sentinel {
6293 continue;
6294 }
6295 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6296 &row_ids, column_id, snapshot, query_now, context,
6297 )? {
6298 cells.entry(row_id).or_default().insert(column_id, value);
6299 }
6300 }
6301 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
6302 ranked.drain(..)
6303 {
6304 let Some(row_cells) = cells.remove(&row_id) else {
6305 continue;
6306 };
6307 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
6308 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
6309 out.push(SearchHit {
6310 row_id,
6311 cells: projection
6312 .iter()
6313 .filter_map(|column_id| {
6314 row_cells
6315 .get(column_id)
6316 .cloned()
6317 .map(|value| (*column_id, value))
6318 })
6319 .collect(),
6320 components,
6321 fused_score,
6322 exact_rerank_score,
6323 final_score,
6324 final_rank,
6325 });
6326 if out.len() == request.limit {
6327 break;
6328 }
6329 }
6330 ranked.append(&mut remainder);
6331 }
6332 crate::trace::QueryTrace::record(|trace| {
6333 trace.union_size = union_size;
6334 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
6335 trace.projection_nanos = trace
6336 .projection_nanos
6337 .saturating_add(projection_started.elapsed().as_nanos() as u64);
6338 trace.total_nanos = trace
6339 .total_nanos
6340 .saturating_add(total_started.elapsed().as_nanos() as u64);
6341 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
6342 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
6343 if let Some(context) = context {
6344 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
6345 }
6346 });
6347 Ok(out)
6348 }
6349
6350 pub fn set_similarity(
6353 &mut self,
6354 request: &crate::query::SetSimilarityRequest,
6355 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6356 self.set_similarity_with_allowed(request, None)
6357 }
6358
6359 pub fn set_similarity_at(
6360 &mut self,
6361 request: &crate::query::SetSimilarityRequest,
6362 snapshot: Snapshot,
6363 allowed: Option<&std::collections::HashSet<RowId>>,
6364 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6365 self.set_similarity_explained_at(request, snapshot, allowed)
6366 .map(|(hits, _)| hits)
6367 }
6368
6369 pub fn ann_rerank(
6371 &mut self,
6372 request: &crate::query::AnnRerankRequest,
6373 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6374 self.ann_rerank_with_allowed(request, None)
6375 }
6376
6377 pub fn ann_rerank_with_allowed(
6378 &mut self,
6379 request: &crate::query::AnnRerankRequest,
6380 allowed: Option<&std::collections::HashSet<RowId>>,
6381 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6382 self.ann_rerank_at(request, self.snapshot(), allowed)
6383 }
6384
6385 pub fn ann_rerank_at(
6386 &mut self,
6387 request: &crate::query::AnnRerankRequest,
6388 snapshot: Snapshot,
6389 allowed: Option<&std::collections::HashSet<RowId>>,
6390 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6391 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
6392 }
6393
6394 pub fn ann_rerank_at_with_context(
6395 &mut self,
6396 request: &crate::query::AnnRerankRequest,
6397 snapshot: Snapshot,
6398 allowed: Option<&std::collections::HashSet<RowId>>,
6399 context: Option<&crate::query::AiExecutionContext>,
6400 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6401 self.ensure_indexes_complete()?;
6402 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6403 }
6404
6405 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6406 &mut self,
6407 request: &crate::query::AnnRerankRequest,
6408 snapshot: Snapshot,
6409 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6410 context: Option<&crate::query::AiExecutionContext>,
6411 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6412 self.ensure_indexes_complete()?;
6413 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6414 }
6415
6416 #[doc(hidden)]
6417 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6418 &self,
6419 request: &crate::query::AnnRerankRequest,
6420 snapshot: Snapshot,
6421 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6422 context: Option<&crate::query::AiExecutionContext>,
6423 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6424 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6425 }
6426
6427 fn ann_rerank_at_with_filters_and_context(
6428 &self,
6429 request: &crate::query::AnnRerankRequest,
6430 snapshot: Snapshot,
6431 allowed: Option<&std::collections::HashSet<RowId>>,
6432 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6433 context: Option<&crate::query::AiExecutionContext>,
6434 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6435 use crate::query::{
6436 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
6437 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6438 };
6439 if request.candidate_k == 0 || request.limit == 0 {
6440 return Err(MongrelError::InvalidArgument(
6441 "candidate_k and limit must be > 0".into(),
6442 ));
6443 }
6444 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6445 return Err(MongrelError::InvalidArgument(format!(
6446 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6447 )));
6448 }
6449 let retriever = Retriever::Ann {
6450 column_id: request.column_id,
6451 query: request.query.clone(),
6452 k: request.candidate_k,
6453 };
6454 self.require_select()?;
6455 self.validate_retriever(&retriever)?;
6456 let hits = self.retrieve_filtered(
6457 &retriever,
6458 snapshot,
6459 None,
6460 allowed,
6461 candidate_authorization,
6462 context,
6463 )?;
6464 let distances: std::collections::HashMap<_, _> = hits
6465 .iter()
6466 .filter_map(|hit| match hit.score {
6467 RetrieverScore::AnnHammingDistance(distance) => {
6468 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
6469 }
6470 RetrieverScore::AnnCosineDistance(distance) => {
6471 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
6472 }
6473 _ => None,
6474 })
6475 .collect();
6476 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6477 if let Some(context) = context {
6478 context.consume(row_ids.len())?;
6479 }
6480 let gather_started = std::time::Instant::now();
6481 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6482 let values = self.values_for_rids_batch_at_with_context(
6483 &row_ids,
6484 request.column_id,
6485 snapshot,
6486 query_now,
6487 context,
6488 )?;
6489 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6490 let score_started = std::time::Instant::now();
6491 let vector_work =
6492 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6493 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6494 if let Some(context) = context {
6495 context.consume(vector_work)?;
6496 }
6497 request
6498 .query
6499 .iter()
6500 .map(|value| f64::from(*value).powi(2))
6501 .sum::<f64>()
6502 .sqrt()
6503 } else {
6504 0.0
6505 };
6506 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6507 for (row_id, value) in values {
6508 let Some(vector) = value.as_embedding() else {
6509 continue;
6510 };
6511 let exact_score = match request.metric {
6512 VectorMetric::DotProduct => {
6513 if let Some(context) = context {
6514 context.consume(vector_work)?;
6515 }
6516 request
6517 .query
6518 .iter()
6519 .zip(vector)
6520 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6521 .sum::<f64>()
6522 }
6523 VectorMetric::Cosine => {
6524 if let Some(context) = context {
6525 context.consume(vector_work.saturating_mul(2))?;
6526 }
6527 let dot = request
6528 .query
6529 .iter()
6530 .zip(vector)
6531 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6532 .sum::<f64>();
6533 let norm = vector
6534 .iter()
6535 .map(|value| f64::from(*value).powi(2))
6536 .sum::<f64>()
6537 .sqrt();
6538 if query_norm == 0.0 || norm == 0.0 {
6539 0.0
6540 } else {
6541 dot / (query_norm * norm)
6542 }
6543 }
6544 VectorMetric::Euclidean => {
6545 if let Some(context) = context {
6546 context.consume(vector_work)?;
6547 }
6548 request
6549 .query
6550 .iter()
6551 .zip(vector)
6552 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6553 .sum::<f64>()
6554 .sqrt()
6555 }
6556 };
6557 let exact_score = exact_score as f32;
6558 if !exact_score.is_finite() {
6559 return Err(MongrelError::InvalidArgument(
6560 "exact ANN score must be finite".into(),
6561 ));
6562 }
6563 let Some(candidate_distance) = distances.get(&row_id).copied() else {
6564 continue;
6565 };
6566 reranked.push(AnnRerankHit {
6567 row_id,
6568 candidate_distance,
6569 exact_score,
6570 });
6571 }
6572 reranked.sort_by(|left, right| {
6573 let score = match request.metric {
6574 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6575 VectorMetric::Cosine | VectorMetric::DotProduct => {
6576 right.exact_score.total_cmp(&left.exact_score)
6577 }
6578 };
6579 score.then_with(|| left.row_id.cmp(&right.row_id))
6580 });
6581 reranked.truncate(request.limit);
6582 crate::trace::QueryTrace::record(|trace| {
6583 trace.exact_vector_gather_nanos =
6584 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6585 trace.exact_vector_score_nanos = trace
6586 .exact_vector_score_nanos
6587 .saturating_add(score_started.elapsed().as_nanos() as u64);
6588 });
6589 Ok(reranked)
6590 }
6591
6592 pub fn set_similarity_with_allowed(
6593 &mut self,
6594 request: &crate::query::SetSimilarityRequest,
6595 allowed: Option<&std::collections::HashSet<RowId>>,
6596 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6597 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6598 .map(|(hits, _)| hits)
6599 }
6600
6601 pub fn set_similarity_explained(
6602 &mut self,
6603 request: &crate::query::SetSimilarityRequest,
6604 ) -> Result<(
6605 Vec<crate::query::SetSimilarityHit>,
6606 crate::query::SetSimilarityTrace,
6607 )> {
6608 self.set_similarity_explained_at(request, self.snapshot(), None)
6609 }
6610
6611 fn set_similarity_explained_at(
6612 &mut self,
6613 request: &crate::query::SetSimilarityRequest,
6614 snapshot: Snapshot,
6615 allowed: Option<&std::collections::HashSet<RowId>>,
6616 ) -> Result<(
6617 Vec<crate::query::SetSimilarityHit>,
6618 crate::query::SetSimilarityTrace,
6619 )> {
6620 self.ensure_indexes_complete()?;
6621 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6622 }
6623
6624 pub fn set_similarity_at_with_context(
6625 &mut self,
6626 request: &crate::query::SetSimilarityRequest,
6627 snapshot: Snapshot,
6628 allowed: Option<&std::collections::HashSet<RowId>>,
6629 context: Option<&crate::query::AiExecutionContext>,
6630 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6631 self.ensure_indexes_complete()?;
6632 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6633 .map(|(hits, _)| hits)
6634 }
6635
6636 pub fn set_similarity_at_with_candidate_authorization_and_context(
6637 &mut self,
6638 request: &crate::query::SetSimilarityRequest,
6639 snapshot: Snapshot,
6640 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6641 context: Option<&crate::query::AiExecutionContext>,
6642 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6643 self.ensure_indexes_complete()?;
6644 self.set_similarity_explained_at_with_context(
6645 request,
6646 snapshot,
6647 None,
6648 authorization,
6649 context,
6650 )
6651 .map(|(hits, _)| hits)
6652 }
6653
6654 #[doc(hidden)]
6655 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6656 &self,
6657 request: &crate::query::SetSimilarityRequest,
6658 snapshot: Snapshot,
6659 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6660 context: Option<&crate::query::AiExecutionContext>,
6661 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6662 self.set_similarity_explained_at_with_context(
6663 request,
6664 snapshot,
6665 None,
6666 authorization,
6667 context,
6668 )
6669 .map(|(hits, _)| hits)
6670 }
6671
6672 fn set_similarity_explained_at_with_context(
6673 &self,
6674 request: &crate::query::SetSimilarityRequest,
6675 snapshot: Snapshot,
6676 allowed: Option<&std::collections::HashSet<RowId>>,
6677 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6678 context: Option<&crate::query::AiExecutionContext>,
6679 ) -> Result<(
6680 Vec<crate::query::SetSimilarityHit>,
6681 crate::query::SetSimilarityTrace,
6682 )> {
6683 use crate::query::{
6684 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6685 MAX_SET_MEMBERS,
6686 };
6687 let mut trace = crate::query::SetSimilarityTrace::default();
6688 if request.members.is_empty() {
6689 return Ok((Vec::new(), trace));
6690 }
6691 if request.candidate_k == 0 || request.limit == 0 {
6692 return Err(MongrelError::InvalidArgument(
6693 "candidate_k and limit must be > 0".into(),
6694 ));
6695 }
6696 if request.candidate_k > MAX_RETRIEVER_K
6697 || request.limit > MAX_FINAL_LIMIT
6698 || request.members.len() > MAX_SET_MEMBERS
6699 {
6700 return Err(MongrelError::InvalidArgument(format!(
6701 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6702 )));
6703 }
6704 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6705 return Err(MongrelError::InvalidArgument(
6706 "min_jaccard must be finite and between 0 and 1".into(),
6707 ));
6708 }
6709 let started = std::time::Instant::now();
6710 let retriever = Retriever::MinHash {
6711 column_id: request.column_id,
6712 members: request.members.clone(),
6713 k: request.candidate_k,
6714 };
6715 self.require_select()?;
6716 self.validate_retriever(&retriever)?;
6717 let hits = self.retrieve_filtered(
6718 &retriever,
6719 snapshot,
6720 None,
6721 allowed,
6722 candidate_authorization,
6723 context,
6724 )?;
6725 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6726 trace.candidate_count = hits.len();
6727 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6728 if let Some(context) = context {
6729 context.consume(row_ids.len())?;
6730 }
6731 let started = std::time::Instant::now();
6732 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6733 let values = self.values_for_rids_batch_at_with_context(
6734 &row_ids,
6735 request.column_id,
6736 snapshot,
6737 query_now,
6738 context,
6739 )?;
6740 trace.gather_us = started.elapsed().as_micros() as u64;
6741 if let Some(context) = context {
6742 context.consume(request.members.len())?;
6743 }
6744 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6745 let estimates: std::collections::HashMap<_, _> = hits
6746 .into_iter()
6747 .filter_map(|hit| match hit.score {
6748 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6749 _ => None,
6750 })
6751 .collect();
6752 let started = std::time::Instant::now();
6753 let mut parsed = Vec::with_capacity(values.len());
6754 for (row_id, value) in values {
6755 let Value::Bytes(bytes) = value else {
6756 continue;
6757 };
6758 if let Some(context) = context {
6759 context.consume(crate::query::work_units(
6760 bytes.len(),
6761 crate::query::PARSE_WORK_QUANTUM,
6762 ))?;
6763 }
6764 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6765 continue;
6766 };
6767 if let Some(context) = context {
6768 context.consume(members.len())?;
6769 }
6770 let stored = members
6771 .into_iter()
6772 .filter_map(|member| match member {
6773 serde_json::Value::String(value) => {
6774 Some(crate::query::SetMember::String(value))
6775 }
6776 serde_json::Value::Number(value) => {
6777 Some(crate::query::SetMember::Number(value))
6778 }
6779 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6780 _ => None,
6781 })
6782 .collect::<std::collections::HashSet<_>>();
6783 parsed.push((row_id, stored));
6784 }
6785 trace.parse_us = started.elapsed().as_micros() as u64;
6786 trace.verified_count = parsed.len();
6787 let started = std::time::Instant::now();
6788 let mut exact = Vec::new();
6789 for (row_id, stored) in parsed {
6790 if let Some(context) = context {
6791 context.consume(query.len().saturating_add(stored.len()))?;
6792 }
6793 let union = query.union(&stored).count();
6794 let score = if union == 0 {
6795 1.0
6796 } else {
6797 query.intersection(&stored).count() as f32 / union as f32
6798 };
6799 if score >= request.min_jaccard {
6800 exact.push(SetSimilarityHit {
6801 row_id,
6802 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6803 exact_jaccard: score,
6804 });
6805 }
6806 }
6807 exact.sort_by(|a, b| {
6808 b.exact_jaccard
6809 .total_cmp(&a.exact_jaccard)
6810 .then_with(|| a.row_id.cmp(&b.row_id))
6811 });
6812 exact.truncate(request.limit);
6813 trace.score_us = started.elapsed().as_micros() as u64;
6814 crate::trace::QueryTrace::record(|query_trace| {
6815 query_trace.exact_set_gather_nanos = query_trace
6816 .exact_set_gather_nanos
6817 .saturating_add(trace.gather_us.saturating_mul(1_000));
6818 query_trace.exact_set_parse_nanos = query_trace
6819 .exact_set_parse_nanos
6820 .saturating_add(trace.parse_us.saturating_mul(1_000));
6821 query_trace.exact_set_score_nanos = query_trace
6822 .exact_set_score_nanos
6823 .saturating_add(trace.score_us.saturating_mul(1_000));
6824 });
6825 Ok((exact, trace))
6826 }
6827
6828 fn values_for_rids_batch_at(
6830 &self,
6831 row_ids: &[u64],
6832 column_id: u16,
6833 snapshot: Snapshot,
6834 now: i64,
6835 ) -> Result<Vec<(RowId, Value)>> {
6836 if self.ttl.is_none()
6837 && self.memtable.is_empty()
6838 && self.mutable_run.is_empty()
6839 && self.run_refs.len() == 1
6840 {
6841 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6842 if row_ids.len().saturating_mul(24) < reader.row_count() {
6847 let mut values = Vec::with_capacity(row_ids.len());
6848 for &raw_row_id in row_ids {
6849 let row_id = RowId(raw_row_id);
6850 if let Some((_, false, Some(value))) =
6851 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6852 {
6853 values.push((row_id, value));
6854 }
6855 }
6856 return Ok(values);
6857 }
6858 let (positions, visible_row_ids) =
6859 reader.visible_positions_with_rids(snapshot.epoch)?;
6860 let requested: Vec<(RowId, usize)> = row_ids
6861 .iter()
6862 .filter_map(|raw| {
6863 visible_row_ids
6864 .binary_search(&(*raw as i64))
6865 .ok()
6866 .map(|index| (RowId(*raw), positions[index]))
6867 })
6868 .collect();
6869 let values = reader.gather_column(
6870 column_id,
6871 &requested
6872 .iter()
6873 .map(|(_, position)| *position)
6874 .collect::<Vec<_>>(),
6875 )?;
6876 return Ok(requested
6877 .into_iter()
6878 .zip(values)
6879 .map(|((row_id, _), value)| (row_id, value))
6880 .collect());
6881 }
6882 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6883 }
6884
6885 fn values_for_rids_batch_at_with_context(
6886 &self,
6887 row_ids: &[u64],
6888 column_id: u16,
6889 snapshot: Snapshot,
6890 now: i64,
6891 context: Option<&crate::query::AiExecutionContext>,
6892 ) -> Result<Vec<(RowId, Value)>> {
6893 let Some(context) = context else {
6894 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6895 };
6896 let mut values = Vec::with_capacity(row_ids.len());
6897 for chunk in row_ids.chunks(256) {
6898 context.checkpoint()?;
6899 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6900 }
6901 Ok(values)
6902 }
6903
6904 fn values_for_rids_at(
6906 &self,
6907 row_ids: &[u64],
6908 column_id: u16,
6909 snapshot: Snapshot,
6910 now: i64,
6911 ) -> Result<Vec<(RowId, Value)>> {
6912 let mut readers: Vec<_> = self
6913 .run_refs
6914 .iter()
6915 .map(|run| self.open_reader(run.run_id))
6916 .collect::<Result<_>>()?;
6917 let mut out = Vec::with_capacity(row_ids.len());
6918 for &raw_row_id in row_ids {
6919 let row_id = RowId(raw_row_id);
6920 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6921 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6922 let overlay = match (mem, mutable) {
6923 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6924 (a_epoch, a)
6925 } else {
6926 (b_epoch, b)
6927 }),
6928 (Some(value), None) | (None, Some(value)) => Some(value),
6929 (None, None) => None,
6930 };
6931 if let Some((_, row)) = overlay {
6932 if !row.deleted && !self.row_expired_at(&row, now) {
6933 if let Some(value) = row.columns.get(&column_id) {
6934 out.push((row_id, value.clone()));
6935 }
6936 }
6937 continue;
6938 }
6939
6940 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6941 for (index, reader) in readers.iter_mut().enumerate() {
6942 if let Some((epoch, deleted, value)) =
6943 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6944 {
6945 if best
6946 .as_ref()
6947 .map(|(best_epoch, ..)| epoch > *best_epoch)
6948 .unwrap_or(true)
6949 {
6950 best = Some((epoch, deleted, value, index));
6951 }
6952 }
6953 }
6954 let Some((_, false, Some(value), reader_index)) = best else {
6955 continue;
6956 };
6957 if let Some(ttl) = self.ttl {
6958 if ttl.column_id != column_id {
6959 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6960 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6961 {
6962 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6963 continue;
6964 }
6965 }
6966 } else if let Value::Int64(timestamp) = value {
6967 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6968 continue;
6969 }
6970 }
6971 }
6972 out.push((row_id, value));
6973 }
6974 Ok(out)
6975 }
6976
6977 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6982 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
6983 }
6984
6985 pub fn rows_for_rids_with_context(
6986 &self,
6987 rids: &[u64],
6988 snapshot: Snapshot,
6989 context: &crate::query::AiExecutionContext,
6990 ) -> Result<Vec<Row>> {
6991 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6992 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
6993 }
6994
6995 fn rows_for_rids_at_time(
6996 &self,
6997 rids: &[u64],
6998 snapshot: Snapshot,
6999 ttl_now: i64,
7000 control: Option<&crate::ExecutionControl>,
7001 ) -> Result<Vec<Row>> {
7002 use std::collections::HashMap;
7003 let mut rows = Vec::with_capacity(rids.len());
7004 let tier_size = self.memtable.len() + self.mutable_run.len();
7021 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
7022 if rids.len().saturating_mul(24) < tier_size {
7023 for &rid in rids {
7024 if overlay.len() & 255 == 0 {
7025 control
7026 .map(crate::ExecutionControl::checkpoint)
7027 .transpose()?;
7028 }
7029 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
7030 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
7031 let newest = match (mem, mrun) {
7032 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
7033 (Some((_, mr)), None) => Some(mr),
7034 (None, Some((_, rr))) => Some(rr),
7035 (None, None) => None,
7036 };
7037 if let Some(row) = newest {
7038 overlay.insert(rid, row);
7039 }
7040 }
7041 } else {
7042 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
7043 overlay
7044 .entry(row.row_id.0)
7045 .and_modify(|e| {
7046 if row.committed_epoch > e.committed_epoch {
7047 *e = row.clone();
7048 }
7049 })
7050 .or_insert(row);
7051 };
7052 for (index, row) in self
7053 .memtable
7054 .visible_versions(snapshot.epoch)
7055 .into_iter()
7056 .enumerate()
7057 {
7058 if index & 255 == 0 {
7059 control
7060 .map(crate::ExecutionControl::checkpoint)
7061 .transpose()?;
7062 }
7063 fold_newest(row, &mut overlay);
7064 }
7065 for (index, row) in self
7066 .mutable_run
7067 .visible_versions(snapshot.epoch)
7068 .into_iter()
7069 .enumerate()
7070 {
7071 if index & 255 == 0 {
7072 control
7073 .map(crate::ExecutionControl::checkpoint)
7074 .transpose()?;
7075 }
7076 fold_newest(row, &mut overlay);
7077 }
7078 }
7079 if self.run_refs.len() == 1 {
7080 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
7081 if rids.len().saturating_mul(24) < reader.row_count() {
7089 for (index, &rid) in rids.iter().enumerate() {
7090 if index & 255 == 0 {
7091 control
7092 .map(crate::ExecutionControl::checkpoint)
7093 .transpose()?;
7094 }
7095 if let Some(r) = overlay.get(&rid) {
7096 if !r.deleted {
7097 rows.push(r.clone());
7098 }
7099 continue;
7100 }
7101 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
7102 if !row.deleted {
7103 rows.push(row);
7104 }
7105 }
7106 }
7107 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7108 return Ok(rows);
7109 }
7110 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
7119 enum Src {
7122 Overlay,
7123 Run,
7124 }
7125 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
7126 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
7127 for (index, rid) in rids.iter().enumerate() {
7128 if index & 255 == 0 {
7129 control
7130 .map(crate::ExecutionControl::checkpoint)
7131 .transpose()?;
7132 }
7133 if overlay.contains_key(rid) {
7134 plan.push(Src::Overlay);
7135 continue;
7136 }
7137 match vis_rids.binary_search(&(*rid as i64)) {
7138 Ok(i) => {
7139 plan.push(Src::Run);
7140 fetch.push(positions[i]);
7141 }
7142 Err(_) => { }
7143 }
7144 }
7145 let fetched = reader.materialize_batch(&fetch)?;
7146 let mut fetched_iter = fetched.into_iter();
7147 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
7148 if index & 255 == 0 {
7149 control
7150 .map(crate::ExecutionControl::checkpoint)
7151 .transpose()?;
7152 }
7153 match src {
7154 Src::Overlay => {
7155 if let Some(r) = overlay.get(rid) {
7156 if !r.deleted {
7157 rows.push(r.clone());
7158 }
7159 }
7160 }
7161 Src::Run => {
7162 if let Some(row) = fetched_iter.next() {
7163 if !row.deleted {
7164 rows.push(row);
7165 }
7166 }
7167 }
7168 }
7169 }
7170 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7171 return Ok(rows);
7172 }
7173 let mut readers: Vec<_> = self
7177 .run_refs
7178 .iter()
7179 .map(|rr| self.open_reader(rr.run_id))
7180 .collect::<Result<Vec<_>>>()?;
7181 for (index, rid) in rids.iter().enumerate() {
7182 if index & 255 == 0 {
7183 control
7184 .map(crate::ExecutionControl::checkpoint)
7185 .transpose()?;
7186 }
7187 if let Some(r) = overlay.get(rid) {
7188 if !r.deleted {
7189 rows.push(r.clone());
7190 }
7191 continue;
7192 }
7193 let mut best: Option<(Epoch, Row)> = None;
7194 for reader in readers.iter_mut() {
7195 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
7196 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
7197 best = Some((epoch, row));
7198 }
7199 }
7200 }
7201 if let Some((_, r)) = best {
7202 if !r.deleted {
7203 rows.push(r);
7204 }
7205 }
7206 }
7207 rows.retain(|row| !self.row_expired_at(row, ttl_now));
7208 Ok(rows)
7209 }
7210
7211 pub fn indexes_complete(&self) -> bool {
7221 self.indexes_complete
7222 }
7223
7224 pub fn index_build_policy(&self) -> IndexBuildPolicy {
7226 self.index_build_policy
7227 }
7228
7229 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
7233 self.index_build_policy = policy;
7234 }
7235
7236 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
7241 if !self.indexes_complete {
7245 return None;
7246 }
7247 let b = self.bitmap.get(&column_id)?;
7248 let result: Vec<Vec<u8>> = b
7249 .keys()
7250 .into_iter()
7251 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
7252 .collect();
7253 Some(result)
7254 }
7255
7256 pub fn fk_join_row_ids(
7257 &self,
7258 fk_column_id: u16,
7259 pk_values: &[Vec<u8>],
7260 fk_conditions: &[crate::query::Condition],
7261 snapshot: Snapshot,
7262 ) -> Result<Vec<u64>> {
7263 let Some(b) = self.bitmap.get(&fk_column_id) else {
7264 return Ok(Vec::new());
7265 };
7266 let mut join_set = {
7267 let mut acc = roaring::RoaringBitmap::new();
7268 for v in pk_values {
7269 acc |= b.get(v);
7270 }
7271 RowIdSet::from_roaring(acc)
7272 };
7273 if !fk_conditions.is_empty() {
7274 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7275 sets.push(join_set);
7276 for c in fk_conditions {
7277 sets.push(self.resolve_condition(c, snapshot)?);
7278 }
7279 join_set = RowIdSet::intersect_many(sets);
7280 }
7281 Ok(join_set.into_sorted_vec())
7282 }
7283
7284 pub fn fk_join_count(
7290 &self,
7291 fk_column_id: u16,
7292 pk_values: &[Vec<u8>],
7293 fk_conditions: &[crate::query::Condition],
7294 snapshot: Snapshot,
7295 ) -> Result<u64> {
7296 let Some(b) = self.bitmap.get(&fk_column_id) else {
7297 return Ok(0);
7298 };
7299 let mut acc = roaring::RoaringBitmap::new();
7300 for v in pk_values {
7301 acc |= b.get(v);
7302 }
7303 if fk_conditions.is_empty() {
7304 return Ok(acc.len());
7305 }
7306 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
7307 sets.push(RowIdSet::from_roaring(acc));
7308 for c in fk_conditions {
7309 sets.push(self.resolve_condition(c, snapshot)?);
7310 }
7311 Ok(RowIdSet::intersect_many(sets).len() as u64)
7312 }
7313
7314 fn resolve_condition(
7319 &self,
7320 c: &crate::query::Condition,
7321 snapshot: Snapshot,
7322 ) -> Result<RowIdSet> {
7323 self.resolve_condition_with_allowed(c, snapshot, None)
7324 }
7325
7326 fn resolve_condition_with_allowed(
7327 &self,
7328 c: &crate::query::Condition,
7329 snapshot: Snapshot,
7330 allowed: Option<&std::collections::HashSet<RowId>>,
7331 ) -> Result<RowIdSet> {
7332 use crate::query::Condition;
7333 self.validate_condition(c)?;
7334 Ok(match c {
7335 Condition::Pk(key) => {
7336 let lookup = self
7337 .schema
7338 .primary_key()
7339 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
7340 .unwrap_or_else(|| key.clone());
7341 self.hot
7342 .get(&lookup)
7343 .map(|r| RowIdSet::one(r.0))
7344 .unwrap_or_else(RowIdSet::empty)
7345 }
7346 Condition::BitmapEq { column_id, value } => {
7347 let lookup = self.index_lookup_key_bytes(*column_id, value);
7348 self.bitmap
7349 .get(column_id)
7350 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
7351 .unwrap_or_else(RowIdSet::empty)
7352 }
7353 Condition::BitmapIn { column_id, values } => {
7354 let bm = self.bitmap.get(column_id);
7355 let mut acc = roaring::RoaringBitmap::new();
7356 if let Some(b) = bm {
7357 for v in values {
7358 let lookup = self.index_lookup_key_bytes(*column_id, v);
7359 acc |= b.get(&lookup);
7360 }
7361 }
7362 RowIdSet::from_roaring(acc)
7363 }
7364 Condition::BytesPrefix { column_id, prefix } => {
7365 if let Some(b) = self.bitmap.get(column_id) {
7370 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
7371 let mut acc = roaring::RoaringBitmap::new();
7372 for key in b.keys() {
7373 if key.starts_with(&lookup_prefix) {
7374 acc |= b.get(&key);
7375 }
7376 }
7377 RowIdSet::from_roaring(acc)
7378 } else {
7379 RowIdSet::empty()
7380 }
7381 }
7382 Condition::FmContains { column_id, pattern } => self
7383 .fm
7384 .get(column_id)
7385 .map(|f| {
7386 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
7387 })
7388 .unwrap_or_else(RowIdSet::empty),
7389 Condition::FmContainsAll {
7390 column_id,
7391 patterns,
7392 } => {
7393 if let Some(f) = self.fm.get(column_id) {
7396 let sets: Vec<RowIdSet> = patterns
7397 .iter()
7398 .map(|pat| {
7399 RowIdSet::from_unsorted(
7400 f.locate(pat).into_iter().map(|r| r.0).collect(),
7401 )
7402 })
7403 .collect();
7404 RowIdSet::intersect_many(sets)
7405 } else {
7406 RowIdSet::empty()
7407 }
7408 }
7409 Condition::Ann {
7410 column_id,
7411 query,
7412 k,
7413 } => RowIdSet::from_unsorted(
7414 self.retrieve_filtered(
7415 &crate::query::Retriever::Ann {
7416 column_id: *column_id,
7417 query: query.clone(),
7418 k: *k,
7419 },
7420 snapshot,
7421 None,
7422 allowed,
7423 None,
7424 None,
7425 )?
7426 .into_iter()
7427 .map(|hit| hit.row_id.0)
7428 .collect(),
7429 ),
7430 Condition::SparseMatch {
7431 column_id,
7432 query,
7433 k,
7434 } => RowIdSet::from_unsorted(
7435 self.retrieve_filtered(
7436 &crate::query::Retriever::Sparse {
7437 column_id: *column_id,
7438 query: query.clone(),
7439 k: *k,
7440 },
7441 snapshot,
7442 None,
7443 allowed,
7444 None,
7445 None,
7446 )?
7447 .into_iter()
7448 .map(|hit| hit.row_id.0)
7449 .collect(),
7450 ),
7451 Condition::MinHashSimilar {
7452 column_id,
7453 query,
7454 k,
7455 } => match self.minhash.get(column_id) {
7456 Some(index) => {
7457 let candidates = index.candidate_row_ids(query);
7458 let eligible =
7459 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7460 RowIdSet::from_unsorted(
7461 index
7462 .search_filtered(query, *k, |row_id| {
7463 eligible.contains(&row_id)
7464 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7465 })
7466 .into_iter()
7467 .map(|(row_id, _)| row_id.0)
7468 .collect(),
7469 )
7470 }
7471 None => RowIdSet::empty(),
7472 },
7473 Condition::Range { column_id, lo, hi } => {
7474 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7483 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7484 } else if self.run_refs.len() == 1 {
7485 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7486 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7487 } else {
7488 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7489 };
7490 set.remove_many(self.overlay_rid_set(snapshot));
7491 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7492 set
7493 }
7494 Condition::RangeF64 {
7495 column_id,
7496 lo,
7497 lo_inclusive,
7498 hi,
7499 hi_inclusive,
7500 } => {
7501 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7504 RowIdSet::from_unsorted(
7505 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7506 .into_iter()
7507 .collect(),
7508 )
7509 } else if self.run_refs.len() == 1 {
7510 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7511 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7512 } else {
7513 return self.range_scan_f64(
7514 *column_id,
7515 *lo,
7516 *lo_inclusive,
7517 *hi,
7518 *hi_inclusive,
7519 snapshot,
7520 );
7521 };
7522 set.remove_many(self.overlay_rid_set(snapshot));
7523 self.range_scan_overlay_f64(
7524 &mut set,
7525 *column_id,
7526 *lo,
7527 *lo_inclusive,
7528 *hi,
7529 *hi_inclusive,
7530 snapshot,
7531 );
7532 set
7533 }
7534 Condition::IsNull { column_id } => {
7535 let mut set = if self.run_refs.len() == 1 {
7536 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7537 r.null_row_id_set(*column_id, true)?
7538 } else {
7539 return self.null_scan(*column_id, true, snapshot);
7540 };
7541 set.remove_many(self.overlay_rid_set(snapshot));
7542 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7543 set
7544 }
7545 Condition::IsNotNull { column_id } => {
7546 let mut set = if self.run_refs.len() == 1 {
7547 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7548 r.null_row_id_set(*column_id, false)?
7549 } else {
7550 return self.null_scan(*column_id, false, snapshot);
7551 };
7552 set.remove_many(self.overlay_rid_set(snapshot));
7553 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7554 set
7555 }
7556 })
7557 }
7558
7559 fn range_scan_i64(
7567 &self,
7568 column_id: u16,
7569 lo: i64,
7570 hi: i64,
7571 snapshot: Snapshot,
7572 ) -> Result<RowIdSet> {
7573 let mut row_ids = Vec::new();
7574 let overlay_rids = self.overlay_rid_set(snapshot);
7575 for rr in &self.run_refs {
7576 let mut reader = self.open_reader(rr.run_id)?;
7577 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7578 for rid in matched {
7579 if !overlay_rids.contains(&rid) {
7580 row_ids.push(rid);
7581 }
7582 }
7583 }
7584 let mut s = RowIdSet::from_unsorted(row_ids);
7585 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7586 Ok(s)
7587 }
7588
7589 fn range_scan_f64(
7592 &self,
7593 column_id: u16,
7594 lo: f64,
7595 lo_inclusive: bool,
7596 hi: f64,
7597 hi_inclusive: bool,
7598 snapshot: Snapshot,
7599 ) -> Result<RowIdSet> {
7600 let mut row_ids = Vec::new();
7601 let overlay_rids = self.overlay_rid_set(snapshot);
7602 for rr in &self.run_refs {
7603 let mut reader = self.open_reader(rr.run_id)?;
7604 let matched = reader.range_row_ids_visible_f64(
7605 column_id,
7606 lo,
7607 lo_inclusive,
7608 hi,
7609 hi_inclusive,
7610 snapshot.epoch,
7611 )?;
7612 for rid in matched {
7613 if !overlay_rids.contains(&rid) {
7614 row_ids.push(rid);
7615 }
7616 }
7617 }
7618 let mut s = RowIdSet::from_unsorted(row_ids);
7619 self.range_scan_overlay_f64(
7620 &mut s,
7621 column_id,
7622 lo,
7623 lo_inclusive,
7624 hi,
7625 hi_inclusive,
7626 snapshot,
7627 );
7628 Ok(s)
7629 }
7630
7631 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7633 let mut s = HashSet::new();
7634 for row in self.memtable.visible_versions(snapshot.epoch) {
7635 s.insert(row.row_id.0);
7636 }
7637 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7638 s.insert(row.row_id.0);
7639 }
7640 s
7641 }
7642
7643 fn range_scan_overlay_i64(
7644 &self,
7645 s: &mut RowIdSet,
7646 column_id: u16,
7647 lo: i64,
7648 hi: i64,
7649 snapshot: Snapshot,
7650 ) {
7651 let mut newest: HashMap<u64, &Row> = HashMap::new();
7656 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7657 let memtable = self.memtable.visible_versions(snapshot.epoch);
7658 for r in &mutable {
7659 newest.entry(r.row_id.0).or_insert(r);
7660 }
7661 for r in &memtable {
7662 newest.insert(r.row_id.0, r);
7663 }
7664 for row in newest.values() {
7665 if !row.deleted {
7666 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7667 if *v >= lo && *v <= hi {
7668 s.insert(row.row_id.0);
7669 }
7670 }
7671 }
7672 }
7673 }
7674
7675 #[allow(clippy::too_many_arguments)]
7676 fn range_scan_overlay_f64(
7677 &self,
7678 s: &mut RowIdSet,
7679 column_id: u16,
7680 lo: f64,
7681 lo_inclusive: bool,
7682 hi: f64,
7683 hi_inclusive: bool,
7684 snapshot: Snapshot,
7685 ) {
7686 let mut newest: HashMap<u64, &Row> = HashMap::new();
7689 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7690 let memtable = self.memtable.visible_versions(snapshot.epoch);
7691 for r in &mutable {
7692 newest.entry(r.row_id.0).or_insert(r);
7693 }
7694 for r in &memtable {
7695 newest.insert(r.row_id.0, r);
7696 }
7697 for row in newest.values() {
7698 if !row.deleted {
7699 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7700 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7701 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7702 if ok_lo && ok_hi {
7703 s.insert(row.row_id.0);
7704 }
7705 }
7706 }
7707 }
7708 }
7709
7710 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7713 let mut row_ids = Vec::new();
7714 let overlay_rids = self.overlay_rid_set(snapshot);
7715 for rr in &self.run_refs {
7716 let mut reader = self.open_reader(rr.run_id)?;
7717 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7718 for rid in matched {
7719 if !overlay_rids.contains(&rid) {
7720 row_ids.push(rid);
7721 }
7722 }
7723 }
7724 let mut s = RowIdSet::from_unsorted(row_ids);
7725 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7726 Ok(s)
7727 }
7728
7729 fn null_scan_overlay(
7733 &self,
7734 s: &mut RowIdSet,
7735 column_id: u16,
7736 want_nulls: bool,
7737 snapshot: Snapshot,
7738 ) {
7739 let mut newest: HashMap<u64, &Row> = HashMap::new();
7740 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7741 let memtable = self.memtable.visible_versions(snapshot.epoch);
7742 for r in &mutable {
7743 newest.entry(r.row_id.0).or_insert(r);
7744 }
7745 for r in &memtable {
7746 newest.insert(r.row_id.0, r);
7747 }
7748 for row in newest.values() {
7749 if row.deleted {
7750 continue;
7751 }
7752 let is_null = !row.columns.contains_key(&column_id)
7753 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7754 if is_null == want_nulls {
7755 s.insert(row.row_id.0);
7756 }
7757 }
7758 }
7759
7760 pub fn snapshot(&self) -> Snapshot {
7761 Snapshot::at(self.epoch.visible())
7762 }
7763
7764 pub fn data_generation(&self) -> u64 {
7766 self.data_generation
7767 }
7768
7769 pub(crate) fn bump_data_generation(&mut self) {
7770 self.data_generation = self.data_generation.wrapping_add(1);
7771 }
7772
7773 pub fn table_id(&self) -> u64 {
7775 self.table_id
7776 }
7777
7778 fn seal_generations(&mut self) {
7784 self.memtable.seal();
7785 self.mutable_run.seal();
7786 self.hot.seal();
7787 for index in self.bitmap.values_mut() {
7788 index.seal();
7789 }
7790 for index in self.ann.values_mut() {
7791 index.seal();
7792 }
7793 for index in self.fm.values_mut() {
7794 index.seal();
7795 }
7796 for index in self.sparse.values_mut() {
7797 index.seal();
7798 }
7799 for index in self.minhash.values_mut() {
7800 index.seal();
7801 }
7802 self.pk_by_row.seal();
7803 }
7804
7805 fn capture_read_generation(&self) -> ReadGeneration {
7810 let visible_through = self.current_epoch();
7811 ReadGeneration {
7812 schema: Arc::new(self.schema.clone()),
7813 base_runs: Arc::new(self.run_refs.clone()),
7814 deltas: TableDeltas {
7815 memtable: self.memtable.clone(),
7816 mutable_run: self.mutable_run.clone(),
7817 hot: self.hot.clone(),
7818 pk_by_row: self.pk_by_row.clone(),
7819 },
7820 indexes: Arc::new(IndexGeneration::capture(
7821 &self.bitmap,
7822 &self.learned_range,
7823 &self.fm,
7824 &self.ann,
7825 &self.sparse,
7826 &self.minhash,
7827 visible_through,
7828 )),
7829 visible_through,
7830 }
7831 }
7832
7833 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
7838 self.ensure_indexes_complete()?;
7839 self.seal_generations();
7840 let view = Arc::new(self.capture_read_generation());
7841 self.published.store(Arc::clone(&view));
7842 Ok(view)
7843 }
7844
7845 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
7851 self.published.load_full()
7852 }
7853
7854 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
7856 &self.pins
7857 }
7858
7859 pub fn version_gc_floor(&self) -> Epoch {
7864 self.min_active_snapshot()
7865 .unwrap_or_else(|| self.current_epoch())
7866 }
7867
7868 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
7875 let mut report = self.pins.report();
7876 let transaction_floor = [
7877 self.pinned.keys().next().copied(),
7878 self.snapshots.min_pinned(),
7879 ]
7880 .into_iter()
7881 .flatten()
7882 .min();
7883 if let Some(epoch) = transaction_floor {
7884 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
7885 }
7886 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
7887 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
7888 }
7889 report
7890 }
7891
7892 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7893 self.publish_read_generation()?;
7894 let mut generation = self.clone();
7895 generation.read_only = true;
7896 generation.wal = WalSink::ReadOnly;
7897 generation.pending_delete_rids.clear();
7898 generation.pending_put_cols.clear();
7899 generation.pending_rows.clear();
7900 generation.pending_rows_auto_inc.clear();
7901 generation.pending_dels.clear();
7902 generation.pending_truncate = None;
7903 generation.agg_cache = Arc::new(HashMap::new());
7904 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
7907 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
7910 crate::retention::PinSource::ReadGeneration,
7911 self.current_epoch(),
7912 )));
7913 Ok(generation)
7914 }
7915
7916 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7917 (std::mem::size_of::<Self>() as u64)
7918 .saturating_add(self.memtable.approx_bytes())
7919 .saturating_add(self.mutable_run.approx_bytes())
7920 .saturating_add(self.live_count.saturating_mul(64))
7921 }
7922
7923 pub fn pin_snapshot(&mut self) -> Snapshot {
7926 let e = self.epoch.visible();
7927 *self.pinned.entry(e).or_insert(0) += 1;
7928 Snapshot::at(e)
7929 }
7930
7931 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7933 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7934 *count -= 1;
7935 if *count == 0 {
7936 self.pinned.remove(&snap.epoch);
7937 }
7938 }
7939 }
7940
7941 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7953 let local = self.pinned.keys().next().copied();
7954 let global = self.snapshots.min_pinned();
7955 let history = self.snapshots.history_floor(self.current_epoch());
7956 let pinned = self.pins.oldest_pinned();
7957 [local, global, history, pinned].into_iter().flatten().min()
7958 }
7959
7960 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7964 self.ensure_writable()?;
7965 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7966 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7967 }
7968
7969 pub fn clear_ttl(&mut self) -> Result<()> {
7970 self.ensure_writable()?;
7971 self.apply_ttl_policy_at(None, self.current_epoch())
7972 }
7973
7974 pub fn ttl(&self) -> Option<TtlPolicy> {
7975 self.ttl
7976 }
7977
7978 pub(crate) fn prepare_ttl_policy(
7979 &self,
7980 column_name: &str,
7981 duration_nanos: u64,
7982 ) -> Result<TtlPolicy> {
7983 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7984 return Err(MongrelError::InvalidArgument(
7985 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7986 ));
7987 }
7988 let column = self
7989 .schema
7990 .columns
7991 .iter()
7992 .find(|column| column.name == column_name)
7993 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7994 if column.ty != TypeId::TimestampNanos {
7995 return Err(MongrelError::Schema(format!(
7996 "TTL column {column_name} must be TimestampNanos, is {:?}",
7997 column.ty
7998 )));
7999 }
8000 Ok(TtlPolicy {
8001 column_id: column.id,
8002 duration_nanos,
8003 })
8004 }
8005
8006 pub(crate) fn apply_ttl_policy_at(
8007 &mut self,
8008 policy: Option<TtlPolicy>,
8009 epoch: Epoch,
8010 ) -> Result<()> {
8011 if let Some(policy) = policy {
8012 let column = self
8013 .schema
8014 .columns
8015 .iter()
8016 .find(|column| column.id == policy.column_id)
8017 .ok_or_else(|| {
8018 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
8019 })?;
8020 if column.ty != TypeId::TimestampNanos
8021 || policy.duration_nanos == 0
8022 || policy.duration_nanos > i64::MAX as u64
8023 {
8024 return Err(MongrelError::Schema("invalid TTL policy".into()));
8025 }
8026 }
8027 self.ttl = policy;
8028 self.agg_cache = Arc::new(HashMap::new());
8029 self.clear_result_cache();
8030 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
8031 self.persist_manifest(epoch)
8032 }
8033
8034 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
8035 let Some(policy) = self.ttl else {
8036 return false;
8037 };
8038 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
8039 return false;
8040 };
8041 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
8042 }
8043
8044 pub fn current_epoch(&self) -> Epoch {
8045 self.epoch.visible()
8046 }
8047
8048 pub fn memtable_len(&self) -> usize {
8049 self.memtable.len()
8050 }
8051
8052 pub fn count(&self) -> u64 {
8055 if self.ttl.is_none()
8056 && self.pending_put_cols.is_empty()
8057 && self.pending_delete_rids.is_empty()
8058 && self.pending_rows.is_empty()
8059 && self.pending_dels.is_empty()
8060 && self.pending_truncate.is_none()
8061 {
8062 self.live_count
8063 } else {
8064 self.visible_rows(self.snapshot())
8065 .map(|rows| rows.len() as u64)
8066 .unwrap_or(self.live_count)
8067 }
8068 }
8069
8070 pub fn count_conditions(
8074 &mut self,
8075 conditions: &[crate::query::Condition],
8076 snapshot: Snapshot,
8077 ) -> Result<Option<u64>> {
8078 use crate::query::Condition;
8079 if self.ttl.is_some() {
8080 if conditions.is_empty() {
8081 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
8082 }
8083 let mut sets = Vec::with_capacity(conditions.len());
8084 for condition in conditions {
8085 sets.push(self.resolve_condition(condition, snapshot)?);
8086 }
8087 let survivors = RowIdSet::intersect_many(sets);
8088 let rows = self.visible_rows(snapshot)?;
8089 return Ok(Some(
8090 rows.into_iter()
8091 .filter(|row| survivors.contains(row.row_id.0))
8092 .count() as u64,
8093 ));
8094 }
8095 if conditions.is_empty() {
8096 return Ok(Some(self.count()));
8097 }
8098 let served = |c: &Condition| {
8099 matches!(
8100 c,
8101 Condition::Pk(_)
8102 | Condition::BitmapEq { .. }
8103 | Condition::BitmapIn { .. }
8104 | Condition::BytesPrefix { .. }
8105 | Condition::FmContains { .. }
8106 | Condition::FmContainsAll { .. }
8107 | Condition::Ann { .. }
8108 | Condition::Range { .. }
8109 | Condition::RangeF64 { .. }
8110 | Condition::SparseMatch { .. }
8111 | Condition::MinHashSimilar { .. }
8112 | Condition::IsNull { .. }
8113 | Condition::IsNotNull { .. }
8114 )
8115 };
8116 if !conditions.iter().all(served) {
8117 return Ok(None);
8118 }
8119 self.ensure_indexes_complete()?;
8120 if !self.pending_put_cols.is_empty()
8121 || !self.pending_delete_rids.is_empty()
8122 || !self.pending_rows.is_empty()
8123 || !self.pending_dels.is_empty()
8124 || self.pending_truncate.is_some()
8125 {
8126 let mut sets = Vec::with_capacity(conditions.len());
8127 for condition in conditions {
8128 sets.push(self.resolve_condition(condition, snapshot)?);
8129 }
8130 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8131 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
8132 }
8133 let mut sets = Vec::with_capacity(conditions.len());
8134 for condition in conditions {
8135 sets.push(self.resolve_condition(condition, snapshot)?);
8136 }
8137 let mut rids = RowIdSet::intersect_many(sets);
8138 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8148 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
8149 }
8150 let count = rids.len() as u64;
8151 crate::trace::QueryTrace::record(|t| {
8152 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
8153 t.survivor_count = Some(count as usize);
8154 t.conditions_pushed = conditions.len();
8155 });
8156 Ok(Some(count))
8157 }
8158
8159 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
8164 let mut out = Vec::new();
8165 for row in self.memtable.visible_versions(snapshot.epoch) {
8166 if row.deleted {
8167 out.push(row.row_id.0);
8168 }
8169 }
8170 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8171 if row.deleted {
8172 out.push(row.row_id.0);
8173 }
8174 }
8175 out
8176 }
8177
8178 pub fn bulk_load_columns(
8187 &mut self,
8188 user_columns: Vec<(u16, columnar::NativeColumn)>,
8189 ) -> Result<Epoch> {
8190 self.bulk_load_columns_with(user_columns, 3, false, true)
8191 }
8192
8193 pub fn bulk_load_fast(
8200 &mut self,
8201 user_columns: Vec<(u16, columnar::NativeColumn)>,
8202 ) -> Result<Epoch> {
8203 self.bulk_load_columns_with(user_columns, -1, true, false)
8204 }
8205
8206 fn bulk_load_columns_with(
8207 &mut self,
8208 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
8209 zstd_level: i32,
8210 force_plain: bool,
8211 lz4: bool,
8212 ) -> Result<Epoch> {
8213 self.ensure_writable()?;
8214 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
8215 if n == 0 {
8216 return Ok(self.current_epoch());
8217 }
8218 let epoch = self.commit_new_epoch()?;
8219 let live_before = self.live_count;
8220 self.spill_mutable_run(epoch)?;
8222 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
8223 && self.indexes_complete
8224 && self.run_refs.is_empty()
8225 && self.memtable.is_empty()
8226 && self.mutable_run.is_empty();
8227 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
8230 self.validate_columns_not_null(&user_columns, n)?;
8231 let winner_idx = self
8232 .bulk_pk_winner_indices(&user_columns, n)
8233 .filter(|idx| idx.len() != n);
8234 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
8235 match winner_idx.as_deref() {
8236 Some(idx) => {
8237 let compacted = user_columns
8238 .iter()
8239 .map(|(id, c)| (*id, c.gather(idx)))
8240 .collect();
8241 (compacted, idx.len())
8242 }
8243 None => (user_columns, n),
8244 };
8245 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
8246 let first = self.allocator.alloc_range(write_n as u64)?.0;
8247 for rid in first..first + write_n as u64 {
8248 self.reservoir.offer(rid);
8249 }
8250 let run_id = self.alloc_run_id()?;
8251 let path = self.run_path(run_id);
8252 let mut writer =
8253 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
8254 if force_plain {
8255 writer = writer.with_plain();
8256 } else if lz4 {
8257 writer = writer.with_lz4();
8260 } else {
8261 writer = writer.with_zstd_level(zstd_level);
8262 }
8263 if let Some(kek) = &self.kek {
8264 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
8265 }
8266 let header = match self.create_run_file(run_id)? {
8267 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
8268 None => writer.write_native(&path, &write_columns, write_n, first)?,
8269 };
8270 self.run_refs.push(RunRef {
8271 run_id: run_id as u128,
8272 level: 0,
8273 epoch_created: epoch.0,
8274 row_count: header.row_count,
8275 });
8276 self.live_count = self.live_count.saturating_add(write_n as u64);
8277 if eager_index_build {
8278 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
8279 self.index_columns_bulk(&write_columns, &row_ids);
8280 self.indexes_complete = true;
8281 self.build_learned_ranges()?;
8282 } else {
8283 self.indexes_complete = false;
8287 }
8288 self.mark_flushed(epoch)?;
8289 self.persist_manifest(epoch)?;
8290 if eager_index_build {
8291 self.checkpoint_indexes(epoch);
8292 }
8293 self.clear_result_cache();
8294 self.data_generation = self.data_generation.wrapping_add(1);
8295 Ok(epoch)
8296 }
8297
8298 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
8316 let n = row_ids.len();
8317 if n == 0 {
8318 return;
8319 }
8320 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
8321 columns.iter().map(|(id, c)| (*id, c)).collect();
8322 let ty_of: std::collections::HashMap<u16, TypeId> = self
8323 .schema
8324 .columns
8325 .iter()
8326 .map(|c| (c.id, c.ty.clone()))
8327 .collect();
8328 let pk_id = self.schema.primary_key().map(|c| c.id);
8329
8330 for (i, &rid) in row_ids.iter().enumerate() {
8331 let row_id = RowId(rid);
8332 if let Some(pid) = pk_id {
8333 if let Some(col) = by_id.get(&pid) {
8334 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
8335 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
8336 self.insert_hot_pk(key, row_id);
8337 }
8338 }
8339 }
8340 for idef in &self.schema.indexes {
8341 let Some(col) = by_id.get(&idef.column_id) else {
8342 continue;
8343 };
8344 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
8345 match idef.kind {
8346 IndexKind::Bitmap => {
8347 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
8348 if let Some(key) =
8349 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
8350 {
8351 b.insert(key, row_id);
8352 }
8353 }
8354 }
8355 IndexKind::FmIndex => {
8356 if let Some(f) = self.fm.get_mut(&idef.column_id) {
8357 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8358 f.insert(bytes.to_vec(), row_id);
8359 }
8360 }
8361 }
8362 IndexKind::Sparse => {
8363 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
8364 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8365 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
8366 s.insert(&terms, row_id);
8367 }
8368 }
8369 }
8370 }
8371 IndexKind::MinHash => {
8372 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
8373 if let Some(bytes) = columnar::native_bytes_at(col, i) {
8374 let tokens = crate::index::token_hashes_from_bytes(bytes);
8375 mh.insert(&tokens, row_id);
8376 }
8377 }
8378 }
8379 _ => {}
8380 }
8381 }
8382 }
8383 }
8384
8385 pub fn visible_columns_native(
8390 &self,
8391 snapshot: Snapshot,
8392 projection: Option<&[u16]>,
8393 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8394 self.visible_columns_native_inner(snapshot, projection, None)
8395 }
8396
8397 pub fn visible_columns_native_with_control(
8398 &self,
8399 snapshot: Snapshot,
8400 projection: Option<&[u16]>,
8401 control: &crate::ExecutionControl,
8402 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8403 self.visible_columns_native_inner(snapshot, projection, Some(control))
8404 }
8405
8406 fn visible_columns_native_inner(
8407 &self,
8408 snapshot: Snapshot,
8409 projection: Option<&[u16]>,
8410 control: Option<&crate::ExecutionControl>,
8411 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
8412 execution_checkpoint(control, 0)?;
8413 let wanted: Vec<u16> = match projection {
8414 Some(p) => p.to_vec(),
8415 None => self.schema.columns.iter().map(|c| c.id).collect(),
8416 };
8417 if self.ttl.is_none()
8418 && self.memtable.is_empty()
8419 && self.mutable_run.is_empty()
8420 && self.run_refs.len() == 1
8421 {
8422 let rr = self.run_refs[0].clone();
8423 let mut reader = self.open_reader(rr.run_id)?;
8424 let idxs = reader.visible_indices_native(snapshot.epoch)?;
8425 execution_checkpoint(control, 0)?;
8426 let all_visible = idxs.len() == reader.row_count();
8427 if reader.has_mmap() && control.is_none() {
8433 use rayon::prelude::*;
8434 let valid: Vec<u16> = wanted
8437 .iter()
8438 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
8439 .copied()
8440 .collect();
8441 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
8443 .par_iter()
8444 .filter_map(|cid| {
8445 reader
8446 .column_native_shared(*cid)
8447 .ok()
8448 .map(|col| (*cid, col))
8449 })
8450 .collect();
8451 let cols = decoded
8452 .into_iter()
8453 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
8454 .collect();
8455 return Ok(cols);
8456 }
8457 let mut cols = Vec::with_capacity(wanted.len());
8458 for (index, cid) in wanted.iter().enumerate() {
8459 execution_checkpoint(control, index)?;
8460 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
8461 Some(c) => c,
8462 None => continue,
8463 };
8464 let col = reader.column_native(cdef.id)?;
8465 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
8466 }
8467 return Ok(cols);
8468 }
8469 let vcols = self.visible_columns(snapshot)?;
8470 execution_checkpoint(control, 0)?;
8471 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
8472 let out: Vec<(u16, columnar::NativeColumn)> = vcols
8473 .into_iter()
8474 .filter(|(id, _)| want_set.contains(id))
8475 .map(|(id, vals)| {
8476 let ty = self
8477 .schema
8478 .columns
8479 .iter()
8480 .find(|c| c.id == id)
8481 .map(|c| c.ty.clone())
8482 .unwrap_or(TypeId::Bytes);
8483 (id, columnar::values_to_native(ty, &vals))
8484 })
8485 .collect();
8486 Ok(out)
8487 }
8488
8489 pub fn run_count(&self) -> usize {
8490 self.run_refs.len()
8491 }
8492
8493 pub fn memtable_is_empty(&self) -> bool {
8495 self.memtable.is_empty()
8496 }
8497
8498 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
8502 self.page_cache.stats()
8503 }
8504
8505 pub fn reset_page_cache_stats(&self) {
8507 self.page_cache.reset_stats();
8508 }
8509
8510 pub fn run_ids(&self) -> Vec<u128> {
8513 self.run_refs.iter().map(|r| r.run_id).collect()
8514 }
8515
8516 pub fn single_run_is_clean(&self) -> bool {
8520 if self.ttl.is_some() || self.run_refs.len() != 1 {
8521 return false;
8522 }
8523 self.open_reader(self.run_refs[0].run_id)
8524 .map(|r| r.is_clean())
8525 .unwrap_or(false)
8526 }
8527
8528 fn resolve_footprint(
8535 &self,
8536 conditions: &[crate::query::Condition],
8537 snapshot: Snapshot,
8538 ) -> roaring::RoaringBitmap {
8539 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
8540 return roaring::RoaringBitmap::new();
8541 }
8542 if self.run_refs.is_empty() {
8543 return roaring::RoaringBitmap::new();
8544 }
8545 if self.run_refs.len() == 1 {
8547 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
8548 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
8549 return rids.to_roaring_lossy();
8550 }
8551 }
8552 }
8553 roaring::RoaringBitmap::new()
8554 }
8555
8556 pub fn query_columns_native_cached(
8567 &mut self,
8568 conditions: &[crate::query::Condition],
8569 projection: Option<&[u16]>,
8570 snapshot: Snapshot,
8571 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8572 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8573 }
8574
8575 pub fn query_columns_native_cached_with_control(
8576 &mut self,
8577 conditions: &[crate::query::Condition],
8578 projection: Option<&[u16]>,
8579 snapshot: Snapshot,
8580 control: &crate::ExecutionControl,
8581 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8582 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8583 }
8584
8585 fn query_columns_native_cached_inner(
8586 &mut self,
8587 conditions: &[crate::query::Condition],
8588 projection: Option<&[u16]>,
8589 snapshot: Snapshot,
8590 control: Option<&crate::ExecutionControl>,
8591 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8592 execution_checkpoint(control, 0)?;
8593 if self.ttl.is_some() {
8596 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8597 }
8598 if conditions.is_empty() {
8599 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8600 }
8601 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8605 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8606 crate::trace::QueryTrace::record(|t| {
8607 t.result_cache_hit = true;
8608 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8609 });
8610 return Ok(Some((*hit).clone()));
8611 }
8612 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8613 execution_checkpoint(control, 0)?;
8614 if let Some(cols) = &res {
8615 let footprint = self.resolve_footprint(conditions, snapshot);
8616 let condition_cols = crate::query::condition_columns(conditions);
8617 execution_checkpoint(control, 0)?;
8618 self.result_cache.lock().insert(
8619 key,
8620 CachedEntry {
8621 data: CachedData::Columns(Arc::new(cols.clone())),
8622 footprint,
8623 condition_cols,
8624 },
8625 );
8626 }
8627 Ok(res)
8628 }
8629
8630 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8635 if self.ttl.is_some() {
8636 return self.query(q);
8637 }
8638 if q.conditions.is_empty() {
8639 return self.query(q);
8640 }
8641 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8642 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8643 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8644 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8645 crate::trace::QueryTrace::record(|t| {
8646 t.result_cache_hit = true;
8647 t.scan_mode = crate::trace::ScanMode::Materialized;
8648 });
8649 return Ok((*hit).clone());
8650 }
8651 let rows = self.query(q)?;
8652 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8653 let condition_cols = crate::query::condition_columns(&q.conditions);
8654 self.result_cache.lock().insert(
8655 key,
8656 CachedEntry {
8657 data: CachedData::Rows(Arc::new(rows.clone())),
8658 footprint,
8659 condition_cols,
8660 },
8661 );
8662 Ok(rows)
8663 }
8664
8665 #[allow(clippy::type_complexity)]
8680 pub fn query_columns_native_traced(
8681 &mut self,
8682 conditions: &[crate::query::Condition],
8683 projection: Option<&[u16]>,
8684 snapshot: Snapshot,
8685 ) -> Result<(
8686 Option<Vec<(u16, columnar::NativeColumn)>>,
8687 crate::trace::QueryTrace,
8688 )> {
8689 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8690 self.query_columns_native(conditions, projection, snapshot)
8691 });
8692 Ok((result?, trace))
8693 }
8694
8695 #[allow(clippy::type_complexity)]
8698 pub fn query_columns_native_cached_traced(
8699 &mut self,
8700 conditions: &[crate::query::Condition],
8701 projection: Option<&[u16]>,
8702 snapshot: Snapshot,
8703 ) -> Result<(
8704 Option<Vec<(u16, columnar::NativeColumn)>>,
8705 crate::trace::QueryTrace,
8706 )> {
8707 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8708 self.query_columns_native_cached(conditions, projection, snapshot)
8709 });
8710 Ok((result?, trace))
8711 }
8712
8713 pub fn native_page_cursor_traced(
8715 &self,
8716 snapshot: Snapshot,
8717 projection: Vec<(u16, TypeId)>,
8718 conditions: &[crate::query::Condition],
8719 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8720 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8721 self.native_page_cursor(snapshot, projection, conditions)
8722 });
8723 Ok((result?, trace))
8724 }
8725
8726 pub fn native_multi_run_cursor_traced(
8728 &self,
8729 snapshot: Snapshot,
8730 projection: Vec<(u16, TypeId)>,
8731 conditions: &[crate::query::Condition],
8732 ) -> Result<(
8733 Option<crate::cursor::MultiRunCursor>,
8734 crate::trace::QueryTrace,
8735 )> {
8736 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8737 self.native_multi_run_cursor(snapshot, projection, conditions)
8738 });
8739 Ok((result?, trace))
8740 }
8741
8742 pub fn count_conditions_traced(
8744 &mut self,
8745 conditions: &[crate::query::Condition],
8746 snapshot: Snapshot,
8747 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8748 let (result, trace) =
8749 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8750 Ok((result?, trace))
8751 }
8752
8753 pub fn query_traced(
8755 &mut self,
8756 q: &crate::query::Query,
8757 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8758 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8759 Ok((result?, trace))
8760 }
8761
8762 pub fn query_columns_native(
8767 &mut self,
8768 conditions: &[crate::query::Condition],
8769 projection: Option<&[u16]>,
8770 snapshot: Snapshot,
8771 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8772 self.query_columns_native_inner(conditions, projection, snapshot, None)
8773 }
8774
8775 pub fn query_columns_native_with_control(
8776 &mut self,
8777 conditions: &[crate::query::Condition],
8778 projection: Option<&[u16]>,
8779 snapshot: Snapshot,
8780 control: &crate::ExecutionControl,
8781 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8782 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8783 }
8784
8785 fn query_columns_native_inner(
8786 &mut self,
8787 conditions: &[crate::query::Condition],
8788 projection: Option<&[u16]>,
8789 snapshot: Snapshot,
8790 control: Option<&crate::ExecutionControl>,
8791 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8792 use crate::query::Condition;
8793 execution_checkpoint(control, 0)?;
8794 if self.ttl.is_some() {
8797 return Ok(None);
8798 }
8799 if conditions.is_empty() {
8800 return Ok(None);
8801 }
8802 self.ensure_indexes_complete()?;
8803
8804 let served = |c: &Condition| {
8809 matches!(
8810 c,
8811 Condition::Pk(_)
8812 | Condition::BitmapEq { .. }
8813 | Condition::BitmapIn { .. }
8814 | Condition::BytesPrefix { .. }
8815 | Condition::FmContains { .. }
8816 | Condition::FmContainsAll { .. }
8817 | Condition::Ann { .. }
8818 | Condition::Range { .. }
8819 | Condition::RangeF64 { .. }
8820 | Condition::SparseMatch { .. }
8821 | Condition::MinHashSimilar { .. }
8822 | Condition::IsNull { .. }
8823 | Condition::IsNotNull { .. }
8824 )
8825 };
8826 if !conditions.iter().all(served) {
8827 return Ok(None);
8828 }
8829 let fast_path =
8830 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8831 crate::trace::QueryTrace::record(|t| {
8832 t.run_count = self.run_refs.len();
8833 t.memtable_rows = self.memtable.len();
8834 t.mutable_run_rows = self.mutable_run.len();
8835 t.conditions_pushed = conditions.len();
8836 t.learned_range_used = conditions.iter().any(|c| match c {
8837 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8838 self.learned_range.contains_key(column_id)
8839 }
8840 _ => false,
8841 });
8842 });
8843 let col_ids: Vec<u16> = projection
8845 .map(|p| p.to_vec())
8846 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8847 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8848 .iter()
8849 .map(|&cid| {
8850 let ty = self
8851 .schema
8852 .columns
8853 .iter()
8854 .find(|c| c.id == cid)
8855 .map(|c| c.ty.clone())
8856 .unwrap_or(TypeId::Bytes);
8857 (cid, ty)
8858 })
8859 .collect();
8860
8861 if fast_path {
8867 let needs_column = conditions.iter().any(|c| match c {
8870 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8871 Condition::RangeF64 { column_id, .. } => {
8872 !self.learned_range.contains_key(column_id)
8873 }
8874 _ => false,
8875 });
8876 let mut reader_opt: Option<RunReader> = if needs_column {
8877 Some(self.open_reader(self.run_refs[0].run_id)?)
8878 } else {
8879 None
8880 };
8881 let mut sets: Vec<RowIdSet> = Vec::new();
8882 for (index, c) in conditions.iter().enumerate() {
8883 execution_checkpoint(control, index)?;
8884 let s = match c {
8885 Condition::Range { column_id, lo, hi }
8886 if !self.learned_range.contains_key(column_id) =>
8887 {
8888 if reader_opt.is_none() {
8889 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8890 }
8891 reader_opt
8892 .as_mut()
8893 .expect("reader opened for range")
8894 .range_row_id_set_i64(*column_id, *lo, *hi)?
8895 }
8896 Condition::RangeF64 {
8897 column_id,
8898 lo,
8899 lo_inclusive,
8900 hi,
8901 hi_inclusive,
8902 } if !self.learned_range.contains_key(column_id) => {
8903 if reader_opt.is_none() {
8904 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8905 }
8906 reader_opt
8907 .as_mut()
8908 .expect("reader opened for range")
8909 .range_row_id_set_f64(
8910 *column_id,
8911 *lo,
8912 *lo_inclusive,
8913 *hi,
8914 *hi_inclusive,
8915 )?
8916 }
8917 _ => self.resolve_condition(c, snapshot)?,
8918 };
8919 sets.push(s);
8920 }
8921 let candidates = RowIdSet::intersect_many(sets);
8922 crate::trace::QueryTrace::record(|t| {
8923 t.survivor_count = Some(candidates.len());
8924 });
8925 if candidates.is_empty() {
8926 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8927 .iter()
8928 .map(|&id| {
8929 (
8930 id,
8931 columnar::null_native(
8932 proj_pairs
8933 .iter()
8934 .find(|(c, _)| c == &id)
8935 .map(|(_, t)| t.clone())
8936 .unwrap_or(TypeId::Bytes),
8937 0,
8938 ),
8939 )
8940 })
8941 .collect();
8942 return Ok(Some(cols));
8943 }
8944 let mut reader = match reader_opt.take() {
8945 Some(r) => r,
8946 None => self.open_reader(self.run_refs[0].run_id)?,
8947 };
8948 let candidate_ids = candidates.into_sorted_vec();
8949 let (positions, fast_rid) = if let Some(positions) =
8950 reader.positions_for_row_ids_fast(&candidate_ids)
8951 {
8952 (positions, true)
8953 } else {
8954 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8955 match col {
8956 columnar::NativeColumn::Int64 { data, .. } => {
8957 let mut p = Vec::with_capacity(candidate_ids.len());
8958 for (index, rid) in candidate_ids.iter().enumerate() {
8959 execution_checkpoint(control, index)?;
8960 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8961 p.push(position);
8962 }
8963 }
8964 p.sort_unstable();
8965 (p, false)
8966 }
8967 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8968 }
8969 };
8970 crate::trace::QueryTrace::record(|t| {
8971 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8972 t.fast_row_id_map = fast_rid;
8973 });
8974 let mut cols = Vec::with_capacity(col_ids.len());
8975 for (index, cid) in col_ids.iter().enumerate() {
8976 execution_checkpoint(control, index)?;
8977 let col = reader.column_native(*cid)?;
8978 cols.push((*cid, col.gather(&positions)));
8979 }
8980 return Ok(Some(cols));
8981 }
8982
8983 if !self.run_refs.is_empty() {
8996 use crate::cursor::{
8997 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8998 };
8999 let remaining: usize;
9000 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
9001 let c = self
9002 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
9003 .expect("single-run cursor should build when run_refs.len() == 1");
9004 remaining = c.remaining_rows();
9005 Box::new(c)
9006 } else {
9007 let c = self
9008 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
9009 .expect("multi-run cursor should build when run_refs.len() >= 1");
9010 remaining = c.remaining_rows();
9011 Box::new(c)
9012 };
9013 crate::trace::QueryTrace::record(|t| {
9014 if t.survivor_count.is_none() {
9015 t.survivor_count = Some(remaining);
9016 }
9017 });
9018 let cols = match control {
9019 Some(control) => {
9020 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
9021 }
9022 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
9023 };
9024 return Ok(Some(cols));
9025 }
9026
9027 crate::trace::QueryTrace::record(|t| {
9032 t.scan_mode = crate::trace::ScanMode::Materialized;
9033 t.row_materialized = true;
9034 });
9035 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9036 for (index, c) in conditions.iter().enumerate() {
9037 execution_checkpoint(control, index)?;
9038 sets.push(self.resolve_condition(c, snapshot)?);
9039 }
9040 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9041 let rows = self.rows_for_rids(&rids, snapshot)?;
9042 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
9043 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
9044 execution_checkpoint(control, index)?;
9045 let vals: Vec<Value> = rows
9046 .iter()
9047 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9048 .collect();
9049 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
9050 }
9051 Ok(Some(cols))
9052 }
9053
9054 pub fn native_page_cursor(
9069 &self,
9070 snapshot: Snapshot,
9071 projection: Vec<(u16, TypeId)>,
9072 conditions: &[crate::query::Condition],
9073 ) -> Result<Option<NativePageCursor>> {
9074 use crate::cursor::build_page_plans;
9075 if self.ttl.is_some() {
9076 return Ok(None);
9077 }
9078 if !conditions.is_empty() && !self.indexes_complete {
9081 return Ok(None);
9082 }
9083 if self.run_refs.len() != 1 {
9084 return Ok(None);
9085 }
9086 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9087 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
9088
9089 let overlay_rids: HashSet<u64> = {
9092 let mut s = HashSet::new();
9093 for row in self.memtable.visible_versions(snapshot.epoch) {
9094 s.insert(row.row_id.0);
9095 }
9096 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9097 s.insert(row.row_id.0);
9098 }
9099 s
9100 };
9101
9102 let survivors = if conditions.is_empty() {
9106 None
9107 } else {
9108 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
9109 };
9110
9111 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
9118 survivors.clone()
9119 } else if let Some(s) = &survivors {
9120 let mut run_set = s.clone();
9121 run_set.remove_many(overlay_rids.iter().copied());
9122 Some(run_set)
9123 } else {
9124 Some(RowIdSet::from_unsorted(
9125 rids.iter()
9126 .map(|&r| r as u64)
9127 .filter(|r| !overlay_rids.contains(r))
9128 .collect(),
9129 ))
9130 };
9131
9132 let overlay_rows = if overlay_rids.is_empty() {
9133 Vec::new()
9134 } else {
9135 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9136 self.overlay_visible_rows(snapshot, bound)
9137 };
9138
9139 let plans = if positions.is_empty() {
9141 Vec::new()
9142 } else {
9143 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
9144 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
9145 };
9146
9147 let overlay = if overlay_rows.is_empty() {
9149 None
9150 } else {
9151 let filtered =
9152 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9153 if filtered.is_empty() {
9154 None
9155 } else {
9156 Some(self.materialize_overlay(&filtered, &projection))
9157 }
9158 };
9159
9160 let overlay_row_count = overlay
9161 .as_ref()
9162 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9163 .unwrap_or(0);
9164 crate::trace::QueryTrace::record(|t| {
9165 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
9166 t.run_count = self.run_refs.len();
9167 t.memtable_rows = self.memtable.len();
9168 t.mutable_run_rows = self.mutable_run.len();
9169 t.overlay_rows = overlay_row_count;
9170 t.conditions_pushed = conditions.len();
9171 t.pages_decoded = plans
9172 .iter()
9173 .map(|p| p.positions.len())
9174 .sum::<usize>()
9175 .min(1);
9176 });
9177
9178 Ok(Some(NativePageCursor::new_with_overlay(
9179 reader, projection, plans, overlay,
9180 )))
9181 }
9182 #[allow(clippy::type_complexity)]
9192 pub fn native_multi_run_cursor(
9193 &self,
9194 snapshot: Snapshot,
9195 projection: Vec<(u16, TypeId)>,
9196 conditions: &[crate::query::Condition],
9197 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
9198 use crate::cursor::{MultiRunCursor, RunStream};
9199 use crate::sorted_run::SYS_ROW_ID;
9200 use std::collections::{BinaryHeap, HashMap, HashSet};
9201 if self.ttl.is_some() {
9202 return Ok(None);
9203 }
9204 if !conditions.is_empty() && !self.indexes_complete {
9207 return Ok(None);
9208 }
9209 if self.run_refs.is_empty() {
9210 return Ok(None);
9211 }
9212
9213 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
9215 Vec::with_capacity(self.run_refs.len());
9216 for rr in &self.run_refs {
9217 let mut reader = self.open_reader(rr.run_id)?;
9218 let (rids, eps, del) = reader.system_columns_native()?;
9219 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
9220 run_meta.push((reader, rids, eps, del, page_rows));
9221 }
9222
9223 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
9227 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
9228 for i in 0..rids.len() {
9229 let rid = rids[i] as u64;
9230 let e = eps[i] as u64;
9231 if e > snapshot.epoch.0 {
9232 continue;
9233 }
9234 let is_del = del[i] != 0;
9235 best.entry(rid)
9236 .and_modify(|cur| {
9237 if e > cur.0 {
9238 *cur = (e, run_idx, i, is_del);
9239 }
9240 })
9241 .or_insert((e, run_idx, i, is_del));
9242 }
9243 }
9244
9245 let overlay_rids: HashSet<u64> = {
9247 let mut s = HashSet::new();
9248 for row in self.memtable.visible_versions(snapshot.epoch) {
9249 s.insert(row.row_id.0);
9250 }
9251 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9252 s.insert(row.row_id.0);
9253 }
9254 s
9255 };
9256
9257 let survivors: Option<RowIdSet> = if conditions.is_empty() {
9259 None
9260 } else {
9261 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9262 for c in conditions {
9263 sets.push(self.resolve_condition(c, snapshot)?);
9264 }
9265 Some(RowIdSet::intersect_many(sets))
9266 };
9267
9268 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
9272 for (rid, (_, run_idx, pos, deleted)) in &best {
9273 if *deleted {
9274 continue;
9275 }
9276 if overlay_rids.contains(rid) {
9277 continue;
9278 }
9279 if let Some(s) = &survivors {
9280 if !s.contains(*rid) {
9281 continue;
9282 }
9283 }
9284 per_run[*run_idx].push((*rid, *pos));
9285 }
9286 for v in per_run.iter_mut() {
9287 v.sort_unstable_by_key(|&(rid, _)| rid);
9288 }
9289
9290 let mut streams = Vec::with_capacity(run_meta.len());
9292 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
9293 let mut total = 0usize;
9294 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
9295 let mut starts = Vec::with_capacity(page_rows.len());
9296 let mut acc = 0usize;
9297 for &r in &page_rows {
9298 starts.push(acc);
9299 acc += r;
9300 }
9301 let mut survivors_vec: Vec<(u64, usize, usize)> =
9302 Vec::with_capacity(per_run[run_idx].len());
9303 for &(rid, pos) in &per_run[run_idx] {
9304 let page_seq = match starts.partition_point(|&s| s <= pos) {
9305 0 => continue,
9306 p => p - 1,
9307 };
9308 let within = pos - starts[page_seq];
9309 survivors_vec.push((rid, page_seq, within));
9310 }
9311 total += survivors_vec.len();
9312 if let Some(&(rid, _, _)) = survivors_vec.first() {
9313 heap.push(std::cmp::Reverse((rid, run_idx)));
9314 }
9315 streams.push(RunStream::new(reader, survivors_vec, page_rows));
9316 }
9317
9318 let overlay_rows = if overlay_rids.is_empty() {
9320 Vec::new()
9321 } else {
9322 let bound = Self::overlay_materialization_bound(conditions, &survivors);
9323 self.overlay_visible_rows(snapshot, bound)
9324 };
9325 let overlay = if overlay_rows.is_empty() {
9326 None
9327 } else {
9328 let filtered =
9329 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
9330 if filtered.is_empty() {
9331 None
9332 } else {
9333 Some(self.materialize_overlay(&filtered, &projection))
9334 }
9335 };
9336
9337 let overlay_row_count = overlay
9338 .as_ref()
9339 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
9340 .unwrap_or(0);
9341 crate::trace::QueryTrace::record(|t| {
9342 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
9343 t.run_count = self.run_refs.len();
9344 t.memtable_rows = self.memtable.len();
9345 t.mutable_run_rows = self.mutable_run.len();
9346 t.overlay_rows = overlay_row_count;
9347 t.conditions_pushed = conditions.len();
9348 t.survivor_count = Some(total);
9349 });
9350
9351 Ok(Some(MultiRunCursor::new(
9352 streams, projection, heap, total, overlay,
9353 )))
9354 }
9355
9356 fn overlay_materialization_bound<'a>(
9368 conditions: &[crate::query::Condition],
9369 survivors: &'a Option<RowIdSet>,
9370 ) -> Option<&'a RowIdSet> {
9371 use crate::query::Condition;
9372 let has_range = conditions
9373 .iter()
9374 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9375 if has_range {
9376 None
9377 } else {
9378 survivors.as_ref()
9379 }
9380 }
9381
9382 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
9394 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
9395 let mut fold = |row: Row| {
9396 if let Some(b) = bound {
9397 if !b.contains(row.row_id.0) {
9398 return;
9399 }
9400 }
9401 best.entry(row.row_id.0)
9402 .and_modify(|(be, br)| {
9403 if row.committed_epoch > *be {
9404 *be = row.committed_epoch;
9405 *br = row.clone();
9406 }
9407 })
9408 .or_insert_with(|| (row.committed_epoch, row));
9409 };
9410 for row in self.memtable.visible_versions(snapshot.epoch) {
9411 fold(row);
9412 }
9413 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9414 fold(row);
9415 }
9416 let mut out: Vec<Row> = best
9417 .into_values()
9418 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
9419 .collect();
9420 out.sort_by_key(|r| r.row_id);
9421 out
9422 }
9423
9424 fn filter_overlay_rows(
9432 &self,
9433 rows: Vec<Row>,
9434 conditions: &[crate::query::Condition],
9435 survivors: Option<&RowIdSet>,
9436 snapshot: Snapshot,
9437 ) -> Result<Vec<Row>> {
9438 if conditions.is_empty() {
9439 return Ok(rows);
9440 }
9441 use crate::query::Condition;
9442 let all_index_served = !conditions
9446 .iter()
9447 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
9448 if all_index_served {
9449 return Ok(rows
9450 .into_iter()
9451 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
9452 .collect());
9453 }
9454 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
9457 for c in conditions {
9458 let s = match c {
9459 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
9460 _ => self.resolve_condition(c, snapshot)?,
9461 };
9462 per_cond_sets.push(s);
9463 }
9464 Ok(rows
9465 .into_iter()
9466 .filter(|row| {
9467 conditions.iter().enumerate().all(|(i, c)| match c {
9468 Condition::Range { column_id, lo, hi } => {
9469 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
9470 }
9471 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
9472 match row.columns.get(column_id) {
9473 Some(Value::Float64(v)) => {
9474 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
9475 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
9476 lo_ok && hi_ok
9477 }
9478 _ => false,
9479 }
9480 }
9481 _ => per_cond_sets[i].contains(row.row_id.0),
9482 })
9483 })
9484 .collect())
9485 }
9486
9487 fn materialize_overlay(
9490 &self,
9491 rows: &[Row],
9492 projection: &[(u16, TypeId)],
9493 ) -> Vec<columnar::NativeColumn> {
9494 if projection.is_empty() {
9495 return vec![columnar::null_native(TypeId::Int64, rows.len())];
9496 }
9497 let mut cols = Vec::with_capacity(projection.len());
9498 for (cid, ty) in projection {
9499 let vals: Vec<Value> = rows
9500 .iter()
9501 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
9502 .collect();
9503 cols.push(columnar::values_to_native(ty.clone(), &vals));
9504 }
9505 cols
9506 }
9507
9508 fn resolve_survivor_rids(
9513 &self,
9514 conditions: &[crate::query::Condition],
9515 reader: &mut RunReader,
9516 snapshot: Snapshot,
9517 ) -> Result<RowIdSet> {
9518 use crate::query::Condition;
9519 let mut sets: Vec<RowIdSet> = Vec::new();
9520 for c in conditions {
9521 self.validate_condition(c)?;
9522 let s: RowIdSet = match c {
9523 Condition::Pk(key) => {
9524 let lookup = self
9525 .schema
9526 .primary_key()
9527 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9528 .unwrap_or_else(|| key.clone());
9529 self.hot
9530 .get(&lookup)
9531 .map(|r| RowIdSet::one(r.0))
9532 .unwrap_or_else(RowIdSet::empty)
9533 }
9534 Condition::BitmapEq { column_id, value } => {
9535 let lookup = self.index_lookup_key_bytes(*column_id, value);
9536 self.bitmap
9537 .get(column_id)
9538 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9539 .unwrap_or_else(RowIdSet::empty)
9540 }
9541 Condition::BitmapIn { column_id, values } => {
9542 let bm = self.bitmap.get(column_id);
9543 let mut acc = roaring::RoaringBitmap::new();
9544 if let Some(b) = bm {
9545 for v in values {
9546 let lookup = self.index_lookup_key_bytes(*column_id, v);
9547 acc |= b.get(&lookup);
9548 }
9549 }
9550 RowIdSet::from_roaring(acc)
9551 }
9552 Condition::BytesPrefix { column_id, prefix } => {
9553 if let Some(b) = self.bitmap.get(column_id) {
9554 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9555 let mut acc = roaring::RoaringBitmap::new();
9556 for key in b.keys() {
9557 if key.starts_with(&lookup_prefix) {
9558 acc |= b.get(&key);
9559 }
9560 }
9561 RowIdSet::from_roaring(acc)
9562 } else {
9563 RowIdSet::empty()
9564 }
9565 }
9566 Condition::FmContains { column_id, pattern } => self
9567 .fm
9568 .get(column_id)
9569 .map(|f| {
9570 RowIdSet::from_unsorted(
9571 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9572 )
9573 })
9574 .unwrap_or_else(RowIdSet::empty),
9575 Condition::FmContainsAll {
9576 column_id,
9577 patterns,
9578 } => {
9579 if let Some(f) = self.fm.get(column_id) {
9580 let sets: Vec<RowIdSet> = patterns
9581 .iter()
9582 .map(|pat| {
9583 RowIdSet::from_unsorted(
9584 f.locate(pat).into_iter().map(|r| r.0).collect(),
9585 )
9586 })
9587 .collect();
9588 RowIdSet::intersect_many(sets)
9589 } else {
9590 RowIdSet::empty()
9591 }
9592 }
9593 Condition::Ann {
9594 column_id,
9595 query,
9596 k,
9597 } => RowIdSet::from_unsorted(
9598 self.retrieve_filtered(
9599 &crate::query::Retriever::Ann {
9600 column_id: *column_id,
9601 query: query.clone(),
9602 k: *k,
9603 },
9604 snapshot,
9605 None,
9606 None,
9607 None,
9608 None,
9609 )?
9610 .into_iter()
9611 .map(|hit| hit.row_id.0)
9612 .collect(),
9613 ),
9614 Condition::SparseMatch {
9615 column_id,
9616 query,
9617 k,
9618 } => RowIdSet::from_unsorted(
9619 self.retrieve_filtered(
9620 &crate::query::Retriever::Sparse {
9621 column_id: *column_id,
9622 query: query.clone(),
9623 k: *k,
9624 },
9625 snapshot,
9626 None,
9627 None,
9628 None,
9629 None,
9630 )?
9631 .into_iter()
9632 .map(|hit| hit.row_id.0)
9633 .collect(),
9634 ),
9635 Condition::MinHashSimilar {
9636 column_id,
9637 query,
9638 k,
9639 } => match self.minhash.get(column_id) {
9640 Some(index) => {
9641 let candidates = index.candidate_row_ids(query);
9642 let eligible =
9643 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9644 RowIdSet::from_unsorted(
9645 index
9646 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9647 .into_iter()
9648 .map(|(row_id, _)| row_id.0)
9649 .collect(),
9650 )
9651 }
9652 None => RowIdSet::empty(),
9653 },
9654 Condition::Range { column_id, lo, hi } => {
9655 if let Some(li) = self.learned_range.get(column_id) {
9656 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9657 } else {
9658 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9659 }
9660 }
9661 Condition::RangeF64 {
9662 column_id,
9663 lo,
9664 lo_inclusive,
9665 hi,
9666 hi_inclusive,
9667 } => {
9668 if let Some(li) = self.learned_range.get(column_id) {
9669 RowIdSet::from_unsorted(
9670 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9671 .into_iter()
9672 .collect(),
9673 )
9674 } else {
9675 reader.range_row_id_set_f64(
9676 *column_id,
9677 *lo,
9678 *lo_inclusive,
9679 *hi,
9680 *hi_inclusive,
9681 )?
9682 }
9683 }
9684 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9685 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9686 };
9687 sets.push(s);
9688 }
9689 Ok(RowIdSet::intersect_many(sets))
9690 }
9691
9692 pub fn scan_cursor(
9713 &self,
9714 snapshot: Snapshot,
9715 projection: Vec<(u16, TypeId)>,
9716 conditions: &[crate::query::Condition],
9717 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9718 if self.ttl.is_some() {
9719 return Ok(None);
9720 }
9721 if !conditions.is_empty() && !self.indexes_complete {
9727 return Ok(None);
9728 }
9729 if self.run_refs.len() == 1 {
9730 Ok(self
9731 .native_page_cursor(snapshot, projection, conditions)?
9732 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9733 } else {
9734 Ok(self
9735 .native_multi_run_cursor(snapshot, projection, conditions)?
9736 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9737 }
9738 }
9739
9740 pub fn aggregate_native(
9754 &self,
9755 snapshot: Snapshot,
9756 column: Option<u16>,
9757 conditions: &[crate::query::Condition],
9758 agg: NativeAgg,
9759 ) -> Result<Option<NativeAggResult>> {
9760 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9761 }
9762
9763 pub fn aggregate_native_with_control(
9764 &self,
9765 snapshot: Snapshot,
9766 column: Option<u16>,
9767 conditions: &[crate::query::Condition],
9768 agg: NativeAgg,
9769 control: &crate::ExecutionControl,
9770 ) -> Result<Option<NativeAggResult>> {
9771 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9772 }
9773
9774 fn aggregate_native_inner(
9775 &self,
9776 snapshot: Snapshot,
9777 column: Option<u16>,
9778 conditions: &[crate::query::Condition],
9779 agg: NativeAgg,
9780 control: Option<&crate::ExecutionControl>,
9781 ) -> Result<Option<NativeAggResult>> {
9782 execution_checkpoint(control, 0)?;
9783 if self.ttl.is_some() {
9784 return Ok(None);
9785 }
9786 if self.run_refs.len() == 1 && conditions.is_empty() {
9788 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9789 return Ok(Some(res));
9790 }
9791 }
9792 if matches!(agg, NativeAgg::Count) && column.is_none() {
9796 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
9797 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
9798 }
9799 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9800 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
9801 }
9802 let cid = match column {
9805 Some(c) => c,
9806 None => return Ok(None),
9807 };
9808 let ty = self.column_type(cid);
9809 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
9810 execution_checkpoint(control, 0)?;
9811 return match ty {
9812 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9813 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9814 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9815 }
9816 TypeId::Float64 => {
9817 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9818 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9819 }
9820 _ => Ok(None),
9821 };
9822 }
9823 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
9825 execution_checkpoint(control, 0)?;
9826 match ty {
9827 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9828 let mut count = 0u64;
9829 let mut sum = 0i128;
9830 let mut mn = i64::MAX;
9831 let mut mx = i64::MIN;
9832 for row in &rows {
9833 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
9834 count += 1;
9835 sum += i128::from(*v);
9836 mn = mn.min(*v);
9837 mx = mx.max(*v);
9838 }
9839 }
9840 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9841 }
9842 TypeId::Float64 => {
9843 let mut count = 0u64;
9844 let mut sum = 0.0f64;
9845 let mut mn = f64::INFINITY;
9846 let mut mx = f64::NEG_INFINITY;
9847 for row in &rows {
9848 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
9849 count += 1;
9850 sum += *v;
9851 mn = mn.min(*v);
9852 mx = mx.max(*v);
9853 }
9854 }
9855 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9856 }
9857 _ => Ok(None),
9858 }
9859 }
9860
9861 fn visible_rows_filtered(
9863 &self,
9864 snapshot: Snapshot,
9865 conditions: &[crate::query::Condition],
9866 control: Option<&crate::ExecutionControl>,
9867 ) -> Result<Vec<Row>> {
9868 let rows = if let Some(control) = control {
9869 self.visible_rows_controlled(snapshot, control)?
9870 } else {
9871 self.visible_rows(snapshot)?
9872 };
9873 if conditions.is_empty() {
9874 return Ok(rows);
9875 }
9876 Ok(rows
9877 .into_iter()
9878 .filter(|row| {
9879 conditions
9880 .iter()
9881 .all(|cond| condition_matches_row(cond, row, &self.schema))
9882 })
9883 .collect())
9884 }
9885
9886 fn aggregate_from_stats(
9894 &self,
9895 snapshot: Snapshot,
9896 column: Option<u16>,
9897 agg: NativeAgg,
9898 ) -> Result<Option<NativeAggResult>> {
9899 let cid = match (agg, column) {
9900 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9901 _ => return Ok(None), };
9903 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9904 return Ok(None);
9905 };
9906 let Some(cs) = stats.get(&cid) else {
9907 return Ok(None);
9908 };
9909 match agg {
9910 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9912 self.live_count.saturating_sub(cs.null_count),
9913 ))),
9914 NativeAgg::Min | NativeAgg::Max => {
9915 let bound = if agg == NativeAgg::Min {
9916 &cs.min
9917 } else {
9918 &cs.max
9919 };
9920 match bound {
9921 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9922 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9923 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9928 None => Ok(None),
9929 }
9930 }
9931 _ => Ok(None),
9932 }
9933 }
9934
9935 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9944 if self.ttl.is_some() {
9945 return Ok(None);
9946 }
9947 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9948 return Ok(None);
9949 }
9950 self.ensure_indexes_complete()?;
9953 let reader = self.open_reader(self.run_refs[0].run_id)?;
9954 if self.live_count != reader.row_count() as u64 {
9955 return Ok(None);
9956 }
9957 let Some(bm) = self.bitmap.get(&column_id) else {
9958 return Ok(None); };
9960 let mut distinct = bm.value_count() as u64;
9961 if !bm.get(&Value::Null.encode_key()).is_empty() {
9964 distinct = distinct.saturating_sub(1);
9965 }
9966 Ok(Some(distinct))
9967 }
9968
9969 pub fn aggregate_incremental(
9981 &mut self,
9982 cache_key: u64,
9983 conditions: &[crate::query::Condition],
9984 column: Option<u16>,
9985 agg: NativeAgg,
9986 ) -> Result<IncrementalAggResult> {
9987 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9988 }
9989
9990 pub fn aggregate_incremental_with_control(
9991 &mut self,
9992 cache_key: u64,
9993 conditions: &[crate::query::Condition],
9994 column: Option<u16>,
9995 agg: NativeAgg,
9996 control: &crate::ExecutionControl,
9997 ) -> Result<IncrementalAggResult> {
9998 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9999 }
10000
10001 fn aggregate_incremental_inner(
10002 &mut self,
10003 cache_key: u64,
10004 conditions: &[crate::query::Condition],
10005 column: Option<u16>,
10006 agg: NativeAgg,
10007 control: Option<&crate::ExecutionControl>,
10008 ) -> Result<IncrementalAggResult> {
10009 execution_checkpoint(control, 0)?;
10010 let snap = self.snapshot();
10011 let cur_wm = self.allocator.current().0;
10012 let cur_epoch = snap.epoch.0;
10013 let incremental_ok = self.ttl.is_none()
10020 && !self.had_deletes
10021 && self.memtable.is_empty()
10022 && self.mutable_run.is_empty();
10023
10024 if incremental_ok {
10027 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
10028 if cached.epoch == cur_epoch {
10029 return Ok(IncrementalAggResult {
10030 state: cached.state,
10031 incremental: true,
10032 delta_rows: 0,
10033 });
10034 }
10035 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
10036 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
10037 let mut delta_rids = Vec::with_capacity(delta_len);
10038 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
10039 execution_checkpoint(control, index)?;
10040 delta_rids.push(row_id);
10041 }
10042 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
10043 execution_checkpoint(control, 0)?;
10044 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10045 let delta_state = agg_state_from_rows(
10046 &delta_rows,
10047 conditions,
10048 &index_sets,
10049 column,
10050 agg,
10051 &self.schema,
10052 control,
10053 )?;
10054 let merged = cached.state.merge(delta_state);
10055 let delta_n = delta_rids.len() as u64;
10056 Arc::make_mut(&mut self.agg_cache).insert(
10057 cache_key,
10058 CachedAgg {
10059 state: merged.clone(),
10060 watermark: cur_wm,
10061 epoch: cur_epoch,
10062 },
10063 );
10064 return Ok(IncrementalAggResult {
10065 state: merged,
10066 incremental: true,
10067 delta_rows: delta_n,
10068 });
10069 }
10070 }
10071 }
10072
10073 let cursor_ok =
10078 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10079 let state = if cursor_ok && agg != NativeAgg::Avg {
10080 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
10081 Some(result) => {
10082 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
10083 }
10084 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
10085 }
10086 } else {
10087 self.agg_state_full_scan(conditions, column, agg, snap, control)?
10088 };
10089 if incremental_ok {
10091 Arc::make_mut(&mut self.agg_cache).insert(
10092 cache_key,
10093 CachedAgg {
10094 state: state.clone(),
10095 watermark: cur_wm,
10096 epoch: cur_epoch,
10097 },
10098 );
10099 }
10100 Ok(IncrementalAggResult {
10101 state,
10102 incremental: false,
10103 delta_rows: 0,
10104 })
10105 }
10106
10107 fn agg_state_full_scan(
10110 &self,
10111 conditions: &[crate::query::Condition],
10112 column: Option<u16>,
10113 agg: NativeAgg,
10114 snap: Snapshot,
10115 control: Option<&crate::ExecutionControl>,
10116 ) -> Result<AggState> {
10117 execution_checkpoint(control, 0)?;
10118 let rows = self.visible_rows(snap)?;
10119 execution_checkpoint(control, 0)?;
10120 let index_sets = self.resolve_index_conditions(conditions, snap)?;
10121 agg_state_from_rows(
10122 &rows,
10123 conditions,
10124 &index_sets,
10125 column,
10126 agg,
10127 &self.schema,
10128 control,
10129 )
10130 }
10131
10132 fn resolve_index_conditions(
10135 &self,
10136 conditions: &[crate::query::Condition],
10137 snapshot: Snapshot,
10138 ) -> Result<Vec<RowIdSet>> {
10139 use crate::query::Condition;
10140 let mut sets = Vec::new();
10141 for c in conditions {
10142 if matches!(
10143 c,
10144 Condition::Ann { .. }
10145 | Condition::SparseMatch { .. }
10146 | Condition::MinHashSimilar { .. }
10147 ) {
10148 sets.push(self.resolve_condition(c, snapshot)?);
10149 }
10150 }
10151 Ok(sets)
10152 }
10153
10154 fn column_type(&self, cid: u16) -> TypeId {
10155 self.schema
10156 .columns
10157 .iter()
10158 .find(|c| c.id == cid)
10159 .map(|c| c.ty.clone())
10160 .unwrap_or(TypeId::Bytes)
10161 }
10162
10163 pub fn approx_aggregate(
10172 &mut self,
10173 conditions: &[crate::query::Condition],
10174 column: Option<u16>,
10175 agg: ApproxAgg,
10176 z: f64,
10177 ) -> Result<Option<ApproxResult>> {
10178 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
10179 }
10180
10181 pub fn approx_aggregate_with_candidate_authorization(
10184 &mut self,
10185 conditions: &[crate::query::Condition],
10186 column: Option<u16>,
10187 agg: ApproxAgg,
10188 z: f64,
10189 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
10190 ) -> Result<Option<ApproxResult>> {
10191 use crate::query::Condition;
10192 self.ensure_reservoir_complete()?;
10193 let snapshot = self.snapshot();
10194 let n_pop = self.count();
10195 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
10196 if sample_rids.is_empty() {
10197 return Ok(None);
10198 }
10199 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
10201 let s = live_sample.len();
10202 if s == 0 {
10203 return Ok(None);
10204 }
10205 let authorized = authorization
10206 .map(|authorization| {
10207 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
10208 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
10209 })
10210 .transpose()?;
10211
10212 let mut index_sets: Vec<RowIdSet> = Vec::new();
10215 for c in conditions {
10216 if matches!(
10217 c,
10218 Condition::Ann { .. }
10219 | Condition::SparseMatch { .. }
10220 | Condition::MinHashSimilar { .. }
10221 ) {
10222 index_sets.push(self.resolve_condition(c, snapshot)?);
10223 }
10224 }
10225
10226 let cid = match (agg, column) {
10228 (ApproxAgg::Count, _) => None,
10229 (_, Some(c)) => Some(c),
10230 _ => return Ok(None),
10231 };
10232 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
10233 for r in &live_sample {
10234 if authorized
10235 .as_ref()
10236 .is_some_and(|authorized| !authorized.contains(&r.row_id))
10237 {
10238 continue;
10239 }
10240 if !conditions
10242 .iter()
10243 .all(|c| condition_matches_row(c, r, &self.schema))
10244 {
10245 continue;
10246 }
10247 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
10249 continue;
10250 }
10251 if let Some(cid) = cid {
10252 let mut cells = r
10253 .columns
10254 .get(&cid)
10255 .cloned()
10256 .map(|value| vec![(cid, value)])
10257 .unwrap_or_default();
10258 if let Some(authorization) = authorization {
10259 authorization.security.apply_masks_to_cells(
10260 authorization.table,
10261 &mut cells,
10262 authorization.principal,
10263 );
10264 }
10265 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
10266 passing_vals.push(v);
10267 } } else {
10269 passing_vals.push(0.0); }
10271 }
10272 let m = passing_vals.len();
10273
10274 let (point, half) = match agg {
10275 ApproxAgg::Count => {
10276 let p = m as f64 / s as f64;
10278 let point = n_pop as f64 * p;
10279 let var = if s > 1 {
10280 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
10281 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10282 } else {
10283 0.0
10284 };
10285 (point, z * var.sqrt())
10286 }
10287 ApproxAgg::Sum => {
10288 let y: Vec<f64> = live_sample
10290 .iter()
10291 .map(|r| {
10292 let passes_row = authorized
10293 .as_ref()
10294 .is_none_or(|authorized| authorized.contains(&r.row_id))
10295 && conditions
10296 .iter()
10297 .all(|c| condition_matches_row(c, r, &self.schema))
10298 && index_sets.iter().all(|set| set.contains(r.row_id.0));
10299 if passes_row {
10300 cid.and_then(|cid| {
10301 let mut cells = r
10302 .columns
10303 .get(&cid)
10304 .cloned()
10305 .map(|value| vec![(cid, value)])
10306 .unwrap_or_default();
10307 if let Some(authorization) = authorization {
10308 authorization.security.apply_masks_to_cells(
10309 authorization.table,
10310 &mut cells,
10311 authorization.principal,
10312 );
10313 }
10314 as_f64(cells.first().map(|(_, value)| value))
10315 })
10316 .unwrap_or(0.0)
10317 } else {
10318 0.0
10319 }
10320 })
10321 .collect();
10322 let mean_y = y.iter().sum::<f64>() / s as f64;
10323 let point = n_pop as f64 * mean_y;
10324 let var = if s > 1 {
10325 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
10326 let var_y = ss / (s - 1) as f64;
10327 n_pop as f64 * n_pop as f64 * var_y / s as f64
10328 * (1.0 - s as f64 / n_pop as f64).max(0.0)
10329 } else {
10330 0.0
10331 };
10332 (point, z * var.sqrt())
10333 }
10334 ApproxAgg::Avg => {
10335 if m == 0 {
10336 return Ok(Some(ApproxResult {
10337 point: 0.0,
10338 ci_low: 0.0,
10339 ci_high: 0.0,
10340 n_population: n_pop,
10341 n_sample_live: s,
10342 n_passing: 0,
10343 }));
10344 }
10345 let mean = passing_vals.iter().sum::<f64>() / m as f64;
10346 let half = if m > 1 {
10347 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
10348 let sd = (ss / (m - 1) as f64).sqrt();
10349 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
10350 z * sd / (m as f64).sqrt() * fpc.sqrt()
10351 } else {
10352 0.0
10353 };
10354 (mean, half)
10355 }
10356 };
10357
10358 Ok(Some(ApproxResult {
10359 point,
10360 ci_low: point - half,
10361 ci_high: point + half,
10362 n_population: n_pop,
10363 n_sample_live: s,
10364 n_passing: m,
10365 }))
10366 }
10367
10368 pub fn exact_column_stats(
10376 &self,
10377 _snapshot: Snapshot,
10378 projection: &[u16],
10379 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
10380 if self.ttl.is_some()
10381 || !(self.memtable.is_empty()
10382 && self.mutable_run.is_empty()
10383 && self.run_refs.len() == 1)
10384 {
10385 return Ok(None);
10386 }
10387 let reader = self.open_reader(self.run_refs[0].run_id)?;
10388 if self.live_count != reader.row_count() as u64 {
10389 return Ok(None);
10390 }
10391 let mut out = HashMap::new();
10392 for &cid in projection {
10393 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
10394 Some(c) => c,
10395 None => continue,
10396 };
10397 let Some(stats) = reader.column_page_stats(cid) else {
10399 out.insert(
10400 cid,
10401 ColumnStat {
10402 min: None,
10403 max: None,
10404 null_count: self.live_count,
10405 },
10406 );
10407 continue;
10408 };
10409 let stat = match cdef.ty {
10410 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
10411 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
10412 min: mn.map(Value::Int64),
10413 max: mx.map(Value::Int64),
10414 null_count: n,
10415 })
10416 }
10417 TypeId::Float64 => {
10418 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
10419 min: mn.map(Value::Float64),
10420 max: mx.map(Value::Float64),
10421 null_count: n,
10422 })
10423 }
10424 _ => None,
10425 };
10426 if let Some(s) = stat {
10427 out.insert(cid, s);
10428 }
10429 }
10430 Ok(Some(out))
10431 }
10432
10433 pub fn dir(&self) -> &Path {
10434 &self.dir
10435 }
10436
10437 pub fn schema(&self) -> &Schema {
10438 &self.schema
10439 }
10440
10441 pub(crate) fn set_catalog_name(&mut self, name: String) {
10442 self.name = name;
10443 }
10444
10445 pub(crate) fn prepare_alter_column(
10446 &mut self,
10447 column_name: &str,
10448 change: &AlterColumn,
10449 ) -> Result<(ColumnDef, Option<Schema>)> {
10450 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
10451 return Err(MongrelError::InvalidArgument(
10452 "ALTER COLUMN requires committing staged writes first".into(),
10453 ));
10454 }
10455 let old = self
10456 .schema
10457 .columns
10458 .iter()
10459 .find(|c| c.name == column_name)
10460 .cloned()
10461 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10462 let mut next = old.clone();
10463
10464 if let Some(name) = &change.name {
10465 let trimmed = name.trim();
10466 if trimmed.is_empty() {
10467 return Err(MongrelError::InvalidArgument(
10468 "ALTER COLUMN name must not be empty".into(),
10469 ));
10470 }
10471 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
10472 return Err(MongrelError::Schema(format!(
10473 "column {trimmed} already exists"
10474 )));
10475 }
10476 next.name = trimmed.to_string();
10477 }
10478
10479 if let Some(ty) = &change.ty {
10480 next.ty = ty.clone();
10481 }
10482 if let Some(flags) = change.flags {
10483 validate_alter_column_flags(old.flags, flags)?;
10484 next.flags = flags;
10485 }
10486
10487 if let Some(default_change) = &change.default_value {
10488 next.default_value = default_change.clone();
10489 }
10490 if let Some(source_change) = &change.embedding_source {
10491 next.embedding_source = source_change.clone();
10492 }
10493
10494 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
10495 if old.flags.contains(ColumnFlags::NULLABLE)
10496 && !next.flags.contains(ColumnFlags::NULLABLE)
10497 && self.column_has_nulls(old.id)?
10498 {
10499 return Err(MongrelError::InvalidArgument(format!(
10500 "column '{}' contains NULL values",
10501 old.name
10502 )));
10503 }
10504 if next == old {
10505 return Ok((next, None));
10506 }
10507 let mut schema = self.schema.clone();
10508 let index = schema
10509 .columns
10510 .iter()
10511 .position(|column| column.id == next.id)
10512 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
10513 schema.columns[index] = next.clone();
10514 schema.schema_id = schema
10515 .schema_id
10516 .checked_add(1)
10517 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10518 schema.validate_auto_increment()?;
10519 schema.validate_defaults()?;
10520 Ok((next, Some(schema)))
10521 }
10522
10523 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
10524 self.schema = schema;
10525 self.auto_inc = resolve_auto_inc(&self.schema);
10526 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
10527 self.clear_result_cache();
10528 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10529 }
10530
10531 pub(crate) fn publish_index_schema_change(
10535 &mut self,
10536 schema: Schema,
10537 artifact: SecondaryIndexArtifact,
10538 ) -> Result<()> {
10539 self.apply_altered_schema_prepared(schema);
10540 self.prune_index_maps_to_schema();
10541 match artifact {
10542 SecondaryIndexArtifact::Bitmap(column_id, index) => {
10543 self.bitmap.insert(column_id, index);
10544 }
10545 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
10546 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
10547 }
10548 SecondaryIndexArtifact::Fm(column_id, index) => {
10549 self.fm.insert(column_id, *index);
10550 }
10551 SecondaryIndexArtifact::Ann(column_id, index) => {
10552 self.ann.insert(column_id, index);
10553 }
10554 SecondaryIndexArtifact::Sparse(column_id, index) => {
10555 self.sparse.insert(column_id, index);
10556 }
10557 SecondaryIndexArtifact::MinHash(column_id, index) => {
10558 self.minhash.insert(column_id, index);
10559 }
10560 }
10561 self.indexes_complete = true;
10562 self.seal_generations();
10563 let view = Arc::new(self.capture_read_generation());
10564 self.published.store(Arc::clone(&view));
10565 checkpoint_current_schema(self)?;
10566 self.invalidate_index_checkpoint();
10569 Ok(())
10570 }
10571
10572 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
10578 self.apply_altered_schema_prepared(schema);
10579 self.prune_index_maps_to_schema();
10580 self.indexes_complete = true;
10581 self.seal_generations();
10582 let view = Arc::new(self.capture_read_generation());
10583 self.published.store(Arc::clone(&view));
10584 checkpoint_current_schema(self)?;
10585 self.invalidate_index_checkpoint();
10587 Ok(())
10588 }
10589
10590 fn prune_index_maps_to_schema(&mut self) {
10591 let keep_bitmap: std::collections::HashSet<u16> = self
10592 .schema
10593 .indexes
10594 .iter()
10595 .filter(|index| index.kind == IndexKind::Bitmap)
10596 .map(|index| index.column_id)
10597 .collect();
10598 let keep_ann: std::collections::HashSet<u16> = self
10599 .schema
10600 .indexes
10601 .iter()
10602 .filter(|index| index.kind == IndexKind::Ann)
10603 .map(|index| index.column_id)
10604 .collect();
10605 let keep_fm: std::collections::HashSet<u16> = self
10606 .schema
10607 .indexes
10608 .iter()
10609 .filter(|index| index.kind == IndexKind::FmIndex)
10610 .map(|index| index.column_id)
10611 .collect();
10612 let keep_sparse: std::collections::HashSet<u16> = self
10613 .schema
10614 .indexes
10615 .iter()
10616 .filter(|index| index.kind == IndexKind::Sparse)
10617 .map(|index| index.column_id)
10618 .collect();
10619 let keep_minhash: std::collections::HashSet<u16> = self
10620 .schema
10621 .indexes
10622 .iter()
10623 .filter(|index| index.kind == IndexKind::MinHash)
10624 .map(|index| index.column_id)
10625 .collect();
10626 let keep_learned: std::collections::HashSet<u16> = self
10627 .schema
10628 .indexes
10629 .iter()
10630 .filter(|index| index.kind == IndexKind::LearnedRange)
10631 .map(|index| index.column_id)
10632 .collect();
10633 self.bitmap
10634 .retain(|column_id, _| keep_bitmap.contains(column_id));
10635 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
10636 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
10637 self.sparse
10638 .retain(|column_id, _| keep_sparse.contains(column_id));
10639 self.minhash
10640 .retain(|column_id, _| keep_minhash.contains(column_id));
10641 {
10642 let learned = Arc::make_mut(&mut self.learned_range);
10643 learned.retain(|column_id, _| keep_learned.contains(column_id));
10644 }
10645 }
10646
10647 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
10648 checkpoint_current_schema(self)
10649 }
10650
10651 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
10652 self.ensure_writable()?;
10653 let previous_schema = self.schema.clone();
10654 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
10655 if let Some(schema) = schema {
10656 self.apply_altered_schema_prepared(schema);
10657 self.checkpoint_standalone_schema_change(previous_schema)?;
10658 }
10659 Ok(column)
10660 }
10661
10662 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
10663 if self.live_count == 0 {
10664 return Ok(false);
10665 }
10666 let snap = self.snapshot();
10667 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
10668 Ok(columns
10669 .first()
10670 .map(|(_, col)| col.null_count(col.len()) != 0)
10671 .unwrap_or(true))
10672 }
10673
10674 fn has_stored_versions(&self) -> bool {
10675 !self.memtable.is_empty()
10676 || !self.mutable_run.is_empty()
10677 || self.run_refs.iter().any(|r| r.row_count > 0)
10678 || !self.retiring.is_empty()
10679 }
10680
10681 pub fn add_column(
10686 &mut self,
10687 name: &str,
10688 ty: TypeId,
10689 flags: ColumnFlags,
10690 default_value: Option<crate::schema::DefaultExpr>,
10691 ) -> Result<u16> {
10692 self.add_column_with_id(name, ty, flags, default_value, None)
10693 }
10694
10695 pub fn add_column_with_id(
10696 &mut self,
10697 name: &str,
10698 ty: TypeId,
10699 flags: ColumnFlags,
10700 default_value: Option<crate::schema::DefaultExpr>,
10701 requested_id: Option<u16>,
10702 ) -> Result<u16> {
10703 self.ensure_writable()?;
10704 let previous_schema = self.schema.clone();
10705 let (column, schema) =
10706 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
10707 self.apply_altered_schema_prepared(schema);
10708 self.checkpoint_standalone_schema_change(previous_schema)?;
10709 Ok(column.id)
10710 }
10711
10712 pub(crate) fn prepare_add_column(
10713 &mut self,
10714 name: &str,
10715 ty: TypeId,
10716 flags: ColumnFlags,
10717 default_value: Option<crate::schema::DefaultExpr>,
10718 requested_id: Option<u16>,
10719 ) -> Result<(ColumnDef, Schema)> {
10720 self.ensure_writable()?;
10721 if self.schema.columns.iter().any(|c| c.name == name) {
10722 return Err(MongrelError::Schema(format!(
10723 "column {name} already exists"
10724 )));
10725 }
10726 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
10727 if self.schema.columns.iter().any(|c| c.id == id) {
10728 return Err(MongrelError::Schema(format!(
10729 "column id {id} already exists"
10730 )));
10731 }
10732 id
10733 } else {
10734 self.schema
10735 .columns
10736 .iter()
10737 .map(|c| c.id)
10738 .max()
10739 .unwrap_or(0)
10740 .checked_add(1)
10741 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
10742 };
10743 let column = ColumnDef {
10744 id,
10745 name: name.to_string(),
10746 ty,
10747 flags,
10748 default_value,
10749 embedding_source: None,
10750 };
10751 let mut next_schema = self.schema.clone();
10752 next_schema.columns.push(column.clone());
10753 next_schema.schema_id = next_schema
10754 .schema_id
10755 .checked_add(1)
10756 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10757 next_schema.validate_auto_increment()?;
10758 next_schema.validate_defaults()?;
10759 Ok((column, next_schema))
10760 }
10761
10762 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10771 self.ensure_writable()?;
10772 let cid = self
10773 .schema
10774 .columns
10775 .iter()
10776 .find(|c| c.name == column_name)
10777 .map(|c| c.id)
10778 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10779 let ty = self
10780 .schema
10781 .columns
10782 .iter()
10783 .find(|c| c.id == cid)
10784 .map(|c| c.ty.clone())
10785 .unwrap_or(TypeId::Int64);
10786 if !matches!(
10787 ty,
10788 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10789 ) {
10790 return Err(MongrelError::Schema(format!(
10791 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10792 )));
10793 }
10794 if self
10795 .schema
10796 .indexes
10797 .iter()
10798 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10799 {
10800 return Ok(()); }
10802 let previous_schema = self.schema.clone();
10803 let previous_learned_range = Arc::clone(&self.learned_range);
10804 let mut next_schema = previous_schema.clone();
10805 next_schema.indexes.push(IndexDef {
10806 name: format!("{}_learned_range", column_name),
10807 column_id: cid,
10808 kind: IndexKind::LearnedRange,
10809 predicate: None,
10810 options: Default::default(),
10811 });
10812 next_schema.schema_id = next_schema
10813 .schema_id
10814 .checked_add(1)
10815 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10816 self.apply_altered_schema_prepared(next_schema);
10817 if let Err(error) = self.build_learned_ranges() {
10818 self.apply_altered_schema_prepared(previous_schema);
10819 self.learned_range = previous_learned_range;
10820 return Err(error);
10821 }
10822 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10823 if !matches!(
10824 &error,
10825 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10826 ) {
10827 self.learned_range = previous_learned_range;
10828 }
10829 return Err(error);
10830 }
10831 Ok(())
10832 }
10833
10834 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10835 let mut schema_published = false;
10836 let schema_result = match self._root_guard.as_deref() {
10837 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10838 schema_published = true;
10839 }),
10840 None => write_schema_with_after(&self.dir, &self.schema, || {
10841 schema_published = true;
10842 }),
10843 };
10844 if schema_result.is_err() && !schema_published {
10845 self.apply_altered_schema_prepared(previous_schema);
10846 return schema_result;
10847 }
10848
10849 let manifest_result = self.persist_manifest(self.current_epoch());
10850 match (schema_result, manifest_result) {
10851 (_, Ok(())) => Ok(()),
10852 (Ok(()), Err(error)) => {
10853 self.poison_after_maintenance_publish_failure();
10854 Err(MongrelError::DurableCommit {
10855 epoch: self.current_epoch().0,
10856 message: format!(
10857 "schema is durable but matching manifest publication failed: {error}"
10858 ),
10859 })
10860 }
10861 (Err(schema_error), Err(manifest_error)) => {
10862 self.poison_after_maintenance_publish_failure();
10863 Err(MongrelError::CommitOutcomeUnknown {
10864 epoch: self.current_epoch().0,
10865 message: format!(
10866 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10867 ),
10868 })
10869 }
10870 }
10871 }
10872
10873 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10876 self.sync_byte_threshold = threshold;
10877 if let WalSink::Private(w) = &mut self.wal {
10878 w.set_sync_byte_threshold(threshold);
10879 }
10880 }
10881
10882 pub fn page_cache_flush(&self) {
10886 self.page_cache.flush_to_disk();
10887 }
10888
10889 pub fn page_cache_len(&self) -> usize {
10891 self.page_cache.len()
10892 }
10893
10894 pub fn decoded_cache_len(&self) -> usize {
10897 self.decoded_cache.len()
10898 }
10899
10900 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10903 self.memtable.drain_sorted()
10904 }
10905
10906 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10907 self.runs_dir().join(format!("r-{run_id}.sr"))
10908 }
10909
10910 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10911 match self.runs_root.as_deref() {
10912 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10913 None => Ok(None),
10914 }
10915 }
10916
10917 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10918 match self.runs_root.as_deref() {
10919 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10920 None => Ok(None),
10921 }
10922 }
10923
10924 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10925 match self.runs_root.as_deref() {
10926 Some(root) => match root.remove_file(name) {
10927 Ok(()) => Ok(()),
10928 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10929 Err(error) => Err(error.into()),
10930 },
10931 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10932 Ok(()) => Ok(()),
10933 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10934 Err(error) => Err(error.into()),
10935 },
10936 }
10937 }
10938
10939 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10940 match self.runs_root.as_deref() {
10941 Some(root) => root
10942 .rename_file_new(source, destination)
10943 .map_err(Into::into),
10944 None => crate::durable_file::rename(
10945 &self.runs_dir().join(source),
10946 &self.runs_dir().join(destination),
10947 )
10948 .map_err(Into::into),
10949 }
10950 }
10951
10952 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10953 self.run_refs.iter().map(|run| run.run_id)
10954 }
10955
10956 pub(crate) fn table_dir(&self) -> &Path {
10957 &self.dir
10958 }
10959
10960 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10961 &self.schema
10962 }
10963
10964 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10965 let id = self.next_run_id;
10966 self.next_run_id = self
10967 .next_run_id
10968 .checked_add(1)
10969 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10970 Ok(id)
10971 }
10972
10973 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10974 self.run_refs.push(run_ref);
10975 }
10976
10977 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10987 self.retiring.push(crate::manifest::RetiredRun {
10988 run_id,
10989 retire_epoch,
10990 });
10991 }
10992
10993 pub(crate) fn reap_retiring(
10997 &mut self,
10998 min_active: Epoch,
10999 backup_pinned: &std::collections::HashSet<u128>,
11000 ) -> Result<usize> {
11001 if self.retiring.is_empty() {
11002 return Ok(0);
11003 }
11004 let mut reaped = 0;
11005 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
11006 for r in std::mem::take(&mut self.retiring) {
11012 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
11013 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
11014 reaped += 1;
11015 } else {
11016 kept.push(r);
11017 }
11018 }
11019 self.retiring = kept;
11020 if reaped > 0 {
11021 self.persist_manifest(self.current_epoch())?;
11022 }
11023 Ok(reaped)
11024 }
11025
11026 pub(crate) fn has_reapable_retiring(
11027 &self,
11028 min_active: Epoch,
11029 backup_pinned: &std::collections::HashSet<u128>,
11030 ) -> bool {
11031 self.retiring
11032 .iter()
11033 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
11034 }
11035
11036 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
11037 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
11038 return false;
11039 }
11040 self.live_count = self.live_count.saturating_add(run_ref.row_count);
11041 self.run_refs.push(run_ref);
11042 self.indexes_complete = false;
11043 true
11044 }
11045
11046 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
11047 self.kek.as_ref()
11048 }
11049
11050 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
11051 let mut reader = match self.runs_root.as_deref() {
11052 Some(root) => RunReader::open_file_with_cache(
11053 root.open_regular(format!("r-{run_id}.sr"))?,
11054 self.schema.clone(),
11055 self.kek.clone(),
11056 Some(self.page_cache.clone()),
11057 Some(self.decoded_cache.clone()),
11058 self.table_id,
11059 Some(&self.verified_runs),
11060 None,
11061 )?,
11062 None => RunReader::open_with_cache(
11063 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
11064 self.schema.clone(),
11065 self.kek.clone(),
11066 Some(self.page_cache.clone()),
11067 Some(self.decoded_cache.clone()),
11068 self.table_id,
11069 Some(&self.verified_runs),
11070 )?,
11071 };
11072 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
11076 reader.set_uniform_epoch(Epoch(rr.epoch_created));
11077 }
11078 Ok(reader)
11079 }
11080
11081 pub(crate) fn run_refs(&self) -> &[RunRef] {
11082 &self.run_refs
11083 }
11084
11085 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
11086 self.retiring.iter().map(|run| run.run_id)
11087 }
11088
11089 pub(crate) fn runs_dir(&self) -> PathBuf {
11090 self.runs_root
11091 .as_deref()
11092 .and_then(|root| root.io_path().ok())
11093 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
11094 }
11095
11096 pub(crate) fn wal_dir(&self) -> PathBuf {
11097 self.dir.join(WAL_DIR)
11098 }
11099
11100 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
11101 self.run_refs = refs;
11102 }
11103
11104 pub(crate) fn compaction_zstd_level(&self) -> i32 {
11105 self.compaction_zstd_level
11106 }
11107
11108 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
11109 self.kek.clone()
11110 }
11111
11112 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
11116 self.kek.as_ref().map(|k| k.derive_idx_key())
11117 }
11118
11119 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
11123 self.kek.as_ref().map(|k| *k.derive_meta_key())
11124 }
11125
11126 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
11129 self.column_keys
11130 .iter()
11131 .map(|(&id, &(_, scheme))| (id, scheme))
11132 .collect()
11133 }
11134
11135 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
11140 self.tokenize_value_enc(column_id, v)
11141 }
11142
11143 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
11144 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11145 let (key, scheme) = self.column_keys.get(&column_id)?;
11146 let token: Vec<u8> = match (*scheme, v) {
11147 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
11148 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
11149 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
11150 _ => hmac_token(key, &v.encode_key()).to_vec(),
11151 };
11152 Some(Value::Bytes(token))
11153 }
11154
11155 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
11157 self.index_lookup_key_bytes(column_id, &v.encode_key())
11158 }
11159
11160 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
11163 {
11164 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
11165 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
11166 if *scheme == SCHEME_HMAC_EQ {
11167 return hmac_token(key, encoded).to_vec();
11168 }
11169 }
11170 }
11171 let _ = column_id;
11172 encoded.to_vec()
11173 }
11174}
11175
11176fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
11177 let columnar::NativeColumn::Int64 { data, validity } = col else {
11178 return false;
11179 };
11180 if data.len() < n || !columnar::all_non_null(validity, n) {
11181 return false;
11182 }
11183 data.iter()
11184 .take(n)
11185 .zip(data.iter().skip(1))
11186 .all(|(a, b)| a < b)
11187}
11188
11189#[derive(Debug, Clone)]
11193pub struct ColumnStat {
11194 pub min: Option<Value>,
11195 pub max: Option<Value>,
11196 pub null_count: u64,
11197}
11198
11199#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11201pub enum NativeAgg {
11202 Count,
11203 Sum,
11204 Min,
11205 Max,
11206 Avg,
11207}
11208
11209#[derive(Debug, Clone, PartialEq)]
11211pub enum NativeAggResult {
11212 Count(u64),
11213 Int(i64),
11214 Float(f64),
11215 Null,
11217}
11218
11219#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11221pub enum ApproxAgg {
11222 Count,
11223 Sum,
11224 Avg,
11225}
11226
11227#[derive(Debug, Clone)]
11231pub struct ApproxResult {
11232 pub point: f64,
11234 pub ci_low: f64,
11236 pub ci_high: f64,
11238 pub n_population: u64,
11240 pub n_sample_live: usize,
11242 pub n_passing: usize,
11244}
11245
11246#[derive(Debug, Clone, PartialEq)]
11251pub enum AggState {
11252 Count(u64),
11254 SumI {
11256 sum: i128,
11257 count: u64,
11258 },
11259 SumF {
11261 sum: f64,
11262 count: u64,
11263 },
11264 AvgI {
11266 sum: i128,
11267 count: u64,
11268 },
11269 AvgF {
11271 sum: f64,
11272 count: u64,
11273 },
11274 MinI(i64),
11276 MaxI(i64),
11277 MinF(f64),
11279 MaxF(f64),
11280 Empty,
11282}
11283
11284impl AggState {
11285 pub fn merge(self, other: AggState) -> AggState {
11287 use AggState::*;
11288 match (self, other) {
11289 (Empty, x) | (x, Empty) => x,
11290 (Count(a), Count(b)) => Count(a + b),
11291 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
11292 sum: sa + sb,
11293 count: ca + cb,
11294 },
11295 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
11296 sum: sa + sb,
11297 count: ca + cb,
11298 },
11299 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
11300 sum: sa + sb,
11301 count: ca + cb,
11302 },
11303 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
11304 sum: sa + sb,
11305 count: ca + cb,
11306 },
11307 (MinI(a), MinI(b)) => MinI(a.min(b)),
11308 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
11309 (MinF(a), MinF(b)) => MinF(a.min(b)),
11310 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
11311 _ => Empty, }
11313 }
11314
11315 pub fn point(&self) -> Option<f64> {
11317 match self {
11318 AggState::Count(n) => Some(*n as f64),
11319 AggState::SumI { sum, .. } => Some(*sum as f64),
11320 AggState::SumF { sum, .. } => Some(*sum),
11321 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
11322 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
11323 AggState::MinI(n) => Some(*n as f64),
11324 AggState::MaxI(n) => Some(*n as f64),
11325 AggState::MinF(n) => Some(*n),
11326 AggState::MaxF(n) => Some(*n),
11327 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
11328 }
11329 }
11330
11331 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
11335 let is_float = matches!(ty, Some(TypeId::Float64));
11336 match (agg, result) {
11337 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
11338 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
11339 sum: x as i128,
11340 count: 1, },
11342 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
11343 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
11344 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
11345 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
11346 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
11347 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
11348 (NativeAgg::Count, _) => AggState::Empty,
11349 (_, NativeAggResult::Null) => AggState::Empty,
11350 _ => {
11351 let _ = is_float;
11352 AggState::Empty
11353 }
11354 }
11355 }
11356}
11357
11358#[derive(Debug, Clone)]
11361pub struct CachedAgg {
11362 pub state: AggState,
11363 pub watermark: u64,
11364 pub epoch: u64,
11365}
11366
11367#[derive(Debug, Clone)]
11369pub struct IncrementalAggResult {
11370 pub state: AggState,
11372 pub incremental: bool,
11375 pub delta_rows: u64,
11377}
11378
11379fn agg_state_from_rows(
11383 rows: &[Row],
11384 conditions: &[crate::query::Condition],
11385 index_sets: &[RowIdSet],
11386 column: Option<u16>,
11387 agg: NativeAgg,
11388 schema: &Schema,
11389 control: Option<&crate::ExecutionControl>,
11390) -> Result<AggState> {
11391 let mut count: u64 = 0;
11392 let mut sum_i: i128 = 0;
11393 let mut sum_f: f64 = 0.0;
11394 let mut mn_i: i64 = i64::MAX;
11395 let mut mx_i: i64 = i64::MIN;
11396 let mut mn_f: f64 = f64::INFINITY;
11397 let mut mx_f: f64 = f64::NEG_INFINITY;
11398 let mut saw_int = false;
11399 let mut saw_float = false;
11400 for (index, r) in rows.iter().enumerate() {
11401 execution_checkpoint(control, index)?;
11402 if !conditions
11403 .iter()
11404 .all(|c| condition_matches_row(c, r, schema))
11405 {
11406 continue;
11407 }
11408 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
11409 continue;
11410 }
11411 match agg {
11412 NativeAgg::Count => match column {
11413 None => count += 1,
11415 Some(cid) => match r.columns.get(&cid) {
11418 None | Some(Value::Null) => {}
11419 Some(_) => count += 1,
11420 },
11421 },
11422 _ => match column.and_then(|cid| r.columns.get(&cid)) {
11423 Some(Value::Int64(n)) => {
11424 count += 1;
11425 sum_i += *n as i128;
11426 mn_i = mn_i.min(*n);
11427 mx_i = mx_i.max(*n);
11428 saw_int = true;
11429 }
11430 Some(Value::Float64(f)) => {
11431 count += 1;
11432 sum_f += f;
11433 mn_f = mn_f.min(*f);
11434 mx_f = mx_f.max(*f);
11435 saw_float = true;
11436 }
11437 _ => {}
11438 },
11439 }
11440 }
11441 Ok(match agg {
11442 NativeAgg::Count => {
11443 if count == 0 {
11444 AggState::Empty
11445 } else {
11446 AggState::Count(count)
11447 }
11448 }
11449 NativeAgg::Sum => {
11450 if count == 0 {
11451 AggState::Empty
11452 } else if saw_int {
11453 AggState::SumI { sum: sum_i, count }
11454 } else {
11455 AggState::SumF { sum: sum_f, count }
11456 }
11457 }
11458 NativeAgg::Avg => {
11459 if count == 0 {
11460 AggState::Empty
11461 } else if saw_int {
11462 AggState::AvgI { sum: sum_i, count }
11463 } else {
11464 AggState::AvgF { sum: sum_f, count }
11465 }
11466 }
11467 NativeAgg::Min => {
11468 if !saw_int && !saw_float {
11469 AggState::Empty
11470 } else if saw_int {
11471 AggState::MinI(mn_i)
11472 } else {
11473 AggState::MinF(mn_f)
11474 }
11475 }
11476 NativeAgg::Max => {
11477 if !saw_int && !saw_float {
11478 AggState::Empty
11479 } else if saw_int {
11480 AggState::MaxI(mx_i)
11481 } else {
11482 AggState::MaxF(mx_f)
11483 }
11484 }
11485 })
11486}
11487
11488fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
11492 use crate::query::Condition;
11493 match c {
11494 Condition::Pk(key) => match schema.primary_key() {
11495 Some(pk) => row
11496 .columns
11497 .get(&pk.id)
11498 .map(|v| v.encode_key() == *key)
11499 .unwrap_or(false),
11500 None => false,
11501 },
11502 Condition::BitmapEq { column_id, value } => row
11503 .columns
11504 .get(column_id)
11505 .map(|v| v.encode_key() == *value)
11506 .unwrap_or(false),
11507 Condition::BitmapIn { column_id, values } => {
11508 let key = row.columns.get(column_id).map(|v| v.encode_key());
11509 match key {
11510 Some(k) => values.contains(&k),
11511 None => false,
11512 }
11513 }
11514 Condition::BytesPrefix { column_id, prefix } => row
11515 .columns
11516 .get(column_id)
11517 .map(|v| v.encode_key().starts_with(prefix))
11518 .unwrap_or(false),
11519 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
11520 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
11521 _ => false,
11522 },
11523 Condition::RangeF64 {
11524 column_id,
11525 lo,
11526 lo_inclusive,
11527 hi,
11528 hi_inclusive,
11529 } => match row.columns.get(column_id) {
11530 Some(Value::Float64(n)) => {
11531 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
11532 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
11533 lo_ok && hi_ok
11534 }
11535 _ => false,
11536 },
11537 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
11538 Some(Value::Bytes(b)) => {
11539 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
11540 }
11541 _ => false,
11542 },
11543 Condition::FmContainsAll {
11544 column_id,
11545 patterns,
11546 } => match row.columns.get(column_id) {
11547 Some(Value::Bytes(b)) => patterns
11548 .iter()
11549 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
11550 _ => false,
11551 },
11552 Condition::Ann { .. }
11553 | Condition::SparseMatch { .. }
11554 | Condition::MinHashSimilar { .. } => true,
11555 Condition::IsNull { column_id } => {
11556 matches!(row.columns.get(column_id), Some(Value::Null) | None)
11557 }
11558 Condition::IsNotNull { column_id } => {
11559 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
11560 }
11561 }
11562}
11563
11564fn as_f64(v: Option<&Value>) -> Option<f64> {
11566 match v {
11567 Some(Value::Int64(n)) => Some(*n as f64),
11568 Some(Value::Float64(f)) => Some(*f),
11569 _ => None,
11570 }
11571}
11572
11573fn accumulate_int(
11577 cursor: &mut dyn crate::cursor::Cursor,
11578 control: Option<&crate::ExecutionControl>,
11579) -> Result<(u64, i128, i64, i64)> {
11580 let mut count: u64 = 0;
11581 let mut sum: i128 = 0;
11582 let mut mn: i64 = i64::MAX;
11583 let mut mx: i64 = i64::MIN;
11584 while let Some(cols) = cursor.next_batch()? {
11585 execution_checkpoint(control, 0)?;
11586 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
11587 if crate::columnar::all_non_null(validity, data.len()) {
11588 count += data.len() as u64;
11590 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11591 execution_checkpoint(control, chunk_index * 1024)?;
11592 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
11593 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
11594 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
11595 }
11596 } else {
11597 for (i, &v) in data.iter().enumerate() {
11598 execution_checkpoint(control, i)?;
11599 if crate::columnar::validity_bit(validity, i) {
11600 count += 1;
11601 sum += v as i128;
11602 mn = mn.min(v);
11603 mx = mx.max(v);
11604 }
11605 }
11606 }
11607 }
11608 }
11609 Ok((count, sum, mn, mx))
11610}
11611
11612fn accumulate_float(
11614 cursor: &mut dyn crate::cursor::Cursor,
11615 control: Option<&crate::ExecutionControl>,
11616) -> Result<(u64, f64, f64, f64)> {
11617 let mut count: u64 = 0;
11618 let mut sum: f64 = 0.0;
11619 let mut mn: f64 = f64::INFINITY;
11620 let mut mx: f64 = f64::NEG_INFINITY;
11621 while let Some(cols) = cursor.next_batch()? {
11622 execution_checkpoint(control, 0)?;
11623 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
11624 if crate::columnar::all_non_null(validity, data.len()) {
11625 count += data.len() as u64;
11626 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
11627 execution_checkpoint(control, chunk_index * 1024)?;
11628 sum += chunk.iter().sum::<f64>();
11629 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
11630 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
11631 }
11632 } else {
11633 for (i, &v) in data.iter().enumerate() {
11634 execution_checkpoint(control, i)?;
11635 if crate::columnar::validity_bit(validity, i) {
11636 count += 1;
11637 sum += v;
11638 mn = mn.min(v);
11639 mx = mx.max(v);
11640 }
11641 }
11642 }
11643 }
11644 }
11645 Ok((count, sum, mn, mx))
11646}
11647
11648#[inline]
11649fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
11650 if index.is_multiple_of(256) {
11651 control
11652 .map(crate::ExecutionControl::checkpoint)
11653 .transpose()?;
11654 }
11655 Ok(())
11656}
11657
11658fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
11659 if count == 0 && !matches!(agg, NativeAgg::Count) {
11660 return NativeAggResult::Null;
11661 }
11662 match agg {
11663 NativeAgg::Count => NativeAggResult::Count(count),
11664 NativeAgg::Sum => match sum.try_into() {
11667 Ok(v) => NativeAggResult::Int(v),
11668 Err(_) => NativeAggResult::Null,
11669 },
11670 NativeAgg::Min => NativeAggResult::Int(mn),
11671 NativeAgg::Max => NativeAggResult::Int(mx),
11672 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
11673 }
11674}
11675
11676fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
11677 if count == 0 && !matches!(agg, NativeAgg::Count) {
11678 return NativeAggResult::Null;
11679 }
11680 match agg {
11681 NativeAgg::Count => NativeAggResult::Count(count),
11682 NativeAgg::Sum => NativeAggResult::Float(sum),
11683 NativeAgg::Min => NativeAggResult::Float(mn),
11684 NativeAgg::Max => NativeAggResult::Float(mx),
11685 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
11686 }
11687}
11688
11689fn agg_int(
11692 stats: &[crate::page::PageStat],
11693 decode: fn(Option<&[u8]>) -> Option<i64>,
11694) -> Option<(Option<i64>, Option<i64>, u64)> {
11695 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
11696 let mut any = false;
11697 for s in stats {
11698 if let Some(v) = decode(s.min.as_deref()) {
11699 mn = mn.min(v);
11700 any = true;
11701 }
11702 if let Some(v) = decode(s.max.as_deref()) {
11703 mx = mx.max(v);
11704 any = true;
11705 }
11706 nulls += s.null_count;
11707 }
11708 any.then_some((Some(mn), Some(mx), nulls))
11709}
11710
11711fn agg_float(
11713 stats: &[crate::page::PageStat],
11714 decode: fn(Option<&[u8]>) -> Option<f64>,
11715) -> Option<(Option<f64>, Option<f64>, u64)> {
11716 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
11717 let mut any = false;
11718 for s in stats {
11719 if let Some(v) = decode(s.min.as_deref()) {
11720 mn = mn.min(v);
11721 any = true;
11722 }
11723 if let Some(v) = decode(s.max.as_deref()) {
11724 mx = mx.max(v);
11725 any = true;
11726 }
11727 nulls += s.null_count;
11728 }
11729 any.then_some((Some(mn), Some(mx), nulls))
11730}
11731
11732type SecondaryIndexes = (
11734 HashMap<u16, BitmapIndex>,
11735 HashMap<u16, AnnIndex>,
11736 HashMap<u16, FmIndex>,
11737 HashMap<u16, SparseIndex>,
11738 HashMap<u16, MinHashIndex>,
11739);
11740
11741fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
11742 let mut bitmap = HashMap::new();
11743 let mut ann = HashMap::new();
11744 let mut fm = HashMap::new();
11745 let mut sparse = HashMap::new();
11746 let mut minhash = HashMap::new();
11747 for idef in &schema.indexes {
11748 match idef.kind {
11749 IndexKind::Bitmap => {
11750 bitmap.insert(idef.column_id, BitmapIndex::new());
11751 }
11752 IndexKind::Ann => {
11753 let dim = schema
11754 .columns
11755 .iter()
11756 .find(|c| c.id == idef.column_id)
11757 .and_then(|c| match c.ty {
11758 TypeId::Embedding { dim } => Some(dim as usize),
11759 _ => None,
11760 })
11761 .unwrap_or(0);
11762 let options = idef.options.ann.clone().unwrap_or_default();
11763 ann.insert(
11764 idef.column_id,
11765 AnnIndex::with_full_options(
11766 dim,
11767 options.m,
11768 options.ef_construction,
11769 options.ef_search,
11770 &options,
11771 ),
11772 );
11773 }
11774 IndexKind::FmIndex => {
11775 fm.insert(idef.column_id, FmIndex::new());
11776 }
11777 IndexKind::Sparse => {
11778 sparse.insert(idef.column_id, SparseIndex::new());
11779 }
11780 IndexKind::MinHash => {
11781 let options = idef.options.minhash.clone().unwrap_or_default();
11782 minhash.insert(
11783 idef.column_id,
11784 MinHashIndex::with_options(options.permutations, options.bands),
11785 );
11786 }
11787 _ => {}
11788 }
11789 }
11790 (bitmap, ann, fm, sparse, minhash)
11791}
11792
11793const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11794 | ColumnFlags::AUTO_INCREMENT
11795 | ColumnFlags::ENCRYPTED
11796 | ColumnFlags::ENCRYPTED_INDEXABLE
11797 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11798
11799fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11800 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11801 return Err(MongrelError::Schema(
11802 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11803 ));
11804 }
11805 Ok(())
11806}
11807
11808fn validate_alter_column_type(
11809 schema: &Schema,
11810 old: &ColumnDef,
11811 next: &ColumnDef,
11812 has_stored_versions: bool,
11813) -> Result<()> {
11814 if old.ty == next.ty {
11815 return Ok(());
11816 }
11817 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11818 return Err(MongrelError::Schema(format!(
11819 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11820 old.name
11821 )));
11822 }
11823 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11824 return Ok(());
11825 }
11826 Err(MongrelError::Schema(format!(
11827 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11828 old.ty, next.ty
11829 )))
11830}
11831
11832fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11833 matches!(
11834 (old, new),
11835 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11836 )
11837}
11838
11839fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11845 let mut prev: Option<i64> = None;
11846 for r in rows {
11847 match r.columns.get(&pk_id) {
11848 Some(Value::Int64(v)) => {
11849 if prev.is_some_and(|p| p >= *v) {
11850 return false;
11851 }
11852 prev = Some(*v);
11853 }
11854 _ => return false,
11855 }
11856 }
11857 true
11858}
11859
11860#[allow(clippy::too_many_arguments)]
11861fn index_into(
11862 schema: &Schema,
11863 row: &Row,
11864 hot: &mut HotIndex,
11865 bitmap: &mut HashMap<u16, BitmapIndex>,
11866 ann: &mut HashMap<u16, AnnIndex>,
11867 fm: &mut HashMap<u16, FmIndex>,
11868 sparse: &mut HashMap<u16, SparseIndex>,
11869 minhash: &mut HashMap<u16, MinHashIndex>,
11870) {
11871 for idef in &schema.indexes {
11872 let Some(val) = row.columns.get(&idef.column_id) else {
11873 continue;
11874 };
11875 match idef.kind {
11876 IndexKind::Bitmap => {
11877 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11878 b.insert(val.encode_key(), row.row_id);
11879 }
11880 }
11881 IndexKind::Ann => {
11882 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
11883 a.insert_validated(v, row.row_id);
11884 }
11885 }
11886 IndexKind::FmIndex => {
11887 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11888 f.insert(b.clone(), row.row_id);
11889 }
11890 }
11891 IndexKind::Sparse => {
11892 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11893 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11896 s.insert(&terms, row.row_id);
11897 }
11898 }
11899 }
11900 IndexKind::MinHash => {
11901 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11902 let tokens = crate::index::token_hashes_from_bytes(b);
11905 mh.insert(&tokens, row.row_id);
11906 }
11907 }
11908 _ => {}
11909 }
11910 }
11911 if let Some(pk_col) = schema.primary_key() {
11912 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11913 hot.insert(pk_val.encode_key(), row.row_id);
11914 }
11915 }
11916}
11917
11918#[allow(clippy::too_many_arguments)]
11921fn index_into_single(
11922 idef: &IndexDef,
11923 _schema: &Schema,
11924 row: &Row,
11925 _hot: &mut HotIndex,
11926 bitmap: &mut HashMap<u16, BitmapIndex>,
11927 ann: &mut HashMap<u16, AnnIndex>,
11928 fm: &mut HashMap<u16, FmIndex>,
11929 sparse: &mut HashMap<u16, SparseIndex>,
11930 minhash: &mut HashMap<u16, MinHashIndex>,
11931) {
11932 let Some(val) = row.columns.get(&idef.column_id) else {
11933 return;
11934 };
11935 match idef.kind {
11936 IndexKind::Bitmap => {
11937 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11938 b.insert(val.encode_key(), row.row_id);
11939 }
11940 }
11941 IndexKind::Ann => {
11942 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
11943 a.insert_validated(v, row.row_id);
11944 }
11945 }
11946 IndexKind::FmIndex => {
11947 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11948 f.insert(b.clone(), row.row_id);
11949 }
11950 }
11951 IndexKind::Sparse => {
11952 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11953 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11954 s.insert(&terms, row.row_id);
11955 }
11956 }
11957 }
11958 IndexKind::MinHash => {
11959 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11960 let tokens = crate::index::token_hashes_from_bytes(b);
11961 mh.insert(&tokens, row.row_id);
11962 }
11963 }
11964 _ => {}
11965 }
11966}
11967
11968fn eval_partial_predicate(
11974 pred: &str,
11975 columns_map: &HashMap<u16, &Value>,
11976 name_to_id: &HashMap<&str, u16>,
11977) -> bool {
11978 let lower = pred.trim().to_ascii_lowercase();
11979 if let Some(rest) = lower.strip_suffix(" is not null") {
11981 let col_name = rest.trim();
11982 if let Some(col_id) = name_to_id.get(col_name) {
11983 return columns_map
11984 .get(col_id)
11985 .is_some_and(|v| !matches!(v, Value::Null));
11986 }
11987 }
11988 if let Some(rest) = lower.strip_suffix(" is null") {
11990 let col_name = rest.trim();
11991 if let Some(col_id) = name_to_id.get(col_name) {
11992 return columns_map
11993 .get(col_id)
11994 .is_none_or(|v| matches!(v, Value::Null));
11995 }
11996 }
11997 true
12000}
12001
12002#[allow(dead_code)]
12008fn bulk_index_key(
12009 column_keys: &HashMap<u16, ([u8; 32], u8)>,
12010 column_id: u16,
12011 ty: TypeId,
12012 col: &columnar::NativeColumn,
12013 i: usize,
12014) -> Option<Vec<u8>> {
12015 let encoded = columnar::encode_key_native(ty, col, i)?;
12016 {
12017 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12018 if let Some((key, scheme)) = column_keys.get(&column_id) {
12019 return Some(match (*scheme, col) {
12020 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
12021 (_, columnar::NativeColumn::Int64 { data, .. }) => {
12022 ope_token_i64(key, data[i]).to_vec()
12023 }
12024 (_, columnar::NativeColumn::Float64 { data, .. }) => {
12025 ope_token_f64(key, data[i]).to_vec()
12026 }
12027 _ => hmac_token(key, &encoded).to_vec(),
12028 });
12029 }
12030 }
12031 Some(encoded)
12032}
12033
12034pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
12035 write_schema_with_after(dir, schema, || {})
12036}
12037
12038pub(crate) fn write_schema_durable(
12039 root: &crate::durable_file::DurableRoot,
12040 schema: &Schema,
12041) -> Result<()> {
12042 write_schema_durable_with_after(root, schema, || {})
12043}
12044
12045fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
12046where
12047 F: FnOnce(),
12048{
12049 let json = serde_json::to_string_pretty(schema)
12050 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
12051 crate::durable_file::write_atomic_with_after(
12052 &dir.join(SCHEMA_FILENAME),
12053 json.as_bytes(),
12054 after_publish,
12055 )?;
12056 Ok(())
12057}
12058
12059fn write_schema_durable_with_after<F>(
12060 root: &crate::durable_file::DurableRoot,
12061 schema: &Schema,
12062 after_publish: F,
12063) -> Result<()>
12064where
12065 F: FnOnce(),
12066{
12067 let json = serde_json::to_string_pretty(schema)
12068 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
12069 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
12070 Ok(())
12071}
12072
12073fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
12074 let mut schema_published = false;
12075 let schema_result = match table._root_guard.as_deref() {
12076 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
12077 schema_published = true;
12078 }),
12079 None => write_schema_with_after(&table.dir, &table.schema, || {
12080 schema_published = true;
12081 }),
12082 };
12083 if schema_result.is_err() && !schema_published {
12084 return schema_result;
12085 }
12086 match table.persist_manifest(table.current_epoch()) {
12087 Ok(()) => Ok(()),
12088 Err(manifest_error) => Err(match schema_result {
12089 Ok(()) => manifest_error,
12090 Err(schema_error) => MongrelError::Other(format!(
12091 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12092 )),
12093 }),
12094 }
12095}
12096
12097fn read_schema(dir: &Path) -> Result<Schema> {
12098 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
12099 read_schema_file(file)
12100}
12101
12102fn read_schema_file(file: std::fs::File) -> Result<Schema> {
12103 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
12104 use std::io::Read;
12105
12106 let length = file.metadata()?.len();
12107 if length > MAX_SCHEMA_BYTES {
12108 return Err(MongrelError::ResourceLimitExceeded {
12109 resource: "schema bytes",
12110 requested: usize::try_from(length).unwrap_or(usize::MAX),
12111 limit: MAX_SCHEMA_BYTES as usize,
12112 });
12113 }
12114 let mut bytes = Vec::with_capacity(length as usize);
12115 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
12116 if bytes.len() as u64 != length {
12117 return Err(MongrelError::Schema(
12118 "schema length changed while reading".into(),
12119 ));
12120 }
12121 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
12122}
12123
12124fn preflight_standalone_open(
12125 dir: &Path,
12126 runs_root: Option<&crate::durable_file::DurableRoot>,
12127 idx_root: Option<&crate::durable_file::DurableRoot>,
12128 manifest: &Manifest,
12129 schema: &Schema,
12130 records: &[crate::wal::Record],
12131 kek: Option<Arc<Kek>>,
12132) -> Result<()> {
12133 crate::wal::validate_shared_transaction_framing(records)?;
12134 if manifest.schema_id > schema.schema_id
12135 || manifest.flushed_epoch > manifest.current_epoch
12136 || manifest.global_idx_epoch > manifest.current_epoch
12137 || manifest.next_row_id == u64::MAX
12138 || manifest.auto_inc_next < 0
12139 || manifest.auto_inc_next == i64::MAX
12140 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
12141 {
12142 return Err(MongrelError::InvalidArgument(
12143 "manifest counters or schema identity are invalid".into(),
12144 ));
12145 }
12146 let mut run_ids = HashSet::new();
12147 let mut maximum_row_id = None::<u64>;
12148 for run in &manifest.runs {
12149 if run.run_id >= u64::MAX as u128
12150 || !run_ids.insert(run.run_id)
12151 || run.epoch_created > manifest.current_epoch
12152 {
12153 return Err(MongrelError::InvalidArgument(
12154 "manifest contains an invalid or duplicate active run".into(),
12155 ));
12156 }
12157 let mut reader = match runs_root {
12158 Some(root) => RunReader::open_file(
12159 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
12160 schema.clone(),
12161 kek.clone(),
12162 )?,
12163 None => RunReader::open(
12164 dir.join(RUNS_DIR)
12165 .join(format!("r-{}.sr", run.run_id as u64)),
12166 schema.clone(),
12167 kek.clone(),
12168 )?,
12169 };
12170 let header = reader.header();
12171 if header.run_id != run.run_id
12172 || header.level != run.level
12173 || header.row_count != run.row_count
12174 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
12175 || header.is_uniform_epoch() && header.epoch_created != 0
12176 || header.schema_id > schema.schema_id
12177 {
12178 return Err(MongrelError::InvalidArgument(format!(
12179 "run {} differs from its manifest",
12180 run.run_id
12181 )));
12182 }
12183 if header.row_count != 0 {
12184 maximum_row_id = Some(
12185 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
12186 );
12187 }
12188 reader.validate_all_pages()?;
12189 }
12190 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
12191 return Err(MongrelError::InvalidArgument(
12192 "manifest next_row_id does not advance beyond persisted rows".into(),
12193 ));
12194 }
12195 for run in &manifest.retiring {
12196 if run.run_id >= u64::MAX as u128
12197 || run.retire_epoch > manifest.current_epoch
12198 || !run_ids.insert(run.run_id)
12199 {
12200 return Err(MongrelError::InvalidArgument(
12201 "manifest contains an invalid or duplicate retired run".into(),
12202 ));
12203 }
12204 }
12205 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
12206 match idx_root {
12207 Some(root) => {
12208 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
12209 }
12210 None => {
12211 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
12212 }
12213 }
12214
12215 let committed = records
12216 .iter()
12217 .filter_map(|record| match record.op {
12218 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
12219 _ => None,
12220 })
12221 .collect::<HashMap<_, _>>();
12222 for record in records {
12223 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
12224 continue;
12225 };
12226 match &record.op {
12227 Op::Put { table_id, rows } => {
12228 if *table_id != manifest.table_id {
12229 return Err(MongrelError::CorruptWal {
12230 offset: record.seq.0,
12231 reason: format!(
12232 "private WAL record references table {table_id}, expected {}",
12233 manifest.table_id
12234 ),
12235 });
12236 }
12237 let rows: Vec<Row> =
12238 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
12239 offset: record.seq.0,
12240 reason: format!("committed Put payload could not be decoded: {error}"),
12241 })?;
12242 for row in rows {
12243 if row.deleted || row.row_id.0 == u64::MAX {
12244 return Err(MongrelError::CorruptWal {
12245 offset: record.seq.0,
12246 reason: "committed Put contains an invalid row identity".into(),
12247 });
12248 }
12249 let cells = row.columns.into_iter().collect::<Vec<_>>();
12250 schema
12251 .validate_values(&cells)
12252 .map_err(|error| MongrelError::CorruptWal {
12253 offset: record.seq.0,
12254 reason: format!("committed Put violates table schema: {error}"),
12255 })?;
12256 if schema.auto_increment_column().is_some_and(|column| {
12257 matches!(
12258 cells.iter().find(|(id, _)| *id == column.id),
12259 Some((_, Value::Int64(value))) if *value == i64::MAX
12260 )
12261 }) {
12262 return Err(MongrelError::CorruptWal {
12263 offset: record.seq.0,
12264 reason: "committed Put exhausts AUTO_INCREMENT".into(),
12265 });
12266 }
12267 }
12268 }
12269 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
12270 if *table_id != manifest.table_id =>
12271 {
12272 return Err(MongrelError::CorruptWal {
12273 offset: record.seq.0,
12274 reason: format!(
12275 "private WAL record references table {table_id}, expected {}",
12276 manifest.table_id
12277 ),
12278 });
12279 }
12280 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
12281 return Err(MongrelError::CorruptWal {
12282 offset: record.seq.0,
12283 reason: "private WAL contains shared spilled-run metadata".into(),
12284 });
12285 }
12286 _ => {}
12287 }
12288 }
12289 Ok(())
12290}
12291
12292fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
12293 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
12294}
12295
12296fn wal_segment_number(path: &Path) -> Option<u64> {
12297 path.file_stem()
12298 .and_then(|stem| stem.to_str())
12299 .and_then(|stem| stem.strip_prefix("seg-"))
12300 .and_then(|number| number.parse().ok())
12301}
12302
12303fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
12304 let n = list_wal_numbers(wal_dir)?;
12305 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
12306}
12307
12308fn next_wal_number(wal_dir: &Path) -> Result<u32> {
12309 list_wal_numbers(wal_dir)?
12310 .map(|maximum| {
12311 maximum
12312 .checked_add(1)
12313 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
12314 })
12315 .unwrap_or(Ok(0))
12316}
12317
12318fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
12319 let mut max_n = None;
12320 let entries = match std::fs::read_dir(wal_dir) {
12321 Ok(entries) => entries,
12322 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
12323 Err(error) => return Err(error.into()),
12324 };
12325 for entry in entries {
12326 let entry = entry?;
12327 let fname = entry.file_name();
12328 let Some(s) = fname.to_str() else {
12329 continue;
12330 };
12331 let Some(stripped) = s.strip_prefix("seg-") else {
12332 continue;
12333 };
12334 let Some(number) = stripped.strip_suffix(".wal") else {
12335 return Err(MongrelError::CorruptWal {
12336 offset: 0,
12337 reason: format!("malformed WAL segment name {s:?}"),
12338 });
12339 };
12340 let n = number
12341 .parse::<u32>()
12342 .map_err(|_| MongrelError::CorruptWal {
12343 offset: 0,
12344 reason: format!("malformed WAL segment name {s:?}"),
12345 })?;
12346 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
12347 return Err(MongrelError::CorruptWal {
12348 offset: n as u64,
12349 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
12350 });
12351 }
12352 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
12353 }
12354 Ok(max_n)
12355}